use std::collections::{BTreeMap, VecDeque}; use std::future::Future; use std::io::Error as IoError; use std::sync::{ atomic::{AtomicBool, AtomicU64, Ordering}, Arc, }; use std::time::{Duration, Instant}; use aether_contracts::{ ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StandardizedUsage, StreamFrame, StreamFramePayload, }; use aether_data_contracts::repository::candidates::RequestCandidateStatus; use aether_data_contracts::repository::usage::UsageBodyCaptureState; use aether_scheduler_core::{ parse_request_candidate_report_context, SchedulerRequestCandidateStatusUpdate, }; use aether_usage_runtime::{ build_lifecycle_usage_seed, build_stream_terminal_usage_payload_seed, build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed, build_usage_event_data_seed, LifecycleUsageSeed, UsageEvent, UsageEventType, DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, }; use async_stream::stream; use axum::body::{Body, Bytes}; use axum::http::Response; use base64::Engine as _; use futures_util::stream::BoxStream; use futures_util::{StreamExt, TryStreamExt}; use serde_json::{json, Value}; use tokio::sync::mpsc; use tokio::time::MissedTickBehavior; use tokio_util::codec::{FramedRead, LinesCodec}; use tokio_util::io::StreamReader; use tracing::{debug, info, warn}; use super::error::{ build_synthetic_non_success_stream_error_body, collect_error_body, decode_stream_error_body, inspect_prefetched_stream_body, read_next_frame, should_synthesize_non_success_stream_error_body, stream_client_error_status_code_for_upstream_status, synthetic_error_response_headers, StreamPrefetchInspection, }; #[path = "execution_failures.rs"] mod execution_failures; use self::execution_failures::{ build_stream_failure_from_execution_error, build_stream_failure_from_provider_error_body, build_stream_failure_report, handle_prefetch_provider_private_stream_error, handle_prefetch_stream_failure, submit_midstream_stream_failure, StreamFailureReport, }; use crate::ai_serving::api::{ extract_provider_private_stream_error_body, maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer, maybe_build_stream_response_rewriter, normalize_provider_private_report_context, StreamingStandardTerminalObserver, }; use crate::ai_serving::is_openai_responses_family_format; use crate::api::response::{ attach_control_metadata_headers, build_client_response, build_client_response_from_parts, }; use crate::clock::current_unix_ms as current_request_candidate_unix_ms; use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER}; use crate::control::GatewayControlDecision; use crate::execution_runtime::build_direct_execution_frame_stream; use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_stream; use crate::execution_runtime::grok::maybe_execute_grok_stream; use crate::execution_runtime::kiro_cache::{ billed_input_tokens as kiro_billed_input_tokens, 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::kiro_web_search::maybe_execute_kiro_web_search_stream; use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry; #[cfg(test)] use crate::execution_runtime::remote_compat::post_stream_plan_to_remote_execution_runtime; use crate::execution_runtime::submission::{ resolve_core_error_background_report_kind, resolve_local_sync_error_status_code, strip_utf8_bom_and_ws, submit_local_core_error_or_sync_finalize, }; use crate::execution_runtime::transport::{ execute_stream_plan_via_local_tunnel, format_upstream_request_error, format_wreq_upstream_request_error, record_manual_proxy_request_failure, record_manual_proxy_request_success, record_manual_proxy_stream_error, stream_first_byte_timeout_message, DirectSyncExecutionRuntime, DirectUpstreamResponse, DirectUpstreamStreamExecution, ExecutionRuntimeTransportError, }; use crate::execution_runtime::windsurf::maybe_execute_windsurf_stream; use crate::execution_runtime::{ apply_endpoint_response_header_rules, attach_provider_response_headers_to_report_context, local_failover_response_text, resolve_core_stream_direct_finalize_report_kind, resolve_core_stream_error_finalize_report_kind, resolve_local_candidate_failover_analysis_stream, should_fallback_to_control_stream, should_retry_next_local_candidate_stream, LocalFailoverDecision, }; use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES}; use crate::log_ids::short_request_id; use crate::orchestration::{ apply_local_execution_effect, build_local_error_flow_metadata, trace_upstream_response_body, with_error_flow_report_context, with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect, }; use crate::provider_pool_demand::{ acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard, }; use crate::request_candidate_runtime::{ ensure_execution_request_candidate_slot, record_local_request_candidate_status, record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status, }; use crate::request_diagnostics::{ attach_current_request_diagnostics_to_report_context, attach_request_diagnostics_to_report_context, current_request_diagnostics, RequestDiagnostics, }; use crate::stage_metrics::{ attach_stage_trace_to_report_context, observe_gateway_stage_ms, observe_gateway_stage_trace_ms, RequestStageTrace, }; use crate::usage::submit_stream_report; use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest}; use crate::{ AppState, GatewayError, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, }; const OPENAI_IMAGE_STREAM_PLAN_KIND: &str = "openai_image_stream"; const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15); const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n"; const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024; const SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES: usize = 1024 * 1024; const STREAM_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60); const STREAM_IDLE_LOG_INTERVAL_MS: u64 = 60_000; const REWRITTEN_STREAM_PREFETCH_TIMEOUT: Duration = Duration::from_millis(750); struct StageElapsedGuard { stage: &'static str, started_at: Instant, } #[derive(Debug)] enum InProcessStreamExecutionError { Transport(ExecutionRuntimeTransportError), Gateway(GatewayError), } impl From for InProcessStreamExecutionError { fn from(error: ExecutionRuntimeTransportError) -> Self { Self::Transport(error) } } impl From for InProcessStreamExecutionError { fn from(error: GatewayError) -> Self { Self::Gateway(error) } } impl StageElapsedGuard { fn from_started_at(stage: &'static str, started_at: Instant) -> Self { Self { stage, started_at } } } fn report_context_with_stage_trace( report_context: Option, mut stage_trace: RequestStageTrace, stream_started_at: Instant, terminal_telemetry: Option<&ExecutionTelemetry>, ) -> Option { stage_trace.observe("stream_total", stream_elapsed_ms_since(stream_started_at)); let fallback_elapsed_ms = terminal_telemetry.and_then(|telemetry| telemetry.ttfb_ms); attach_stage_trace_to_report_context( report_context, stage_trace.into_metadata_value(fallback_elapsed_ms), ) } fn report_context_with_request_diagnostics( report_context: Option, diagnostics: Option<&Arc>, ) -> Option { attach_request_diagnostics_to_report_context(report_context, diagnostics) } impl Drop for StageElapsedGuard { fn drop(&mut self) { observe_gateway_stage_ms(self.stage, self.started_at.elapsed().as_millis() as u64); } } fn record_sync_terminal_usage( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, ) { let report_context_with_diagnostics = attach_current_request_diagnostics_to_report_context(report_context); 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.data.as_ref(), context_seed, payload_seed); } fn build_stream_sync_payload( trace_id: &str, report_kind: String, report_context: Option, status_code: u16, headers: BTreeMap, body_json: Option, body_base64: Option, telemetry: Option, ) -> GatewaySyncReportRequest { GatewaySyncReportRequest { trace_id: trace_id.to_string(), report_kind, report_context, status_code, headers, body_json, client_body_json: None, body_base64, telemetry, } } #[allow(clippy::too_many_arguments)] fn build_stream_error_sync_payload( trace_id: &str, report_kind: String, report_context: Option, upstream_status_code: u16, provider_headers: BTreeMap, provider_body_json: Option, provider_body_base64: Option, client_headers: BTreeMap, client_body_json: Option, telemetry: Option, ) -> GatewaySyncReportRequest { let client_status_code = stream_client_error_status_code_for_upstream_status(upstream_status_code); let mut report_context = report_context; if client_status_code != upstream_status_code || client_headers != provider_headers { let mut object = match report_context { Some(Value::Object(object)) => object, Some(other) => serde_json::Map::from_iter([("seed".to_string(), other)]), None => serde_json::Map::new(), }; object.insert( "client_response_status_code".to_string(), Value::from(client_status_code), ); object.insert( "client_response_headers".to_string(), serde_json::to_value(client_headers).unwrap_or(Value::Null), ); report_context = Some(Value::Object(object)); } GatewaySyncReportRequest { trace_id: trace_id.to_string(), report_kind, report_context, status_code: upstream_status_code, headers: provider_headers, body_json: provider_body_json, client_body_json, body_base64: provider_body_base64, telemetry, } } fn record_stream_terminal_usage( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewayStreamReportRequest, cancelled: bool, ) { let context_seed = build_terminal_usage_context_seed(plan, report_context); let payload_seed = build_stream_terminal_usage_payload_seed(payload); state.usage_runtime.record_stream_terminal( state.data.as_ref(), context_seed, payload_seed, cancelled, ); } async fn record_stream_admission_timeout_terminal_state( state: &AppState, plan: &ExecutionPlan, report_context: Option<&Value>, candidate_started_unix_ms: u64, error: &GatewayError, ) { let status_code = 429; let error_type = "gateway_admission_timeout"; let error_message = match error { GatewayError::AdmissionTimeout { gate, queue_budget_ms, .. } => format!("gateway admission gate {gate} timed out after {queue_budget_ms}ms"), other => format!("{other:?}"), }; let terminal_unix_ms = current_request_candidate_unix_ms(); let latency_ms = terminal_unix_ms.saturating_sub(candidate_started_unix_ms); record_local_request_candidate_status( state, plan, report_context, SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, 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); usage_data.status_code = Some(status_code); usage_data.error_message = Some(error_message.clone()); usage_data.error_category = Some("client_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.submit_terminal_event( state.data.as_ref(), UsageEvent::new(UsageEventType::Failed, plan.request_id.clone(), usage_data), ); } fn build_stream_body_capture( body: &[u8], truncated: bool, ) -> (Option, Option) { let body_base64 = (!body.is_empty()).then(|| base64::engine::general_purpose::STANDARD.encode(body)); let body_state = Some(if truncated { UsageBodyCaptureState::Truncated } else if body.is_empty() { UsageBodyCaptureState::None } else { UsageBodyCaptureState::Inline }); (body_base64, body_state) } fn wrap_non_json_binary_stream_error_for_client( plan_kind: &str, headers: &BTreeMap, error_body: &[u8], ) -> Result, GatewayError> { let content_type = headers .get("content-type") .map(|value| value.to_ascii_lowercase()) .unwrap_or_default(); if content_type.starts_with("application/json") { return Ok(None); } let body = match plan_kind { GEMINI_FILES_DOWNLOAD_PLAN_KIND => json!({ "error": String::from_utf8_lossy(error_body).to_string(), }), OPENAI_VIDEO_CONTENT_PLAN_KIND => json!({ "error": { "type": "upstream_error", "message": "Video not available", } }), _ => return Ok(None), }; Ok(Some(body)) } fn with_stream_error_trace_context( report_context: Option<&Value>, status_code: u16, headers: &BTreeMap, body_json: Option<&Value>, body_bytes: &[u8], response_text: Option<&str>, local_failover_analysis: crate::orchestration::LocalFailoverAnalysis, ) -> Option { 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), ) } #[allow(clippy::too_many_arguments)] // stream report payload assembly mirrors runtime state fn build_stream_usage_payload( trace_id: String, report_kind: String, report_context: Option, status_code: u16, headers: BTreeMap, provider_body: &[u8], provider_body_truncated: bool, client_body: &[u8], client_body_truncated: bool, terminal_summary: Option, telemetry: Option, ) -> GatewayStreamReportRequest { let (provider_body_base64, provider_body_state) = build_stream_body_capture(provider_body, provider_body_truncated); let (client_body_base64, client_body_state) = build_stream_body_capture(client_body, client_body_truncated); GatewayStreamReportRequest { trace_id, report_kind, report_context, status_code, headers, provider_body_base64, provider_body_state, client_body_base64, client_body_state, terminal_summary, telemetry, } } fn seed_kiro_report_context_input_tokens(plan: &ExecutionPlan, report_context: &mut Option) { 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_simulated_cache_enabled( state: &AppState, plan: &ExecutionPlan, report_context: &mut Option, ) { 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_report_context_prompt_cache_usage( state: &AppState, plan: &ExecutionPlan, report_context: &mut Option, ) { 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_stream_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_stream_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, ) -> Option { 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 kiro_cache_usage_from_report_context(report_context: &Value) -> Option { report_context .as_object() .and_then(kiro_cache_usage_from_context_object) } async fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary( state: &AppState, plan: &ExecutionPlan, report_context: Option<&Value>, summary: &mut Option, ) { if !plan .provider_name .as_deref() .is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro")) { return; } let Some(report_context) = report_context else { return; }; let Some(original_request_body) = report_context.get("original_request_body") else { return; }; let simulated_cache_enabled = kiro_simulated_cache_enabled_from_report_context(Some(report_context)); let summary = summary.get_or_insert_with(ExecutionStreamTerminalSummary::default); let usage = summary .standardized_usage .get_or_insert_with(StandardizedUsage::new); let estimated_input_tokens = report_context .get("input_tokens") .and_then(Value::as_u64) .filter(|value| *value > 0) .unwrap_or_else(|| { let estimated_input_tokens = estimate_kiro_prompt_input_tokens(original_request_body); if estimated_input_tokens > 0 { estimated_input_tokens } else { usage.input_tokens.max(0) as u64 } }); if !simulated_cache_enabled { usage.cache_creation_tokens = 0; usage.cache_read_tokens = 0; if usage.input_tokens <= 0 { usage.input_tokens = estimated_input_tokens as i64; } return; } if let Some(cache_usage) = kiro_cache_usage_from_report_context(report_context) { usage.input_tokens = kiro_billed_input_tokens(estimated_input_tokens, cache_usage) as i64; usage.cache_creation_tokens = cache_usage.cache_creation_input_tokens as i64; usage.cache_read_tokens = cache_usage.cache_read_input_tokens as i64; return; } if usage.cache_creation_tokens > 0 || usage.cache_read_tokens > 0 { if usage.input_tokens <= 0 { usage.input_tokens = kiro_billed_input_tokens( estimated_input_tokens, KiroPromptCacheUsage { cache_creation_input_tokens: usage.cache_creation_tokens.max(0) as u64, cache_read_input_tokens: usage.cache_read_tokens.max(0) as u64, }, ) as i64; } return; } if usage.input_tokens <= 0 { usage.input_tokens = estimated_input_tokens as i64; } let Some(profile) = build_kiro_prompt_cache_profile(original_request_body, estimated_input_tokens) else { return; }; let cache_usage = compute_kiro_prompt_cache_usage( state.runtime_state(), kiro_stream_cache_credential_id(plan), &profile, ) .await; if cache_usage.cache_creation_input_tokens == 0 && cache_usage.cache_read_input_tokens == 0 { return; } let billed_input_tokens = kiro_billed_input_tokens(estimated_input_tokens, cache_usage); usage.input_tokens = billed_input_tokens as i64; usage.cache_creation_tokens = cache_usage.cache_creation_input_tokens as i64; usage.cache_read_tokens = cache_usage.cache_read_input_tokens as i64; } fn append_stream_capture_bytes( buffer: &mut Vec, chunk: &[u8], max_bytes: usize, truncated: &mut bool, ) { if chunk.is_empty() || max_bytes == 0 { return; } if buffer.len() >= max_bytes { *truncated = true; return; } let remaining = max_bytes - buffer.len(); let keep_len = remaining.min(chunk.len()); buffer.extend_from_slice(&chunk[..keep_len]); if keep_len < chunk.len() { *truncated = true; } } fn observe_stream_usage_bytes( observer: &mut StreamingStandardTerminalObserver, report_context: &Value, buffered: &mut Vec, chunk: &[u8], ) { if chunk.is_empty() { return; } buffered.extend_from_slice(chunk); while let Some(line_end) = buffered.iter().position(|byte| *byte == b'\n') { let line = buffered.drain(..=line_end).collect::>(); if let Err(err) = observer.push_line(report_context, line) { observer.disable_with_error(err.to_string()); buffered.clear(); break; } } } fn finalize_stream_usage_observer( observer: &mut Option, report_context: Option<&Value>, buffered: &mut Vec, ) -> Option { let (Some(observer), Some(report_context)) = (observer.as_mut(), report_context) else { return None; }; if !buffered.is_empty() { let line = std::mem::take(buffered); if let Err(err) = observer.push_line(report_context, line) { observer.disable_with_error(err.to_string()); } } match observer.finish(report_context) { Ok(summary) => summary, Err(err) => { observer.disable_with_error(err.to_string()); observer.latest_summary().cloned() } } } fn merge_stream_terminal_summary( mut current: Option, observed: Option, ) -> Option { let Some(observed) = observed else { return current; }; let Some(current_summary) = current.as_mut() else { return Some(observed); }; if should_replace_stream_usage( current_summary.standardized_usage.as_ref(), observed.standardized_usage.as_ref(), ) { current_summary.standardized_usage = observed.standardized_usage; } if current_summary.finish_reason.is_none() { current_summary.finish_reason = observed.finish_reason; } if current_summary.response_id.is_none() { current_summary.response_id = observed.response_id; } if current_summary.model.is_none() { current_summary.model = observed.model; } current_summary.observed_finish |= observed.observed_finish; current_summary.unknown_event_count = current_summary .unknown_event_count .saturating_add(observed.unknown_event_count); if current_summary.parser_error.is_none() { current_summary.parser_error = observed.parser_error; } current } fn should_replace_stream_usage( current: Option<&aether_contracts::StandardizedUsage>, observed: Option<&aether_contracts::StandardizedUsage>, ) -> bool { let Some(observed) = observed else { return false; }; let Some(current) = current else { return true; }; observed.is_more_complete_than(current) } fn stream_terminal_summary_missing_observed_finish( summary: Option<&ExecutionStreamTerminalSummary>, ) -> bool { summary.is_some_and(|summary| { !summary.observed_finish && !summary .standardized_usage .as_ref() .is_some_and(StandardizedUsage::has_token_signal) }) } fn stream_report_context_format_field<'a>( report_context: Option<&'a Value>, field: &str, ) -> Option<&'a str> { report_context .and_then(Value::as_object) .and_then(|object| object.get(field)) .and_then(Value::as_str) .map(str::trim) .filter(|value| !value.is_empty()) } fn stream_requires_observed_terminal_event( provider_api_format: &str, report_context: Option<&Value>, ) -> bool { is_openai_responses_family_format(provider_api_format) || [ "provider_stream_event_api_format", "provider_stream_api_format", "provider_api_format", ] .into_iter() .filter_map(|field| stream_report_context_format_field(report_context, field)) .any(is_openai_responses_family_format) } fn stream_terminal_summary_missing_observed_finish_with_requirement( summary: Option<&ExecutionStreamTerminalSummary>, requires_observed_terminal_event: bool, ) -> bool { if !requires_observed_terminal_event { return stream_terminal_summary_missing_observed_finish(summary); } summary.is_some_and(|summary| !summary.observed_finish) } fn ensure_stream_terminal_summary_for_missing_observed_finish( summary: &mut Option, requires_observed_terminal_event: bool, ) { if !requires_observed_terminal_event { return; } let summary = summary.get_or_insert_with(ExecutionStreamTerminalSummary::default); if !summary.observed_finish && summary.parser_error.is_none() { summary.parser_error = Some("execution runtime stream ended before provider terminal event".to_string()); } } fn stream_terminal_summary_represents_failure_with_requirement( summary: Option<&ExecutionStreamTerminalSummary>, requires_observed_terminal_event: bool, ) -> bool { summary.is_some_and(|summary| { summary.parser_error.is_some() || stream_terminal_summary_missing_observed_finish_with_requirement( Some(summary), requires_observed_terminal_event, ) }) } async fn execute_in_process_stream( state: &AppState, plan: &ExecutionPlan, trace_id: &str, ) -> Result { if let Some(execution) = execute_stream_plan_via_local_tunnel(state, plan).await? { return Ok(execution); } let upstream_target_permit = state .upstream_target_admission .acquire(plan, trace_id) .await?; match DirectSyncExecutionRuntime::new().execute_stream(plan).await { Ok(mut execution) => { execution.upstream_target_permit = upstream_target_permit; record_manual_proxy_request_success(state, plan).await; Ok(execution) } Err(error) => { record_manual_proxy_request_failure(state, plan).await; Err(error.into()) } } } async fn execute_in_process_stream_with_oauth_retry( state: &AppState, plan: &mut ExecutionPlan, trace_id: &str, report_context: Option<&Value>, ) -> Result { let mut execution = execute_in_process_stream(state, plan, trace_id).await?; apply_stream_summary_report_context(&mut execution, report_context); if execution.status_code >= 400 && refresh_oauth_plan_auth_for_retry(state, plan, execution.status_code, None, trace_id) .await { drop(execution); execution = execute_in_process_stream(state, plan, trace_id).await?; apply_stream_summary_report_context(&mut execution, report_context); } Ok(execution) } fn should_use_direct_sse_passthrough( plan: &ExecutionPlan, plan_kind: &str, report_context: Option<&Value>, execution: &DirectUpstreamStreamExecution, ) -> bool { if !(200..300).contains(&execution.status_code) { return false; } if plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND { return false; } if !response_headers_indicate_sse(&execution.headers) { return false; } if !plan .provider_api_format .eq_ignore_ascii_case(plan.client_api_format.as_str()) { return false; } if client_format_allows_proxy_generated_sse_control_blocks(plan) { return false; } if maybe_build_provider_private_stream_normalizer(report_context).is_some() { return false; } let normalized_stream_report_context = normalize_provider_private_report_context(report_context); if maybe_build_stream_response_rewriter(normalized_stream_report_context.as_ref()).is_some() { return false; } let direct_stream_finalize_kind = resolve_core_stream_direct_finalize_report_kind(plan_kind); should_skip_direct_finalize_prefetch( direct_stream_finalize_kind.as_deref(), execution.headers.get("content-type").map(String::as_str), plan.provider_api_format.as_str(), plan.client_api_format.as_str(), false, false, ) } fn direct_upstream_response_byte_stream( response: DirectUpstreamResponse, ) -> BoxStream<'static, Result> { match response { DirectUpstreamResponse::Reqwest(response) => response .bytes_stream() .map(|item| item.map_err(|err| format_upstream_request_error(&err))) .boxed(), DirectUpstreamResponse::BrowserWreq(response) => response .bytes_stream() .map(|item| item.map_err(|err| format_wreq_upstream_request_error(&err))) .boxed(), DirectUpstreamResponse::LocalTunnel(mut response) => stream! { loop { match response.next_chunk().await { Ok(Some(chunk)) => yield Ok(chunk), Ok(None) => break, Err(err) => { yield Err(err); break; } } } } .boxed(), } } async fn await_direct_passthrough_first_item( future: F, started_at: Instant, timeout: Option, ) -> Result where F: Future, { let Some(timeout) = timeout else { return Ok(future.await); }; let Some(remaining) = timeout.checked_sub(started_at.elapsed()) else { return Err(timeout); }; if remaining.is_zero() { return Err(timeout); } tokio::time::timeout(remaining, future) .await .map_err(|_| timeout) } #[allow(clippy::too_many_arguments)] async fn forward_direct_passthrough_client_chunk( tx: &mpsc::Sender>, chunk: Bytes, downstream_dropped: &mut bool, client_visible_stream_completed: &mut bool, client_stream_completion_tracker: &mut ClientVisibleStreamCompletionTracker, client_stream_bytes: &mut u64, buffered_body: &mut Vec, client_body_truncated: &mut bool, max_stream_body_buffer_bytes: usize, stream_started_at: Instant, last_client_chunk_elapsed_ms: &mut u64, trace_id: &str, request_id_for_log: &str, candidate_id: Option<&str>, ) -> bool { if chunk.is_empty() { return false; } append_stream_capture_bytes( buffered_body, chunk.as_ref(), max_stream_body_buffer_bytes, client_body_truncated, ); if *downstream_dropped { return false; } let chunk_len = u64::try_from(chunk.len()).unwrap_or(u64::MAX); let send_started_at = Instant::now(); if tx.send(Ok(chunk.clone())).await.is_err() { debug!( event_name = "direct_passthrough_downstream_disconnected", log_type = "ops", trace_id = %trace_id, request_id = %request_id_for_log, candidate_id = ?candidate_id, "gateway direct passthrough downstream dropped; cancelling upstream stream" ); *downstream_dropped = true; return false; } observe_gateway_stage_ms( "direct_passthrough_body_send_wait", send_started_at.elapsed().as_millis() as u64, ); *client_visible_stream_completed |= client_stream_completion_tracker.observe_chunk(chunk.as_ref()); *client_stream_bytes = client_stream_bytes.saturating_add(chunk_len); *last_client_chunk_elapsed_ms = stream_started_at .elapsed() .as_millis() .min(u128::from(u64::MAX)) as u64; true } #[allow(clippy::too_many_arguments)] async fn execute_stream_from_direct_passthrough( state: &AppState, plan: ExecutionPlan, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, report_kind: Option, report_context: Option, candidate_started_unix_secs: u64, stream_started_at: Instant, mut stage_trace: RequestStageTrace, execution: DirectUpstreamStreamExecution, in_flight_guard: Option, ) -> Result>, GatewayError> { let DirectUpstreamStreamExecution { request_id: _, candidate_id: _, status_code, mut headers, provider_api_format: _, stream_summary_report_context: _, response, started_at: upstream_started_at, stream_first_byte_timeout, upstream_target_permit, } = execution; let request_id = plan.request_id.clone(); let candidate_id = plan.candidate_id.clone(); let request_id_for_log = short_request_id(request_id.as_str()); let mut report_context = attach_provider_response_headers_to_report_context(report_context, &headers); if status_code == 200 { seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await; if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { seed_kiro_report_context_input_tokens(&plan, &mut report_context); } seed_kiro_report_context_prompt_cache_usage(state, &plan, &mut report_context).await; } let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); state.usage_runtime.record_stream_started( state.data.as_ref(), &lifecycle_seed, status_code, None, ); let request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); if let Some(snapshot) = request_candidate_status_snapshot { let state_bg = state.clone(); tokio::spawn(async move { record_local_request_candidate_status_snapshot( &state_bg, &snapshot, SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Streaming, status_code: Some(status_code), error_type: None, error_message: None, latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: None, }, ) .await; }); } apply_endpoint_response_header_rules(state, &plan, &mut headers, None).await?; let headers_for_report = headers.clone(); headers.insert(CONTROL_REQUEST_ID_HEADER.to_string(), request_id.clone()); if let Some(candidate_id) = candidate_id .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) { headers.insert( CONTROL_CANDIDATE_ID_HEADER.to_string(), candidate_id.to_string(), ); } headers.remove("content-length"); let (tx, rx) = mpsc::channel::>(128); let state_for_report = state.clone(); let plan_for_report = plan; let trace_id_owned = trace_id.to_string(); let report_kind_owned = report_kind; let report_context_owned = report_context; let lifecycle_seed_for_report = lifecycle_seed; let direct_stream_finalize_kind_owned = resolve_core_stream_direct_finalize_report_kind(plan_kind); let normalized_stream_report_context_owned = normalize_provider_private_report_context(report_context_owned.as_ref()); let stream_started_at_for_report = stream_started_at; observe_gateway_stage_trace_ms( &mut stage_trace, "stream_response_ready", stream_elapsed_ms_since(stream_started_at), ); let stage_trace_for_report = stage_trace; let request_diagnostics_for_report = current_request_diagnostics(); let request_id_for_report = request_id.clone(); let request_id_for_report_log = request_id_for_log.clone(); let candidate_id_for_report = candidate_id.clone(); let provider_pool_in_flight_guard_for_report = in_flight_guard; tokio::spawn(async move { let mut stage_trace_for_report = stage_trace_for_report; let _stream_total_guard = StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report); let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report; let _upstream_target_permit = upstream_target_permit; let max_stream_body_buffer_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; let stream_usage_report_context = normalized_stream_report_context_owned.clone().or_else(|| { Some(serde_json::json!({ "provider_api_format": plan_for_report.provider_api_format.as_str(), "client_api_format": plan_for_report.client_api_format.as_str(), })) }); let mut stream_usage_observer = stream_usage_report_context .as_ref() .map(|_| StreamingStandardTerminalObserver::default()); let mut stream_usage_observer_buffered = Vec::new(); let mut provider_buffered_body = Vec::new(); let mut buffered_body = Vec::new(); let mut provider_body_truncated = false; let mut client_body_truncated = false; let mut upstream_control_filter = Some(SseControlBlockFilter::default()); let mut client_stream_completion_tracker = ClientVisibleStreamCompletionTracker::default(); let mut client_visible_stream_completed = false; let mut usage_stream_telemetry: Option = None; let telemetry: Option = None; let mut provider_stream_bytes = 0u64; let mut client_stream_bytes = 0u64; let mut last_client_chunk_elapsed_ms = 0u64; let mut downstream_dropped = false; let mut terminal_failure: Option = None; let mut upstream = direct_upstream_response_byte_stream(response); let mut observed_first_upstream_body = false; let mut observed_first_client_send = false; loop { if downstream_dropped { break; } let item = if usage_stream_telemetry .as_ref() .and_then(|telemetry| telemetry.ttfb_ms) .is_none() { tokio::select! { biased; _ = tx.closed(), if !downstream_dropped => { downstream_dropped = true; break; } result = await_direct_passthrough_first_item( upstream.next(), upstream_started_at, stream_first_byte_timeout, ) => { match result { Ok(item) => item, Err(timeout) => { terminal_failure = Some(build_stream_failure_report( "first_byte_timeout", stream_first_byte_timeout_message(timeout), 504, )); break; } } } } } else { tokio::select! { biased; _ = tx.closed(), if !downstream_dropped => { downstream_dropped = true; break; } item = upstream.next() => item, } }; let Some(item) = item else { break; }; let chunk = match item { Ok(chunk) => chunk, Err(message) => { warn!( event_name = "direct_passthrough_body_read_error", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), upstream_bytes = provider_stream_bytes, error = %message, "gateway direct passthrough upstream body read failed" ); terminal_failure = Some(build_stream_failure_report( "execution_runtime_stream_read_error", message, 502, )); break; } }; if chunk.is_empty() { continue; } let observed_at = Instant::now(); if !observed_first_upstream_body { observed_first_upstream_body = true; observe_gateway_stage_trace_ms( &mut stage_trace_for_report, "direct_passthrough_upstream_body_first", stream_elapsed_ms_at(stream_started_at_for_report, observed_at), ); } maybe_record_first_stream_event_started( &state_for_report, &lifecycle_seed_for_report, status_code, stream_started_at_for_report, observed_at, telemetry.as_ref(), &mut usage_stream_telemetry, ); if usage_stream_telemetry .as_ref() .and_then(|telemetry| telemetry.ttfb_ms) .is_some() && provider_stream_bytes == 0 { observe_gateway_stage_trace_ms( &mut stage_trace_for_report, "stream_first_data", stream_elapsed_ms_at(stream_started_at_for_report, observed_at), ); } let provider_chunk = chunk.clone(); if let Some(client_chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { let sent_client_chunk = forward_direct_passthrough_client_chunk( &tx, client_chunk, &mut downstream_dropped, &mut client_visible_stream_completed, &mut client_stream_completion_tracker, &mut client_stream_bytes, &mut buffered_body, &mut client_body_truncated, max_stream_body_buffer_bytes, stream_started_at_for_report, &mut last_client_chunk_elapsed_ms, trace_id_owned.as_str(), request_id_for_report_log.as_str(), candidate_id_for_report.as_deref(), ) .await; if sent_client_chunk && !observed_first_client_send { observed_first_client_send = true; observe_gateway_stage_trace_ms( &mut stage_trace_for_report, "direct_passthrough_first_client_send", stream_elapsed_ms_since(stream_started_at_for_report), ); } } provider_stream_bytes = provider_stream_bytes .saturating_add(u64::try_from(provider_chunk.len()).unwrap_or(u64::MAX)); append_stream_capture_bytes( &mut provider_buffered_body, provider_chunk.as_ref(), max_stream_body_buffer_bytes, &mut provider_body_truncated, ); if let (Some(observer), Some(report_context)) = ( stream_usage_observer.as_mut(), stream_usage_report_context.as_ref(), ) { observe_stream_usage_bytes( observer, report_context, &mut stream_usage_observer_buffered, provider_chunk.as_ref(), ); } let provider_private_error_body_json = extract_provider_private_stream_error_body( stream_usage_report_context.as_ref(), provider_chunk.as_ref(), ); if let Some(error_body_json) = provider_private_error_body_json { let error_status_code = resolve_local_sync_error_status_code(status_code, &error_body_json); terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, )); break; } } if terminal_failure.is_none() { if let Some(client_chunk) = flush_upstream_sse_control_filter(&mut upstream_control_filter) { let _ = forward_direct_passthrough_client_chunk( &tx, client_chunk, &mut downstream_dropped, &mut client_visible_stream_completed, &mut client_stream_completion_tracker, &mut client_stream_bytes, &mut buffered_body, &mut client_body_truncated, max_stream_body_buffer_bytes, stream_started_at_for_report, &mut last_client_chunk_elapsed_ms, trace_id_owned.as_str(), request_id_for_report_log.as_str(), candidate_id_for_report.as_deref(), ) .await; } } if let Some(failure) = terminal_failure.as_ref().filter(|_| !downstream_dropped) { match encode_terminal_sse_error_event(failure) { Ok(error_event) => { let _ = forward_direct_passthrough_client_chunk( &tx, error_event, &mut downstream_dropped, &mut client_visible_stream_completed, &mut client_stream_completion_tracker, &mut client_stream_bytes, &mut buffered_body, &mut client_body_truncated, max_stream_body_buffer_bytes, stream_started_at_for_report, &mut last_client_chunk_elapsed_ms, trace_id_owned.as_str(), request_id_for_report_log.as_str(), candidate_id_for_report.as_deref(), ) .await; } Err(err) => { warn!( event_name = "direct_passthrough_terminal_error_event_encode_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway direct passthrough failed to encode terminal SSE error event" ); } } } drop(tx); let mut stream_terminal_summary = finalize_stream_usage_observer( &mut stream_usage_observer, stream_usage_report_context.as_ref(), &mut stream_usage_observer_buffered, ); if downstream_dropped && client_visible_stream_completed && terminal_failure.is_none() { debug!( event_name = "direct_passthrough_downstream_closed_after_done", log_type = "debug", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), "gateway treats direct passthrough downstream close after terminal SSE event as completed" ); downstream_dropped = false; } if downstream_dropped { let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at_for_report, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes, )); let report_context_for_payload = report_context_with_stage_trace( report_context_owned, stage_trace_for_report, stream_started_at_for_report, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id_owned, report_kind_owned.unwrap_or_default(), report_context_for_payload, 499, headers_for_report, &provider_buffered_body, provider_body_truncated, &buffered_body, client_body_truncated, stream_terminal_summary, terminal_telemetry, ); record_stream_terminal_usage( &state_for_report, &plan_for_report, usage_payload.report_context.as_ref(), &usage_payload, true, ); record_local_request_candidate_status( &state_for_report, &plan_for_report, usage_payload.report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Cancelled, status_code: Some(499), error_type: Some("downstream_disconnect".to_string()), error_message: Some("client disconnected before stream completion".to_string()), latency_ms: usage_payload .telemetry .as_ref() .and_then(|value| value.elapsed_ms), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(current_request_candidate_unix_ms()), }, ) .await; return; } if let Some(failure) = terminal_failure { record_manual_proxy_stream_error(&state_for_report, &plan_for_report).await; let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at_for_report, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes, )); let report_context_for_payload = report_context_with_stage_trace( report_context_owned, stage_trace_for_report, stream_started_at_for_report, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), ); submit_midstream_stream_failure( &state_for_report, &trace_id_owned, &plan_for_report, direct_stream_finalize_kind_owned.as_deref(), report_context_for_payload, headers_for_report, terminal_telemetry, &provider_buffered_body, candidate_started_unix_secs, failure, ) .await; return; } maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state_for_report, &plan_for_report, report_context_owned.as_ref(), &mut stream_terminal_summary, ) .await; let requires_observed_terminal_event = stream_requires_observed_terminal_event( plan_for_report.provider_api_format.as_str(), stream_usage_report_context.as_ref(), ); ensure_stream_terminal_summary_for_missing_observed_finish( &mut stream_terminal_summary, requires_observed_terminal_event, ); let missing_observed_finish = stream_terminal_summary_missing_observed_finish_with_requirement( stream_terminal_summary.as_ref(), requires_observed_terminal_event, ); let stream_failed = stream_terminal_summary_represents_failure_with_requirement( stream_terminal_summary.as_ref(), requires_observed_terminal_event, ); let stream_terminal_error_message = stream_terminal_summary .as_ref() .and_then(|summary| summary.parser_error.clone()) .or_else(|| { missing_observed_finish.then(|| { "execution runtime stream ended before provider terminal event".to_string() }) }); let should_submit_report = report_kind_owned.is_some(); let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at_for_report, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes, )); let report_context_for_payload = report_context_with_stage_trace( report_context_owned, stage_trace_for_report, stream_started_at_for_report, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id_owned.clone(), report_kind_owned.unwrap_or_default(), report_context_for_payload, status_code, headers_for_report, &provider_buffered_body, provider_body_truncated, &buffered_body, client_body_truncated, stream_terminal_summary, terminal_telemetry, ); if stream_failed { warn!( event_name = "direct_passthrough_stream_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), status_code, error_message = stream_terminal_error_message.as_deref().unwrap_or_default(), "gateway direct passthrough stream ended with a failed terminal state" ); } else { apply_local_execution_effect( &state_for_report, LocalExecutionEffectContext { plan: &plan_for_report, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect), ) .await; apply_local_execution_effect( &state_for_report, LocalExecutionEffectContext { plan: &plan_for_report, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect), ) .await; apply_local_execution_effect( &state_for_report, LocalExecutionEffectContext { plan: &plan_for_report, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::PoolSuccessStream { payload: &usage_payload, }, ) .await; } record_stream_terminal_usage( &state_for_report, &plan_for_report, usage_payload.report_context.as_ref(), &usage_payload, false, ); record_local_request_candidate_status( &state_for_report, &plan_for_report, usage_payload.report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: if stream_failed { RequestCandidateStatus::Failed } else { RequestCandidateStatus::Success }, status_code: Some(status_code), error_type: if stream_failed { if missing_observed_finish { Some("stream_missing_terminal_event".to_string()) } else { Some("stream_terminal_error".to_string()) } } else { None }, error_message: stream_failed .then_some(stream_terminal_error_message) .flatten(), latency_ms: usage_payload .telemetry .as_ref() .and_then(|value| value.elapsed_ms), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(current_request_candidate_unix_ms()), }, ) .await; if should_submit_report { if let Err(err) = submit_stream_report(&state_for_report, usage_payload).await { warn!( event_name = "execution_report_submit_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), report_scope = "direct_passthrough_stream", error = ?err, "gateway failed to submit direct passthrough stream execution report" ); } } }); let body_stream = build_sse_body_stream(Vec::new(), rx, false, false, SSE_KEEPALIVE_INTERVAL); Ok(Some(build_client_response_from_parts( status_code, &headers, Body::from_stream(body_stream), trace_id, Some(decision), )?)) } #[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection pub(crate) async fn execute_execution_runtime_stream( state: &AppState, mut plan: ExecutionPlan, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, report_kind: Option, mut report_context: Option, ) -> Result>, GatewayError> { let stream_started_at = Instant::now(); let mut stage_trace = RequestStageTrace::from_env(); let candidate_slot_started_at = Instant::now(); ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await; observe_gateway_stage_trace_ms( &mut stage_trace, "stream_candidate_slot", candidate_slot_started_at.elapsed().as_millis() as u64, ); let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); let request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); let usage_pending_started_at = Instant::now(); state .usage_runtime .record_pending(state.data.as_ref(), lifecycle_seed.clone()); observe_gateway_stage_trace_ms( &mut stage_trace, "stream_usage_pending", usage_pending_started_at.elapsed().as_millis() as u64, ); let candidate_started_unix_secs = current_request_candidate_unix_ms(); if let Some(snapshot) = request_candidate_status_snapshot.clone() { let state_bg = state.clone(); tokio::spawn(async move { record_local_request_candidate_status_snapshot( &state_bg, &snapshot, 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 plan_request_id_for_log = short_request_id(plan.request_id.as_str()); 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 provider_in_flight_started_at = Instant::now(); let mut provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard( state.runtime_state.clone(), &plan.provider_id, plan.request_id.as_str(), plan.candidate_id.as_deref(), key_id.as_str(), ) .await; observe_gateway_stage_trace_ms( &mut stage_trace, "stream_provider_in_flight", provider_in_flight_started_at.elapsed().as_millis() as u64, ); match maybe_execute_grok_stream(&plan, report_context.as_ref()).await { Ok(Some(grok_stream)) => { return execute_stream_from_frame_stream( state, plan, trace_id, decision, plan_kind, report_kind, grok_stream.report_context.or(report_context), candidate_started_unix_secs, stream_started_at, stage_trace, grok_stream.frame_stream, provider_pool_in_flight_guard.take(), ) .await; } Ok(None) => {} Err(err) => { info!( 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 = provider_name.as_str(), endpoint_id = %endpoint_id, key_id = %key_id, model_name = model_name.as_str(), candidate_index = candidate_index.as_str(), error = %err, "gateway Grok stream 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(format!("{err:?}")), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } } match maybe_execute_windsurf_stream(state, &plan, report_context.as_ref()).await { Ok(Some(windsurf_stream)) => { return execute_stream_from_frame_stream( state, plan, trace_id, decision, plan_kind, report_kind, windsurf_stream.report_context.or(report_context), candidate_started_unix_secs, stream_started_at, stage_trace, windsurf_stream.frame_stream, provider_pool_in_flight_guard.take(), ) .await; } Ok(None) => {} Err(err) => { info!( event_name = "windsurf_native_execution_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan.candidate_id, provider_name = provider_name.as_str(), endpoint_id = %endpoint_id, key_id = %key_id, model_name = model_name.as_str(), candidate_index = candidate_index.as_str(), error = %err, "gateway native Windsurf stream 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("windsurf_native_execution_unavailable".to_string()), error_message: Some(err.to_string()), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } } match maybe_execute_kiro_web_search_stream(state, &plan, report_context.as_ref()).await { Ok(Some(kiro_web_search)) => { return execute_stream_from_frame_stream( state, plan, trace_id, decision, plan_kind, report_kind, kiro_web_search.report_context.or(report_context), candidate_started_unix_secs, stream_started_at, stage_trace, kiro_web_search.frame_stream, provider_pool_in_flight_guard.take(), ) .await; } Ok(None) => {} Err(err) => { info!( event_name = "kiro_web_search_mcp_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan.candidate_id, provider_name = provider_name.as_str(), endpoint_id = %endpoint_id, key_id = %key_id, model_name = model_name.as_str(), candidate_index = candidate_index.as_str(), error = %err, "gateway Kiro web_search MCP 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("kiro_web_search_mcp_unavailable".to_string()), error_message: Some(format!("{err:?}")), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } } match maybe_execute_chatgpt_web_image_stream(state, &plan, report_context.as_ref()).await { Ok(Some(chatgpt_web_image)) => { return execute_stream_from_frame_stream( state, plan, trace_id, decision, plan_kind, report_kind, chatgpt_web_image.report_context.or(report_context), candidate_started_unix_secs, stream_started_at, stage_trace, chatgpt_web_image.frame_stream, provider_pool_in_flight_guard.take(), ) .await; } Ok(None) => {} Err(err) => { info!( 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 = provider_name.as_str(), endpoint_id = %endpoint_id, key_id = %key_id, model_name = model_name.as_str(), candidate_index = candidate_index.as_str(), error = %err, "gateway ChatGPT-Web image stream 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(format!("{err:?}")), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } } #[cfg(not(test))] { let upstream_headers_started_at = Instant::now(); let execution = match execute_in_process_stream_with_oauth_retry( state, &mut plan, trace_id, report_context.as_ref(), ) .await { Ok(execution) => execution, Err(InProcessStreamExecutionError::Gateway(err)) => { if matches!(err, GatewayError::AdmissionTimeout { .. }) { record_stream_admission_timeout_terminal_state( state, &plan, report_context.as_ref(), candidate_started_unix_secs, &err, ) .await; } return Err(err); } Err(InProcessStreamExecutionError::Transport(err)) => { info!( event_name = "stream_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 = %err, "gateway in-process stream 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("execution_runtime_unavailable".to_string()), error_message: Some(format!("{err:?}")), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } }; observe_gateway_stage_trace_ms( &mut stage_trace, "stream_upstream_headers", upstream_headers_started_at.elapsed().as_millis() as u64, ); if should_use_direct_sse_passthrough(&plan, plan_kind, report_context.as_ref(), &execution) { return execute_stream_from_direct_passthrough( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, execution, provider_pool_in_flight_guard.take(), ) .await; } let frame_stream = build_direct_execution_frame_stream(execution).boxed(); return execute_stream_from_frame_stream( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, frame_stream, provider_pool_in_flight_guard.take(), ) .await; } #[cfg(test)] { let remote_execution_runtime_base_url = state .execution_runtime_override_base_url() .unwrap_or_default(); if remote_execution_runtime_base_url.trim().is_empty() { let upstream_headers_started_at = Instant::now(); let execution = match execute_in_process_stream_with_oauth_retry( state, &mut plan, trace_id, report_context.as_ref(), ) .await { Ok(execution) => execution, Err(InProcessStreamExecutionError::Gateway(err)) => { if matches!(err, GatewayError::AdmissionTimeout { .. }) { record_stream_admission_timeout_terminal_state( state, &plan, report_context.as_ref(), candidate_started_unix_secs, &err, ) .await; } return Err(err); } Err(InProcessStreamExecutionError::Transport(err)) => { info!( event_name = "stream_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 = %err, "gateway in-process stream 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("execution_runtime_unavailable".to_string()), error_message: Some(err.to_string()), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } }; observe_gateway_stage_trace_ms( &mut stage_trace, "stream_upstream_headers", upstream_headers_started_at.elapsed().as_millis() as u64, ); if should_use_direct_sse_passthrough( &plan, plan_kind, report_context.as_ref(), &execution, ) { return execute_stream_from_direct_passthrough( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, execution, provider_pool_in_flight_guard.take(), ) .await; } let frame_stream = build_direct_execution_frame_stream(execution).boxed(); return execute_stream_from_frame_stream( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, frame_stream, provider_pool_in_flight_guard.take(), ) .await; } let response = match post_stream_plan_to_remote_execution_runtime( state, remote_execution_runtime_base_url, Some(trace_id), &plan, ) .await { Ok(response) => response, Err(err) => { warn!( event_name = "stream_execution_runtime_remote_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan.candidate_id, error = ?err, "gateway remote execution runtime stream 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("execution_runtime_unavailable".to_string()), error_message: Some(format!("{err:?}")), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } }; if response.status() != http::StatusCode::OK { 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: 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: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(Some(attach_control_metadata_headers( build_client_response(response, trace_id, Some(decision))?, Some(plan.request_id.as_str()), plan.candidate_id.as_deref(), )?)); } let frame_stream = response .bytes_stream() .map_err(|err| IoError::other(err.to_string())) .boxed(); return execute_stream_from_frame_stream( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, frame_stream, provider_pool_in_flight_guard.take(), ) .await; } } fn decode_stream_data_chunk( chunk_b64: Option<&str>, text: Option<&str>, ) -> Result, GatewayError> { if let Some(chunk_b64) = chunk_b64 { return base64::engine::general_purpose::STANDARD .decode(chunk_b64) .map_err(|err| GatewayError::Internal(err.to_string())); } Ok(text.unwrap_or_default().as_bytes().to_vec()) } fn response_headers_indicate_sse(headers: &BTreeMap) -> bool { headers .get("content-type") .map(String::as_str) .map(str::trim) .filter(|value| !value.is_empty()) .is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream")) } fn parse_prefetched_sync_json_body(body: &[u8]) -> Option { let stripped = strip_utf8_bom_and_ws(body); serde_json::from_slice::(stripped).ok() } fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result { let payload = failure .to_json_string() .map_err(|err| IoError::other(err.to_string()))?; let mut event = String::new(); for line in payload.lines() { event.push_str("data: "); event.push_str(line); event.push('\n'); } event.push_str("\ndata: [DONE]\n\n"); Ok(Bytes::from(event)) } fn image_stream_failed_event_name(report_context: Option<&Value>) -> &'static str { let operation = report_context .and_then(|value| value.get("image_request")) .and_then(|value| value.get("operation")) .and_then(Value::as_str) .map(str::trim) .unwrap_or_default(); if operation == "edit" { "image_edit.failed" } else { "image_generation.failed" } } fn encode_openai_image_failed_event( report_context: Option<&Value>, failure: &StreamFailureReport, ) -> Result { let event_name = image_stream_failed_event_name(report_context); let failure_body = failure .to_json_string() .map_err(|err| IoError::other(err.to_string()))?; let failure_json: Value = serde_json::from_str(&failure_body).map_err(|err| IoError::other(err.to_string()))?; let error = failure_json.get("error").cloned().unwrap_or_else(|| { serde_json::json!({ "type": failure.error_type.as_str(), "message": failure.error_message.as_str(), "code": failure.status_code, }) }); let payload = serde_json::json!({ "type": event_name, "error": error, }); let payload = serde_json::to_string(&payload).map_err(|err| IoError::other(err.to_string()))?; let mut event = format!("event: {event_name}\n"); for line in payload.lines() { event.push_str("data: "); event.push_str(line); event.push('\n'); } event.push('\n'); Ok(Bytes::from(event)) } fn should_limit_direct_finalize_prefetch(plan_kind: &str, has_local_stream_rewriter: bool) -> bool { plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND || has_local_stream_rewriter } fn client_format_allows_proxy_generated_sse_control_blocks(plan: &ExecutionPlan) -> bool { // OpenAI-compatible clients commonly parse every client-visible SSE event as // an OpenAI JSON payload or [DONE]. Keep the downstream wire format strict: // do not inject proxy-generated comments, pings, or keepalives for openai:*. !plan .client_api_format .trim() .to_ascii_lowercase() .starts_with("openai:") } fn build_sse_body_stream( prefetched_chunks_for_body: Vec, mut rx: mpsc::Receiver>, filter_control_blocks: bool, emit_keepalive: bool, keepalive_interval: Duration, ) -> impl futures_util::Stream> + Send + 'static { stream! { let mut upstream_control_filter = filter_control_blocks.then(SseControlBlockFilter::default); let mut sent_prefetched_chunk = false; for chunk in prefetched_chunks_for_body { if let Some(chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { sent_prefetched_chunk = true; yield Ok(chunk); } } if emit_keepalive { if !sent_prefetched_chunk { yield Ok(Bytes::from_static(SSE_KEEPALIVE_BYTES)); } let mut keepalive = tokio::time::interval(keepalive_interval); keepalive.set_missed_tick_behavior(MissedTickBehavior::Delay); keepalive.tick().await; loop { tokio::select! { biased; item = rx.recv() => { let Some(item) = item else { break; }; match item { Ok(chunk) => { if let Some(chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { yield Ok(chunk); } } Err(err) => yield Err(err), } } _ = keepalive.tick() => { yield Ok(Bytes::from_static(SSE_KEEPALIVE_BYTES)); } } } if let Some(chunk) = flush_upstream_sse_control_filter(&mut upstream_control_filter) { yield Ok(chunk); } } else { while let Some(item) = rx.recv().await { match item { Ok(chunk) => { if let Some(chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { yield Ok(chunk); } } Err(err) => yield Err(err), } } if let Some(chunk) = flush_upstream_sse_control_filter(&mut upstream_control_filter) { yield Ok(chunk); } } } } #[derive(Default)] struct SseControlBlockFilter { buffered: Vec, emitted_len: usize, passthrough_current_block: bool, } impl SseControlBlockFilter { fn push_chunk(&mut self, chunk: &[u8]) -> Vec { if chunk.is_empty() { return Vec::new(); } self.buffered.extend_from_slice(chunk); let mut output = Vec::new(); while let Some((block_end, separator_len)) = find_sse_block_boundary(&self.buffered) { let block_len = block_end + separator_len; let block = self.buffered.drain(..block_len).collect::>(); if self.passthrough_current_block { let emitted_len = self.emitted_len.min(block.len()); output.extend_from_slice(&block[emitted_len..]); } else if sse_block_has_data_line(&block) { output.extend_from_slice(&block); } self.emitted_len = 0; self.passthrough_current_block = false; } if self.passthrough_current_block { if self.buffered.len() > self.emitted_len { output.extend_from_slice(&self.buffered[self.emitted_len..]); self.emitted_len = self.buffered.len(); } } else if sse_buffer_has_data_line(&self.buffered) { self.passthrough_current_block = true; output.extend_from_slice(&self.buffered); self.emitted_len = self.buffered.len(); } if self.buffered.len() > SSE_CONTROL_FILTER_MAX_BUFFER_BYTES { let buffered = std::mem::take(&mut self.buffered); if self.passthrough_current_block { let emitted_len = self.emitted_len.min(buffered.len()); output.extend_from_slice(&buffered[emitted_len..]); } else { output.extend(buffered); } self.emitted_len = 0; self.passthrough_current_block = false; } output } fn finish(&mut self) -> Vec { if self.buffered.is_empty() { return Vec::new(); } let block = std::mem::take(&mut self.buffered); let emitted_len = self.emitted_len.min(block.len()); let passthrough_current_block = self.passthrough_current_block; self.emitted_len = 0; self.passthrough_current_block = false; if passthrough_current_block { block[emitted_len..].to_vec() } else if sse_block_has_data_line(&block) { block } else { Vec::new() } } } fn filter_upstream_sse_control_chunk( filter: &mut Option, chunk: Bytes, ) -> Option { let Some(filter) = filter.as_mut() else { return Some(chunk); }; let filtered = filter.push_chunk(chunk.as_ref()); (!filtered.is_empty()).then(|| Bytes::from(filtered)) } fn flush_upstream_sse_control_filter(filter: &mut Option) -> Option { let filtered = filter.as_mut()?.finish(); (!filtered.is_empty()).then(|| Bytes::from(filtered)) } fn find_sse_block_boundary(buffer: &[u8]) -> Option<(usize, usize)> { let lf = buffer .windows(2) .position(|window| window == b"\n\n") .map(|index| (index, 2)); let crlf = buffer .windows(4) .position(|window| window == b"\r\n\r\n") .map(|index| (index, 4)); match (lf, crlf) { (Some(lf), Some(crlf)) => Some(if lf.0 <= crlf.0 { lf } else { crlf }), (Some(lf), None) => Some(lf), (None, Some(crlf)) => Some(crlf), (None, None) => None, } } fn sse_block_has_data_line(block: &[u8]) -> bool { let Ok(text) = std::str::from_utf8(block) else { return true; }; text.lines() .any(|line| line.trim_start().starts_with("data:")) } fn sse_buffer_has_data_line(buffer: &[u8]) -> bool { let Ok(text) = std::str::from_utf8(buffer) else { return true; }; text.lines() .any(|line| line.trim_start().starts_with("data:")) } #[derive(Default)] struct ClientVisibleStreamCompletionTracker { line_buffer: Vec, event_type: Option, data_payload: String, has_data_payload: bool, skip_next_lf: bool, completed: bool, } impl ClientVisibleStreamCompletionTracker { fn observe_chunk(&mut self, chunk: &[u8]) -> bool { if self.completed { return true; } if chunk.is_empty() { return false; } for byte in chunk { if self.skip_next_lf { self.skip_next_lf = false; if *byte == b'\n' { continue; } } match *byte { b'\n' => self.finish_line(), b'\r' => { self.finish_line(); self.skip_next_lf = true; } _ => { self.line_buffer.push(*byte); if self.line_buffer.len() > SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES { self.line_buffer.clear(); } } } if self.completed { break; } } self.completed } fn finish_line(&mut self) { let line = std::mem::take(&mut self.line_buffer); let Ok(line) = std::str::from_utf8(&line) else { self.reset_current_event(); return; }; let line = line.trim(); if line.is_empty() { self.completed = self.current_event_is_terminal(); self.reset_current_event(); return; } if let Some(event_type) = line.strip_prefix("event:").map(str::trim) { self.event_type = Some(event_type.to_string()); return; } if let Some(data) = line.strip_prefix("data:").map(str::trim) { if data.is_empty() { return; } if self.has_data_payload { self.data_payload.push('\n'); } self.data_payload.push_str(data); self.has_data_payload = true; } } fn current_event_is_terminal(&self) -> bool { self.event_type .as_deref() .is_some_and(is_terminal_sse_event_type) || (self.has_data_payload && sse_data_payload_is_terminal(&self.data_payload)) } fn reset_current_event(&mut self) { self.event_type = None; self.data_payload.clear(); self.has_data_payload = false; } } fn is_terminal_sse_event_type(event_type: &str) -> bool { matches!( event_type, "message_stop" | "response.completed" | "response.failed" | "response.incomplete" | "error" ) } fn sse_data_payload_is_terminal(data: &str) -> bool { data == "[DONE]" || serde_json::from_str::(data).is_ok_and(|value| { value .get("type") .and_then(serde_json::Value::as_str) .is_some_and(is_terminal_sse_event_type) }) } fn stream_chunk_contains_sse_done(chunk: &[u8]) -> bool { let mut tracker = ClientVisibleStreamCompletionTracker::default(); tracker.observe_chunk(chunk) } struct ObservedStreamFrame { frame: StreamFrame, observed_at: Instant, } async fn read_next_observed_stream_frame( lines: &mut FramedRead, ) -> Result, GatewayError> where R: tokio::io::AsyncRead + Unpin, { Ok(read_next_frame(lines) .await? .map(|frame| ObservedStreamFrame { frame, observed_at: Instant::now(), })) } async fn next_stream_frame( buffered_frames: &mut VecDeque, lines: &mut FramedRead, ) -> Result, GatewayError> where R: tokio::io::AsyncRead + Unpin, { if let Some(frame) = buffered_frames.pop_front() { return Ok(Some(frame)); } read_next_observed_stream_frame(lines).await } fn should_refresh_stream_usage_telemetry( previous: Option<&ExecutionTelemetry>, next: &ExecutionTelemetry, ) -> bool { let previous_ttfb = previous.and_then(|telemetry| telemetry.ttfb_ms); let previous_elapsed = previous.and_then(|telemetry| telemetry.elapsed_ms); let next_ttfb = next.ttfb_ms; let next_elapsed = next.elapsed_ms; (next_ttfb.is_some() && next_ttfb != previous_ttfb) || (next_elapsed.is_some() && next_elapsed != previous_elapsed) } fn stream_elapsed_ms_since(started_at: Instant) -> u64 { started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 } fn stream_elapsed_ms_at(started_at: Instant, observed_at: Instant) -> u64 { observed_at .saturating_duration_since(started_at) .as_millis() .min(u128::from(u64::MAX)) as u64 } fn first_stream_event_telemetry( stream_started_at: Instant, event_observed_at: Instant, upstream_telemetry: Option<&ExecutionTelemetry>, ) -> ExecutionTelemetry { let elapsed_ms = stream_elapsed_ms_at(stream_started_at, event_observed_at); ExecutionTelemetry { ttfb_ms: Some(elapsed_ms), elapsed_ms: Some(elapsed_ms), upstream_bytes: upstream_telemetry.and_then(|telemetry| telemetry.upstream_bytes), } } fn maybe_capture_first_stream_event_telemetry( stream_started_at: Instant, event_observed_at: Instant, upstream_telemetry: Option<&ExecutionTelemetry>, usage_stream_telemetry: &mut Option, ) -> bool { if usage_stream_telemetry .as_ref() .and_then(|telemetry| telemetry.ttfb_ms) .is_some() { return false; } *usage_stream_telemetry = Some(first_stream_event_telemetry( stream_started_at, event_observed_at, upstream_telemetry, )); true } fn usage_refresh_telemetry( upstream_telemetry: &ExecutionTelemetry, usage_stream_telemetry: Option<&ExecutionTelemetry>, ) -> ExecutionTelemetry { ExecutionTelemetry { ttfb_ms: usage_stream_telemetry.and_then(|telemetry| telemetry.ttfb_ms), elapsed_ms: upstream_telemetry.elapsed_ms, upstream_bytes: upstream_telemetry.upstream_bytes, } } fn maybe_record_first_stream_event_started( state: &AppState, lifecycle_seed: &LifecycleUsageSeed, status_code: u16, stream_started_at: Instant, event_observed_at: Instant, upstream_telemetry: Option<&ExecutionTelemetry>, usage_stream_telemetry: &mut Option, ) { if !maybe_capture_first_stream_event_telemetry( stream_started_at, event_observed_at, upstream_telemetry, usage_stream_telemetry, ) { return; } let Some(telemetry) = usage_stream_telemetry.as_ref() else { return; }; state.usage_runtime.record_stream_started( state.data.as_ref(), lifecycle_seed, status_code, Some(telemetry), ); } fn build_terminal_stream_telemetry( stream_started_at: Instant, telemetry: Option<&ExecutionTelemetry>, usage_stream_telemetry: Option<&ExecutionTelemetry>, upstream_bytes: u64, ) -> ExecutionTelemetry { let current_elapsed_ms = stream_elapsed_ms_since(stream_started_at); let ttfb_ms = usage_stream_telemetry.and_then(|telemetry| telemetry.ttfb_ms); let prior_elapsed_ms = telemetry .and_then(|telemetry| telemetry.elapsed_ms) .or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.elapsed_ms)) .unwrap_or(0); let elapsed_ms = current_elapsed_ms .max(prior_elapsed_ms) .max(ttfb_ms.unwrap_or(0)); ExecutionTelemetry { ttfb_ms, elapsed_ms: Some(elapsed_ms), upstream_bytes: Some(upstream_bytes), } } fn should_skip_direct_finalize_prefetch( direct_stream_finalize_kind: Option<&str>, content_type: Option<&str>, provider_api_format: &str, client_api_format: &str, has_private_stream_normalizer: bool, has_local_stream_rewriter: bool, ) -> bool { if direct_stream_finalize_kind.is_none() { return false; } let content_type = content_type .map(str::trim) .filter(|value| !value.is_empty()) .unwrap_or_default() .to_ascii_lowercase(); if content_type.contains("text/event-stream") { return true; } if has_private_stream_normalizer || has_local_stream_rewriter { return false; } if !provider_api_format.eq_ignore_ascii_case(client_api_format) { return false; } if content_type.is_empty() { return true; } !(content_type.contains("json") || content_type.ends_with("+json")) } fn should_probe_success_failover_before_stream(headers: &BTreeMap) -> bool { let content_type = headers .get("content-type") .map(String::as_str) .map(str::trim) .filter(|value| !value.is_empty()) .unwrap_or_default() .to_ascii_lowercase(); content_type.contains("json") || content_type.ends_with("+json") } async fn probe_local_stream_success_failover_text( buffered_frames: &mut VecDeque, lines: &mut FramedRead, ) -> Result, GatewayError> where R: tokio::io::AsyncRead + Unpin, { while let Some(observed_frame) = read_next_observed_stream_frame(lines).await? { let probe_text = match &observed_frame.frame.payload { StreamFramePayload::Data { chunk_b64, text } => { match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) { Ok(chunk) if !chunk.is_empty() => { Some(String::from_utf8_lossy(&chunk).into_owned()) } Ok(_) | Err(_) => None, } } StreamFramePayload::Error { .. } | StreamFramePayload::Eof { .. } => None, StreamFramePayload::Headers { .. } | StreamFramePayload::Telemetry { .. } => None, }; buffered_frames.push_back(observed_frame); if probe_text.is_some() { return Ok(probe_text); } } Ok(None) } async fn execute_stream_from_frame_stream( state: &AppState, plan: ExecutionPlan, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, report_kind: Option, report_context: Option, candidate_started_unix_secs: u64, stream_started_at: Instant, mut stage_trace: RequestStageTrace, frame_stream: BoxStream<'static, Result>, in_flight_guard: Option, ) -> Result>, GatewayError> { let request_id = plan.request_id.as_str(); let request_id_for_log = short_request_id(request_id); let candidate_id = plan.candidate_id.as_deref(); let provider_name = plan.provider_name.as_deref().unwrap_or("-"); let model_name = plan.model_name.as_deref().unwrap_or("-"); let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); let request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); 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 reader = StreamReader::new(frame_stream); let mut lines = FramedRead::new(reader, LinesCodec::new()); let first_frame_started_at = Instant::now(); let first_frame = read_next_frame(&mut lines).await?.ok_or_else(|| { GatewayError::Internal("execution runtime stream ended before headers frame".to_string()) })?; observe_gateway_stage_trace_ms( &mut stage_trace, "stream_first_frame", first_frame_started_at.elapsed().as_millis() as u64, ); let StreamFramePayload::Headers { status_code, mut headers, } = first_frame.payload else { return Err(GatewayError::Internal( "execution runtime stream must start with headers frame".to_string(), )); }; let mut report_context = attach_provider_response_headers_to_report_context(report_context, &headers); if status_code == 200 { seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await; if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { seed_kiro_report_context_input_tokens(&plan, &mut report_context); } seed_kiro_report_context_prompt_cache_usage(state, &plan, &mut report_context).await; } let mut buffered_frames = VecDeque::new(); let mut stream_terminal_summary: Option = None; if status_code == 200 && should_probe_success_failover_before_stream(&headers) { let success_probe_text = probe_local_stream_success_failover_text(&mut buffered_frames, &mut lines).await?; if should_retry_next_local_candidate_stream( state, &plan, plan_kind, report_context.as_ref(), status_code, success_probe_text.as_deref(), ) .await { 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: Some(status_code), error_type: Some("success_failover_pattern".to_string()), error_message: Some( "execution runtime stream matched provider success failover rule" .to_string(), ), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; warn!( event_name = "local_stream_candidate_retry_scheduled", log_type = "event", trace_id = %trace_id, request_id = %request_id_for_log, status_code, provider_name = provider_name, endpoint_id = %plan.endpoint_id, key_id = %plan.key_id, model_name, candidate_index = candidate_index.as_str(), "gateway local stream decision retrying next candidate after success failover rule match" ); return Ok(None); } } let stream_error_finalize_kind = resolve_core_stream_error_finalize_report_kind(plan_kind, status_code); if !(200..300).contains(&status_code) { let provider_error_body = collect_error_body(&mut lines).await?; let private_error_body_json = extract_provider_private_stream_error_body( report_context.as_ref(), &provider_error_body, ); let provider_private_error_decoded = private_error_body_json.is_some(); let synthetic_body_json = (!provider_private_error_decoded && should_synthesize_non_success_stream_error_body(status_code, &provider_error_body)) .then(|| build_synthetic_non_success_stream_error_body(status_code, &headers)); let (provider_body_json, provider_body_base64) = if let Some(error_body_json) = private_error_body_json { (Some(error_body_json), None) } else { decode_stream_error_body(&headers, &provider_error_body) }; let client_status_code = stream_client_error_status_code_for_upstream_status(status_code); let wrapped_binary_body_json = if provider_private_error_decoded { None } else { wrap_non_json_binary_stream_error_for_client(plan_kind, &headers, &provider_error_body)? }; let (client_body_json, client_error_body, payload_client_body_json) = if let Some(body_json) = synthetic_body_json.or(wrapped_binary_body_json) { let body_bytes = serde_json::to_vec(&body_json) .map_err(|err| GatewayError::Internal(err.to_string()))?; (Some(body_json.clone()), body_bytes, Some(body_json)) } else if provider_private_error_decoded { let body_json = provider_body_json.clone().ok_or_else(|| { GatewayError::Internal( "decoded provider private stream error body is missing".to_string(), ) })?; let body_bytes = serde_json::to_vec(&body_json) .map_err(|err| GatewayError::Internal(err.to_string()))?; (Some(body_json), body_bytes, None) } else { ( provider_body_json.clone(), provider_error_body.clone(), provider_body_json.clone(), ) }; let error_response_text = local_failover_response_text(client_body_json.as_ref(), &client_error_body, None); let failover_analysis = resolve_local_candidate_failover_analysis_stream( state, &plan, report_context.as_ref(), status_code, error_response_text.as_deref(), ) .await; apply_local_execution_effect( state, LocalExecutionEffectContext { plan: &plan, report_context: report_context.as_ref(), }, LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect { status_code, classification: failover_analysis.classification, }), ) .await; apply_local_execution_effect( state, LocalExecutionEffectContext { plan: &plan, report_context: report_context.as_ref(), }, LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect { status_code, classification: 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, classification: failover_analysis.classification, }), ) .await; apply_local_execution_effect( state, LocalExecutionEffectContext { plan: &plan, report_context: report_context.as_ref(), }, LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect { status_code, response_text: error_response_text.as_deref(), }), ) .await; apply_local_execution_effect( state, LocalExecutionEffectContext { plan: &plan, report_context: report_context.as_ref(), }, LocalExecutionEffect::PoolError(LocalPoolErrorEffect { status_code, classification: failover_analysis.classification, headers: &headers, error_body: error_response_text.as_deref(), }), ) .await; let failover_decision = failover_analysis.decision; debug!( event_name = "execution_runtime_stream_failover_decided", log_type = "debug", trace_id = %trace_id, request_id = %request_id_for_log, candidate_id = ?candidate_id, plan_kind, status_code, provider_name, endpoint_id = %plan.endpoint_id, key_id = %plan.key_id, model_name, candidate_index = candidate_index.as_str(), failover_decision = failover_decision.as_str(), "gateway resolved execution runtime stream failover decision" ); if matches!(failover_decision, LocalFailoverDecision::RetryNextCandidate) { let terminal_unix_secs = current_request_candidate_unix_ms(); let error_trace_report_context = with_stream_error_trace_context( report_context.as_ref(), status_code, &headers, provider_body_json.as_ref(), &provider_error_body, error_response_text.as_deref(), 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(status_code), error_type: Some("retryable_upstream_status".to_string()), error_message: Some(format!( "execution runtime stream returned retryable status {status_code}" )), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; warn!( event_name = "local_stream_candidate_retry_scheduled", log_type = "event", trace_id = %trace_id, request_id = %request_id_for_log, status_code, provider_name = provider_name, endpoint_id = %plan.endpoint_id, key_id = %plan.key_id, model_name, candidate_index = candidate_index.as_str(), "gateway local stream decision retrying next candidate after retryable execution runtime status" ); return Ok(None); } if !matches!(failover_decision, LocalFailoverDecision::StopLocalFailover) && should_fallback_to_control_stream( plan_kind, status_code, stream_error_finalize_kind.is_some(), ) { let terminal_unix_secs = current_request_candidate_unix_ms(); let error_trace_report_context = with_stream_error_trace_context( report_context.as_ref(), status_code, &headers, provider_body_json.as_ref(), &provider_error_body, error_response_text.as_deref(), 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(status_code), error_type: Some("control_fallback".to_string()), error_message: Some(format!( "stream decision fell back to control after status {status_code}" )), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } let mut client_headers = if (300..400).contains(&status_code) { let mut headers = synthetic_error_response_headers(headers.clone()); headers.insert( "x-aether-upstream-status".to_string(), status_code.to_string(), ); headers } else { headers.clone() }; if provider_private_error_decoded { client_headers.remove("content-encoding"); client_headers.remove("content-length"); client_headers.insert("content-type".to_string(), "application/json".to_string()); } apply_endpoint_response_header_rules( state, &plan, &mut client_headers, client_body_json.as_ref(), ) .await?; let client_response_headers = client_headers.clone(); let error_trace_report_context = with_stream_error_trace_context( report_context.as_ref(), status_code, &headers, provider_body_json.as_ref(), &provider_error_body, error_response_text.as_deref(), failover_analysis, ); let payload = build_stream_error_sync_payload( trace_id, stream_error_finalize_kind .as_deref() .or(report_kind.as_deref()) .unwrap_or_default() .to_string(), error_trace_report_context.or(report_context), status_code, headers.clone(), provider_body_json, provider_body_base64, client_headers, payload_client_body_json, None, ); record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload); let terminal_unix_secs = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, &plan, payload.report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, status_code: Some(status_code), error_type: Some("execution_runtime_stream_non_success_status".to_string()), error_message: Some(format!( "execution runtime stream returned non-success status {status_code}" )), latency_ms: None, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; if stream_error_finalize_kind.is_some() { let response = submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload) .await?; return Ok(Some(attach_control_metadata_headers( response, Some(request_id), candidate_id, )?)); } return Ok(Some(attach_control_metadata_headers( build_client_response_from_parts( client_status_code, &client_response_headers, Body::from(client_error_body), trace_id, Some(decision), )?, Some(request_id), candidate_id, )?)); } let direct_stream_finalize_kind = resolve_core_stream_direct_finalize_report_kind(plan_kind); let normalized_stream_report_context = normalize_provider_private_report_context(report_context.as_ref()); let upstream_headers = headers.clone(); let mut private_stream_normalizer = maybe_build_provider_private_stream_normalizer(report_context.as_ref()); let mut local_stream_rewriter = maybe_build_stream_response_rewriter(normalized_stream_report_context.as_ref()); if private_stream_normalizer.is_some() || local_stream_rewriter.is_some() { headers.remove("content-encoding"); headers.remove("content-length"); headers.insert("content-type".to_string(), "text/event-stream".to_string()); } let upstream_content_type = upstream_headers.get("content-type").map(String::as_str); let skip_direct_finalize_prefetch = should_skip_direct_finalize_prefetch( direct_stream_finalize_kind.as_deref(), upstream_content_type, plan.provider_api_format.as_str(), plan.client_api_format.as_str(), private_stream_normalizer.is_some(), local_stream_rewriter.is_some(), ); let limit_direct_finalize_prefetch = should_limit_direct_finalize_prefetch(plan_kind, local_stream_rewriter.is_some()); let mut prefetched_chunks: Vec = Vec::new(); let mut provider_prefetched_body = Vec::new(); let mut prefetched_body = Vec::new(); let mut prefetched_inspection_body = Vec::new(); let mut prefetched_telemetry: Option = None; let mut prefetched_usage_telemetry: Option = None; let mut reached_eof = false; let mut sync_json_stream_bridge_active = false; if skip_direct_finalize_prefetch { debug!( event_name = "execution_runtime_stream_prefetch_skipped", log_type = "debug", trace_id = %trace_id, request_id = %request_id_for_log, candidate_id = ?candidate_id, plan_kind, provider_name, endpoint_id = %plan.endpoint_id, key_id = %plan.key_id, model_name, candidate_index = candidate_index.as_str(), content_type = upstream_content_type.unwrap_or("-"), provider_api_format = plan.provider_api_format.as_str(), client_api_format = plan.client_api_format.as_str(), "gateway skipped direct finalize prefetch for same-format passthrough stream" ); } if let Some(report_kind) = direct_stream_finalize_kind .as_ref() .filter(|_| !skip_direct_finalize_prefetch) { while prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES && prefetched_inspection_body.len() < MAX_STREAM_PREFETCH_BYTES { let next_frame_result = if limit_direct_finalize_prefetch { match tokio::time::timeout( REWRITTEN_STREAM_PREFETCH_TIMEOUT, next_stream_frame(&mut buffered_frames, &mut lines), ) .await { Ok(result) => result, Err(_) => { debug!( event_name = "execution_runtime_stream_prefetch_limited", log_type = "debug", trace_id = %trace_id, request_id = %request_id_for_log, candidate_id = ?candidate_id, plan_kind, report_kind, provider_name, endpoint_id = %plan.endpoint_id, key_id = %plan.key_id, model_name, candidate_index = candidate_index.as_str(), timeout_ms = REWRITTEN_STREAM_PREFETCH_TIMEOUT.as_millis() as u64, "gateway stopped rewritten stream prefetch before client-visible body" ); break; } } } else { next_stream_frame(&mut buffered_frames, &mut lines).await }; let Some(observed_frame) = (match next_frame_result { Ok(frame) => frame, Err(err) => { let failure = build_stream_failure_report( "execution_runtime_stream_frame_decode_error", format!("failed to decode execution runtime stream frame: {err:?}"), 502, ); return handle_prefetch_stream_failure( state, trace_id, decision, &plan, report_context, request_id, candidate_id, report_kind, headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, failure, ) .await; } }) else { reached_eof = true; break; }; let frame_observed_at = observed_frame.observed_at; match observed_frame.frame.payload { StreamFramePayload::Data { chunk_b64, text } => { if maybe_capture_first_stream_event_telemetry( stream_started_at, frame_observed_at, prefetched_telemetry.as_ref(), &mut prefetched_usage_telemetry, ) { observe_gateway_stage_trace_ms( &mut stage_trace, "stream_first_data", stream_elapsed_ms_at(stream_started_at, frame_observed_at), ); } let chunk = match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) { Ok(chunk) => chunk, Err(err) => { let failure = build_stream_failure_report( "execution_runtime_stream_chunk_decode_error", format!( "failed to decode execution runtime stream chunk: {err:?}" ), 502, ); return handle_prefetch_stream_failure( state, trace_id, decision, &plan, report_context, request_id, candidate_id, report_kind, headers, prefetched_usage_telemetry.clone(), &prefetched_body, failure, ) .await; } }; if chunk.is_empty() { continue; } provider_prefetched_body.extend_from_slice(&chunk); prefetched_inspection_body.extend_from_slice(&chunk); if let Some(error_body_json) = extract_provider_private_stream_error_body( report_context.as_ref(), &prefetched_inspection_body, ) { let error_status_code = resolve_local_sync_error_status_code(status_code, &error_body_json); return handle_prefetch_provider_private_stream_error( state, trace_id, decision, &plan, report_context, request_id, candidate_id, report_kind, headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, error_status_code, error_body_json, ) .await; } let inspection = inspect_prefetched_stream_body( &upstream_headers, &prefetched_inspection_body, ); match inspection { StreamPrefetchInspection::EmbeddedError(body_json) => { debug!( event_name = "execution_runtime_stream_prefetch_embedded_error_detected", log_type = "debug", trace_id = %trace_id, request_id = %request_id_for_log, candidate_id = ?candidate_id, plan_kind, report_kind, provider_name, endpoint_id = %plan.endpoint_id, key_id = %plan.key_id, model_name, candidate_index = candidate_index.as_str(), provider_prefetched_body_bytes = provider_prefetched_body.len(), "gateway detected embedded error while prefetching execution runtime stream" ); let payload = build_stream_sync_payload( trace_id, report_kind.clone(), report_context, status_code, headers, Some(body_json), None, prefetched_usage_telemetry.clone(), ); record_sync_terminal_usage( state, &plan, payload.report_context.as_ref(), &payload, ); let response = submit_local_core_error_or_sync_finalize( state, trace_id, decision, payload, ) .await?; return Ok(Some(attach_control_metadata_headers( response, Some(request_id), candidate_id, )?)); } StreamPrefetchInspection::NeedMore => {} StreamPrefetchInspection::NonError => {} } if !response_headers_indicate_sse(&upstream_headers) && (200..300).contains(&status_code) { if let Some(body_json) = parse_prefetched_sync_json_body(&prefetched_inspection_body) { match maybe_bridge_standard_sync_json_to_stream( &body_json, plan.provider_api_format.as_str(), plan.client_api_format.as_str(), report_context.as_ref(), ) { Ok(Some(outcome)) => { headers.remove("content-encoding"); headers.remove("content-length"); headers.insert( "content-type".to_string(), "text/event-stream".to_string(), ); stream_terminal_summary = outcome.terminal_summary; prefetched_body.extend_from_slice(&outcome.sse_body); prefetched_chunks.push(Bytes::from(outcome.sse_body)); sync_json_stream_bridge_active = true; break; } Ok(None) => {} Err(err) => { let failure = build_stream_failure_report( "execution_runtime_sync_json_stream_bridge_error", format!( "failed to bridge execution runtime sync json to stream: {err:?}" ), 502, ); return handle_prefetch_stream_failure( state, trace_id, decision, &plan, report_context, request_id, candidate_id, report_kind, headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, failure, ) .await; } } } } let normalized_chunk = if let Some(normalizer) = private_stream_normalizer.as_mut() { match normalizer.push_chunk(&chunk) { Ok(normalized_chunk) => normalized_chunk, Err(err) => { let failure = build_stream_failure_report( "execution_runtime_stream_rewrite_error", format!( "failed to normalize execution runtime stream chunk: {err:?}" ), 502, ); return handle_prefetch_stream_failure( state, trace_id, decision, &plan, report_context, request_id, candidate_id, report_kind, headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, failure, ) .await; } } } else { chunk }; let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut() { match rewriter.push_chunk(&normalized_chunk) { Ok(rewritten_chunk) => rewritten_chunk, Err(err) => { let failure = build_stream_failure_report( "execution_runtime_stream_rewrite_error", format!( "failed to rewrite execution runtime stream chunk: {err:?}" ), 502, ); return handle_prefetch_stream_failure( state, trace_id, decision, &plan, report_context, request_id, candidate_id, report_kind, headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, failure, ) .await; } } } else { normalized_chunk }; if !rewritten_chunk.is_empty() { prefetched_body.extend_from_slice(&rewritten_chunk); prefetched_chunks.push(Bytes::from(rewritten_chunk)); } if matches!(inspection, StreamPrefetchInspection::NonError) { break; } } StreamFramePayload::Telemetry { telemetry: frame_telemetry, } => { prefetched_telemetry = Some(frame_telemetry); } StreamFramePayload::Eof { summary } => { if summary.is_some() { stream_terminal_summary = summary; } reached_eof = true; break; } StreamFramePayload::Error { error } => { warn!( event_name = "stream_execution_prefetch_error_frame", log_type = "ops", trace_id = %trace_id, request_id, candidate_id = ?candidate_id, error = %error.message, "execution runtime stream emitted error frame during prefetch" ); return handle_prefetch_stream_failure( state, trace_id, decision, &plan, report_context, request_id, candidate_id, report_kind, headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, build_stream_failure_from_execution_error(&error), ) .await; } StreamFramePayload::Headers { .. } => {} } } } drop(private_stream_normalizer); drop(local_stream_rewriter); let initial_usage_telemetry = prefetched_usage_telemetry.clone().or_else(|| { prefetched_telemetry .as_ref() .map(|telemetry| usage_refresh_telemetry(telemetry, None)) }); state.usage_runtime.record_stream_started( state.data.as_ref(), &lifecycle_seed, status_code, initial_usage_telemetry.as_ref(), ); if let Some(snapshot) = request_candidate_status_snapshot { let state_bg = state.clone(); let latency_ms = prefetched_telemetry .as_ref() .and_then(|telemetry| telemetry.elapsed_ms); tokio::spawn(async move { record_local_request_candidate_status_snapshot( &state_bg, &snapshot, SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Streaming, status_code: Some(status_code), error_type: None, error_message: None, latency_ms, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: None, }, ) .await; }); } apply_endpoint_response_header_rules(state, &plan, &mut headers, None).await?; let request_id = request_id.to_string(); let candidate_id = candidate_id.map(ToOwned::to_owned); let (tx, mut rx) = mpsc::channel::>(16); let state_for_report = state.clone(); let trace_id_owned = trace_id.to_string(); let headers_for_report = headers.clone(); let report_kind_owned = report_kind; let report_context_owned = report_context; let normalized_stream_report_context_owned = normalized_stream_report_context; let lifecycle_seed_for_report = lifecycle_seed; let provider_prefetched_body_for_report = provider_prefetched_body; let prefetched_body_for_report = prefetched_body; let prefetched_chunks_for_body = prefetched_chunks; let sync_json_stream_bridge_active_for_report = sync_json_stream_bridge_active; let initial_telemetry = prefetched_telemetry; let initial_reached_eof = reached_eof; let direct_stream_finalize_kind_owned = direct_stream_finalize_kind; let candidate_started_unix_secs_for_report = candidate_started_unix_secs; let request_id_for_report = request_id.clone(); let request_id_for_report_log = short_request_id(&request_id); let candidate_id_for_report = candidate_id.clone(); let candidate_index_for_report = candidate_index.clone(); let is_openai_image_stream_for_report = plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND; let response_headers_are_sse = response_headers_indicate_sse(&headers); let emit_proxy_generated_sse_control_blocks = response_headers_are_sse && client_format_allows_proxy_generated_sse_control_blocks(&plan); let plan_for_report = plan; let emit_passthrough_sse_terminal_error = skip_direct_finalize_prefetch && response_headers_indicate_sse(&upstream_headers) && !is_openai_image_stream_for_report; let plan_kind_for_report = plan_kind.to_string(); let stream_started_at_for_report = stream_started_at; observe_gateway_stage_trace_ms( &mut stage_trace, "stream_response_ready", stream_elapsed_ms_since(stream_started_at), ); let stage_trace_for_report = stage_trace; let request_diagnostics_for_report = current_request_diagnostics(); let provider_pool_in_flight_guard_for_report = in_flight_guard; tokio::spawn(async move { let mut stage_trace_for_report = stage_trace_for_report; let _stream_total_guard = StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report); let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report; let max_stream_body_buffer_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; let mut provider_buffered_body = Vec::new(); let mut buffered_body = Vec::new(); let mut provider_body_truncated = false; let mut client_body_truncated = false; let mut private_stream_normalizer = if sync_json_stream_bridge_active_for_report { None } else { maybe_build_provider_private_stream_normalizer(report_context_owned.as_ref()) }; let mut local_stream_rewriter = if sync_json_stream_bridge_active_for_report { None } else { maybe_build_stream_response_rewriter(normalized_stream_report_context_owned.as_ref()) }; let stream_usage_report_context = normalized_stream_report_context_owned.clone().or_else(|| { Some(serde_json::json!({ "provider_api_format": plan_for_report.provider_api_format.as_str(), "client_api_format": plan_for_report.client_api_format.as_str(), })) }); let mut stream_usage_observer = stream_usage_report_context .as_ref() .filter(|_| !sync_json_stream_bridge_active_for_report) .map(|_| StreamingStandardTerminalObserver::default()); let mut stream_usage_observer_buffered = Vec::new(); append_stream_capture_bytes( &mut provider_buffered_body, &provider_prefetched_body_for_report, max_stream_body_buffer_bytes, &mut provider_body_truncated, ); append_stream_capture_bytes( &mut buffered_body, &prefetched_body_for_report, max_stream_body_buffer_bytes, &mut client_body_truncated, ); let mut client_stream_completion_tracker = ClientVisibleStreamCompletionTracker::default(); let mut client_visible_stream_completed = client_stream_completion_tracker.observe_chunk(&prefetched_body_for_report); let mut usage_stream_telemetry: Option = initial_usage_telemetry; let mut telemetry: Option = initial_telemetry; let reached_eof = initial_reached_eof; let mut downstream_dropped = false; let mut terminal_failure: Option = None; let initial_elapsed_ms = stream_started_at_for_report .elapsed() .as_millis() .min(u128::from(u64::MAX)) as u64; let last_upstream_frame_elapsed_ms = Arc::new(AtomicU64::new(initial_elapsed_ms)); let last_client_chunk_elapsed_ms = Arc::new(AtomicU64::new(if prefetched_body_for_report.is_empty() { 0 } else { initial_elapsed_ms })); let provider_stream_bytes = Arc::new(AtomicU64::new( u64::try_from(provider_prefetched_body_for_report.len()).unwrap_or(u64::MAX), )); let client_stream_bytes = Arc::new(AtomicU64::new( u64::try_from(prefetched_body_for_report.len()).unwrap_or(u64::MAX), )); let idle_monitor_done = Arc::new(AtomicBool::new(false)); let idle_monitor_handle = { let done = Arc::clone(&idle_monitor_done); let last_upstream = Arc::clone(&last_upstream_frame_elapsed_ms); let last_client = Arc::clone(&last_client_chunk_elapsed_ms); let provider_bytes = Arc::clone(&provider_stream_bytes); let client_bytes = Arc::clone(&client_stream_bytes); let trace_id_for_idle = trace_id_owned.clone(); let request_id_for_idle = request_id_for_report_log.clone(); let candidate_id_for_idle = candidate_id_for_report.clone(); let candidate_index_for_idle = candidate_index_for_report.clone(); let plan_kind_for_idle = plan_kind_for_report.clone(); let provider_name_for_idle = plan_for_report .provider_name .clone() .unwrap_or_else(|| "-".to_string()); let endpoint_id_for_idle = plan_for_report.endpoint_id.clone(); let key_id_for_idle = plan_for_report.key_id.clone(); let model_name_for_idle = plan_for_report .model_name .clone() .unwrap_or_else(|| "-".to_string()); let has_local_stream_rewriter_for_idle = local_stream_rewriter.is_some(); tokio::spawn(async move { let mut interval = tokio::time::interval(STREAM_IDLE_LOG_INTERVAL); interval.set_missed_tick_behavior(MissedTickBehavior::Delay); interval.tick().await; loop { interval.tick().await; if done.load(Ordering::Relaxed) { break; } let elapsed_ms = stream_started_at_for_report .elapsed() .as_millis() .min(u128::from(u64::MAX)) as u64; let last_upstream_frame_elapsed_ms = last_upstream.load(Ordering::Relaxed); let last_client_chunk_elapsed_ms = last_client.load(Ordering::Relaxed); let upstream_idle_ms = elapsed_ms.saturating_sub(last_upstream_frame_elapsed_ms); let client_idle_ms = if last_client_chunk_elapsed_ms == 0 { elapsed_ms } else { elapsed_ms.saturating_sub(last_client_chunk_elapsed_ms) }; if upstream_idle_ms >= STREAM_IDLE_LOG_INTERVAL_MS { warn!( event_name = "stream_execution_upstream_idle", log_type = "ops", trace_id = %trace_id_for_idle, request_id = %request_id_for_idle, candidate_id = ?candidate_id_for_idle.as_deref(), candidate_index = candidate_index_for_idle.as_str(), plan_kind = plan_kind_for_idle.as_str(), provider_name = provider_name_for_idle.as_str(), endpoint_id = %endpoint_id_for_idle, key_id = %key_id_for_idle, model_name = model_name_for_idle.as_str(), elapsed_ms, provider_bytes = provider_bytes.load(Ordering::Relaxed), client_bytes = client_bytes.load(Ordering::Relaxed), last_upstream_frame_elapsed_ms, last_client_chunk_elapsed_ms, "gateway stream has not received an upstream frame within the idle window" ); } else if client_idle_ms >= STREAM_IDLE_LOG_INTERVAL_MS && last_upstream_frame_elapsed_ms >= last_client_chunk_elapsed_ms { warn!( event_name = "stream_execution_client_visible_idle", log_type = "ops", trace_id = %trace_id_for_idle, request_id = %request_id_for_idle, candidate_id = ?candidate_id_for_idle.as_deref(), candidate_index = candidate_index_for_idle.as_str(), plan_kind = plan_kind_for_idle.as_str(), provider_name = provider_name_for_idle.as_str(), endpoint_id = %endpoint_id_for_idle, key_id = %key_id_for_idle, model_name = model_name_for_idle.as_str(), elapsed_ms, provider_bytes = provider_bytes.load(Ordering::Relaxed), client_bytes = client_bytes.load(Ordering::Relaxed), last_upstream_frame_elapsed_ms, last_client_chunk_elapsed_ms, local_stream_rewriter = has_local_stream_rewriter_for_idle, "gateway stream received upstream frames but has no recent client-visible chunk" ); } } }) }; if !provider_prefetched_body_for_report.is_empty() { let normalized_prefetched_chunk = if let Some(normalizer) = private_stream_normalizer.as_mut() { match normalizer.push_chunk(&provider_prefetched_body_for_report) { Ok(normalized_chunk) => Some(normalized_chunk), Err(err) => { warn!( event_name = "stream_execution_prefetch_normalize_restore_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to restore private stream normalization state after prefetch" ); terminal_failure = Some(build_stream_failure_report( "execution_runtime_stream_rewrite_error", format!( "failed to restore private stream normalization state after prefetch: {err:?}" ), 502, )); None } } } else { None }; let replay_chunk = normalized_prefetched_chunk .as_deref() .unwrap_or(provider_prefetched_body_for_report.as_slice()); if let (Some(observer), Some(report_context)) = ( stream_usage_observer.as_mut(), stream_usage_report_context.as_ref(), ) { observe_stream_usage_bytes( observer, report_context, &mut stream_usage_observer_buffered, replay_chunk, ); } if terminal_failure.is_none() { if let Some(rewriter) = local_stream_rewriter.as_mut() { if let Err(err) = rewriter.push_chunk(replay_chunk) { warn!( event_name = "stream_execution_prefetch_rewrite_restore_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to restore local stream rewrite state after prefetch" ); terminal_failure = Some(build_stream_failure_report( "execution_runtime_stream_rewrite_error", format!( "failed to restore local stream rewrite state after prefetch: {err:?}" ), 502, )); } } } } if terminal_failure.is_none() && !reached_eof { loop { let next_frame_result = tokio::select! { biased; _ = tx.closed(), if client_visible_stream_completed => { downstream_dropped = true; break; } result = next_stream_frame(&mut buffered_frames, &mut lines) => result, }; let next_frame = match next_frame_result { Ok(frame) => frame, Err(err) => { warn!( event_name = "stream_execution_frame_decode_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to decode execution runtime stream frame" ); terminal_failure = Some(build_stream_failure_report( "execution_runtime_stream_frame_decode_error", format!("failed to decode execution runtime stream frame: {err:?}"), 502, )); break; } }; let Some(observed_frame) = next_frame else { if tx.is_closed() { downstream_dropped = true; } break; }; let frame_observed_at = observed_frame.observed_at; let frame_elapsed_ms = stream_elapsed_ms_at(stream_started_at_for_report, frame_observed_at); last_upstream_frame_elapsed_ms.store(frame_elapsed_ms, Ordering::Relaxed); match observed_frame.frame.payload { StreamFramePayload::Data { chunk_b64, text } => { let first_data_before = usage_stream_telemetry .as_ref() .and_then(|telemetry| telemetry.ttfb_ms) .is_some(); maybe_record_first_stream_event_started( &state_for_report, &lifecycle_seed_for_report, status_code, stream_started_at_for_report, frame_observed_at, telemetry.as_ref(), &mut usage_stream_telemetry, ); let first_data_after = usage_stream_telemetry .as_ref() .and_then(|telemetry| telemetry.ttfb_ms) .is_some(); if !first_data_before && first_data_after { observe_gateway_stage_trace_ms( &mut stage_trace_for_report, "stream_first_data", stream_elapsed_ms_at( stream_started_at_for_report, frame_observed_at, ), ); } if sync_json_stream_bridge_active_for_report { continue; } let chunk = match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) { Ok(chunk) => chunk, Err(err) => { warn!( event_name = "stream_execution_chunk_decode_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to decode execution runtime chunk" ); terminal_failure = Some(build_stream_failure_report( "execution_runtime_stream_chunk_decode_error", format!( "failed to decode execution runtime stream chunk: {err:?}" ), 502, )); break; } }; if chunk.is_empty() { continue; } provider_stream_bytes.fetch_add( u64::try_from(chunk.len()).unwrap_or(u64::MAX), Ordering::Relaxed, ); append_stream_capture_bytes( &mut provider_buffered_body, &chunk, max_stream_body_buffer_bytes, &mut provider_body_truncated, ); let normalized_chunk = if let Some(normalizer) = private_stream_normalizer.as_mut() { match normalizer.push_chunk(&chunk) { Ok(normalized_chunk) => normalized_chunk, Err(err) => { warn!( event_name = "stream_execution_chunk_normalize_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to normalize execution runtime stream chunk" ); terminal_failure = Some(build_stream_failure_report( "execution_runtime_stream_rewrite_error", format!("failed to normalize execution runtime stream chunk: {err:?}"), 502, )); break; } } } else { chunk }; let provider_private_error_body_json = extract_provider_private_stream_error_body( stream_usage_report_context.as_ref(), &normalized_chunk, ); if let (Some(observer), Some(report_context)) = ( stream_usage_observer.as_mut(), stream_usage_report_context.as_ref(), ) { observe_stream_usage_bytes( observer, report_context, &mut stream_usage_observer_buffered, &normalized_chunk, ); } let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut() { match rewriter.push_chunk(&normalized_chunk) { Ok(rewritten_chunk) => rewritten_chunk, Err(err) => { warn!( event_name = "stream_execution_chunk_rewrite_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to rewrite execution runtime stream chunk" ); terminal_failure = Some(build_stream_failure_report( "execution_runtime_stream_rewrite_error", format!("failed to rewrite execution runtime stream chunk: {err:?}"), 502, )); break; } } } else { normalized_chunk }; if rewritten_chunk.is_empty() { if let Some(error_body_json) = provider_private_error_body_json { let error_status_code = resolve_local_sync_error_status_code( status_code, &error_body_json, ); terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, )); break; } continue; } append_stream_capture_bytes( &mut buffered_body, &rewritten_chunk, max_stream_body_buffer_bytes, &mut client_body_truncated, ); let rewritten_chunk_len = u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX); if downstream_dropped { continue; } let rewritten_chunk = Bytes::from(rewritten_chunk); if tx.send(Ok(rewritten_chunk.clone())).await.is_err() { debug!( event_name = "stream_execution_downstream_disconnected", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), "gateway stream downstream dropped; continuing to drain execution runtime stream" ); downstream_dropped = true; } else { client_visible_stream_completed |= client_stream_completion_tracker .observe_chunk(rewritten_chunk.as_ref()); client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed); last_client_chunk_elapsed_ms.store( stream_started_at_for_report .elapsed() .as_millis() .min(u128::from(u64::MAX)) as u64, Ordering::Relaxed, ); } if let Some(error_body_json) = provider_private_error_body_json { let error_status_code = resolve_local_sync_error_status_code(status_code, &error_body_json); terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, )); break; } } StreamFramePayload::Telemetry { telemetry: frame_telemetry, } => { let usage_frame_telemetry = usage_refresh_telemetry( &frame_telemetry, usage_stream_telemetry.as_ref(), ); let should_refresh_stream_usage = should_refresh_stream_usage_telemetry( usage_stream_telemetry.as_ref(), &usage_frame_telemetry, ); if should_refresh_stream_usage { state_for_report.usage_runtime.record_stream_started( state_for_report.data.as_ref(), &lifecycle_seed_for_report, status_code, Some(&usage_frame_telemetry), ); usage_stream_telemetry = Some(usage_frame_telemetry); } telemetry = Some(frame_telemetry); } StreamFramePayload::Eof { summary } => { if summary.is_some() { stream_terminal_summary = summary; } break; } StreamFramePayload::Error { error } => { warn!( event_name = "stream_execution_error_frame", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = %error.message, "execution runtime stream emitted error frame" ); terminal_failure = Some(build_stream_failure_from_execution_error(&error)); break; } StreamFramePayload::Headers { .. } => {} } } } if downstream_dropped { debug!( event_name = "execution_runtime_stream_client_flush_skipped", log_type = "debug", debug_context = "redacted", stream_status = "downstream_disconnected", trace_id = %trace_id_owned, "gateway skipped client stream flush after downstream disconnect" ); } // Buffered stream state is partial after a terminal failure; normal // finish paths may synthesize successful terminal events. let should_finish_stream_rewriters = terminal_failure.is_none(); if let Some(normalizer) = private_stream_normalizer .as_mut() .filter(|_| should_finish_stream_rewriters) { match normalizer.finish() { Ok(normalized_chunk) if !normalized_chunk.is_empty() => { let provider_private_error_body_json = extract_provider_private_stream_error_body( stream_usage_report_context.as_ref(), &normalized_chunk, ); if let (Some(observer), Some(report_context)) = ( stream_usage_observer.as_mut(), stream_usage_report_context.as_ref(), ) { observe_stream_usage_bytes( observer, report_context, &mut stream_usage_observer_buffered, &normalized_chunk, ); } if !downstream_dropped { let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut() { match rewriter.push_chunk(&normalized_chunk) { Ok(rewritten_chunk) => rewritten_chunk, Err(err) => { warn!( event_name = "stream_execution_normalized_flush_rewrite_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to rewrite normalized private stream chunk during flush" ); let failure = build_stream_failure_report( "execution_runtime_stream_rewrite_flush_error", format!("failed to rewrite normalized private stream chunk during flush: {err:?}"), 502, ); terminal_failure.get_or_insert(failure); Vec::new() } } } else { normalized_chunk }; if !rewritten_chunk.is_empty() { append_stream_capture_bytes( &mut buffered_body, &rewritten_chunk, max_stream_body_buffer_bytes, &mut client_body_truncated, ); let rewritten_chunk_len = u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX); let rewritten_chunk = Bytes::from(rewritten_chunk); if tx.send(Ok(rewritten_chunk.clone())).await.is_err() { warn!( event_name = "stream_execution_downstream_flush_disconnected", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), "gateway stream downstream dropped while flushing private stream normalization" ); downstream_dropped = true; } else { client_visible_stream_completed |= client_stream_completion_tracker .observe_chunk(rewritten_chunk.as_ref()); client_stream_bytes .fetch_add(rewritten_chunk_len, Ordering::Relaxed); last_client_chunk_elapsed_ms.store( stream_started_at_for_report .elapsed() .as_millis() .min(u128::from(u64::MAX)) as u64, Ordering::Relaxed, ); } } if let Some(error_body_json) = provider_private_error_body_json { let error_status_code = resolve_local_sync_error_status_code(status_code, &error_body_json); terminal_failure.get_or_insert_with(|| { build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, ) }); } } } Ok(_) => {} Err(err) => { warn!( event_name = "stream_execution_normalization_flush_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to flush private stream normalization" ); terminal_failure.get_or_insert_with(|| { build_stream_failure_report( "execution_runtime_stream_rewrite_flush_error", format!("failed to flush private stream normalization: {err:?}"), 502, ) }); } } } if !downstream_dropped && terminal_failure.is_none() { if let Some(rewriter) = local_stream_rewriter.as_mut() { match rewriter.finish() { Ok(flushed_chunk) if !flushed_chunk.is_empty() => { append_stream_capture_bytes( &mut buffered_body, &flushed_chunk, max_stream_body_buffer_bytes, &mut client_body_truncated, ); let flushed_chunk_len = u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX); let flushed_chunk = Bytes::from(flushed_chunk); if tx.send(Ok(flushed_chunk.clone())).await.is_err() { warn!( event_name = "stream_execution_downstream_rewrite_flush_disconnected", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), "gateway stream downstream dropped while flushing local stream rewrite" ); downstream_dropped = true; } else { client_visible_stream_completed |= client_stream_completion_tracker .observe_chunk(flushed_chunk.as_ref()); client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed); last_client_chunk_elapsed_ms.store( stream_started_at_for_report .elapsed() .as_millis() .min(u128::from(u64::MAX)) as u64, Ordering::Relaxed, ); } } Ok(_) => {} Err(err) => { warn!( event_name = "stream_execution_rewrite_flush_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to flush local stream rewrite" ); terminal_failure.get_or_insert_with(|| { build_stream_failure_report( "execution_runtime_stream_rewrite_flush_error", format!("failed to flush local stream rewrite: {err:?}"), 502, ) }); } } } } if !downstream_dropped { if let Some(failure) = terminal_failure.as_ref() { let terminal_event = if is_openai_image_stream_for_report { Some(encode_openai_image_failed_event( report_context_owned.as_ref(), failure, )) } else if emit_passthrough_sse_terminal_error { Some(encode_terminal_sse_error_event(failure)) } else { None }; if let Some(terminal_event) = terminal_event { match terminal_event { Ok(error_event) => { let error_event_len = u64::try_from(error_event.len()).unwrap_or(u64::MAX); append_stream_capture_bytes( &mut buffered_body, error_event.as_ref(), max_stream_body_buffer_bytes, &mut client_body_truncated, ); if tx.send(Ok(error_event)).await.is_err() { warn!( event_name = "stream_execution_downstream_terminal_error_disconnected", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), "gateway stream downstream dropped while sending terminal SSE error event" ); downstream_dropped = true; } else { client_stream_bytes.fetch_add(error_event_len, Ordering::Relaxed); last_client_chunk_elapsed_ms.store( stream_started_at_for_report .elapsed() .as_millis() .min(u128::from(u64::MAX)) as u64, Ordering::Relaxed, ); } } Err(err) => { warn!( event_name = "stream_execution_terminal_error_event_encode_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error = ?err, "gateway failed to encode terminal SSE error event" ); } } } } } drop(tx); idle_monitor_done.store(true, Ordering::Relaxed); idle_monitor_handle.abort(); stream_terminal_summary = merge_stream_terminal_summary( stream_terminal_summary, finalize_stream_usage_observer( &mut stream_usage_observer, stream_usage_report_context.as_ref(), &mut stream_usage_observer_buffered, ), ); if downstream_dropped && client_visible_stream_completed && terminal_failure.is_none() { debug!( event_name = "execution_runtime_stream_downstream_closed_after_done", log_type = "debug", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), "gateway treats downstream close after client-visible SSE DONE as completed" ); downstream_dropped = false; } if downstream_dropped { debug!( event_name = "execution_runtime_stream_report_skipped", log_type = "debug", debug_context = "redacted", stream_status = "downstream_disconnected", status_code = 499_u16, trace_id = %trace_id_owned, "gateway skipped stream report because downstream disconnected before completion" ); let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at_for_report, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes.load(Ordering::Relaxed), )); let report_context_for_payload = report_context_with_stage_trace( report_context_owned, stage_trace_for_report, stream_started_at_for_report, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id_owned, report_kind_owned.unwrap_or_default(), report_context_for_payload, 499, headers_for_report, &provider_buffered_body, provider_body_truncated, &buffered_body, client_body_truncated, stream_terminal_summary, terminal_telemetry, ); record_stream_terminal_usage( &state_for_report, &plan_for_report, usage_payload.report_context.as_ref(), &usage_payload, true, ); record_local_request_candidate_status( &state_for_report, &plan_for_report, usage_payload.report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Cancelled, status_code: Some(499), error_type: Some("downstream_disconnect".to_string()), error_message: Some("client disconnected before stream completion".to_string()), latency_ms: usage_payload .telemetry .as_ref() .and_then(|value| value.elapsed_ms), started_at_unix_ms: Some(candidate_started_unix_secs_for_report), finished_at_unix_ms: Some(current_request_candidate_unix_ms()), }, ) .await; return; } if let Some(failure) = terminal_failure { record_manual_proxy_stream_error(&state_for_report, &plan_for_report).await; let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at_for_report, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes.load(Ordering::Relaxed), )); let report_context_for_payload = report_context_with_stage_trace( report_context_owned, stage_trace_for_report, stream_started_at_for_report, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), ); submit_midstream_stream_failure( &state_for_report, &trace_id_owned, &plan_for_report, direct_stream_finalize_kind_owned.as_deref(), report_context_for_payload, headers_for_report, terminal_telemetry, &provider_buffered_body, candidate_started_unix_secs_for_report, failure, ) .await; return; } maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state_for_report, &plan_for_report, report_context_owned.as_ref(), &mut stream_terminal_summary, ) .await; let requires_observed_terminal_event = stream_requires_observed_terminal_event( plan_for_report.provider_api_format.as_str(), stream_usage_report_context.as_ref(), ); ensure_stream_terminal_summary_for_missing_observed_finish( &mut stream_terminal_summary, requires_observed_terminal_event, ); let missing_observed_finish = stream_terminal_summary_missing_observed_finish_with_requirement( stream_terminal_summary.as_ref(), requires_observed_terminal_event, ); let should_submit_report = report_kind_owned.is_some(); let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at_for_report, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes.load(Ordering::Relaxed), )); let stream_failed = stream_terminal_summary_represents_failure_with_requirement( stream_terminal_summary.as_ref(), requires_observed_terminal_event, ); let stream_terminal_error_message = stream_terminal_summary .as_ref() .and_then(|summary| summary.parser_error.clone()) .or_else(|| { missing_observed_finish.then(|| { "execution runtime stream ended before provider terminal event".to_string() }) }); let report_context_for_payload = report_context_with_stage_trace( report_context_owned, stage_trace_for_report, stream_started_at_for_report, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id_owned.clone(), report_kind_owned.unwrap_or_default(), report_context_for_payload, status_code, headers_for_report, &provider_buffered_body, provider_body_truncated, &buffered_body, client_body_truncated, stream_terminal_summary, terminal_telemetry, ); if stream_failed { warn!( event_name = "execution_runtime_stream_missing_terminal_event", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), status_code, error_message = stream_terminal_error_message.as_deref().unwrap_or_default(), "gateway stream ended with a failed terminal state" ); } else { apply_local_execution_effect( &state_for_report, LocalExecutionEffectContext { plan: &plan_for_report, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect), ) .await; apply_local_execution_effect( &state_for_report, LocalExecutionEffectContext { plan: &plan_for_report, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect), ) .await; apply_local_execution_effect( &state_for_report, LocalExecutionEffectContext { plan: &plan_for_report, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::PoolSuccessStream { payload: &usage_payload, }, ) .await; } record_stream_terminal_usage( &state_for_report, &plan_for_report, usage_payload.report_context.as_ref(), &usage_payload, false, ); record_local_request_candidate_status( &state_for_report, &plan_for_report, usage_payload.report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: if stream_failed { RequestCandidateStatus::Failed } else { RequestCandidateStatus::Success }, status_code: Some(status_code), error_type: if stream_failed { if missing_observed_finish { Some("stream_missing_terminal_event".to_string()) } else { Some("stream_terminal_error".to_string()) } } else { None }, error_message: stream_failed .then_some(stream_terminal_error_message) .flatten(), latency_ms: usage_payload .telemetry .as_ref() .and_then(|value| value.elapsed_ms), started_at_unix_ms: Some(candidate_started_unix_secs_for_report), finished_at_unix_ms: Some(current_request_candidate_unix_ms()), }, ) .await; if should_submit_report { if let Err(err) = submit_stream_report(&state_for_report, usage_payload).await { warn!( event_name = "execution_report_submit_failed", log_type = "ops", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), report_scope = "stream", error = ?err, "gateway failed to submit stream execution report" ); } } }); headers.insert(CONTROL_REQUEST_ID_HEADER.to_string(), request_id.clone()); if let Some(candidate_id) = candidate_id .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) { headers.insert( CONTROL_CANDIDATE_ID_HEADER.to_string(), candidate_id.to_string(), ); } if response_headers_are_sse { headers.remove("content-length"); } let body_stream = build_sse_body_stream( prefetched_chunks_for_body, rx, response_headers_are_sse, emit_proxy_generated_sse_control_blocks, SSE_KEEPALIVE_INTERVAL, ); Ok(Some(build_client_response_from_parts( status_code, &headers, Body::from_stream(body_stream), trace_id, Some(decision), )?)) } fn apply_stream_summary_report_context( execution: &mut DirectUpstreamStreamExecution, report_context: Option<&Value>, ) { if let Some(report_context) = report_context.cloned() { execution.stream_summary_report_context = report_context; } } #[cfg(test)] mod tests { use std::collections::BTreeMap; use std::convert::Infallible; use std::sync::Arc; use std::time::{Duration, Instant}; use aether_contracts::{ ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTimeouts, RequestBody, StandardizedUsage, StreamFrame, StreamFramePayload, StreamFrameType, }; use aether_data::repository::candidates::InMemoryRequestCandidateRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::repository::usage::InMemoryUsageReadRepository; use aether_data_contracts::repository::candidates::{ RequestCandidateReadRepository, RequestCandidateStatus, }; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use aether_data_contracts::repository::usage::UsageReadRepository; use aether_usage_runtime::UsageRuntimeConfig; use async_stream::stream; use axum::body::{to_bytes, Body, Bytes}; use axum::extract::ws::Message; use axum::extract::Request; use axum::routing::any; use axum::{http::header, http::HeaderValue, Router}; use base64::Engine as _; use futures_util::StreamExt as _; use serde_json::{json, Value}; use tokio::sync::{mpsc, watch, Notify}; use super::{ build_sse_body_stream, client_format_allows_proxy_generated_sse_control_blocks, ensure_stream_terminal_summary_for_missing_observed_finish, execute_execution_runtime_stream, execute_stream_from_frame_stream, maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary, should_limit_direct_finalize_prefetch, should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch, stream_chunk_contains_sse_done, stream_requires_observed_terminal_event, stream_terminal_summary_missing_observed_finish, stream_terminal_summary_missing_observed_finish_with_requirement, stream_terminal_summary_represents_failure_with_requirement, ClientVisibleStreamCompletionTracker, }; use crate::control::GatewayControlDecision; use crate::stage_metrics::RequestStageTrace; use crate::tunnel::{tunnel_protocol, TunnelProxyConn}; use crate::AppState; fn provider_catalog_stop_429_for_plan( plan: &ExecutionPlan, ) -> InMemoryProviderCatalogReadRepository { let provider_type = plan.provider_name.as_deref().unwrap_or("custom"); let provider = StoredProviderCatalogProvider::new( plan.provider_id.clone(), plan.provider_id.clone(), Some("https://provider.example".to_string()), provider_type.to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, Some(3), None, None, None, Some(json!({ "failover_rules": { "stop_status_codes": [429] } })), ); let endpoint = StoredProviderCatalogEndpoint::new( plan.endpoint_id.clone(), plan.provider_id.clone(), plan.provider_api_format.clone(), None, None, true, ) .expect("endpoint should build") .with_transport_fields( "https://provider.example".to_string(), None, None, Some(2), None, None, None, None, ) .expect("endpoint transport should build"); let key = StoredProviderCatalogKey::new( plan.key_id.clone(), plan.provider_id.clone(), plan.key_id.clone(), "api_key".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(json!([plan.provider_api_format.clone()])), "plain-upstream-key".to_string(), None, None, Some(json!({ "openai:chat": 1 })), None, None, None, None, ) .expect("key transport should build"); InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![endpoint], vec![key]) } 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) } fn test_state() -> AppState { AppState::new().expect("gateway state should build") } #[test] fn detects_client_visible_sse_terminal_events() { assert!(stream_chunk_contains_sse_done(b"data: [DONE]\n\n")); assert!(stream_chunk_contains_sse_done( b"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" )); assert!(stream_chunk_contains_sse_done( b"event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{}}\n\n" )); assert!(stream_chunk_contains_sse_done( b"event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\"}}\n\n" )); assert!(!stream_chunk_contains_sse_done( b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\n" )); } #[test] fn detects_client_visible_sse_terminal_events_across_chunks() { let mut tracker = ClientVisibleStreamCompletionTracker::default(); assert!(!tracker.observe_chunk(b"data: [DO")); assert!(!tracker.observe_chunk(b"NE]\n")); assert!(tracker.observe_chunk(b"\n")); let mut tracker = ClientVisibleStreamCompletionTracker::default(); assert!(!tracker.observe_chunk(b"event: response.comp")); assert!(!tracker.observe_chunk(b"leted\r\n")); assert!(tracker .observe_chunk(b"data: {\"type\":\"response.completed\",\"response\":{}}\r\n\r\n")); } fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot { aether_contracts::ProxySnapshot { enabled: Some(true), mode: Some("tunnel".into()), node_id: Some("node-1".into()), label: Some("relay-node".into()), url: None, extra: Some(json!({"tunnel_base_url": base_url})), } } fn connect_json_frame(flags: u8, payload: &[u8]) -> Vec { let mut out = Vec::with_capacity(5 + payload.len()); out.push(flags); out.extend_from_slice(&(payload.len() as u32).to_be_bytes()); out.extend_from_slice(payload); out } fn ndjson_frame(frame: StreamFrame) -> Bytes { let mut bytes = serde_json::to_vec(&frame).expect("stream frame should serialize"); bytes.push(b'\n'); Bytes::from(bytes) } #[test] fn merge_stream_terminal_summary_prefers_more_complete_observed_usage() { let mut runtime_usage = StandardizedUsage::new(); runtime_usage.output_tokens = 137; let mut observed_usage = StandardizedUsage::new(); observed_usage.input_tokens = 26; observed_usage.output_tokens = 137; let merged = merge_stream_terminal_summary( Some(ExecutionStreamTerminalSummary { standardized_usage: Some(runtime_usage), model: Some("gpt-5.5".to_string()), unknown_event_count: 1, ..ExecutionStreamTerminalSummary::default() }), Some(ExecutionStreamTerminalSummary { standardized_usage: Some(observed_usage), response_id: Some("resp_123".to_string()), observed_finish: true, unknown_event_count: 2, ..ExecutionStreamTerminalSummary::default() }), ) .expect("summary should merge"); let usage = merged .standardized_usage .expect("merged usage should exist"); assert_eq!(usage.input_tokens, 26); assert_eq!(usage.output_tokens, 137); assert_eq!(merged.model.as_deref(), Some("gpt-5.5")); assert_eq!(merged.response_id.as_deref(), Some("resp_123")); assert!(merged.observed_finish); assert_eq!(merged.unknown_event_count, 3); } #[test] fn detects_missing_observed_finish_only_without_usage_signal() { assert!(stream_terminal_summary_missing_observed_finish(Some( &ExecutionStreamTerminalSummary { response_id: Some("resp_missing_finish".to_string()), model: Some("gpt-5.5".to_string()), observed_finish: false, ..ExecutionStreamTerminalSummary::default() } ))); let mut usage = StandardizedUsage::new(); usage.output_tokens = 12; assert!(!stream_terminal_summary_missing_observed_finish(Some( &ExecutionStreamTerminalSummary { standardized_usage: Some(usage), observed_finish: false, ..ExecutionStreamTerminalSummary::default() } ))); assert!(!stream_terminal_summary_missing_observed_finish(Some( &ExecutionStreamTerminalSummary { observed_finish: true, ..ExecutionStreamTerminalSummary::default() } ))); assert!(!stream_terminal_summary_missing_observed_finish(None)); } #[test] fn requires_terminal_event_for_openai_responses_streams() { assert!(stream_requires_observed_terminal_event( "openai:responses", None )); assert!(stream_requires_observed_terminal_event( "openai:responses:compact", None )); assert!(!stream_requires_observed_terminal_event( "openai:chat", None )); assert!(stream_requires_observed_terminal_event( "openai:chat", Some(&json!({ "provider_stream_event_api_format": "openai:responses" })) )); } #[test] fn synthesizes_missing_terminal_summary_for_openai_responses_empty_stream() { let mut summary = None; ensure_stream_terminal_summary_for_missing_observed_finish(&mut summary, true); let summary = summary.expect("summary should be synthesized"); assert!(!summary.observed_finish); assert_eq!( summary.parser_error.as_deref(), Some("execution runtime stream ended before provider terminal event") ); assert!( stream_terminal_summary_missing_observed_finish_with_requirement(Some(&summary), true) ); assert!(stream_terminal_summary_represents_failure_with_requirement( Some(&summary), true )); } #[test] fn terminal_required_stream_fails_even_with_usage_without_finish() { let mut usage = StandardizedUsage::new(); usage.output_tokens = 12; let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(usage), observed_finish: false, ..ExecutionStreamTerminalSummary::default() }); ensure_stream_terminal_summary_for_missing_observed_finish(&mut summary, true); let summary = summary.as_ref().expect("summary should remain present"); assert!( stream_terminal_summary_missing_observed_finish_with_requirement(Some(summary), true) ); assert!(stream_terminal_summary_represents_failure_with_requirement( Some(summary), true )); } #[tokio::test] async fn kiro_stream_summary_applies_prompt_cache_usage_from_original_request() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ { "type": "text", "text": "cacheable system ".repeat(600), "cache_control": {"type": "ephemeral"} } ], "messages": [ { "role": "user", "content": [ { "type": "text", "text": "cacheable prompt ".repeat(1200), "cache_control": {"type": "ephemeral"} } ] } ] }); let report_context = json!({ "original_request_body": request_body, "kiro_simulated_cache_enabled": true, }); let plan = ExecutionPlan { request_id: "req-kiro-cache-stream".into(), candidate_id: Some("cand-kiro-cache-stream".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-cache-stream".into(), endpoint_id: "endpoint-kiro-cache-stream".into(), key_id: "key-kiro-cache-stream".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-opus-4-7".into()), proxy: None, transport_profile: None, timeouts: None, }; let state = test_state(); let mut first_summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { input_tokens: 6_000, output_tokens: 17, ..StandardizedUsage::new() }), ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, Some(&report_context), &mut first_summary, ) .await; let first_usage = first_summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) .expect("first usage should exist"); assert!(first_usage.cache_creation_tokens > 0); assert_eq!(first_usage.cache_read_tokens, 0); let mut second_summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { input_tokens: 6_000, output_tokens: 19, ..StandardizedUsage::new() }), ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, Some(&report_context), &mut second_summary, ) .await; let second_usage = second_summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) .expect("second usage should exist"); assert!(second_usage.cache_read_tokens > 0); assert_eq!(second_usage.cache_creation_tokens, 0); assert!(second_usage.input_tokens < 6_000); assert_eq!(second_usage.output_tokens, 19); } #[tokio::test] async fn kiro_stream_summary_reads_cached_prefix_within_prompt_cache_lookback_window() { let first_request_body = json!({ "model": "claude-sonnet-4.6", "messages": [{ "role": "user", "content": [{ "type": "text", "text": "shared first turn ".repeat(600), "cache_control": {"type": "ephemeral"} }] }] }); let mut second_messages = vec![json!({ "role": "user", "content": [{ "type": "text", "text": "shared first turn ".repeat(600) }] })]; for index in 0..12 { second_messages.push(json!({ "role": if index % 2 == 0 { "assistant" } else { "user" }, "content": format!("intermediate stream turn {index}") })); } second_messages.push(json!({ "role": "user", "content": [{ "type": "text", "text": "new tail turn ".repeat(600), "cache_control": {"type": "ephemeral"} }] })); let second_request_body = json!({ "model": "claude-sonnet-4.6", "messages": second_messages }); let plan = ExecutionPlan { request_id: "req-kiro-cache-stream-long-tail".into(), candidate_id: Some("cand-kiro-cache-stream-long-tail".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-cache-stream-long-tail".into(), endpoint_id: "endpoint-kiro-cache-stream-long-tail".into(), key_id: "key-kiro-cache-stream-long-tail".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-sonnet-4.6".into()), proxy: None, transport_profile: None, timeouts: None, }; let first_report_context = json!({ "original_request_body": first_request_body, "kiro_simulated_cache_enabled": true, }); let second_report_context = json!({ "original_request_body": second_request_body, "kiro_simulated_cache_enabled": true, }); let state = test_state(); let mut first_summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { input_tokens: 4_000, output_tokens: 17, ..StandardizedUsage::new() }), ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, Some(&first_report_context), &mut first_summary, ) .await; let first_usage = first_summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) .expect("first usage should exist"); assert!(first_usage.cache_creation_tokens > 0); assert_eq!(first_usage.cache_read_tokens, 0); let mut second_summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { input_tokens: 8_000, output_tokens: 19, ..StandardizedUsage::new() }), ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, Some(&second_report_context), &mut second_summary, ) .await; let second_usage = second_summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) .expect("second usage should exist"); assert!( second_usage.cache_read_tokens > 0, "stream summary should reuse the far earlier cached prefix" ); assert!(second_usage.cache_creation_tokens > 0); assert_eq!(second_usage.output_tokens, 19); } #[tokio::test] async fn kiro_stream_summary_seeds_input_tokens_without_cache_control() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ { "type": "text", "text": "non cacheable system ".repeat(400) } ], "messages": [ { "role": "user", "content": [ { "type": "text", "text": "non cacheable prompt ".repeat(800) } ] } ] }); let report_context = json!({ "original_request_body": request_body, "kiro_simulated_cache_enabled": true, }); let plan = ExecutionPlan { request_id: "req-kiro-non-cache".into(), candidate_id: Some("cand-kiro-non-cache".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-non-cache".into(), endpoint_id: "endpoint-kiro-non-cache".into(), key_id: "key-kiro-non-cache".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-opus-4-7".into()), proxy: None, transport_profile: None, timeouts: None, }; let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { input_tokens: 0, output_tokens: 13, ..StandardizedUsage::new() }), ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, Some(&report_context), &mut summary, ) .await; let usage = summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) .expect("usage should exist"); assert!(usage.input_tokens > 0); assert_eq!(usage.cache_creation_tokens, 0); assert_eq!(usage.cache_read_tokens, 0); assert_eq!(usage.output_tokens, 13); } #[tokio::test] async fn kiro_stream_summary_bills_existing_cache_usage_when_input_is_zero() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ { "type": "text", "text": "cached system ".repeat(800) } ], "messages": [ { "role": "user", "content": [ { "type": "text", "text": "cached prompt ".repeat(1400) } ] } ] }); let report_context = json!({ "original_request_body": request_body, "kiro_simulated_cache_enabled": true, }); let plan = ExecutionPlan { request_id: "req-kiro-existing-cache".into(), candidate_id: Some("cand-kiro-existing-cache".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-existing-cache".into(), endpoint_id: "endpoint-kiro-existing-cache".into(), key_id: "key-kiro-existing-cache".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-opus-4-7".into()), proxy: None, transport_profile: None, timeouts: None, }; let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { input_tokens: 0, output_tokens: 23, cache_read_tokens: 200, ..StandardizedUsage::new() }), ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, Some(&report_context), &mut summary, ) .await; let usage = summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) .expect("usage should exist"); assert!(usage.input_tokens > 0); assert_eq!(usage.cache_read_tokens, 200); assert_eq!(usage.cache_creation_tokens, 0); assert_eq!(usage.output_tokens, 23); } #[tokio::test] async fn kiro_stream_summary_clears_cache_usage_when_simulated_cache_disabled() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ { "type": "text", "text": "disabled cache summary system ".repeat(800), "cache_control": {"type": "ephemeral"} } ], "messages": [ { "role": "user", "content": [ { "type": "text", "text": "disabled cache summary prompt ".repeat(1400), "cache_control": {"type": "ephemeral"} } ] } ] }); let report_context = json!({ "original_request_body": request_body, }); let plan = ExecutionPlan { request_id: "req-kiro-summary-cache-disabled".into(), candidate_id: Some("cand-kiro-summary-cache-disabled".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-summary-cache-disabled".into(), endpoint_id: "endpoint-kiro-summary-cache-disabled".into(), key_id: "key-kiro-summary-cache-disabled".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-opus-4-7".into()), proxy: None, transport_profile: None, timeouts: None, }; let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { input_tokens: 0, output_tokens: 23, cache_creation_tokens: 500, cache_read_tokens: 700, ..StandardizedUsage::new() }), ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, Some(&report_context), &mut summary, ) .await; let usage = summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) .expect("usage should exist"); assert!(usage.input_tokens > 0); assert_eq!(usage.cache_creation_tokens, 0); assert_eq!(usage.cache_read_tokens, 0); assert_eq!(usage.output_tokens, 23); } #[tokio::test] async fn kiro_stream_summary_does_not_subtract_cache_from_already_billed_input() { let request_body = json!({ "model": "claude-opus-4-7", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "cached history ".repeat(400), "cache_control": {"type": "ephemeral"} }, { "type": "text", "text": "new user turn" } ] } ] }); let report_context = json!({ "original_request_body": request_body, "input_tokens": 24_770, "cache_creation_input_tokens": 175, "cache_read_input_tokens": 24_463, "kiro_simulated_cache_enabled": true }); let plan = ExecutionPlan { request_id: "req-kiro-billed-input".into(), candidate_id: Some("cand-kiro-billed-input".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-billed-input".into(), endpoint_id: "endpoint-kiro-billed-input".into(), key_id: "key-kiro-billed-input".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-opus-4-7".into()), proxy: None, transport_profile: None, timeouts: None, }; let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { input_tokens: 132, output_tokens: 167, cache_creation_tokens: 175, cache_read_tokens: 24_463, ..StandardizedUsage::new() }), ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, Some(&report_context), &mut summary, ) .await; let usage = summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) .expect("usage should exist"); assert_eq!(usage.input_tokens, 132); assert_eq!(usage.cache_creation_tokens, 175); assert_eq!(usage.cache_read_tokens, 24_463); assert_eq!(usage.output_tokens, 167); } #[tokio::test] async fn kiro_report_context_seeds_input_tokens_from_original_request_body() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ { "type": "text", "text": "seeded system ".repeat(600) } ], "messages": [ { "role": "user", "content": [ { "type": "text", "text": "seeded prompt ".repeat(1200) } ] } ] }); let plan = ExecutionPlan { request_id: "req-kiro-seed".into(), candidate_id: Some("cand-kiro-seed".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-seed".into(), endpoint_id: "endpoint-kiro-seed".into(), key_id: "key-kiro-seed".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-opus-4-7".into()), proxy: None, transport_profile: None, timeouts: None, }; let mut report_context = Some(json!({ "original_request_body": request_body, "kiro_simulated_cache_enabled": true, })); super::seed_kiro_report_context_input_tokens(&plan, &mut report_context); let input_tokens = report_context .as_ref() .and_then(|context| context.get("input_tokens")) .and_then(Value::as_u64) .expect("kiro input tokens should be seeded"); assert!(input_tokens > 0); } #[tokio::test] async fn kiro_report_context_seeds_prompt_cache_usage_before_stream_rewrite() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ { "type": "text", "text": "cache seed system ".repeat(600), "cache_control": {"type": "ephemeral"} } ], "messages": [ { "role": "user", "content": [ { "type": "text", "text": "cache seed prompt ".repeat(1200), "cache_control": {"type": "ephemeral"} } ] } ] }); let plan = ExecutionPlan { request_id: "req-kiro-cache-seed".into(), candidate_id: Some("cand-kiro-cache-seed".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-cache-seed".into(), endpoint_id: "endpoint-kiro-cache-seed".into(), key_id: "key-kiro-cache-seed".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-opus-4-7".into()), proxy: None, transport_profile: None, timeouts: None, }; let mut report_context = Some(json!({ "original_request_body": request_body, "kiro_simulated_cache_enabled": true, })); let state = AppState::new().expect("gateway state should build"); super::seed_kiro_report_context_input_tokens(&plan, &mut report_context); super::seed_kiro_report_context_prompt_cache_usage(&state, &plan, &mut report_context) .await; let context = report_context.as_ref().expect("context should exist"); assert!(context .get("input_tokens") .and_then(Value::as_u64) .is_some_and(|value| value > 0)); assert!(context .get("cache_creation_input_tokens") .and_then(Value::as_u64) .is_some_and(|value| value > 0)); assert_eq!( context .get("cache_read_input_tokens") .and_then(Value::as_u64), Some(0) ); } #[tokio::test] async fn kiro_report_context_skips_prompt_cache_usage_when_disabled() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ { "type": "text", "text": "disabled cache system ".repeat(600), "cache_control": {"type": "ephemeral"} } ], "messages": [ { "role": "user", "content": [ { "type": "text", "text": "disabled cache prompt ".repeat(1200), "cache_control": {"type": "ephemeral"} } ] } ] }); let plan = ExecutionPlan { request_id: "req-kiro-cache-disabled".into(), candidate_id: Some("cand-kiro-cache-disabled".into()), provider_name: Some("Kiro".into()), provider_id: "provider-kiro-cache-disabled".into(), endpoint_id: "endpoint-kiro-cache-disabled".into(), key_id: "key-kiro-cache-disabled".into(), method: "POST".into(), url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"conversationState": {}})), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "claude:messages".into(), model_name: Some("claude-opus-4-7".into()), proxy: None, transport_profile: None, timeouts: None, }; let mut report_context = Some(json!({ "original_request_body": request_body, })); let state = AppState::new().expect("gateway state should build"); super::seed_kiro_report_context_input_tokens(&plan, &mut report_context); super::seed_kiro_report_context_prompt_cache_usage(&state, &plan, &mut report_context) .await; let context = report_context.as_ref().expect("context should exist"); assert!(context .get("input_tokens") .and_then(Value::as_u64) .is_some_and(|value| value > 0)); assert_eq!(context.get("cache_creation_input_tokens"), None); assert_eq!(context.get("cache_read_input_tokens"), None); } #[test] fn skips_prefetch_for_same_format_passthrough_event_streams() { assert!(should_skip_direct_finalize_prefetch( Some("claude_cli_sync_finalize"), Some("text/event-stream"), "claude:messages", "claude:messages", false, false, )); } #[test] fn skips_prefetch_for_same_format_passthrough_streams_without_content_type() { assert!(should_skip_direct_finalize_prefetch( Some("claude_cli_sync_finalize"), None, "claude:messages", "claude:messages", false, false, )); } #[test] fn keeps_prefetch_for_same_format_json_streams() { assert!(!should_skip_direct_finalize_prefetch( Some("claude_cli_sync_finalize"), Some("application/json"), "claude:messages", "claude:messages", false, false, )); } #[test] fn skips_prefetch_for_event_streams_even_when_cross_format_or_rewritten() { assert!(should_skip_direct_finalize_prefetch( Some("claude_cli_sync_finalize"), Some("text/event-stream"), "openai:chat", "claude:messages", false, true, )); } #[test] fn skips_success_failover_probe_for_event_streams() { assert!(!should_probe_success_failover_before_stream( &BTreeMap::from([( "content-type".to_string(), "text/event-stream; charset=utf-8".to_string(), )]) )); assert!(should_probe_success_failover_before_stream( &BTreeMap::from([("content-type".to_string(), "application/json".to_string(),)]) )); } #[test] fn limits_prefetch_for_openai_image_and_rewritten_streams() { assert!(should_limit_direct_finalize_prefetch( "openai_image_stream", false )); assert!(should_limit_direct_finalize_prefetch( "openai_chat_stream", true )); assert!(!should_limit_direct_finalize_prefetch( "openai_chat_stream", false )); } #[test] fn openai_client_formats_disallow_proxy_generated_sse_control_blocks() { let mut plan = ExecutionPlan { request_id: "req-openai-keepalive".into(), candidate_id: Some("cand-openai-keepalive".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://example.com/v1/chat/completions".into(), headers: BTreeMap::new(), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"stream": true})), stream: true, client_api_format: "openai:chat".into(), provider_api_format: "openai:chat".into(), model_name: Some("gpt-5.4".into()), proxy: None, transport_profile: None, timeouts: None, }; assert!(!client_format_allows_proxy_generated_sse_control_blocks( &plan )); plan.client_api_format = "openai:responses".into(); assert!(!client_format_allows_proxy_generated_sse_control_blocks( &plan )); plan.client_api_format = "claude:messages".into(); assert!(client_format_allows_proxy_generated_sse_control_blocks( &plan )); } #[tokio::test] async fn sse_body_stream_emits_initial_and_periodic_keepalive_without_business_chunks() { let (_tx, rx) = mpsc::channel::>(1); let mut body_stream = Box::pin(build_sse_body_stream( Vec::new(), rx, true, true, Duration::from_millis(10), )); let first = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("initial keepalive should be immediate") .expect("stream should yield initial keepalive") .expect("initial keepalive should be ok"); assert_eq!(first.as_ref(), b": aether-keepalive\n\n"); let second = tokio::time::timeout(Duration::from_millis(100), body_stream.next()) .await .expect("periodic keepalive should arrive") .expect("stream should yield periodic keepalive") .expect("periodic keepalive should be ok"); assert_eq!(second.as_ref(), b": aether-keepalive\n\n"); } #[tokio::test] async fn sse_body_stream_filters_control_blocks_without_synthetic_keepalive() { let (tx, rx) = mpsc::channel::>(1); let mut body_stream = Box::pin(build_sse_body_stream( vec![Bytes::from_static(b": upstream-keepalive\n\n")], rx, true, false, Duration::from_millis(10), )); assert!( tokio::time::timeout(Duration::from_millis(30), body_stream.next()) .await .is_err(), "control-only prefetched blocks should not produce client-visible chunks" ); tx.send(Ok(Bytes::from_static( b"data: {\"id\":\"chatcmpl-no-keepalive\"}\n\n", ))) .await .expect("business chunk should send"); let chunk = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("business chunk should arrive") .expect("stream should yield business chunk") .expect("business chunk should be ok"); assert_eq!( chunk.as_ref(), b"data: {\"id\":\"chatcmpl-no-keepalive\"}\n\n" ); } #[tokio::test] async fn sse_body_stream_drops_upstream_control_only_blocks() { let (_tx, rx) = mpsc::channel::>(1); let mut body_stream = Box::pin(build_sse_body_stream( vec![ Bytes::from_static(b": upstream-keepalive\n\n"), Bytes::from_static(b"event: ping\nid: 1\nretry: 1000\n\n"), Bytes::from_static( b"event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\n", ), ], rx, true, true, Duration::from_secs(60), )); let chunk = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("business chunk should arrive") .expect("stream should yield business chunk") .expect("business chunk should be ok"); let text = std::str::from_utf8(chunk.as_ref()).expect("chunk should be utf8"); assert!(text.contains("event: response.output_text.delta")); assert!(text.contains("data: {\"type\":\"response.output_text.delta\"")); assert!(!text.contains("upstream-keepalive")); assert!(!text.contains("event: ping")); assert!(!text.contains("retry: 1000")); } #[tokio::test] async fn sse_body_stream_filters_control_blocks_across_chunk_boundaries() { let (_tx, rx) = mpsc::channel::>(1); let mut body_stream = Box::pin(build_sse_body_stream( vec![ Bytes::from_static(b": upstream-keepalive\n"), Bytes::from_static(b"\n"), Bytes::from_static(b"event: response.created\n"), Bytes::from_static(b"data: {\"type\":\"response.created\"}\n\n"), ], rx, true, true, Duration::from_secs(60), )); let chunk = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("business chunk should arrive") .expect("stream should yield business chunk") .expect("business chunk should be ok"); let text = std::str::from_utf8(chunk.as_ref()).expect("chunk should be utf8"); assert_eq!( text, "event: response.created\ndata: {\"type\":\"response.created\"}\n\n" ); } #[tokio::test] async fn sse_body_stream_forwards_data_line_before_block_boundary() { let (tx, rx) = mpsc::channel::>(4); let mut body_stream = Box::pin(build_sse_body_stream( Vec::new(), rx, true, true, Duration::from_secs(60), )); let keepalive = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("initial keepalive should be immediate") .expect("stream should yield initial keepalive") .expect("initial keepalive should be ok"); assert_eq!(keepalive.as_ref(), b": aether-keepalive\n\n"); tx.send(Ok(Bytes::from_static( b"event: response.output_text.delta\n", ))) .await .expect("event line should send"); assert!( tokio::time::timeout(Duration::from_millis(20), body_stream.next()) .await .is_err(), "event-only partial block should remain buffered" ); tx.send(Ok(Bytes::from_static( b"data: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n", ))) .await .expect("data line should send"); let data_chunk = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("data-bearing block should stream before terminator") .expect("stream should yield data-bearing block") .expect("data-bearing block should be ok"); assert_eq!( data_chunk.as_ref(), b"event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n" ); tx.send(Ok(Bytes::from_static(b"\n"))) .await .expect("terminator should send"); let terminator = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("terminator should stream") .expect("stream should yield terminator") .expect("terminator should be ok"); assert_eq!(terminator.as_ref(), b"\n"); } #[tokio::test] async fn sse_body_stream_uses_local_keepalive_when_prefetched_blocks_are_control_only() { let (_tx, rx) = mpsc::channel::>(1); let mut body_stream = Box::pin(build_sse_body_stream( vec![Bytes::from_static(b": upstream-keepalive\n\n")], rx, true, true, Duration::from_secs(60), )); let first = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("local keepalive should arrive") .expect("stream should yield local keepalive") .expect("local keepalive should be ok"); assert_eq!(first.as_ref(), b": aether-keepalive\n\n"); } #[tokio::test] async fn execute_stream_from_frame_stream_does_not_finalize_rewritten_tool_call_after_midstream_error( ) { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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 plan = ExecutionPlan { request_id: "req-responses-tool-midstream-error".into(), candidate_id: Some("cand-responses-tool-midstream-error".into()), provider_name: Some("openai".into()), provider_id: "provider-openai-responses".into(), endpoint_id: "endpoint-openai-responses".into(), key_id: "key-openai-responses".into(), method: "POST".into(), url: "https://api.openai.com/v1/responses".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.5", "input": [], "stream": true })), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "openai:responses".into(), model_name: Some("gpt-5.5".into()), proxy: None, transport_profile: None, timeouts: None, }; let upstream_chunk = concat!( "event: response.created\n", "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_midstream_error\",\"model\":\"gpt-5.5\",\"status\":\"in_progress\"}}\n\n", "event: response.output_item.added\n", "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"fc_1\",\"call_id\":\"call_1\",\"name\":\"lookup\",\"arguments\":\"\",\"status\":\"in_progress\"}}\n\n", "event: response.function_call_arguments.delta\n", "data: {\"type\":\"response.function_call_arguments.delta\",\"output_index\":0,\"item_id\":\"fc_1\",\"call_id\":\"call_1\",\"delta\":\"{\\\"query\\\":\\\"abc\"}\n\n" ); let frame_stream = stream! { yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Headers, payload: StreamFramePayload::Headers { status_code: 200, headers: BTreeMap::from([( "content-type".to_string(), "text/event-stream".to_string(), )]), }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(upstream_chunk.to_string()), }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Error, payload: StreamFramePayload::Error { error: ExecutionError { kind: ExecutionErrorKind::Internal, phase: ExecutionPhase::StreamRead, message: "error reading a body from connection: stream error received: unexpected internal error encountered".to_string(), upstream_status: Some(200), retryable: false, failover_recommended: false, }, }, })); } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-responses-tool-midstream-error", &test_decision(), "openai_responses_stream", Some("openai_responses_stream_success".to_string()), Some(json!({ "request_id": "req-responses-tool-midstream-error", "candidate_id": "cand-responses-tool-midstream-error", "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:responses", "client_api_format": "claude:messages", "needs_conversion": true, })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), frame_stream, None, ) .await .expect("execution should succeed") .expect("execution should return a client response"); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body_text = String::from_utf8(body.to_vec()).expect("body should be utf8"); assert!(body_text.contains("event: content_block_start")); assert!(body_text.contains("event: content_block_delta")); assert!(body_text.contains("\"type\":\"tool_use\"")); assert!(!body_text.contains("event: content_block_stop")); assert!(!body_text.contains("event: message_delta")); assert!(!body_text.contains("event: message_stop")); assert!(!body_text.contains("\"stop_reason\":\"tool_use\"")); assert!(body_text.contains("\"error\"")); assert!(body_text.contains("unexpected internal error encountered")); assert!(body_text.contains("data: [DONE]")); let candidates = tokio::time::timeout(Duration::from_secs(1), async { loop { let candidates = request_candidate_repository .list_by_request_id("req-responses-tool-midstream-error") .await .expect("request candidates should read"); if candidates .first() .is_some_and(|candidate| candidate.status == RequestCandidateStatus::Failed) { break candidates; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("candidate should be marked failed"); assert_eq!(candidates[0].status_code, Some(200)); assert_eq!(candidates[0].error_type.as_deref(), Some("internal")); } #[tokio::test] async fn openai_image_stream_ignores_plan_total_timeout() { let state = AppState::new().expect("app state should build"); let plan = ExecutionPlan { request_id: "req-image-stream-timeout".into(), candidate_id: Some("cand-image-stream-timeout".into()), provider_name: Some("codex".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://chatgpt.com/backend-api/codex/responses".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-image-1", "prompt": "hello", "stream": true })), stream: true, client_api_format: "openai:image".into(), provider_api_format: "openai:image".into(), model_name: Some("gpt-image-1".into()), proxy: None, transport_profile: None, timeouts: Some(ExecutionTimeouts { total_ms: Some(25), ..ExecutionTimeouts::default() }), }; let decision = GatewayControlDecision::synthetic( "/v1/images/generations", Some("ai_public".to_string()), Some("openai".to_string()), Some("image".to_string()), Some("openai:image".to_string()), ) .with_execution_runtime_candidate(true); let frame_stream = stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", )); std::future::pending::<()>().await; } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-image-stream-timeout", &decision, "openai_image_stream", None, Some(json!({ "provider_api_format": "openai:image", "client_api_format": "openai:image", "image_request": { "operation": "generate" } })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), frame_stream, None, ) .await .expect("execution should succeed") .expect("execution should return a client response"); let mut body_stream = response.into_body().into_data_stream(); let next_chunk = tokio::time::timeout(Duration::from_millis(100), body_stream.next()).await; assert!( next_chunk.is_err(), "stream total_ms must not synthesize a keepalive, image failure, or close the response body" ); } #[tokio::test] async fn execute_stream_from_frame_stream_treats_windsurf_connect_trailer_error_as_failure() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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 plan = ExecutionPlan { request_id: "req-windsurf-connect-error".into(), candidate_id: Some("cand-windsurf-connect-error".into()), provider_name: Some("windsurf".into()), provider_id: "provider-windsurf".into(), endpoint_id: "endpoint-windsurf-chat".into(), key_id: "key-windsurf".into(), method: "POST".into(), url: "https://server.codeium.com/exa.api_server_pb.ApiServerService/GetChatMessage?beta=true".into(), headers: BTreeMap::from([ ("content-type".into(), "application/connect+json".into()), ("accept".into(), "application/connect+json".into()), ]), content_type: Some("application/connect+json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "claude-sonnet-4", "messages": [], "stream": true })), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "openai:chat".into(), model_name: Some("claude-sonnet-4".into()), proxy: None, transport_profile: None, timeouts: None, }; let state = state.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_provider_catalog_reader(Arc::new(provider_catalog_stop_429_for_plan(&plan))) .with_encryption_key_for_tests("development-key"), ); let trailer_error = connect_json_frame( 2, br#"{"error":{"code":"resource_exhausted","message":"an internal error occurred"}}"#, ); let trailer_error_b64 = base64::engine::general_purpose::STANDARD.encode(trailer_error); let frame = format!( "{{\"type\":\"data\",\"payload\":{{\"kind\":\"data\",\"chunk_b64\":\"{trailer_error_b64}\"}}}}\n" ); let frame_stream = stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/connect+json\"}}}\n", )); yield Ok::(Bytes::from(frame)); yield Ok::(Bytes::from_static( b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n", )); } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-windsurf-connect-error", &test_decision(), "claude_chat_stream", Some("claude_chat_stream_success".to_string()), Some(json!({ "request_id": "req-windsurf-connect-error", "candidate_id": "cand-windsurf-connect-error", "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:chat", "client_api_format": "claude:messages", "needs_conversion": true, "has_envelope": true, "envelope_name": "windsurf:GetChatMessage" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), frame_stream, None, ) .await .expect("execution should succeed") .expect("execution should return a client response"); let status = response.status(); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body_json: Value = serde_json::from_slice(&body).expect("response body should decode as json"); assert_eq!(status.as_u16(), 429); assert_eq!(body_json["type"], json!("error")); assert_eq!(body_json["error"]["type"], json!("rate_limit_error")); assert_eq!(body_json["error"]["code"], json!("resource_exhausted")); assert_eq!( body_json["error"]["message"], json!("an internal error occurred") ); let candidates = tokio::time::timeout(Duration::from_secs(1), async { loop { let candidates = request_candidate_repository .list_by_request_id("req-windsurf-connect-error") .await .expect("request candidates should read"); if candidates .first() .is_some_and(|candidate| candidate.status == RequestCandidateStatus::Failed) { break candidates; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("candidate should be marked failed"); assert_eq!(candidates[0].status_code, Some(429)); assert_eq!( candidates[0].error_type.as_deref(), Some("resource_exhausted") ); } #[tokio::test] async fn execute_stream_from_frame_stream_decodes_non_success_windsurf_connect_error_body() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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 plan = ExecutionPlan { request_id: "req-windsurf-connect-429".into(), candidate_id: Some("cand-windsurf-connect-429".into()), provider_name: Some("windsurf".into()), provider_id: "provider-windsurf".into(), endpoint_id: "endpoint-windsurf-chat".into(), key_id: "key-windsurf".into(), method: "POST".into(), url: "https://server.codeium.com/exa.api_server_pb.ApiServerService/GetChatMessage?beta=true".into(), headers: BTreeMap::from([ ("content-type".into(), "application/connect+json".into()), ("accept".into(), "application/connect+json".into()), ]), content_type: Some("application/connect+json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "claude-sonnet-4", "messages": [], "stream": true })), stream: true, client_api_format: "claude:messages".into(), provider_api_format: "openai:chat".into(), model_name: Some("claude-sonnet-4".into()), proxy: None, transport_profile: None, timeouts: None, }; let state = state.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_provider_catalog_reader(Arc::new(provider_catalog_stop_429_for_plan(&plan))) .with_encryption_key_for_tests("development-key"), ); let connect_error = connect_json_frame( 2, br#"{"error":{"code":"resource_exhausted","message":"quota exhausted"}}"#, ); let connect_error_b64 = base64::engine::general_purpose::STANDARD.encode(connect_error); let frame_stream = stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":429,\"headers\":{\"content-type\":\"application/connect+json\"}}}\n", )); yield Ok::(Bytes::from(format!( "{{\"type\":\"data\",\"payload\":{{\"kind\":\"data\",\"chunk_b64\":\"{connect_error_b64}\"}}}}\n" ))); yield Ok::(Bytes::from_static( b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n", )); } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-windsurf-connect-429", &test_decision(), "claude_chat_stream", Some("claude_chat_stream_success".to_string()), Some(json!({ "request_id": "req-windsurf-connect-429", "candidate_id": "cand-windsurf-connect-429", "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:chat", "client_api_format": "claude:messages", "needs_conversion": true, "has_envelope": true, "envelope_name": "windsurf:GetChatMessage" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), frame_stream, None, ) .await .expect("execution should succeed") .expect("execution should return a client response"); assert_eq!(response.status().as_u16(), 429); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body_json: Value = serde_json::from_slice(&body).expect("response body should decode as json"); assert_eq!(body_json["type"], json!("error")); assert_eq!(body_json["error"]["type"], json!("rate_limit_error")); assert_eq!(body_json["error"]["code"], json!("resource_exhausted")); assert_eq!(body_json["error"]["message"], json!("quota exhausted")); let record = tokio::time::timeout(Duration::from_secs(2), async { loop { if let Some(usage) = usage_repository .find_by_request_id("req-windsurf-connect-429") .await .expect("usage should read") .filter(|usage| usage.status == "failed") { break usage; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("usage should be written"); assert_eq!(record.status_code, Some(429)); assert_eq!( record .response_body .as_ref() .and_then(|body| body.get("error")) .and_then(|error| error.get("code")), Some(&json!("resource_exhausted")) ); assert!(record.response_body_ref.is_none()); } #[tokio::test] async fn execute_stream_from_frame_stream_drains_upstream_when_client_drops_body() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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 plan = ExecutionPlan { request_id: "req-client-drop-cancels-upstream".into(), candidate_id: Some("cand-client-drop-cancels-upstream".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://example.com/v1/chat/completions".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.4", "messages": [], "stream": true })), stream: true, client_api_format: "openai:chat".into(), provider_api_format: "openai:chat".into(), model_name: Some("gpt-5.4".into()), proxy: None, transport_profile: None, timeouts: None, }; let release_terminal = Arc::new(Notify::new()); let terminal_frame_drained = Arc::new(Notify::new()); let release_terminal_for_stream = Arc::clone(&release_terminal); let terminal_frame_drained_for_stream = Arc::clone(&terminal_frame_drained); let frame_stream = stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"first\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}]}\\n\\n\"}}\n", )); release_terminal_for_stream.notified().await; yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"terminal\\\",\\\"object\\\":\\\"chat.completion.chunk\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{},\\\"finish_reason\\\":\\\"stop\\\"}],\\\"usage\\\":{\\\"prompt_tokens\\\":7,\\\"completion_tokens\\\":11,\\\"total_tokens\\\":18}}\\n\\ndata: [DONE]\\n\\n\"}}\n", )); terminal_frame_drained_for_stream.notify_one(); } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-client-drop-cancels-upstream", &test_decision(), "openai_chat_stream", None, Some(json!({ "request_id": "req-client-drop-cancels-upstream", "candidate_id": "cand-client-drop-cancels-upstream", "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:chat", "client_api_format": "openai:chat" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), frame_stream, None, ) .await .expect("execution should succeed") .expect("execution should return a client response"); let mut body_stream = response.into_body().into_data_stream(); let first = tokio::time::timeout(Duration::from_secs(1), async { loop { let chunk = body_stream .next() .await .expect("body should yield first chunk") .expect("first chunk should be ok"); if chunk.as_ref() != b": aether-keepalive\n\n" { break chunk; } } }) .await .expect("first business chunk should arrive"); assert_eq!( first.as_ref(), b"data: {\"id\":\"first\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n" ); tokio::time::sleep(Duration::from_millis(30)).await; drop(body_stream); release_terminal.notify_one(); tokio::time::timeout(Duration::from_secs(1), terminal_frame_drained.notified()) .await .expect("upstream frame stream should be drained after client disconnect"); let candidates = tokio::time::timeout(Duration::from_secs(1), async { loop { let candidates = request_candidate_repository .list_by_request_id("req-client-drop-cancels-upstream") .await .expect("request candidates should read"); if candidates .first() .is_some_and(|candidate| candidate.status == RequestCandidateStatus::Cancelled) { break candidates; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("candidate should be marked cancelled"); assert_eq!(candidates[0].status_code, Some(499)); assert_eq!( candidates[0].error_type.as_deref(), Some("downstream_disconnect") ); let stored_usage = tokio::time::timeout(Duration::from_secs(1), async { loop { let usage = usage_repository .find_by_request_id("req-client-drop-cancels-upstream") .await .expect("usage should read"); if usage .as_ref() .is_some_and(|usage| usage.status == "cancelled") { break usage.expect("cancelled usage should exist"); } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("usage should be marked cancelled"); assert_eq!(stored_usage.billing_status, "pending"); assert_eq!(stored_usage.status_code, Some(499)); assert_eq!(stored_usage.input_tokens, 7); assert_eq!(stored_usage.output_tokens, 11); assert_eq!(stored_usage.total_tokens, 18); let first_byte_time_ms = stored_usage .first_byte_time_ms .expect("cancelled stream should retain first byte time"); let response_time_ms = stored_usage .response_time_ms .expect("cancelled stream should record terminal duration"); assert!( response_time_ms > first_byte_time_ms, "terminal duration should include time after the first byte" ); } #[tokio::test] async fn split_done_then_downstream_close_is_recorded_success() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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 plan = ExecutionPlan { request_id: "req-split-done-close-success".into(), candidate_id: Some("cand-split-done-close-success".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://example.com/v1/chat/completions".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.4", "messages": [], "stream": true })), stream: true, client_api_format: "openai:chat".into(), provider_api_format: "openai:chat".into(), model_name: Some("gpt-5.4".into()), proxy: None, transport_profile: None, timeouts: None, }; let release_eof = Arc::new(Notify::new()); let release_eof_for_stream = Arc::clone(&release_eof); let frame_stream = stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"first\\\",\\\"object\\\":\\\"chat.completion.chunk\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hi\\\"},\\\"finish_reason\\\":null}]}\\n\\n\"}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"terminal\\\",\\\"object\\\":\\\"chat.completion.chunk\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{},\\\"finish_reason\\\":\\\"stop\\\"}],\\\"usage\\\":{\\\"prompt_tokens\\\":7,\\\"completion_tokens\\\":11,\\\"total_tokens\\\":18}}\\n\\n\"}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DO\"}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"NE]\\n\\n\"}}\n", )); release_eof_for_stream.notified().await; } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-split-done-close-success", &test_decision(), "openai_chat_stream", None, Some(json!({ "request_id": "req-split-done-close-success", "candidate_id": "cand-split-done-close-success", "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:chat", "client_api_format": "openai:chat" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), frame_stream, None, ) .await .expect("execution should succeed") .expect("execution should return a client response"); let mut body_stream = response.into_body().into_data_stream(); let mut body = Vec::new(); tokio::time::timeout(Duration::from_secs(1), async { while !String::from_utf8_lossy(&body).contains("data: [DONE]") { let chunk = body_stream .next() .await .expect("body should yield until done") .expect("chunk should be ok"); body.extend_from_slice(&chunk); } }) .await .expect("final DONE should arrive"); drop(body_stream); release_eof.notify_one(); let candidates = tokio::time::timeout(Duration::from_secs(1), async { loop { let candidates = request_candidate_repository .list_by_request_id("req-split-done-close-success") .await .expect("request candidates should read"); if candidates .first() .is_some_and(|candidate| candidate.status == RequestCandidateStatus::Success) { break candidates; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("candidate should be marked success"); assert_eq!(candidates[0].status_code, Some(200)); let stored_usage = tokio::time::timeout(Duration::from_secs(1), async { loop { let usage = usage_repository .find_by_request_id("req-split-done-close-success") .await .expect("usage should read"); if usage .as_ref() .is_some_and(|usage| usage.status == "completed") { break usage.expect("completed usage should exist"); } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("usage should be marked completed"); assert_eq!(stored_usage.status_code, Some(200)); assert_eq!(stored_usage.input_tokens, 7); assert_eq!(stored_usage.output_tokens, 11); assert_eq!(stored_usage.total_tokens, 18); } #[tokio::test] async fn image_stream_downstream_close_after_done_is_recorded_success() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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 plan = ExecutionPlan { request_id: "req-image-done-close-success".into(), candidate_id: Some("cand-image-done-close-success".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://example.com/v1/images/generations".into(), headers: BTreeMap::from([("accept".into(), "text/event-stream".into())]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-image-2", "prompt": "draw a small image", "stream": true })), stream: true, client_api_format: "openai:chat".into(), provider_api_format: "openai:image".into(), model_name: Some("gpt-image-2".into()), proxy: None, transport_profile: None, timeouts: None, }; let frame_stream = stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_1\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\"}}\\n\\nevent: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_1\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":null}}\\n\\n\"}}\n", )); std::future::pending::<()>().await; } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-image-done-close-success", &test_decision(), "openai_chat_stream", Some("openai_chat_stream_success".to_string()), Some(json!({ "request_id": "req-image-done-close-success", "candidate_id": "cand-image-done-close-success", "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:image", "client_api_format": "openai:chat", "image_request": { "size": "1024x1024", "quality": "medium" } })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), frame_stream, None, ) .await .expect("execution should succeed") .expect("execution should return a client response"); let mut body_stream = response.into_body().into_data_stream(); let mut body = Vec::new(); tokio::time::timeout(Duration::from_secs(1), async { while !String::from_utf8_lossy(&body).contains("data: [DONE]") { let chunk = body_stream .next() .await .expect("body should yield until done") .expect("chunk should be ok"); body.extend_from_slice(&chunk); } }) .await .expect("final DONE should arrive"); drop(body_stream); let candidates = tokio::time::timeout(Duration::from_secs(1), async { loop { let candidates = request_candidate_repository .list_by_request_id("req-image-done-close-success") .await .expect("request candidates should read"); if candidates .first() .is_some_and(|candidate| candidate.status == RequestCandidateStatus::Success) { break candidates; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("candidate should be marked success"); assert_eq!(candidates[0].status_code, Some(200)); let stored_usage = tokio::time::timeout(Duration::from_secs(1), async { loop { let usage = usage_repository .find_by_request_id("req-image-done-close-success") .await .expect("usage should read"); if usage .as_ref() .is_some_and(|usage| usage.status == "completed") { break usage.expect("completed usage should exist"); } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("usage should be marked completed"); assert_eq!(stored_usage.status_code, Some(200)); assert!(stored_usage.total_tokens > 0); } #[tokio::test] async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry( ) { let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("local addr should resolve"); let first_data_seen = Arc::new(Notify::new()); let release_terminal = Arc::new(Notify::new()); let first_data_seen_for_route = Arc::clone(&first_data_seen); let release_terminal_for_route = Arc::clone(&release_terminal); let server = tokio::spawn(async move { let app = Router::new().route( "/v1/execute/stream", any(move |_request: Request| { let first_data_seen = Arc::clone(&first_data_seen_for_route); let release_terminal = Arc::clone(&release_terminal_for_route); async move { let frames = stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", )); tokio::time::sleep(Duration::from_millis(10)).await; yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_text.delta\\ndata: {\\\"type\\\":\\\"response.output_text.delta\\\",\\\"delta\\\":\\\"hi\\\"}\\n\\n\"}}\n", )); first_data_seen.notify_one(); release_terminal.notified().await; yield Ok::(Bytes::from_static( b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"ttfb_ms\":123,\"elapsed_ms\":456}}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n", )); }; let mut response = axum::http::Response::new(Body::from_stream(frames)); response.headers_mut().insert( header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response } }), ); axum::serve(listener, app) .await .expect("server should start"); }); let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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() }) .with_execution_runtime_override_base_url(format!("http://{addr}")); let plan = ExecutionPlan { request_id: "req-live-stream-first-data".into(), candidate_id: Some("cand-live-stream-first-data".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://chatgpt.com/backend-api/codex/responses".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.4", "input": "hello", "stream": true })), stream: true, client_api_format: "openai:responses".into(), provider_api_format: "openai:responses".into(), model_name: Some("gpt-5.4".into()), proxy: None, transport_profile: None, timeouts: Some(ExecutionTimeouts { connect_ms: Some(5_000), total_ms: Some(5_000), ..ExecutionTimeouts::default() }), }; let decision = GatewayControlDecision::synthetic( "/v1/responses", Some("ai_public".to_string()), Some("openai".to_string()), Some("cli".to_string()), Some("openai:responses".to_string()), ) .with_execution_runtime_candidate(true); let response = execute_execution_runtime_stream( &state, plan, "trace-live-stream-first-data", &decision, "openai_responses_stream", None, Some(json!({ "provider_api_format": "openai:responses", "client_api_format": "openai:responses", })), ) .await .expect("execution should succeed") .expect("execution should return a client response"); first_data_seen.notified().await; let deadline = tokio::time::Instant::now() + Duration::from_secs(15); let live_usage = loop { let usage = usage_repository .find_by_request_id("req-live-stream-first-data") .await .expect("usage should read"); if usage.as_ref().is_some_and(|usage| { usage.status == "streaming" && usage.first_byte_time_ms.is_some() }) { break usage.expect("live usage should exist"); } assert!( tokio::time::Instant::now() < deadline, "usage should record streaming status with first byte before terminal telemetry" ); tokio::time::sleep(Duration::from_millis(10)).await; }; assert_eq!(live_usage.status, "streaming"); assert!(live_usage.first_byte_time_ms.is_some()); release_terminal.notify_one(); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let text = String::from_utf8(body.to_vec()).expect("response body should be utf8"); assert!(text.contains("response.output_text.delta")); server.abort(); } #[tokio::test] async fn execute_execution_runtime_stream_records_first_stream_event_before_visible_text() { let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("local addr should resolve"); let first_event_seen = Arc::new(Notify::new()); let release_text = Arc::new(Notify::new()); let text_seen = Arc::new(Notify::new()); let release_terminal = Arc::new(Notify::new()); let first_event_seen_for_route = Arc::clone(&first_event_seen); let release_text_for_route = Arc::clone(&release_text); let text_seen_for_route = Arc::clone(&text_seen); let release_terminal_for_route = Arc::clone(&release_terminal); let server = tokio::spawn(async move { let app = Router::new().route( "/v1/execute/stream", any(move |_request: Request| { let first_event_seen = Arc::clone(&first_event_seen_for_route); let release_text = Arc::clone(&release_text_for_route); let text_seen = Arc::clone(&text_seen_for_route); let release_terminal = Arc::clone(&release_terminal_for_route); async move { let frames = stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"\"}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"ttfb_ms\":11,\"elapsed_ms\":12}}}\n", )); first_event_seen.notify_one(); release_text.notified().await; yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"choices\\\":[{\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}]}\\n\\n\"}}\n", )); text_seen.notify_one(); release_terminal.notified().await; yield Ok::(Bytes::from_static( b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":50}}}\n", )); yield Ok::(Bytes::from_static( b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n", )); }; let mut response = axum::http::Response::new(Body::from_stream(frames)); response.headers_mut().insert( header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response } }), ); axum::serve(listener, app) .await .expect("server should start"); }); let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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() }) .with_execution_runtime_override_base_url(format!("http://{addr}")); let plan = ExecutionPlan { request_id: "req-live-stream-first-event".into(), candidate_id: Some("cand-live-stream-first-event".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://api.openai.com/v1/chat/completions".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "stream": true })), stream: true, client_api_format: "openai:chat".into(), provider_api_format: "openai:chat".into(), model_name: Some("gpt-5.4".into()), proxy: None, transport_profile: None, timeouts: Some(ExecutionTimeouts { connect_ms: Some(5_000), total_ms: Some(5_000), ..ExecutionTimeouts::default() }), }; let decision = 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); let response = execute_execution_runtime_stream( &state, plan, "trace-live-stream-first-event", &decision, "openai_chat_stream", None, Some(json!({ "provider_api_format": "openai:chat", "client_api_format": "openai:chat", })), ) .await .expect("execution should succeed") .expect("execution should return a client response"); first_event_seen.notified().await; let deadline = tokio::time::Instant::now() + Duration::from_secs(15); let first_event_usage = loop { let usage = usage_repository .find_by_request_id("req-live-stream-first-event") .await .expect("usage should read"); if usage.as_ref().is_some_and(|usage| { usage.status == "streaming" && usage.first_byte_time_ms.is_some() }) { break usage.expect("streaming usage should exist"); } assert!( tokio::time::Instant::now() < deadline, "usage should record first byte on the first upstream stream event" ); tokio::time::sleep(Duration::from_millis(10)).await; }; assert!(first_event_usage.first_byte_time_ms.is_some()); release_text.notify_one(); text_seen.notified().await; release_terminal.notify_one(); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let text = String::from_utf8(body.to_vec()).expect("response body should be utf8"); assert!(text.contains("\"content\":\"hello\"")); server.abort(); } #[tokio::test] async fn execute_execution_runtime_stream_bridges_sync_json_body_from_remote_runtime_to_sse() { let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("local addr should resolve"); let server = tokio::spawn(async move { let app = Router::new().route( "/v1/execute/stream", any(|_request: Request| async move { let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"id\\\":\\\"resp-remote-runtime-sync-json-123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"status\\\":\\\"completed\\\",\\\"output\\\":[{\\\"type\\\":\\\"message\\\",\\\"id\\\":\\\"msg-remote-runtime-sync-json-123\\\",\\\"role\\\":\\\"assistant\\\",\\\"content\\\":[{\\\"type\\\":\\\"output_text\\\",\\\"text\\\":\\\"Hello from remote runtime sync json\\\",\\\"annotations\\\":[]}]}],\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); let mut response = axum::http::Response::new(Body::from(frames)); response.headers_mut().insert( header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response }), ); axum::serve(listener, app) .await .expect("server should start"); }); let state = AppState::new() .expect("app state should build") .with_execution_runtime_override_base_url(format!("http://{addr}")); let plan = ExecutionPlan { request_id: "req-remote-runtime-sync-json-stream".into(), candidate_id: Some("cand-remote-runtime-sync-json-stream".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://chatgpt.com/backend-api/codex/responses".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.4", "input": "hello", "stream": true })), stream: true, client_api_format: "openai:responses".into(), provider_api_format: "openai:responses".into(), model_name: Some("gpt-5.4".into()), proxy: None, transport_profile: None, timeouts: Some(ExecutionTimeouts { connect_ms: Some(5_000), total_ms: Some(5_000), ..ExecutionTimeouts::default() }), }; let decision = GatewayControlDecision::synthetic( "/v1/responses", Some("ai_public".to_string()), Some("openai".to_string()), Some("cli".to_string()), Some("openai:responses".to_string()), ) .with_execution_runtime_candidate(true); let response = execute_execution_runtime_stream( &state, plan, "trace-remote-runtime-sync-json-stream", &decision, "openai_responses_stream", None, Some(json!({ "provider_api_format": "openai:responses", "client_api_format": "openai:responses", })), ) .await .expect("execution should succeed") .expect("execution should return a client response"); assert_eq!( response .headers() .get(header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()), Some("text/event-stream") ); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let text = String::from_utf8(body.to_vec()).expect("response body should be utf8"); assert!(text.contains("event: response.output_text.delta")); assert!(text.contains("Hello from remote runtime sync json")); assert!(text.contains("event: response.completed")); server.abort(); } #[tokio::test] async fn execute_execution_runtime_stream_rewrites_redirect_to_structured_failure() { let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("local addr should resolve"); let server = tokio::spawn(async move { let app = Router::new().route( "/v1/execute/stream", any(|_request: Request| async move { let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":302,\"headers\":{\"location\":\"/\",\"content-type\":\"text/html\",\"content-length\":\"0\"}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); let mut response = axum::http::Response::new(Body::from(frames)); response.headers_mut().insert( header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response }), ); axum::serve(listener, app) .await .expect("server should start"); }); let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("app 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() }) .with_execution_runtime_override_base_url(format!("http://{addr}")); let plan = ExecutionPlan { request_id: "req-remote-runtime-stream-redirect".into(), candidate_id: Some("cand-remote-runtime-stream-redirect".into()), provider_name: Some("ChatGPTWeb".into()), provider_id: "prov-redirect".into(), endpoint_id: "ep-redirect".into(), key_id: "key-redirect".into(), method: "POST".into(), url: "https://chatgpt.com/backend-api/codex/responses".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.4", "input": "hello", "stream": true })), stream: true, client_api_format: "gemini:generate_content".into(), provider_api_format: "openai:responses".into(), model_name: Some("gemini-3.1-flash-image-preview".into()), proxy: None, transport_profile: None, timeouts: Some(ExecutionTimeouts { connect_ms: Some(5_000), total_ms: Some(5_000), ..ExecutionTimeouts::default() }), }; let decision = GatewayControlDecision::synthetic( "/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent", Some("ai_public".to_string()), Some("gemini".to_string()), Some("generate_content".to_string()), Some("gemini:generate_content".to_string()), ) .with_execution_runtime_candidate(true); let response = execute_execution_runtime_stream( &state, plan, "trace-remote-runtime-stream-redirect", &decision, "gemini_chat_stream", None, Some(json!({ "request_id": "req-remote-runtime-stream-redirect", "candidate_id": "cand-remote-runtime-stream-redirect", "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:responses", "client_api_format": "gemini:generate_content", "needs_conversion": true })), ) .await .expect("execution should succeed") .expect("execution should return a client response"); assert_eq!(response.status(), axum::http::StatusCode::BAD_GATEWAY); assert_eq!( response .headers() .get(header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()), Some("application/json") ); assert_eq!( response .headers() .get("x-aether-upstream-status") .and_then(|value| value.to_str().ok()), Some("302") ); assert!( response.headers().get(header::LOCATION).is_none(), "redirect location should not be forwarded to AI clients" ); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body_json: Value = serde_json::from_slice(&body).expect("response body should decode as json"); assert_eq!( body_json["error"]["type"], json!("execution_runtime_non_success_status") ); assert_eq!(body_json["error"]["upstream_status"], json!(302)); assert_eq!(body_json["error"]["location"], json!("/")); assert!(body_json["error"]["message"] .as_str() .is_some_and(|value| value.contains("non-success status 302"))); let usage = tokio::time::timeout(Duration::from_secs(2), async { loop { if let Some(usage) = usage_repository .find_by_request_id("req-remote-runtime-stream-redirect") .await .expect("usage should read") .filter(|usage| usage.status == "failed") { break usage; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("usage should be written"); assert_eq!(usage.status_code, Some(302)); assert_eq!(usage.error_category.as_deref(), Some("redirect")); assert!(usage .error_message .as_deref() .is_some_and(|value| value.contains("non-success status 302"))); assert_eq!( usage .client_response_headers .as_ref() .and_then(|headers| headers.get("x-aether-upstream-status")), Some(&json!("302")) ); assert_eq!( usage .response_headers .as_ref() .and_then(|headers| headers.get("location")), Some(&json!("/")) ); assert!( usage.response_body.is_none(), "upstream redirect did not include a body" ); assert_eq!( usage .client_response_body .as_ref() .and_then(|body| body.pointer("/error/upstream_status")), Some(&json!(302)) ); let candidates = request_candidate_repository .list_by_request_id("req-remote-runtime-stream-redirect") .await .expect("candidate trace should read"); let candidate_extra = candidates .first() .and_then(|candidate| candidate.extra_data.as_ref()) .expect("failed candidate extra_data should exist"); assert_eq!( candidate_extra["upstream_response"]["status_code"], json!(302) ); assert_eq!( candidate_extra["upstream_response"]["headers"]["location"], json!("/") ); assert!(candidate_extra["upstream_response"].get("body").is_none()); assert!(candidate_extra.get("client_response").is_none()); server.abort(); } #[tokio::test] async fn execute_execution_runtime_stream_bridges_openai_image_sync_json_from_remote_runtime_to_image_sse( ) { let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("local addr should resolve"); let server = tokio::spawn(async move { let app = Router::new().route( "/v1/execute/stream", any(|_request: Request| async move { let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"created\\\":1776972364,\\\"data\\\":[{\\\"b64_json\\\":\\\"aGVsbG8=\\\"}],\\\"usage\\\":{\\\"total_tokens\\\":100,\\\"input_tokens\\\":50,\\\"output_tokens\\\":50,\\\"input_tokens_details\\\":{\\\"text_tokens\\\":10,\\\"image_tokens\\\":40}}}\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); let mut response = axum::http::Response::new(Body::from(frames)); response.headers_mut().insert( header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response }), ); axum::serve(listener, app) .await .expect("server should start"); }); let state = AppState::new() .expect("app state should build") .with_execution_runtime_override_base_url(format!("http://{addr}")); let plan = ExecutionPlan { request_id: "req-remote-runtime-image-sync-json-stream".into(), candidate_id: Some("cand-remote-runtime-image-sync-json-stream".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://chatgpt.com/backend-api/codex/responses".into(), headers: BTreeMap::from([ ("content-type".into(), "application/json".into()), ("accept".into(), "text/event-stream".into()), ]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-image-1", "prompt": "hello", "stream": true })), stream: true, client_api_format: "openai:image".into(), provider_api_format: "openai:image".into(), model_name: Some("gpt-image-1".into()), proxy: None, transport_profile: None, timeouts: Some(ExecutionTimeouts { connect_ms: Some(5_000), total_ms: Some(5_000), ..ExecutionTimeouts::default() }), }; let decision = GatewayControlDecision::synthetic( "/v1/images/generations", Some("ai_public".to_string()), Some("openai".to_string()), Some("image".to_string()), Some("openai:image".to_string()), ) .with_execution_runtime_candidate(true); let response = execute_execution_runtime_stream( &state, plan, "trace-remote-runtime-image-sync-json-stream", &decision, "openai_image_stream", None, Some(json!({ "provider_api_format": "openai:image", "client_api_format": "openai:image", "mapped_model": "gpt-image-1", "image_request": { "operation": "generate" } })), ) .await .expect("execution should succeed") .expect("execution should return a client response"); assert_eq!( response .headers() .get(header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()), Some("text/event-stream") ); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let text = String::from_utf8(body.to_vec()).expect("response body should be utf8"); assert!(text.contains("event: image_generation.completed")); assert!(text.contains("\"type\":\"image_generation.completed\"")); assert!(text.contains("\"b64_json\":\"aGVsbG8=\"")); assert!(text.contains("\"total_tokens\":100")); server.abort(); } #[tokio::test] async fn execute_execution_runtime_stream_returns_client_error_with_local_tunnel_message_before_first_data( ) { let state = AppState::new().expect("app state should build"); let tunnel_app = state.tunnel.app_state(); let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8); let (proxy_close_tx, _) = watch::channel(false); tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new( 901, "node-1".to_string(), "Node 1".to_string(), proxy_tx, proxy_close_tx, 16, 2, ))); let plan = ExecutionPlan { request_id: "req-client-stream-error-1".into(), candidate_id: Some("cand-client-stream-error-1".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://example.com/chat".into(), headers: BTreeMap::from([("content-type".into(), "application/json".into())]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"stream": true})), stream: true, client_api_format: "openai:chat".into(), provider_api_format: "openai:chat".into(), model_name: Some("gpt-5".into()), proxy: Some(tunnel_proxy_snapshot("http://127.0.0.1:1".to_string())), transport_profile: None, timeouts: Some(ExecutionTimeouts { connect_ms: Some(5_000), total_ms: Some(5_000), ..ExecutionTimeouts::default() }), }; let decision = test_decision(); let state_for_task = state.clone(); let plan_for_task = plan.clone(); let decision_for_task = decision.clone(); let execution_task = tokio::spawn(async move { execute_execution_runtime_stream( &state_for_task, plan_for_task, "trace-local-stream-client-error", &decision_for_task, "openai_chat_stream", None, Some(json!({ "client_api_format": "openai:chat", "provider_api_format": "openai:chat", })), ) .await }); let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") { Message::Binary(data) => data, other => panic!("unexpected message: {other:?}"), }; let request_header = tunnel_protocol::FrameHeader::parse(&request_headers) .expect("request header frame should parse"); assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS); let request_body = match proxy_rx.recv().await.expect("body frame should arrive") { Message::Binary(data) => data, other => panic!("unexpected message: {other:?}"), }; let request_body_header = tunnel_protocol::FrameHeader::parse(&request_body) .expect("request body frame should parse"); assert_eq!(request_body_header.msg_type, tunnel_protocol::REQUEST_BODY); let response_meta = tunnel_protocol::ResponseMeta { status: 200, // Use a non-SSE content type so direct finalize prefetch stays enabled and the // pre-body tunnel error is surfaced as a client-visible structured error response. headers: vec![("content-type".to_string(), "application/json".to_string())], }; let response_payload = serde_json::to_vec(&response_meta).expect("response meta should serialize"); let mut response_headers_frame = tunnel_protocol::encode_frame( request_header.stream_id, tunnel_protocol::RESPONSE_HEADERS, 0, &response_payload, ); tunnel_app .hub .handle_proxy_frame(901, &mut response_headers_frame) .await; let original_error = "proxy disconnected before first upstream event"; let mut response_error_frame = tunnel_protocol::encode_stream_error(request_header.stream_id, original_error); tunnel_app .hub .handle_proxy_frame(901, &mut response_error_frame) .await; let response = execution_task .await .expect("execution task should complete") .expect("execution should succeed") .expect("execution should return a client response"); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body_json: Value = serde_json::from_slice(&body).expect("response body should decode as json"); let error_message = body_json .get("error") .and_then(|error| error.get("message")) .and_then(Value::as_str) .expect("response body should contain error.message"); assert_eq!(error_message, original_error); assert!( !error_message.contains("unexpected EOF during chunk size line"), "client-facing response should preserve the original local tunnel error" ); } #[tokio::test] async fn execute_execution_runtime_stream_emits_terminal_sse_error_event_after_body_started() { let state = AppState::new().expect("app state should build"); let tunnel_app = state.tunnel.app_state(); let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8); let (proxy_close_tx, _) = watch::channel(false); tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new( 902, "node-1".to_string(), "Node 1".to_string(), proxy_tx, proxy_close_tx, 16, 2, ))); let plan = ExecutionPlan { request_id: "req-client-stream-sse-error-1".into(), candidate_id: Some("cand-client-stream-sse-error-1".into()), provider_name: Some("openai".into()), provider_id: "prov-1".into(), endpoint_id: "ep-1".into(), key_id: "key-1".into(), method: "POST".into(), url: "https://example.com/chat".into(), headers: BTreeMap::from([("content-type".into(), "application/json".into())]), content_type: Some("application/json".into()), content_encoding: None, body: RequestBody::from_json(json!({"stream": true})), stream: true, client_api_format: "openai:chat".into(), provider_api_format: "openai:chat".into(), model_name: Some("gpt-5".into()), proxy: Some(tunnel_proxy_snapshot("http://127.0.0.1:1".to_string())), transport_profile: None, timeouts: Some(ExecutionTimeouts { connect_ms: Some(5_000), total_ms: Some(5_000), ..ExecutionTimeouts::default() }), }; let decision = test_decision(); let state_for_task = state.clone(); let plan_for_task = plan.clone(); let decision_for_task = decision.clone(); let execution_task = tokio::spawn(async move { execute_execution_runtime_stream( &state_for_task, plan_for_task, "trace-local-stream-sse-error", &decision_for_task, "openai_chat_stream", None, Some(json!({ "client_api_format": "openai:chat", "provider_api_format": "openai:chat", })), ) .await }); let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") { Message::Binary(data) => data, other => panic!("unexpected message: {other:?}"), }; let request_header = tunnel_protocol::FrameHeader::parse(&request_headers) .expect("request header frame should parse"); assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS); let request_body = match proxy_rx.recv().await.expect("body frame should arrive") { Message::Binary(data) => data, other => panic!("unexpected message: {other:?}"), }; let request_body_header = tunnel_protocol::FrameHeader::parse(&request_body) .expect("request body frame should parse"); assert_eq!(request_body_header.msg_type, tunnel_protocol::REQUEST_BODY); let response_meta = tunnel_protocol::ResponseMeta { status: 200, headers: vec![("content-type".to_string(), "text/event-stream".to_string())], }; let response_payload = serde_json::to_vec(&response_meta).expect("response meta should serialize"); let mut response_headers_frame = tunnel_protocol::encode_frame( request_header.stream_id, tunnel_protocol::RESPONSE_HEADERS, 0, &response_payload, ); tunnel_app .hub .handle_proxy_frame(902, &mut response_headers_frame) .await; let mut response_body_frame = tunnel_protocol::encode_frame( request_header.stream_id, tunnel_protocol::RESPONSE_BODY, 0, b"data: hello\n\n", ); tunnel_app .hub .handle_proxy_frame(902, &mut response_body_frame) .await; let response = execution_task .await .expect("execution task should complete") .expect("execution should succeed") .expect("execution should return a client response"); assert_eq!( response .headers() .get(axum::http::header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()), Some("text/event-stream") ); let body_task = tokio::spawn(async move { let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); String::from_utf8(body.to_vec()).expect("response body should be utf8") }); let original_error = "proxy disconnected while forwarding upstream body"; let mut response_error_frame = tunnel_protocol::encode_stream_error(request_header.stream_id, original_error); tunnel_app .hub .handle_proxy_frame(902, &mut response_error_frame) .await; let body = body_task.await.expect("body task should complete"); assert!(body.contains("data: hello\n\n")); assert!(body.contains("data: {\"error\":")); assert!(body.contains(original_error)); assert!(body.contains("data: [DONE]\n\n")); assert!( !body.contains("unexpected EOF during chunk size line"), "same-format SSE path should surface the original terminal error event" ); } }