mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-09 20:50:20 +08:00
4430 lines
164 KiB
Rust
4430 lines
164 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, 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;
|
|
pub(crate) use response::{
|
|
maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response,
|
|
maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind,
|
|
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild,
|
|
LocalVideoSyncSuccessOutcome,
|
|
};
|
|
|
|
const OPENAI_IMAGE_SYNC_PLAN_KIND: &str = "openai_image_sync";
|
|
const OPENAI_IMAGE_SYNC_DEFAULT_TOTAL_TIMEOUT_MS: u64 = 900_000;
|
|
const SYNC_EXECUTION_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60);
|
|
const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15);
|
|
const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES: &[u8] = b"\n";
|
|
const OPENAI_IMAGE_SYNC_PROGRESS_WRITE_INTERVAL: Duration = Duration::from_secs(5);
|
|
const INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE: &str = "Provider returned HTTP 200 but the Gemini response did not contain visible model output; refusing to finalize it as a successful response.";
|
|
|
|
fn elapsed_ms_since(started_at: Instant) -> u64 {
|
|
started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64
|
|
}
|
|
|
|
fn calibrated_sync_candidate_first_byte_elapsed_ms(
|
|
candidate_started_at: Instant,
|
|
result: &ExecutionResult,
|
|
) -> Option<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,
|
|
}
|
|
|
|
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(),
|
|
self.report_context.clone(),
|
|
self.request_diagnostics.clone(),
|
|
self.candidate_started_unix_ms,
|
|
self.candidate_started_at,
|
|
UsageEventType::Failed,
|
|
RequestCandidateStatus::Failed,
|
|
StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
|
|
"local_sync_attempt_aborted",
|
|
format!("Local sync attempt failed before terminal finalization: {error:?}"),
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
|
|
impl Drop for SyncAttemptTerminalGuard {
|
|
fn drop(&mut self) {
|
|
if !self.armed {
|
|
return;
|
|
}
|
|
self.armed = false;
|
|
let state = self.state.clone();
|
|
let plan = self.plan.clone();
|
|
let report_context = self.report_context.clone();
|
|
let request_diagnostics = self.request_diagnostics.clone();
|
|
let candidate_started_unix_ms = self.candidate_started_unix_ms;
|
|
let candidate_started_at = self.candidate_started_at;
|
|
if let Ok(handle) = tokio::runtime::Handle::try_current() {
|
|
handle.spawn(async move {
|
|
record_sync_attempt_forced_terminal_state(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
request_diagnostics,
|
|
candidate_started_unix_ms,
|
|
candidate_started_at,
|
|
UsageEventType::Cancelled,
|
|
RequestCandidateStatus::Cancelled,
|
|
499,
|
|
"local_sync_attempt_cancelled",
|
|
"Local sync attempt was dropped before terminal finalization, usually because the client disconnected or the request task was cancelled.",
|
|
)
|
|
.await;
|
|
});
|
|
} else {
|
|
warn!(
|
|
event_name = "local_sync_attempt_terminal_guard_no_runtime",
|
|
log_type = "ops",
|
|
request_id = %short_request_id(self.plan.request_id.as_str()),
|
|
candidate_id = ?self.plan.candidate_id,
|
|
"gateway could not finalize dropped local sync attempt because no Tokio runtime is available"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
async fn record_sync_attempt_forced_terminal_state(
|
|
state: AppState,
|
|
plan: ExecutionPlan,
|
|
report_context: Option<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 terminal_unix_ms = current_request_candidate_unix_ms();
|
|
let latency_ms = elapsed_ms_since(candidate_started_at);
|
|
record_local_request_candidate_status(
|
|
&state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: candidate_status,
|
|
status_code: Some(status_code),
|
|
error_type: Some(error_type.to_string()),
|
|
error_message: Some(error_message.clone()),
|
|
latency_ms: Some(latency_ms),
|
|
started_at_unix_ms: Some(candidate_started_unix_ms),
|
|
finished_at_unix_ms: Some(terminal_unix_ms),
|
|
},
|
|
)
|
|
.await;
|
|
|
|
if !state.usage_runtime.is_enabled() {
|
|
return;
|
|
}
|
|
|
|
let mut usage_data = build_usage_event_data_seed(&plan, report_context.as_ref());
|
|
usage_data.status_code = Some(status_code);
|
|
usage_data.error_message = Some(error_message.clone());
|
|
usage_data.error_category = Some(
|
|
match usage_event_type {
|
|
UsageEventType::Cancelled => "cancelled",
|
|
_ => "server_error",
|
|
}
|
|
.to_string(),
|
|
);
|
|
usage_data.response_time_ms = Some(latency_ms);
|
|
let error_body = json!({
|
|
"error": {
|
|
"type": error_type,
|
|
"message": error_message,
|
|
"code": status_code
|
|
}
|
|
});
|
|
usage_data.response_headers = Some(json!({"content-type": "application/json"}));
|
|
usage_data.response_body = Some(error_body.clone());
|
|
usage_data.client_response_headers = Some(json!({"content-type": "application/json"}));
|
|
usage_data.client_response_body = Some(error_body);
|
|
|
|
state
|
|
.usage_runtime
|
|
.record_terminal_event_direct(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
UsageEvent::new(usage_event_type, plan.request_id.clone(), usage_data),
|
|
)
|
|
.await;
|
|
}
|
|
|
|
impl SyncExecutionFailure {
|
|
fn from_transport(err: ExecutionRuntimeTransportError) -> Self {
|
|
let fallback_kind = match &err {
|
|
ExecutionRuntimeTransportError::UpstreamResponseTooLarge { .. } => {
|
|
Some(SyncExecutionFailureFallbackKind::UpstreamResponseTooLarge)
|
|
}
|
|
ExecutionRuntimeTransportError::UpstreamResponseDecode { .. } => {
|
|
Some(SyncExecutionFailureFallbackKind::UpstreamResponseDecode)
|
|
}
|
|
_ => None,
|
|
};
|
|
Self {
|
|
error_type: fallback_kind
|
|
.map(SyncExecutionFailureFallbackKind::error_type)
|
|
.unwrap_or("execution_runtime_unavailable"),
|
|
message: err.to_string(),
|
|
status_code: fallback_kind.map(|_| StatusCode::BAD_GATEWAY.as_u16()),
|
|
latency_ms: None,
|
|
fallback_kind,
|
|
}
|
|
}
|
|
|
|
fn image_sync_total_timeout(timeout_ms: u64, elapsed_ms: u64) -> Self {
|
|
Self {
|
|
error_type: "image_sync_total_timeout",
|
|
message: format!(
|
|
"OpenAI image sync execution exceeded total timeout of {timeout_ms}ms"
|
|
),
|
|
status_code: Some(StatusCode::GATEWAY_TIMEOUT.as_u16()),
|
|
latency_ms: Some(elapsed_ms),
|
|
fallback_kind: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
fn build_sync_execution_failure_fallback_body(
|
|
client_api_format: &str,
|
|
kind: SyncExecutionFailureFallbackKind,
|
|
) -> Value {
|
|
let message = kind.client_message();
|
|
let error_type = kind.error_type();
|
|
match crate::ai_serving::normalize_api_format_alias(client_api_format).as_str() {
|
|
"claude:messages" => json!({
|
|
"type": "error",
|
|
"error": {
|
|
"type": "upstream_error",
|
|
"message": message,
|
|
}
|
|
}),
|
|
"gemini:generate_content" => json!({
|
|
"error": {
|
|
"code": StatusCode::BAD_GATEWAY.as_u16(),
|
|
"message": message,
|
|
"status": "BAD_GATEWAY",
|
|
}
|
|
}),
|
|
_ => json!({
|
|
"error": {
|
|
"type": "upstream_error",
|
|
"message": message,
|
|
"code": error_type,
|
|
}
|
|
}),
|
|
}
|
|
}
|
|
|
|
fn build_sync_execution_failure_fallback_response(
|
|
failure: &SyncExecutionFailure,
|
|
plan: &ExecutionPlan,
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
) -> Result<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 report_context_with_diagnostics =
|
|
attach_current_request_diagnostics_and_candidate_start_timing_to_report_context(
|
|
report_context,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
);
|
|
let context_seed = build_terminal_usage_context_seed(
|
|
plan,
|
|
report_context_with_diagnostics.as_ref().or(report_context),
|
|
);
|
|
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
|
state
|
|
.usage_runtime
|
|
.record_sync_terminal(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
context_seed,
|
|
payload_seed,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
async fn record_sync_terminal_usage_and_disarm_guard(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&serde_json::Value>,
|
|
payload: &GatewaySyncReportRequest,
|
|
candidate_started_at: Instant,
|
|
candidate_first_byte_elapsed_ms: Option<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;
|
|
}
|
|
self.buffer.extend_from_slice(chunk);
|
|
let mut force_persist = false;
|
|
while let Some(block_end) = find_sse_block_end(&self.buffer) {
|
|
let block = self.buffer.drain(..block_end).collect::<Vec<_>>();
|
|
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(&err.to_string()),
|
|
},
|
|
Ok(None) => openai_image_sync_json_heartbeat_error_body(
|
|
"sync image execution ended without a local response",
|
|
),
|
|
Err(err) => openai_image_sync_json_heartbeat_error_body(&format!("{err:?}")),
|
|
}
|
|
}
|
|
|
|
fn openai_image_sync_json_heartbeat_error_body(message: &str) -> Vec<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.runtime_state(),
|
|
report_context,
|
|
payload
|
|
.client_body_json
|
|
.as_ref()
|
|
.or(payload.body_json.as_ref()),
|
|
)
|
|
.await;
|
|
}
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan,
|
|
report_context,
|
|
},
|
|
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan,
|
|
report_context,
|
|
},
|
|
LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan,
|
|
report_context,
|
|
},
|
|
LocalExecutionEffect::PoolSuccessSync { payload },
|
|
)
|
|
.await;
|
|
}
|
|
|
|
#[cfg(test)]
|
|
enum RemoteSyncFallbackOutcome {
|
|
Executed(ExecutionResult),
|
|
ClientResponse(Response<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)
|
|
.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 = err.to_string();
|
|
warn!(
|
|
event_name = "chatgpt_web_image_execution_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = %err,
|
|
"gateway ChatGPT-Web image execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some(
|
|
"chatgpt_web_image_execution_unavailable".to_string(),
|
|
),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"chatgpt_web_image_execution_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
}
|
|
}
|
|
Err(err) => {
|
|
let transport_error_message = err.to_string();
|
|
warn!(
|
|
event_name = "grok_execution_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = %err,
|
|
"gateway Grok execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some("grok_execution_unavailable".to_string()),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"grok_execution_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
}
|
|
};
|
|
#[cfg(test)]
|
|
let mut result = {
|
|
if let Some(override_fn) = state.execution_runtime_sync_override.as_ref() {
|
|
match (override_fn.0)(&plan) {
|
|
Ok(result) => result,
|
|
Err(err) => {
|
|
let transport_error_message = format!("{err:?}");
|
|
warn!(
|
|
event_name = "sync_execution_runtime_test_override_failed",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = ?err,
|
|
"gateway test sync execution override failed"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some("execution_runtime_unavailable".to_string()),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"execution_runtime_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
}
|
|
} else if state
|
|
.execution_runtime_override_base_url()
|
|
.unwrap_or_default()
|
|
.trim()
|
|
.is_empty()
|
|
{
|
|
match maybe_execute_grok_sync(&plan, report_context.as_ref()).await {
|
|
Ok(Some(result)) => result,
|
|
Ok(None) => match maybe_execute_chatgpt_web_image_sync(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(Some(result)) => result,
|
|
Ok(None) => match execute_direct_sync_runtime_candidate(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
plan_kind,
|
|
candidate_started_unix_secs,
|
|
plan_request_id_for_log.as_str(),
|
|
plan_candidate_id.as_deref(),
|
|
provider_name.as_str(),
|
|
endpoint_id.as_str(),
|
|
key_id.as_str(),
|
|
model_name.as_str(),
|
|
candidate_index.as_str(),
|
|
progress_snapshot.clone(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => result,
|
|
Err(err) => {
|
|
let failure_error_type = err.error_type;
|
|
let failure_message = err.message.clone();
|
|
let failure_latency_ms = err
|
|
.latency_ms
|
|
.unwrap_or_else(|| elapsed_ms_since(candidate_started_at));
|
|
maybe_store_sync_execution_failure_fallback(
|
|
&err,
|
|
&plan,
|
|
trace_id,
|
|
decision,
|
|
&mut retry_scope_out,
|
|
&mut retry_fallback_out,
|
|
)?;
|
|
warn!(
|
|
event_name = "sync_execution_runtime_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error_type = err.error_type,
|
|
error = %err.message,
|
|
"gateway in-process sync execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some(failure_error_type.to_string()),
|
|
error_message: Some(err.message),
|
|
latency_ms: Some(failure_latency_ms),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
failure_error_type,
|
|
failure_message.as_str(),
|
|
failure_latency_ms,
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
},
|
|
Err(err) => {
|
|
let transport_error_message = err.to_string();
|
|
warn!(
|
|
event_name = "chatgpt_web_image_execution_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = %err,
|
|
"gateway ChatGPT-Web image execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some(
|
|
"chatgpt_web_image_execution_unavailable".to_string(),
|
|
),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"chatgpt_web_image_execution_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
},
|
|
Err(err) => {
|
|
let transport_error_message = err.to_string();
|
|
warn!(
|
|
event_name = "grok_execution_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = %err,
|
|
"gateway Grok execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some("grok_execution_unavailable".to_string()),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"grok_execution_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
}
|
|
} else {
|
|
let remote_execution_runtime_base_url = state
|
|
.execution_runtime_override_base_url()
|
|
.unwrap_or_default();
|
|
let remote_outcome = execute_sync_via_remote_execution_runtime(
|
|
state,
|
|
remote_execution_runtime_base_url,
|
|
trace_id,
|
|
decision,
|
|
&plan,
|
|
plan_request_id.as_str(),
|
|
plan_candidate_id.as_deref(),
|
|
report_context.as_ref(),
|
|
candidate_started_unix_secs,
|
|
candidate_started_at,
|
|
)
|
|
.await?;
|
|
match remote_outcome {
|
|
RemoteSyncFallbackOutcome::Executed(result) => result,
|
|
RemoteSyncFallbackOutcome::ClientResponse(response) => return Ok(Some(response)),
|
|
RemoteSyncFallbackOutcome::Unavailable => return Ok(None),
|
|
}
|
|
}
|
|
};
|
|
let mut candidate_first_byte_elapsed_ms =
|
|
calibrated_sync_candidate_first_byte_elapsed_ms(candidate_started_at, &result);
|
|
let initial_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
|
let mut provider_response_observation =
|
|
result
|
|
.response_observation
|
|
.clone()
|
|
.unwrap_or(ExecutionResponseObservation {
|
|
request_started_at_unix_ms: candidate_started_unix_secs,
|
|
response_headers_observed_at_unix_ms: initial_response_observed_at_unix_ms,
|
|
request_order_id: uuid::Uuid::now_v7().to_string(),
|
|
});
|
|
let mut oauth_retry_attempted = false;
|
|
let (
|
|
result_error_type,
|
|
result_error_message,
|
|
result_latency_ms,
|
|
headers,
|
|
body_bytes,
|
|
body_json,
|
|
body_base64,
|
|
local_failover_response_text,
|
|
local_failover_analysis,
|
|
) = loop {
|
|
spawn_local_oauth_success_effect(
|
|
state.clone(),
|
|
&plan,
|
|
report_context.as_ref(),
|
|
LocalOAuthSuccessEffect {
|
|
status_code: result.status_code,
|
|
request_started_at_unix_ms: Some(
|
|
provider_response_observation.request_started_at_unix_ms,
|
|
),
|
|
request_order_id: Some(&provider_response_observation.request_order_id),
|
|
},
|
|
);
|
|
let result_latency_ms = result
|
|
.telemetry
|
|
.as_ref()
|
|
.and_then(|telemetry| telemetry.elapsed_ms);
|
|
let mut headers = std::mem::take(&mut result.headers);
|
|
let (body_bytes, mut body_json, body_base64) =
|
|
decode_execution_result_body(result.body.take(), &mut headers)?;
|
|
if let Some(message) = invalid_gemini_provider_success_message(
|
|
&plan,
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
body_json.as_ref(),
|
|
)
|
|
.or_else(|| {
|
|
invalid_gemini_provider_stream_success_message(
|
|
&plan,
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
body_base64.is_some(),
|
|
)
|
|
}) {
|
|
result.status_code = StatusCode::BAD_GATEWAY.as_u16();
|
|
result.error = Some(invalid_gemini_provider_success_execution_error(message));
|
|
if let Some(error_body) =
|
|
build_invalid_provider_success_body(&plan, report_context.as_ref(), message)
|
|
{
|
|
body_json = Some(error_body);
|
|
headers.insert("content-type".to_string(), "application/json".to_string());
|
|
}
|
|
}
|
|
let (mut result_error_type, mut result_error_message) =
|
|
execution_error_details(result.error.as_ref(), body_json.as_ref());
|
|
if result.status_code < 400 && body_json.is_none() {
|
|
if let Some(error_body_json) =
|
|
extract_provider_private_stream_error_body(report_context.as_ref(), &body_bytes)
|
|
{
|
|
result.status_code =
|
|
resolve_local_sync_error_status_code(result.status_code, &error_body_json);
|
|
let (private_error_type, private_error_message) =
|
|
provider_private_error_details(&error_body_json);
|
|
result_error_type = private_error_type.or(result_error_type);
|
|
result_error_message = private_error_message.or(result_error_message);
|
|
body_json = Some(error_body_json);
|
|
}
|
|
}
|
|
let local_failover_response_text = local_failover_response_text(
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
result.error.as_ref().map(|error| error.message.as_str()),
|
|
);
|
|
|
|
if result.status_code >= 400
|
|
&& !oauth_retry_attempted
|
|
&& refresh_oauth_plan_auth_for_retry(
|
|
state,
|
|
&mut plan,
|
|
result.status_code,
|
|
local_failover_response_text.as_deref(),
|
|
trace_id,
|
|
report_context.as_ref(),
|
|
Some(provider_response_observation.request_started_at_unix_ms),
|
|
Some(&provider_response_observation.request_order_id),
|
|
)
|
|
.await
|
|
{
|
|
oauth_retry_attempted = true;
|
|
let retry_started_at_unix_ms = current_request_candidate_unix_ms();
|
|
let retry_request_order_id = uuid::Uuid::now_v7().to_string();
|
|
match crate::execution_runtime::execute_execution_runtime_sync_plan(
|
|
state,
|
|
Some(trace_id),
|
|
&plan,
|
|
)
|
|
.await
|
|
{
|
|
Ok(retry_result) => {
|
|
let retry_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
|
provider_response_observation = retry_result
|
|
.response_observation
|
|
.clone()
|
|
.unwrap_or(ExecutionResponseObservation {
|
|
request_started_at_unix_ms: retry_started_at_unix_ms,
|
|
response_headers_observed_at_unix_ms:
|
|
retry_response_observed_at_unix_ms,
|
|
request_order_id: retry_request_order_id,
|
|
});
|
|
candidate_first_byte_elapsed_ms =
|
|
calibrated_sync_candidate_first_byte_elapsed_ms(
|
|
candidate_started_at,
|
|
&retry_result,
|
|
);
|
|
result = retry_result;
|
|
continue;
|
|
}
|
|
Err(err) => {
|
|
warn!(
|
|
event_name = "local_sync_oauth_retry_execution_failed",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = ?err,
|
|
"gateway oauth retry sync execution failed"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
let local_failover_analysis = analyze_local_candidate_failover_sync(
|
|
state,
|
|
&plan,
|
|
plan_kind,
|
|
report_context.as_ref(),
|
|
&result,
|
|
local_failover_response_text.as_deref(),
|
|
)
|
|
.await;
|
|
break (
|
|
result_error_type,
|
|
result_error_message,
|
|
result_latency_ms,
|
|
headers,
|
|
body_bytes,
|
|
body_json,
|
|
body_base64,
|
|
local_failover_response_text,
|
|
local_failover_analysis,
|
|
);
|
|
};
|
|
let mut report_context = attach_provider_response_headers_to_report_context(
|
|
report_context,
|
|
&headers,
|
|
provider_response_observation.request_started_at_unix_ms,
|
|
provider_response_observation.response_headers_observed_at_unix_ms,
|
|
&provider_response_observation.request_order_id,
|
|
);
|
|
if result.status_code >= 400 {
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
|
|
status_code: result.status_code,
|
|
classification: local_failover_analysis.classification,
|
|
}),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
|
|
status_code: result.status_code,
|
|
classification: local_failover_analysis.classification,
|
|
headers: Some(&headers),
|
|
}),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
|
status_code: result.status_code,
|
|
classification: local_failover_analysis.classification,
|
|
}),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
|
|
status_code: result.status_code,
|
|
response_text: local_failover_response_text.as_deref(),
|
|
}),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
|
|
status_code: result.status_code,
|
|
classification: local_failover_analysis.classification,
|
|
headers: &headers,
|
|
error_body: local_failover_response_text.as_deref(),
|
|
}),
|
|
)
|
|
.await;
|
|
}
|
|
if matches!(
|
|
local_failover_analysis.decision,
|
|
LocalFailoverDecision::RetryNextCandidate
|
|
) {
|
|
let failure_disposition = crate::orchestration::classify_failure_disposition(
|
|
&plan.provider_api_format,
|
|
local_failover_analysis.classification,
|
|
result.status_code,
|
|
);
|
|
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
|
|
*retry_scope =
|
|
ai_attempt_retry_scope_from_failure_disposition(failure_disposition);
|
|
}
|
|
if failure_disposition.preserve_upstream_error {
|
|
if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() {
|
|
let mut fallback_headers = headers.clone();
|
|
apply_endpoint_response_header_rules(
|
|
state,
|
|
&plan,
|
|
&mut fallback_headers,
|
|
body_json.as_ref(),
|
|
)
|
|
.await?;
|
|
*retry_fallback = Some(attach_control_metadata_headers(
|
|
build_client_response_from_parts(
|
|
result.status_code,
|
|
&fallback_headers,
|
|
Body::from(body_bytes.clone()),
|
|
trace_id,
|
|
Some(decision),
|
|
)?,
|
|
Some(plan.request_id.as_str()),
|
|
plan.candidate_id.as_deref(),
|
|
)?);
|
|
}
|
|
}
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
let error_trace_report_context = with_sync_error_trace_context(
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
&headers,
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
local_failover_response_text.as_deref(),
|
|
local_failover_analysis,
|
|
);
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
error_trace_report_context
|
|
.as_ref()
|
|
.or(report_context.as_ref()),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: Some(result.status_code),
|
|
error_type: result_error_type.clone(),
|
|
error_message: result_error_message.clone(),
|
|
latency_ms: result_latency_ms,
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
warn!(
|
|
event_name = "local_sync_candidate_retry_scheduled",
|
|
log_type = "event",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
status_code = result.status_code,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
"gateway local sync decision retrying next candidate after retryable execution runtime result"
|
|
);
|
|
return Ok(None);
|
|
}
|
|
let status_code = result.status_code;
|
|
let has_body_bytes = body_base64.is_some();
|
|
if (200..300).contains(&status_code) {
|
|
seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await;
|
|
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
|
|
seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context);
|
|
}
|
|
seed_kiro_sync_report_context_prompt_cache_usage(state, &plan, &mut report_context).await;
|
|
}
|
|
let mut client_headers = headers.clone();
|
|
apply_endpoint_response_header_rules(state, &plan, &mut client_headers, body_json.as_ref())
|
|
.await?;
|
|
let explicit_finalize = should_finalize_sync_response(report_kind.as_deref());
|
|
let mapped_error_finalize_kind =
|
|
resolve_core_sync_error_finalize_report_kind(plan_kind, &result, body_json.as_ref());
|
|
let implicit_finalize = if !explicit_finalize && mapped_error_finalize_kind.is_none() {
|
|
maybe_build_implicit_sync_finalize_outcome(
|
|
trace_id,
|
|
decision,
|
|
plan_kind,
|
|
&report_context,
|
|
status_code,
|
|
&client_headers,
|
|
&body_json,
|
|
&body_base64,
|
|
&result.telemetry,
|
|
)?
|
|
} else {
|
|
None
|
|
};
|
|
if !matches!(
|
|
local_failover_analysis.decision,
|
|
LocalFailoverDecision::StopLocalFailover
|
|
) && should_fallback_to_control_sync(
|
|
plan_kind,
|
|
&result,
|
|
body_json.as_ref(),
|
|
has_body_bytes,
|
|
explicit_finalize || implicit_finalize.is_some(),
|
|
mapped_error_finalize_kind.is_some(),
|
|
) {
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
let error_trace_report_context = with_sync_error_trace_context(
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
&headers,
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
local_failover_response_text.as_deref(),
|
|
local_failover_analysis,
|
|
);
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
error_trace_report_context
|
|
.as_ref()
|
|
.or(report_context.as_ref()),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: Some(result.status_code),
|
|
error_type: result_error_type.clone(),
|
|
error_message: result_error_message.clone(),
|
|
latency_ms: result_latency_ms,
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
return Ok(None);
|
|
}
|
|
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
let error_flow_report_context = (result.status_code >= 400)
|
|
.then(|| {
|
|
with_sync_error_trace_context(
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
&headers,
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
local_failover_response_text.as_deref(),
|
|
local_failover_analysis,
|
|
)
|
|
})
|
|
.flatten();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
error_flow_report_context
|
|
.as_ref()
|
|
.or(report_context.as_ref()),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: if result.status_code >= 400 {
|
|
RequestCandidateStatus::Failed
|
|
} else {
|
|
RequestCandidateStatus::Success
|
|
},
|
|
status_code: Some(result.status_code),
|
|
error_type: result_error_type.clone(),
|
|
error_message: result_error_message.clone(),
|
|
latency_ms: result_latency_ms,
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
|
|
let request_id_owned = result.request_id;
|
|
let candidate_id_owned = result.candidate_id;
|
|
let request_id = (!request_id_owned.trim().is_empty())
|
|
.then_some(request_id_owned.as_str())
|
|
.or(Some(plan_request_id.as_str()));
|
|
let request_id_for_log = short_request_id(request_id.unwrap_or("-"));
|
|
let candidate_id = candidate_id_owned
|
|
.as_deref()
|
|
.or(plan_candidate_id.as_deref());
|
|
let report_context = report_context;
|
|
let body_json = body_json;
|
|
let telemetry = result.telemetry;
|
|
|
|
if let Some(implicit_finalize) = implicit_finalize {
|
|
let usage_payload = implicit_finalize
|
|
.outcome
|
|
.background_report
|
|
.as_ref()
|
|
.unwrap_or(&implicit_finalize.payload);
|
|
apply_sync_success_effects(
|
|
state,
|
|
&plan,
|
|
implicit_finalize.payload.report_context.as_ref(),
|
|
usage_payload,
|
|
)
|
|
.await;
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
implicit_finalize.payload.report_context.as_ref(),
|
|
usage_payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
if let Some(report_payload) = implicit_finalize.outcome.background_report {
|
|
spawn_sync_report(state.clone(), report_payload);
|
|
} else {
|
|
warn!(
|
|
event_name = "local_core_finalize_missing_success_report_mapping",
|
|
log_type = "event",
|
|
trace_id = %trace_id,
|
|
report_kind = %implicit_finalize.payload.report_kind,
|
|
"gateway implicit local core finalize produced response without background success report mapping"
|
|
);
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
implicit_finalize.outcome.response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
|
|
let finalize_report_kind = if explicit_finalize {
|
|
report_kind.clone()
|
|
} else {
|
|
mapped_error_finalize_kind
|
|
};
|
|
|
|
if let Some(finalize_report_kind) = finalize_report_kind {
|
|
let mut payload = build_sync_report_payload(
|
|
trace_id,
|
|
finalize_report_kind,
|
|
report_context,
|
|
status_code,
|
|
client_headers,
|
|
body_json,
|
|
body_base64,
|
|
telemetry,
|
|
);
|
|
if let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, decision, &payload)? {
|
|
let usage_payload = outcome.background_report.as_ref().unwrap_or(&payload);
|
|
if status_code < 400 {
|
|
apply_sync_success_effects(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
usage_payload,
|
|
)
|
|
.await;
|
|
}
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
usage_payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
if let Some(report_payload) = outcome.background_report {
|
|
spawn_sync_report(state.clone(), report_payload);
|
|
} else {
|
|
warn!(
|
|
event_name = "local_core_finalize_missing_success_report_mapping",
|
|
log_type = "event",
|
|
trace_id = %trace_id,
|
|
report_kind = %payload.report_kind,
|
|
"gateway local core finalize produced response without background success report mapping"
|
|
);
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
outcome.response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
let mut payload = match maybe_build_local_video_success_outcome(
|
|
trace_id,
|
|
decision,
|
|
payload,
|
|
&state.video_tasks,
|
|
&plan,
|
|
)? {
|
|
LocalVideoSyncSuccessBuild::Handled(outcome) => {
|
|
let LocalVideoSyncSuccessOutcome {
|
|
response,
|
|
report_payload,
|
|
original_report_context,
|
|
report_mode,
|
|
local_task_snapshot,
|
|
} = outcome;
|
|
apply_sync_success_effects(
|
|
state,
|
|
&plan,
|
|
original_report_context.as_ref(),
|
|
&report_payload,
|
|
)
|
|
.await;
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
original_report_context.as_ref(),
|
|
&report_payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
if let Some(snapshot) = local_task_snapshot {
|
|
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
|
|
state.video_tasks.record_snapshot(snapshot);
|
|
}
|
|
match report_mode {
|
|
VideoTaskSyncReportMode::InlineSync => {
|
|
submit_sync_report(state, report_payload).await?;
|
|
}
|
|
VideoTaskSyncReportMode::Background => {
|
|
spawn_sync_report(state.clone(), report_payload);
|
|
}
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
LocalVideoSyncSuccessBuild::NotHandled(payload) => payload,
|
|
};
|
|
if let Some(response) =
|
|
maybe_build_local_sync_finalize_response(trace_id, decision, &payload)?
|
|
{
|
|
let background_success_report_kind =
|
|
resolve_local_sync_success_background_report_kind(payload.report_kind.as_str());
|
|
apply_sync_success_effects(state, &plan, payload.report_context.as_ref(), &payload)
|
|
.await;
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
&payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
state
|
|
.video_tasks
|
|
.apply_finalize_mutation(request_path, payload.report_kind.as_str());
|
|
if let Some(snapshot) = state
|
|
.video_tasks
|
|
.snapshot_for_route(decision.route_family.as_deref(), request_path)
|
|
{
|
|
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
|
|
}
|
|
if let Some(success_report_kind) = background_success_report_kind {
|
|
payload.report_kind = success_report_kind.to_string();
|
|
}
|
|
if background_success_report_kind.is_some() {
|
|
spawn_sync_report(state.clone(), payload);
|
|
} else {
|
|
warn!(
|
|
event_name = "local_video_finalize_missing_success_report_mapping",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %request_id_for_log,
|
|
candidate_id = ?candidate_id,
|
|
report_kind = %payload.report_kind,
|
|
"gateway local video finalize produced response without background success report mapping"
|
|
);
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
if let Some(response) =
|
|
maybe_build_local_video_error_response(trace_id, decision, &payload)?
|
|
{
|
|
let background_error_report_kind =
|
|
resolve_local_sync_error_background_report_kind(payload.report_kind.as_str());
|
|
if let Some(error_report_kind) = background_error_report_kind {
|
|
payload.report_kind = error_report_kind.to_string();
|
|
}
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
&payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
if background_error_report_kind.is_some() {
|
|
spawn_sync_report(state.clone(), payload);
|
|
} else {
|
|
warn!(
|
|
event_name = "local_video_finalize_missing_error_report_mapping",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %request_id_for_log,
|
|
candidate_id = ?candidate_id,
|
|
report_kind = %payload.report_kind,
|
|
"gateway local video finalize produced response without background error report mapping"
|
|
);
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
&payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
let response =
|
|
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
|
|
let usage_payload = build_sync_report_payload(
|
|
trace_id,
|
|
report_kind.unwrap_or_default(),
|
|
report_context,
|
|
status_code,
|
|
client_headers,
|
|
body_json,
|
|
body_base64,
|
|
telemetry,
|
|
);
|
|
if status_code < 400 {
|
|
apply_sync_success_effects(
|
|
state,
|
|
&plan,
|
|
usage_payload.report_context.as_ref(),
|
|
&usage_payload,
|
|
)
|
|
.await;
|
|
}
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
usage_payload.report_context.as_ref(),
|
|
&usage_payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
let response = attach_control_metadata_headers(
|
|
build_client_response_from_parts(
|
|
status_code,
|
|
&usage_payload.headers,
|
|
Body::from(body_bytes),
|
|
trace_id,
|
|
Some(decision),
|
|
)?,
|
|
request_id,
|
|
candidate_id,
|
|
)?;
|
|
if !usage_payload.report_kind.trim().is_empty() {
|
|
if status_code >= 400 {
|
|
let report_kind = usage_payload.report_kind.clone();
|
|
if let Err(err) = submit_sync_report(state, usage_payload).await {
|
|
warn!(
|
|
event_name = "local_sync_error_report_submit_failed",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
report_kind = %report_kind,
|
|
"gateway failed to submit local sync error report before returning response: {err:?}"
|
|
);
|
|
}
|
|
} else {
|
|
spawn_sync_report(state.clone(), usage_payload);
|
|
}
|
|
}
|
|
|
|
Ok(Some(response))
|
|
})
|
|
.await;
|
|
if let Err(error) = result.as_ref() {
|
|
terminal_guard.fail_and_disarm(error).await;
|
|
} else {
|
|
terminal_guard.disarm();
|
|
}
|
|
result
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)] // mirrors sync execution context
|
|
fn maybe_build_implicit_sync_finalize_outcome(
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
plan_kind: &str,
|
|
report_context: &Option<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(format!("{err:?}")),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
return Ok(RemoteSyncFallbackOutcome::Unavailable);
|
|
}
|
|
};
|
|
|
|
if response.status() != http::StatusCode::OK {
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: Some(response.status().as_u16()),
|
|
error_type: Some("execution_runtime_http_error".to_string()),
|
|
error_message: Some(format!(
|
|
"execution runtime returned HTTP {}",
|
|
response.status()
|
|
)),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
return Ok(RemoteSyncFallbackOutcome::ClientResponse(
|
|
attach_control_metadata_headers(
|
|
build_client_response(response, trace_id, Some(decision))?,
|
|
Some(plan_request_id),
|
|
plan_candidate_id,
|
|
)?,
|
|
));
|
|
}
|
|
|
|
let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
|
let mut result = response
|
|
.json::<ExecutionResult>()
|
|
.await
|
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
|
result
|
|
.response_observation
|
|
.get_or_insert(ExecutionResponseObservation {
|
|
request_started_at_unix_ms: remote_request_started_at_unix_ms,
|
|
response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms,
|
|
request_order_id: remote_request_order_id,
|
|
});
|
|
Ok(RemoteSyncFallbackOutcome::Executed(result))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
|
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
|
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
|
|
use aether_data_contracts::repository::usage::UsageReadRepository;
|
|
use aether_usage_runtime::UsageRuntimeConfig;
|
|
use futures_util::{pin_mut, StreamExt as _};
|
|
use serde_json::json;
|
|
use std::collections::BTreeMap;
|
|
use std::sync::Arc;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
|
|
fn test_openai_image_plan(stream: bool) -> ExecutionPlan {
|
|
ExecutionPlan {
|
|
request_id: "req-1".to_string(),
|
|
candidate_id: Some("candidate-1".to_string()),
|
|
provider_name: Some("OpenAI".to_string()),
|
|
provider_id: "provider-1".to_string(),
|
|
endpoint_id: "endpoint-1".to_string(),
|
|
key_id: "key-1".to_string(),
|
|
method: "POST".to_string(),
|
|
url: "https://chatgpt.com/backend-api/codex/responses".to_string(),
|
|
headers: BTreeMap::new(),
|
|
content_type: Some("application/json".to_string()),
|
|
content_encoding: None,
|
|
body: aether_contracts::RequestBody::from_json(json!({"stream": true})),
|
|
stream,
|
|
client_api_format: "openai:image".to_string(),
|
|
provider_api_format: "openai:image".to_string(),
|
|
model_name: Some("gpt-image-2".to_string()),
|
|
proxy: None,
|
|
transport_profile: None,
|
|
timeouts: None,
|
|
}
|
|
}
|
|
|
|
fn test_gemini_chat_plan() -> ExecutionPlan {
|
|
let mut plan = test_openai_image_plan(false);
|
|
plan.client_api_format = "openai:chat".to_string();
|
|
plan.provider_api_format = "gemini:generate_content".to_string();
|
|
plan.model_name = Some("gemini-3-flash-preview".to_string());
|
|
plan
|
|
}
|
|
|
|
fn test_decision() -> GatewayControlDecision {
|
|
GatewayControlDecision::synthetic(
|
|
"/v1/chat/completions",
|
|
Some("ai_public".to_string()),
|
|
Some("openai".to_string()),
|
|
Some("chat".to_string()),
|
|
Some("openai:chat".to_string()),
|
|
)
|
|
.with_execution_runtime_candidate(true)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn 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_error_is_retryable_candidate_failure() {
|
|
let error = invalid_gemini_provider_success_execution_error(
|
|
INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE,
|
|
);
|
|
|
|
assert_eq!(error.kind, ExecutionErrorKind::Upstream5xx);
|
|
assert_eq!(error.phase, ExecutionPhase::Finalize);
|
|
assert_eq!(error.upstream_status, Some(StatusCode::OK.as_u16()));
|
|
assert!(error.retryable);
|
|
assert!(error.failover_recommended);
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_gemini_provider_success_accepts_antigravity_chunks_with_visible_output() {
|
|
let plan = test_gemini_chat_plan();
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"envelope_name": "antigravity:v1internal",
|
|
"provider_api_format": "gemini:generate_content",
|
|
});
|
|
let body = json!({
|
|
"chunks": [{
|
|
"response": {
|
|
"responseId": "resp_antigravity_chunks_123",
|
|
"candidates": [{
|
|
"content": {
|
|
"parts": [{"text": "Hello Gemini"}],
|
|
"role": "model"
|
|
},
|
|
"finishReason": "STOP",
|
|
"index": 0
|
|
}],
|
|
"modelVersion": "gemini-3-flash-agent",
|
|
"usageMetadata": {
|
|
"promptTokenCount": 2,
|
|
"candidatesTokenCount": 2,
|
|
"totalTokenCount": 4
|
|
}
|
|
},
|
|
"traceId": "trace-antigravity-chunks"
|
|
}],
|
|
"metadata": {
|
|
"stream": true,
|
|
"stored_chunks": 1,
|
|
"total_chunks": 1
|
|
}
|
|
});
|
|
|
|
let message = invalid_gemini_provider_success_message(
|
|
&plan,
|
|
Some(&report_context),
|
|
StatusCode::OK.as_u16(),
|
|
Some(&body),
|
|
);
|
|
|
|
assert!(message.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_gemini_provider_success_unwraps_gemini_cli_v1internal_envelope() {
|
|
let plan = test_gemini_chat_plan();
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"envelope_name": "gemini_cli:v1internal",
|
|
"provider_api_format": "gemini:generate_content",
|
|
});
|
|
let body = json!({
|
|
"response": {
|
|
"candidates": [{
|
|
"content": {
|
|
"role": "model",
|
|
"parts": [{"text": "Hello from Gemini CLI"}]
|
|
},
|
|
"finishReason": "STOP"
|
|
}]
|
|
},
|
|
"remainingCredits": 41,
|
|
"consumedCredits": 1,
|
|
"traceId": "trace-upstream-sync-1"
|
|
});
|
|
|
|
let message = invalid_gemini_provider_success_message(
|
|
&plan,
|
|
Some(&report_context),
|
|
StatusCode::OK.as_u16(),
|
|
Some(&body),
|
|
);
|
|
|
|
assert!(message.is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sync_attempt_terminal_guard_marks_dropped_pending_attempt_cancelled() {
|
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
|
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(
|
|
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
|
Arc::clone(&request_candidate_repository),
|
|
Arc::clone(&usage_repository),
|
|
),
|
|
)
|
|
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
|
enabled: true,
|
|
..UsageRuntimeConfig::default()
|
|
});
|
|
let mut plan = test_openai_image_plan(false);
|
|
plan.request_id = "sync-cancel-guard-request".to_string();
|
|
plan.candidate_id = None;
|
|
let mut report_context = Some(json!({
|
|
"candidate_index": 0,
|
|
"retry_index": 0,
|
|
"user_id": "user-cancel",
|
|
"api_key_id": "api-key-cancel",
|
|
"client_api_format": "openai:image",
|
|
"provider_api_format": "openai:image",
|
|
"request_path": "/v1/images/generations",
|
|
"request_path_and_query": "/v1/images/generations",
|
|
"upstream_url": "https://example.test/v1/images/generations",
|
|
"mapped_model": "gpt-image-2",
|
|
}));
|
|
|
|
ensure_execution_request_candidate_slot(&state, &mut plan, &mut report_context).await;
|
|
let candidate_started_at = Instant::now();
|
|
let started_at = current_request_candidate_unix_ms();
|
|
state.usage_runtime.record_pending(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
build_lifecycle_usage_seed(&plan, report_context.as_ref()),
|
|
);
|
|
record_local_request_candidate_status(
|
|
&state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Pending,
|
|
status_code: None,
|
|
error_type: None,
|
|
error_message: None,
|
|
latency_ms: None,
|
|
started_at_unix_ms: Some(started_at),
|
|
finished_at_unix_ms: None,
|
|
},
|
|
)
|
|
.await;
|
|
|
|
crate::request_diagnostics::scope_request_diagnostics(async {
|
|
crate::request_diagnostics::record_request_accepted_at(
|
|
Instant::now() - Duration::from_millis(25),
|
|
);
|
|
let _guard = SyncAttemptTerminalGuard::new(
|
|
&state,
|
|
&plan,
|
|
report_context.clone(),
|
|
started_at,
|
|
candidate_started_at,
|
|
);
|
|
})
|
|
.await;
|
|
|
|
let mut stored_usage = None;
|
|
for _ in 0..50 {
|
|
if let Some(usage) = usage_repository
|
|
.find_by_request_id("sync-cancel-guard-request")
|
|
.await
|
|
.expect("usage should read")
|
|
{
|
|
if usage.status == "cancelled" {
|
|
stored_usage = Some(usage);
|
|
break;
|
|
}
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
}
|
|
let stored_usage = stored_usage.expect("cancelled usage should be recorded");
|
|
assert_eq!(stored_usage.status, "cancelled");
|
|
assert_eq!(stored_usage.billing_status, "void");
|
|
assert_eq!(stored_usage.status_code, Some(499));
|
|
assert_eq!(stored_usage.error_category.as_deref(), Some("cancelled"));
|
|
let request_metadata = stored_usage
|
|
.request_metadata
|
|
.as_ref()
|
|
.expect("cancelled usage should retain request diagnostics");
|
|
assert!(request_metadata
|
|
.get("end_to_end_time_ms")
|
|
.and_then(Value::as_u64)
|
|
.is_some());
|
|
assert!(request_metadata
|
|
.get("end_to_end_first_byte_time_ms")
|
|
.is_none());
|
|
|
|
let stored_candidates = request_candidate_repository
|
|
.list_by_request_id("sync-cancel-guard-request")
|
|
.await
|
|
.expect("request candidates should read");
|
|
assert_eq!(stored_candidates.len(), 1);
|
|
assert_eq!(
|
|
stored_candidates[0].status,
|
|
RequestCandidateStatus::Cancelled
|
|
);
|
|
assert_eq!(stored_candidates[0].status_code, Some(499));
|
|
assert_eq!(
|
|
stored_candidates[0].error_type.as_deref(),
|
|
Some("local_sync_attempt_cancelled")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sync_direct_response_start_marks_usage_and_candidate_active_before_body_finishes() {
|
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
|
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(
|
|
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
|
Arc::clone(&request_candidate_repository),
|
|
Arc::clone(&usage_repository),
|
|
),
|
|
)
|
|
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
|
enabled: true,
|
|
..UsageRuntimeConfig::default()
|
|
});
|
|
|
|
let listener = crate::test_support::bind_loopback_listener()
|
|
.await
|
|
.expect("listener should bind");
|
|
let addr = listener.local_addr().expect("local addr should resolve");
|
|
let (headers_tx, headers_rx) = tokio::sync::oneshot::channel();
|
|
let (body_tx, body_rx) = tokio::sync::oneshot::channel();
|
|
let server = tokio::spawn(async move {
|
|
let (mut socket, _) = listener.accept().await.expect("client should connect");
|
|
let mut request = [0_u8; 4096];
|
|
let _ = socket
|
|
.read(&mut request)
|
|
.await
|
|
.expect("request should read");
|
|
socket
|
|
.write_all(
|
|
b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 11\r\n\r\n",
|
|
)
|
|
.await
|
|
.expect("headers should write");
|
|
socket.flush().await.expect("headers should flush");
|
|
let _ = headers_tx.send(());
|
|
let _ = body_rx.await;
|
|
socket
|
|
.write_all(br#"{"ok":true}"#)
|
|
.await
|
|
.expect("body should write");
|
|
});
|
|
|
|
let mut plan = test_gemini_chat_plan();
|
|
plan.request_id = "sync-response-start-active-request".to_string();
|
|
plan.candidate_id = Some("sync-response-start-active-candidate".to_string());
|
|
plan.url = format!("http://{addr}/chat");
|
|
plan.provider_api_format = "openai:chat".to_string();
|
|
plan.model_name = Some("gpt-5".to_string());
|
|
plan.body = aether_contracts::RequestBody::from_json(json!({
|
|
"model": "gpt-5",
|
|
"messages": [{"role": "user", "content": "slow body"}],
|
|
}));
|
|
let report_context = Some(json!({
|
|
"candidate_index": 0,
|
|
"retry_index": 0,
|
|
"user_id": "user-active",
|
|
"api_key_id": "api-key-active",
|
|
"client_api_format": "openai:chat",
|
|
"provider_api_format": "openai:chat",
|
|
"request_path": "/v1/chat/completions",
|
|
"request_path_and_query": "/v1/chat/completions",
|
|
"upstream_url": plan.url.clone(),
|
|
"mapped_model": "gpt-5",
|
|
}));
|
|
let started_at = current_request_candidate_unix_ms();
|
|
state
|
|
.usage_runtime
|
|
.record_pending_direct(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
build_lifecycle_usage_seed(&plan, report_context.as_ref()),
|
|
)
|
|
.await;
|
|
record_local_request_candidate_status(
|
|
&state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Pending,
|
|
status_code: None,
|
|
error_type: None,
|
|
error_message: None,
|
|
latency_ms: None,
|
|
started_at_unix_ms: Some(started_at),
|
|
finished_at_unix_ms: None,
|
|
},
|
|
)
|
|
.await;
|
|
|
|
let state_for_exec = state.clone();
|
|
let plan_for_exec = plan.clone();
|
|
let report_context_for_exec = report_context.clone();
|
|
let exec = tokio::spawn(async move {
|
|
execute_direct_sync_runtime_candidate(
|
|
&state_for_exec,
|
|
&plan_for_exec,
|
|
report_context_for_exec.as_ref(),
|
|
"trace-response-start-active",
|
|
"openai_chat_sync",
|
|
started_at,
|
|
"sync-response-start-active-request",
|
|
plan_for_exec.candidate_id.as_deref(),
|
|
"openai",
|
|
"endpoint-1",
|
|
"key-1",
|
|
"gpt-5",
|
|
"0",
|
|
None,
|
|
)
|
|
.await
|
|
});
|
|
|
|
headers_rx.await.expect("headers should be written");
|
|
let mut active_usage = None;
|
|
for _ in 0..50 {
|
|
if let Some(usage) = usage_repository
|
|
.find_by_request_id("sync-response-start-active-request")
|
|
.await
|
|
.expect("usage should read")
|
|
{
|
|
if usage.status == "streaming" {
|
|
active_usage = Some(usage);
|
|
break;
|
|
}
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
}
|
|
let active_usage = active_usage.expect("usage should become active before body finishes");
|
|
assert_eq!(active_usage.status_code, Some(200));
|
|
assert!(active_usage.first_byte_time_ms.is_some());
|
|
assert!(active_usage.response_time_ms.is_some());
|
|
|
|
let stored_candidates = request_candidate_repository
|
|
.list_by_request_id("sync-response-start-active-request")
|
|
.await
|
|
.expect("candidate should read");
|
|
let active_candidate = stored_candidates
|
|
.iter()
|
|
.find(|candidate| candidate.id == "sync-response-start-active-candidate")
|
|
.expect("candidate should exist");
|
|
assert_eq!(active_candidate.status, RequestCandidateStatus::Streaming);
|
|
assert_eq!(active_candidate.status_code, Some(200));
|
|
|
|
let _ = body_tx.send(());
|
|
let result = tokio::time::timeout(Duration::from_secs(2), exec)
|
|
.await
|
|
.expect("sync execution should finish")
|
|
.expect("sync execution task should not panic")
|
|
.expect("sync execution should succeed");
|
|
assert_eq!(result.status_code, 200);
|
|
server.abort();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sync_execution_active_marks_usage_before_response_headers() {
|
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
|
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(
|
|
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
|
Arc::clone(&request_candidate_repository),
|
|
Arc::clone(&usage_repository),
|
|
),
|
|
)
|
|
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
|
enabled: true,
|
|
..UsageRuntimeConfig::default()
|
|
});
|
|
|
|
let listener = crate::test_support::bind_loopback_listener()
|
|
.await
|
|
.expect("listener should bind");
|
|
let addr = listener.local_addr().expect("local addr should resolve");
|
|
let (request_seen_tx, request_seen_rx) = tokio::sync::oneshot::channel();
|
|
let (finish_tx, finish_rx) = tokio::sync::oneshot::channel();
|
|
let server = tokio::spawn(async move {
|
|
let (mut socket, _) = listener.accept().await.expect("client should connect");
|
|
let mut request = [0_u8; 4096];
|
|
let _ = socket
|
|
.read(&mut request)
|
|
.await
|
|
.expect("request should read");
|
|
let _ = request_seen_tx.send(());
|
|
let _ = finish_rx.await;
|
|
socket
|
|
.write_all(
|
|
b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 11\r\n\r\n{\"ok\":true}",
|
|
)
|
|
.await
|
|
.expect("response should write");
|
|
});
|
|
|
|
let mut plan = test_gemini_chat_plan();
|
|
plan.request_id = "sync-active-before-headers-request".to_string();
|
|
plan.candidate_id = Some("sync-active-before-headers-candidate".to_string());
|
|
plan.url = format!("http://{addr}/chat");
|
|
plan.provider_name = Some("OpenAI".to_string());
|
|
plan.provider_api_format = "openai:chat".to_string();
|
|
plan.model_name = Some("gpt-5".to_string());
|
|
plan.body = aether_contracts::RequestBody::from_json(json!({
|
|
"model": "gpt-5",
|
|
"messages": [{"role": "user", "content": "slow headers"}],
|
|
}));
|
|
let report_context = Some(json!({
|
|
"candidate_index": 0,
|
|
"retry_index": 0,
|
|
"user_id": "user-active-before-headers",
|
|
"api_key_id": "api-key-active-before-headers",
|
|
"candidate_id": "sync-active-before-headers-candidate",
|
|
"provider_id": "provider-1",
|
|
"endpoint_id": "endpoint-1",
|
|
"key_id": "key-1",
|
|
"provider_name": "OpenAI",
|
|
"client_api_format": "openai:chat",
|
|
"provider_api_format": "openai:chat",
|
|
"request_path": "/v1/chat/completions",
|
|
"request_path_and_query": "/v1/chat/completions",
|
|
"upstream_url": plan.url.clone(),
|
|
"mapped_model": "gpt-5",
|
|
}));
|
|
let state_for_exec = state.clone();
|
|
let plan_for_exec = plan.clone();
|
|
let report_context_for_exec = report_context.clone();
|
|
let exec = tokio::spawn(async move {
|
|
execute_execution_runtime_sync(
|
|
&state_for_exec,
|
|
"/v1/chat/completions",
|
|
plan_for_exec,
|
|
"trace-active-before-headers",
|
|
&test_decision(),
|
|
"openai_chat_sync",
|
|
Some("openai_chat_sync".to_string()),
|
|
report_context_for_exec,
|
|
)
|
|
.await
|
|
});
|
|
|
|
request_seen_rx
|
|
.await
|
|
.expect("upstream request should be observed");
|
|
let mut active_usage = None;
|
|
for _ in 0..50 {
|
|
if let Some(usage) = usage_repository
|
|
.find_by_request_id("sync-active-before-headers-request")
|
|
.await
|
|
.expect("usage should read")
|
|
{
|
|
if usage.status == "streaming" {
|
|
active_usage = Some(usage);
|
|
break;
|
|
}
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
}
|
|
let active_usage =
|
|
active_usage.expect("usage should become active before upstream headers");
|
|
assert_eq!(active_usage.status_code, None);
|
|
assert_eq!(active_usage.first_byte_time_ms, None);
|
|
assert_eq!(active_usage.response_time_ms, None);
|
|
|
|
let stored_candidates = request_candidate_repository
|
|
.list_by_request_id("sync-active-before-headers-request")
|
|
.await
|
|
.expect("candidate should read");
|
|
let active_candidate = stored_candidates
|
|
.iter()
|
|
.find(|candidate| candidate.id == "sync-active-before-headers-candidate")
|
|
.expect("candidate should exist");
|
|
assert_eq!(active_candidate.status, RequestCandidateStatus::Streaming);
|
|
assert_eq!(active_candidate.status_code, None);
|
|
assert!(active_candidate.started_at_unix_ms.is_some());
|
|
assert!(active_candidate.finished_at_unix_ms.is_none());
|
|
|
|
let _ = finish_tx.send(());
|
|
let response = tokio::time::timeout(Duration::from_secs(2), exec)
|
|
.await
|
|
.expect("sync execution should finish")
|
|
.expect("sync execution task should not panic")
|
|
.expect("sync execution should succeed")
|
|
.expect("sync execution should produce a response");
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
server.abort();
|
|
}
|
|
|
|
#[test]
|
|
fn kiro_sync_report_context_seeds_input_tokens_from_original_request_body() {
|
|
let plan = test_kiro_sync_plan();
|
|
let mut report_context = Some(json!({
|
|
"original_request_body": test_kiro_sync_cacheable_request_body(),
|
|
}));
|
|
|
|
seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context);
|
|
|
|
assert!(report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.is_some_and(|tokens| tokens > 0));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn kiro_sync_report_context_applies_prompt_cache_usage_from_tracker() {
|
|
let state = AppState::new().expect("gateway state should build");
|
|
let plan = test_kiro_sync_plan();
|
|
|
|
let mut first_report_context = Some(json!({
|
|
"original_request_body": test_kiro_sync_cacheable_request_body(),
|
|
"kiro_simulated_cache_enabled": true,
|
|
}));
|
|
seed_kiro_sync_report_context_input_tokens(&plan, &mut first_report_context);
|
|
seed_kiro_sync_report_context_prompt_cache_usage(&state, &plan, &mut first_report_context)
|
|
.await;
|
|
let first_creation = first_report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("cache_creation_input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or_default();
|
|
let first_read = first_report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("cache_read_input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or_default();
|
|
assert!(first_creation > 0);
|
|
assert_eq!(first_read, 0);
|
|
|
|
let mut second_report_context = Some(json!({
|
|
"original_request_body": test_kiro_sync_cacheable_request_body(),
|
|
"kiro_simulated_cache_enabled": true,
|
|
}));
|
|
seed_kiro_sync_report_context_input_tokens(&plan, &mut second_report_context);
|
|
seed_kiro_sync_report_context_prompt_cache_usage(&state, &plan, &mut second_report_context)
|
|
.await;
|
|
let second_creation = second_report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("cache_creation_input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or_default();
|
|
let second_read = second_report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("cache_read_input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or_default();
|
|
assert_eq!(second_creation, 0);
|
|
assert!(second_read > 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn json_whitespace_heartbeat_stream_prefixes_final_json() {
|
|
let (tx, rx) = mpsc::channel::<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),
|
|
));
|
|
}
|
|
}
|