use std::collections::{BTreeMap, VecDeque}; use std::future::Future; use std::io::Error as IoError; use std::pin::Pin; use std::sync::{ atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, Arc, }; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope}; use aether_contracts::{ ExecutionPlan, ExecutionResponseObservation, ExecutionStreamTerminalSummary, ExecutionTelemetry, StandardizedUsage, StreamFrame, StreamFramePayload, }; use aether_data_contracts::repository::candidates::{ RequestCandidateStatus, UpsertRequestCandidateRecord, }; use aether_data_contracts::repository::usage::UsageBodyCaptureState; use aether_scheduler_core::{ parse_request_candidate_report_context, SchedulerRequestCandidateStatusUpdate, }; #[cfg(test)] use aether_usage_runtime::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; 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, LifecycleUsageSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed, UsageRequestRecordLevel, }; use async_stream::stream; use axum::body::{Body, Bytes}; use axum::http::Response; use base64::Engine as _; use futures_util::stream::{self as futures_stream, BoxStream}; use futures_util::{Stream, StreamExt, TryStreamExt}; use http_body_util::BodyExt; use serde_json::{json, Value}; use tokio::io::{AsyncRead, ReadBuf}; use tokio::sync::mpsc; use tokio::time::MissedTickBehavior; use tokio_util::codec::{FramedRead, LinesCodec}; use tracing::{debug, info, warn}; use super::commit_policy::{ anthropic_error_status_code, find_sse_record_boundary, StreamCommitGate, StreamCommitPolicy, StreamPrecommitObservation, }; 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, build_stream_transport_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, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, UPSTREAM_IS_STREAM_KEY, }; use crate::ai_serving::is_openai_responses_family_format; use crate::ai_serving::record_local_runtime_candidate_skip_reason; use crate::api::response::{ attach_control_metadata_headers, build_client_response, build_client_response_from_parts, build_client_response_from_parts_with_mutator, }; use crate::clock::current_unix_ms as current_request_candidate_unix_ms; use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER}; use crate::control::GatewayControlDecision; use crate::execution_runtime::attempt_cancellation::AttemptCancellationGuard; 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::{ decode_base64_body_with_limit, execute_stream_plan_via_local_tunnel, format_hyper_error_chain, 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::{ ai_attempt_retry_scope_from_failure_disposition, 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_ERROR_BODY_BYTES, 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, classify_failure_disposition, spawn_local_oauth_success_effect, trace_upstream_response_body, with_error_flow_report_context, with_upstream_response_report_context, FailureDisposition, FailureTokenAction, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis, LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect, }; use crate::provider_pool_demand::{ acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission, ProviderPoolInFlightGuard, }; use crate::request_candidate_runtime::{ ensure_execution_request_candidate_slot, persist_local_request_candidate_status_record, record_local_request_candidate_status, record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status, try_enqueue_local_request_candidate_status_snapshot, LocalRequestCandidateStatusSnapshot, }; use crate::request_diagnostics::{ attach_current_request_diagnostics_to_report_context, attach_request_diagnostics_and_candidate_start_timing_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, record_stream_pre_first_byte_spawn, 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, }; /// Settlement labels for a stream attempt whose future is dropped before the /// transport reaches a terminal state. const STREAM_ATTEMPT_CANCELLED_ERROR_TYPE: &str = "local_stream_attempt_cancelled"; const STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE: &str = "Local stream attempt was dropped before terminal finalization, usually because the client disconnected or the request task was cancelled."; 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 SSE_TERMINAL_DETECTOR_MAX_RECORD_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES; const PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES; const BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES: usize = 5 * 1024 * 1024; const MAX_EXECUTION_STREAM_DATA_CHUNK_BYTES: usize = 64 * 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); const OAUTH_ERROR_PREFETCH_MAX_WAIT: Duration = Duration::from_millis(750); const ANTHROPIC_POST_STOP_DRAIN_MAX_WAIT: Duration = Duration::from_millis(250); const ANTHROPIC_POST_STOP_DRAIN_MAX_FRAMES: usize = 8; const ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES: usize = 64 * 1024; const POST_STOP_FRAME_READ_BUDGET_INACTIVE: usize = usize::MAX; const POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL: usize = 32; const DEFAULT_DIRECT_PASSTHROUGH_CHANNEL_CAPACITY: usize = 16; const MAX_DIRECT_PASSTHROUGH_CHANNEL_CAPACITY: usize = 1024; const DIRECT_PASSTHROUGH_CHANNEL_CAPACITY_ENV: &str = "AETHER_GATEWAY_DIRECT_PASSTHROUGH_CHANNEL_CAPACITY"; const DIRECT_PASSTHROUGH_MODE_ENV: &str = "AETHER_GATEWAY_DIRECT_PASSTHROUGH_MODE"; /// Retains the incomplete tail needed to recognize provider error events split across transport /// chunks without retaining an unbounded copy of the stream. #[derive(Default)] struct ProviderStreamErrorInspection { buffered: Vec, } impl ProviderStreamErrorInspection { fn observe(&mut self, report_context: Option<&Value>, chunk: &[u8]) -> Option { if chunk.is_empty() { return None; } // A transport implementation may deliver a very large chunk. Keep // every parser invocation bounded: inspect the prefix (including the // previous rolling tail for events split across chunks) and suffix, // while retaining only the bounded suffix for the next observation. // The middle of an oversized chunk is deliberately skipped because // this observer is best-effort and must never duplicate the client // stream or turn a single upstream read into an unbounded parse. if chunk.len() > PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES { let prefix_len = PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES .saturating_sub(self.buffered.len()) .min(chunk.len()); let mut boundary = Vec::with_capacity( self.buffered .len() .saturating_add(prefix_len) .min(PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES), ); boundary.extend_from_slice(&self.buffered); boundary.extend_from_slice(&chunk[..prefix_len]); if let Some(error_body) = extract_provider_private_stream_error_body(report_context, &boundary) { return Some(error_body); } let prefix = &chunk[..PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES]; if let Some(error_body) = extract_provider_private_stream_error_body(report_context, prefix) { return Some(error_body); } let suffix_start = chunk.len() - PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES; if let Some(error_body) = extract_provider_private_stream_error_body(report_context, &chunk[suffix_start..]) { return Some(error_body); } } else if let Some(error_body) = extract_provider_private_stream_error_body(report_context, chunk) { return Some(error_body); } self.append_rolling(chunk); let error_body = extract_provider_private_stream_error_body(report_context, &self.buffered); if error_body.is_none() { self.trim_completed_sse_events(); } error_body } fn append_rolling(&mut self, chunk: &[u8]) { if chunk.len() >= PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES { self.buffered.clear(); self.buffered.extend_from_slice( &chunk[chunk.len() - PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES..], ); return; } let overflow = self .buffered .len() .saturating_add(chunk.len()) .saturating_sub(PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES); if overflow > 0 { self.buffered.drain(..overflow); } self.buffered.extend_from_slice(chunk); } fn trim_completed_sse_events(&mut self) { let Ok(text) = std::str::from_utf8(&self.buffered) else { return; }; if !text.lines().any(|line| { let line = line.trim_start(); line.starts_with("event:") || line.starts_with("data:") || line.starts_with(':') }) { return; } let lf_end = self .buffered .windows(2) .rposition(|window| window == b"\n\n") .map(|index| index + 2); let crlf_end = self .buffered .windows(4) .rposition(|window| window == b"\r\n\r\n") .map(|index| index + 4); if let Some(event_end) = lf_end.into_iter().chain(crlf_end).max() { self.buffered.drain(..event_end); } } } 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>, candidate_started_at: Instant, terminal_telemetry: Option<&ExecutionTelemetry>, ) -> Option { attach_request_diagnostics_and_candidate_start_timing_to_report_context( report_context, diagnostics, Some(candidate_started_at), terminal_telemetry.and_then(|telemetry| telemetry.ttfb_ms), ) } fn request_accepted_elapsed_ms(diagnostics: Option<&Arc>) -> Option { diagnostics.and_then(|diagnostics| { diagnostics .request_accepted_at() .map(|accepted_at| accepted_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64) }) } fn observe_request_accepted_stage_trace_ms( trace: &mut RequestStageTrace, diagnostics: Option<&Arc>, stage: &'static str, ) { if let Some(elapsed_ms) = request_accepted_elapsed_ms(diagnostics) { observe_gateway_stage_trace_ms(trace, stage, elapsed_ms); } } impl Drop for StageElapsedGuard { fn drop(&mut self) { observe_gateway_stage_ms(self.stage, self.started_at.elapsed().as_millis() as u64); } } fn direct_passthrough_channel_capacity() -> usize { std::env::var(DIRECT_PASSTHROUGH_CHANNEL_CAPACITY_ENV) .ok() .and_then(|value| value.trim().parse::().ok()) .filter(|value| *value > 0) .unwrap_or(DEFAULT_DIRECT_PASSTHROUGH_CHANNEL_CAPACITY) .clamp(1, MAX_DIRECT_PASSTHROUGH_CHANNEL_CAPACITY) } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum DirectPassthroughMode { Inline, Legacy, } fn direct_passthrough_mode() -> DirectPassthroughMode { std::env::var(DIRECT_PASSTHROUGH_MODE_ENV) .ok() .as_deref() .map(parse_direct_passthrough_mode) .unwrap_or(DirectPassthroughMode::Inline) } fn stream_body_buffer_limit_for_record_level(record_level: UsageRequestRecordLevel) -> usize { match record_level { UsageRequestRecordLevel::Basic => BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES, UsageRequestRecordLevel::Full => crate::execution_runtime::MAX_STREAM_BODY_CAPTURE_BYTES, } } async fn resolve_stream_body_buffer_limit(state: &AppState) -> usize { if !state.usage_runtime.is_enabled() { return BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES; } match state .usage_runtime .body_capture_policy_for(state.usage_lifecycle_data_state().as_ref()) .await { Ok(policy) => stream_body_buffer_limit_for_record_level(policy.record_level), Err(_error) => { warn!( event_name = "stream_body_capture_policy_read_failed", log_type = "ops", error_category = "capture_policy_read_failed", fallback = "basic", "gateway could not resolve stream body capture policy" ); BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES } } } fn parse_direct_passthrough_mode(value: &str) -> DirectPassthroughMode { match value.trim().to_ascii_lowercase().as_str() { "legacy" | "pump" | "mpsc" => DirectPassthroughMode::Legacy, _ => DirectPassthroughMode::Inline, } } fn build_sync_terminal_usage_seeds( plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, ) -> (TerminalUsageContextSeed, SyncTerminalUsagePayloadSeed) { 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); (context_seed, payload_seed) } async fn record_sync_terminal_usage_with_handoff( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, ) { record_sync_terminal_usage_with_handoff_after_spawn( state, plan, report_context, payload, std::future::ready(()), ) .await; } async fn record_sync_terminal_usage_with_handoff_after_spawn( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, before_dispatch: F, ) where F: Future + Send + 'static, { crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started(); // Capture request task-local diagnostics before handing the work to a spawned task. Tokio // task-local values do not propagate across spawn boundaries. let (context_seed, payload_seed) = build_sync_terminal_usage_seeds(plan, report_context, payload); let state = state.clone(); let task = tokio::spawn(async move { before_dispatch.await; state .usage_runtime .record_sync_terminal( state.usage_lifecycle_data_state().as_ref(), context_seed, payload_seed, ) .await; }); if let Err(_err) = task.await { warn!( event_name = "sync_terminal_usage_handoff_failed", log_type = "ops", error_category = "terminal_usage_handoff_failed", "gateway sync terminal usage handoff task failed" ); } } 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, } } async fn record_stream_terminal_usage( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewayStreamReportRequest, cancelled: bool, ) { crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started(); 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.usage_lifecycle_data_state().as_ref(), context_seed, payload_seed, cancelled, ) .await; } async fn record_stream_admission_timeout_candidate_failure( 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; } fn build_stream_body_capture( body: &[u8], truncated: bool, ) -> (Option, Option) { build_stream_body_capture_with_limit( body, truncated, crate::execution_runtime::MAX_STREAM_BODY_CAPTURE_BYTES, ) } fn build_stream_body_capture_with_limit( body: &[u8], truncated: bool, max_bytes: usize, ) -> (Option, Option) { let captured = &body[..body.len().min(max_bytes)]; let body_base64 = (!captured.is_empty()).then(|| base64::engine::general_purpose::STANDARD.encode(captured)); let body_state = Some(if truncated || captured.len() < body.len() { UsageBodyCaptureState::Truncated } else if captured.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": "File download failed", }), OPENAI_VIDEO_CONTENT_PLAN_KIND => json!({ "error": { "type": "upstream_error", "message": "Video not available", } }), _ => json!({ "error": { "type": "upstream_error", "message": "Upstream request failed", } }), }; 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_category = "provider_catalog_read_failed", "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() || observer .latest_summary() .and_then(|summary| summary.parser_error.as_deref()) .is_some() { return; } let mut remaining = chunk; while !remaining.is_empty() { let line_part_len = remaining .iter() .position(|byte| *byte == b'\n') .map_or(remaining.len(), |index| index + 1); if buffered.len().saturating_add(line_part_len) > SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES { observer.disable_with_error(format!( "stream usage event exceeded {SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES} bytes" )); buffered.clear(); return; } buffered.extend_from_slice(&remaining[..line_part_len]); remaining = &remaining[line_part_len..]; if buffered.last() == Some(&b'\n') { let line = std::mem::take(buffered); if let Err(_err) = observer.push_line(report_context, line) { observer.disable_with_error("stream usage parsing failed"); buffered.clear(); return; } } } } 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("stream usage parsing failed"); } } match observer.finish(report_context) { Ok(summary) => summary, Err(_err) => { observer.disable_with_error("stream usage parsing failed"); 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; } if observed.provider_actual_service_tier.is_some() { current_summary.provider_actual_service_tier = observed.provider_actual_service_tier; } 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); let uses_oauth_credential = stream_plan_uses_oauth_credential(state, plan).await; let embedded_oauth_credential = execution.status_code == 200 && plan .provider_api_format .eq_ignore_ascii_case("claude:messages") && uses_oauth_credential; let prefetched_failure = if embedded_oauth_credential { prefetch_direct_anthropic_stream_failure(&mut execution, plan, report_context).await } else { None }; let analyzed_prefetched_failure = match prefetched_failure { Some(failure) => { Some(analyze_prefetched_stream_failure(state, plan, report_context, failure).await) } None => None, }; let response_text = if let Some(failure) = analyzed_prefetched_failure.as_ref() { Some(failure.response_text.clone()) } else if execution.status_code == 403 && uses_oauth_credential { prefetch_direct_stream_error_body(&mut execution).await } else if execution.status_code == 401 && stream_plan_uses_codex_agent_identity(state, plan).await { prefetch_direct_stream_error_body(&mut execution).await } else { None }; let retry_status_code = analyzed_prefetched_failure .as_ref() .map(|failure| failure.status_code) .unwrap_or(execution.status_code); let retry_requested = analyzed_prefetched_failure .as_ref() .map_or(execution.status_code >= 400, |failure| { matches!( failure.disposition.token_action, FailureTokenAction::ForceRefresh ) }); if retry_requested && uses_oauth_credential && refresh_oauth_plan_auth_for_retry( state, plan, retry_status_code, response_text.as_deref(), trace_id, report_context, Some(execution.response_observation.request_started_at_unix_ms), Some(&execution.response_observation.request_order_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) } #[derive(Debug)] struct PrefetchedStreamFailure { status_code: u16, response_text: String, } #[derive(Debug)] struct AnalyzedPrefetchedStreamFailure { status_code: u16, response_text: String, #[allow(dead_code)] analysis: LocalFailoverAnalysis, disposition: FailureDisposition, } async fn analyze_prefetched_stream_failure( state: &AppState, plan: &ExecutionPlan, report_context: Option<&Value>, failure: PrefetchedStreamFailure, ) -> AnalyzedPrefetchedStreamFailure { let analysis = resolve_local_candidate_failover_analysis_stream( state, plan, report_context, failure.status_code, Some(failure.response_text.as_str()), ) .await; let disposition = classify_failure_disposition( plan.provider_api_format.as_str(), analysis.classification, failure.status_code, ); AnalyzedPrefetchedStreamFailure { status_code: failure.status_code, response_text: failure.response_text, analysis, disposition, } } async fn stream_plan_uses_oauth_credential(state: &AppState, plan: &ExecutionPlan) -> bool { state .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id) .await .ok() .flatten() .as_ref() .is_some_and(|transport| { aether_provider_transport::auth::resolve_local_auth_type_for_transport_format(transport) .eq_ignore_ascii_case("oauth") }) } async fn prefetch_direct_anthropic_stream_failure( execution: &mut DirectUpstreamStreamExecution, plan: &ExecutionPlan, report_context: Option<&Value>, ) -> Option { if execution.status_code != 200 { return None; } let normalized_stream_report_context = normalize_provider_private_report_context(report_context); let policy = StreamCommitPolicy::for_response( true, execution.headers.get("content-type").map(String::as_str), plan.provider_api_format.as_str(), plan.client_api_format.as_str(), maybe_build_provider_private_stream_normalizer(report_context).is_some(), maybe_build_stream_response_rewriter(normalized_stream_report_context.as_ref()).is_some(), false, ); if !policy.is_native_anthropic() { return None; } let mut gate = StreamCommitGate::new(policy); let mut semantic_commit_observed = false; let precommit_started_at = Instant::now(); let max_wait = policy.max_precommit_wait()?; let mut observed_first_body = execution .prefetched_body .iter() .any(|item| item.as_ref().is_ok_and(|chunk| !chunk.is_empty())); while gate.is_uncommitted() { let wait = select_direct_anthropic_prefetch_wait( precommit_started_at, max_wait, execution.started_at, execution.stream_first_byte_timeout, observed_first_body, Instant::now(), ); if wait.remaining.is_zero() { if wait.commit_on_timeout { gate.commit(); } break; } let next_chunk = match tokio::time::timeout( wait.remaining, next_direct_upstream_response_chunk(&mut execution.response), ) .await { Ok(result) => result, Err(_) => { if wait.commit_on_timeout { gate.commit(); } break; } }; let chunk = match next_chunk { Ok(Some(chunk)) => chunk, Ok(None) => break, Err(error) => { execution.prefetched_body.push_back(Err(error)); break; } }; if chunk.is_empty() { continue; } observed_first_body = true; execution.prefetched_body.push_back(Ok(chunk.clone())); match gate.observe_provider_bytes(&chunk) { StreamPrecommitObservation::Pending => {} StreamPrecommitObservation::Commit => { semantic_commit_observed = true; break; } StreamPrecommitObservation::UpstreamError { status_code, body_json, } => { let response_text = serde_json::to_string(&body_json) .unwrap_or_else(|_| String::from_utf8_lossy(&chunk).into_owned()); return Some(PrefetchedStreamFailure { status_code, response_text, }); } } } execution.stream_precommit_committed = semantic_commit_observed; None } #[derive(Debug, Clone, Copy, PartialEq, Eq)] struct DirectAnthropicPrefetchWait { remaining: Duration, commit_on_timeout: bool, } fn select_direct_anthropic_prefetch_wait( precommit_started_at: Instant, max_precommit_wait: Duration, upstream_started_at: Instant, first_byte_timeout: Option, observed_first_body: bool, now: Instant, ) -> DirectAnthropicPrefetchWait { let precommit_remaining = max_precommit_wait.saturating_sub(now.saturating_duration_since(precommit_started_at)); if observed_first_body { return DirectAnthropicPrefetchWait { remaining: precommit_remaining, commit_on_timeout: true, }; } let Some(first_byte_timeout) = first_byte_timeout else { return DirectAnthropicPrefetchWait { remaining: precommit_remaining, commit_on_timeout: true, }; }; let first_byte_remaining = first_byte_timeout.saturating_sub(now.saturating_duration_since(upstream_started_at)); if first_byte_remaining <= precommit_remaining { DirectAnthropicPrefetchWait { remaining: first_byte_remaining, commit_on_timeout: false, } } else { DirectAnthropicPrefetchWait { remaining: precommit_remaining, commit_on_timeout: true, } } } async fn stream_plan_uses_codex_agent_identity(state: &AppState, plan: &ExecutionPlan) -> bool { state .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id) .await .ok() .flatten() .as_ref() .is_some_and(aether_provider_transport::is_codex_agent_identity_transport) } async fn next_direct_upstream_response_chunk( response: &mut DirectUpstreamResponse, ) -> Result, String> { match response { DirectUpstreamResponse::Reqwest(response) => response .chunk() .await .map_err(|err| format_upstream_request_error(&err)), DirectUpstreamResponse::HyperH2c(response) => loop { let Some(frame) = response.body_mut().frame().await else { return Ok(None); }; let frame = frame.map_err(|err| format_hyper_error_chain(&err))?; if let Ok(chunk) = frame.into_data() { return Ok(Some(chunk)); } }, DirectUpstreamResponse::BrowserWreq(response) => response .chunk() .await .map_err(|err| format_wreq_upstream_request_error(&err)), DirectUpstreamResponse::LocalTunnel(response) => response.next_chunk().await, } } async fn prefetch_direct_stream_error_body( execution: &mut DirectUpstreamStreamExecution, ) -> Option { let prefetch_started_at = Instant::now(); let mut inspected = Vec::with_capacity(MAX_ERROR_BODY_BYTES); let mut fully_buffered = false; while inspected.len() < MAX_ERROR_BODY_BYTES { let remaining = OAUTH_ERROR_PREFETCH_MAX_WAIT.saturating_sub(prefetch_started_at.elapsed()); if remaining.is_zero() { break; } let next_chunk = if execution.prefetched_body.is_empty() { match tokio::time::timeout( remaining, await_direct_passthrough_first_item( next_direct_upstream_response_chunk(&mut execution.response), execution.started_at, execution.stream_first_byte_timeout, ), ) .await { Ok(Ok(item)) => item, Ok(Err(_)) | Err(_) => break, } } else { match tokio::time::timeout( remaining, next_direct_upstream_response_chunk(&mut execution.response), ) .await { Ok(item) => item, Err(_) => break, } }; let chunk = match next_chunk { Ok(Some(chunk)) => chunk, Ok(None) => { fully_buffered = true; break; } Err(error) => { execution.prefetched_body.push_back(Err(error)); break; } }; if chunk.is_empty() { continue; } let remaining = MAX_ERROR_BODY_BYTES.saturating_sub(inspected.len()); inspected.extend_from_slice(&chunk[..chunk.len().min(remaining)]); execution.prefetched_body.push_back(Ok(chunk)); let response_text = String::from_utf8_lossy(&inspected); if aether_provider_transport::is_codex_agent_identity_invalid_task_response( execution.status_code, Some(response_text.as_ref()), ) { break; } } if inspected.is_empty() { return None; } if fully_buffered { let (body_json, _) = decode_stream_error_body(&execution.headers, &inspected); if let Some(body_json) = body_json { if let Ok(response_text) = serde_json::to_string(&body_json) { return Some(response_text); } } } Some(String::from_utf8_lossy(&inspected).into_owned()) } 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 is_openai_responses_family_format(plan.provider_api_format.as_str()) || is_openai_responses_family_format(plan.client_api_format.as_str()) { 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, false, ) } type DirectUpstreamByteStream = BoxStream<'static, Result>; fn direct_upstream_response_byte_stream( prefetched_body: VecDeque>, response: DirectUpstreamResponse, ) -> DirectUpstreamByteStream { let response_stream = match response { DirectUpstreamResponse::Reqwest(response) => response .bytes_stream() .map(|item| item.map_err(|err| format_upstream_request_error(&err))) .boxed(), DirectUpstreamResponse::HyperH2c(response) => response .into_body() .into_data_stream() .map(|item| item.map_err(|err| format_hyper_error_chain(&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(), }; futures_stream::iter(prefetched_body) .chain(response_stream) .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, observe_stream_completion: bool, 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, first_client_chunk: bool, 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; } let send_wait_ms = send_started_at.elapsed().as_millis() as u64; observe_gateway_stage_ms("direct_passthrough_body_send_wait", send_wait_ms); if first_client_chunk { observe_gateway_stage_ms("direct_passthrough_first_client_send_wait", send_wait_ms); } if observe_stream_completion { *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 } struct DirectPassthroughFinalizer { core: Option, } struct DirectPassthroughFinalizerCore { state: AppState, plan: ExecutionPlan, trace_id: String, report_kind: Option, report_context: Option, lifecycle_seed: LifecycleUsageSeed, direct_stream_finalize_kind: Option, stream_started_at: Instant, stage_trace: RequestStageTrace, request_diagnostics: Option>, request_id_for_log: String, candidate_id: Option, request_candidate_status_snapshot: Option, deferred_request_candidate_status_record: Option, candidate_started_unix_secs: u64, status_code: u16, headers: BTreeMap, stream_usage_report_context: Option, stream_usage_observer: Option, stream_usage_observer_buffered: Vec, provider_error_inspection: ProviderStreamErrorInspection, max_stream_body_buffer_bytes: usize, provider_buffered_body: Vec, buffered_body: Vec, provider_body_truncated: bool, client_body_truncated: bool, client_stream_completion_tracker: ClientVisibleStreamCompletionTracker, requires_anthropic_message_stop: bool, client_visible_stream_completed: bool, usage_stream_telemetry: Option, telemetry: Option, provider_stream_bytes: u64, client_stream_bytes: u64, last_client_chunk_elapsed_ms: u64, pending_recorded: bool, stream_started_recorded: bool, terminal_failure: Option, _provider_pool_in_flight_guard: Option, _upstream_target_permit: Option, } impl DirectPassthroughFinalizer { fn new(core: DirectPassthroughFinalizerCore) -> Self { Self { core: Some(core) } } fn core(&self) -> &DirectPassthroughFinalizerCore { self.core .as_ref() .expect("direct passthrough finalizer core should exist") } fn core_mut(&mut self) -> &mut DirectPassthroughFinalizerCore { self.core .as_mut() .expect("direct passthrough finalizer core should exist") } fn stream_started_at(&self) -> Instant { self.core().stream_started_at } fn ttfb_observed(&self) -> bool { self.core() .usage_stream_telemetry .as_ref() .and_then(|telemetry| telemetry.ttfb_ms) .is_some() } fn terminal_failure(&self) -> Option<&StreamFailureReport> { self.core().terminal_failure.as_ref() } fn set_terminal_failure(&mut self, failure: StreamFailureReport) { self.core_mut().terminal_failure = Some(failure); } fn prepare_upstream_chunk(&mut self, mut chunk: Bytes) -> Option { let core = self.core_mut(); if !core.requires_anthropic_message_stop { return Some(chunk); } if core.client_visible_stream_completed { return None; } if let Some(terminal_end) = core .client_stream_completion_tracker .observe_anthropic_message_stop_terminal_end(chunk.as_ref()) { chunk.truncate(terminal_end); core.client_visible_stream_completed = true; } Some(chunk) } fn fail_if_anthropic_message_stop_missing(&mut self) { let core = self.core_mut(); if core.requires_anthropic_message_stop && !core.client_visible_stream_completed && core.terminal_failure.is_none() { core.terminal_failure = Some(build_anthropic_premature_eof_failure( "upstream Anthropic stream ended before message_stop", )); } } fn completed_native_anthropic_stream(&self) -> bool { let core = self.core(); core.requires_anthropic_message_stop && core.client_visible_stream_completed && core.terminal_failure.is_none() } fn log_terminal_error_event_encode_failed(&self, _err: impl std::fmt::Debug) { let core = self.core(); warn!( event_name = "direct_passthrough_terminal_error_event_encode_failed", log_type = "ops", trace_id = %core.trace_id, request_id = %core.request_id_for_log, candidate_id = ?core.candidate_id.as_deref(), error_category = "terminal_error_event_encode_failed", "gateway direct passthrough failed to encode terminal SSE error event" ); } fn observe_first_body_poll(&mut self) { let core = self.core_mut(); observe_gateway_stage_trace_ms( &mut core.stage_trace, "stream_body_inline_first_poll", stream_elapsed_ms_since(core.stream_started_at), ); let request_diagnostics = core.request_diagnostics.clone(); observe_request_accepted_stage_trace_ms( &mut core.stage_trace, request_diagnostics.as_ref(), "frontdoor_to_stream_body_first_poll", ); } fn observe_upstream_chunk(&mut self, chunk: &Bytes, observed_at: Instant) { let core = self.core_mut(); if core.provider_stream_bytes == 0 { observe_gateway_stage_trace_ms( &mut core.stage_trace, "direct_passthrough_upstream_body_first", stream_elapsed_ms_at(core.stream_started_at, observed_at), ); } let captured_first_stream_event = maybe_capture_first_stream_event_telemetry( core.stream_started_at, observed_at, core.telemetry.as_ref(), &mut core.usage_stream_telemetry, ); if captured_first_stream_event && core.provider_stream_bytes == 0 { observe_gateway_stage_trace_ms( &mut core.stage_trace, "stream_first_data", stream_elapsed_ms_at(core.stream_started_at, observed_at), ); } core.provider_stream_bytes = core .provider_stream_bytes .saturating_add(u64::try_from(chunk.len()).unwrap_or(u64::MAX)); append_stream_capture_bytes( &mut core.provider_buffered_body, chunk.as_ref(), core.max_stream_body_buffer_bytes, &mut core.provider_body_truncated, ); if let (Some(observer), Some(report_context)) = ( core.stream_usage_observer.as_mut(), core.stream_usage_report_context.as_ref(), ) { observe_stream_usage_bytes( observer, report_context, &mut core.stream_usage_observer_buffered, chunk.as_ref(), ); } if let Some(error_body_json) = core .provider_error_inspection .observe(core.stream_usage_report_context.as_ref(), chunk.as_ref()) { let error_status_code = resolve_provider_stream_error_status_code( core.plan.provider_api_format.as_str(), core.status_code, &error_body_json, ); core.terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, )); } } fn observe_client_chunk(&mut self, chunk: &Bytes) { if chunk.is_empty() { return; } let core = self.core_mut(); append_stream_capture_bytes( &mut core.buffered_body, chunk.as_ref(), core.max_stream_body_buffer_bytes, &mut core.client_body_truncated, ); if !core.requires_anthropic_message_stop { core.client_visible_stream_completed |= core .client_stream_completion_tracker .observe_chunk(chunk.as_ref()); } core.client_stream_bytes = core .client_stream_bytes .saturating_add(u64::try_from(chunk.len()).unwrap_or(u64::MAX)); core.last_client_chunk_elapsed_ms = core .stream_started_at .elapsed() .as_millis() .min(u128::from(u64::MAX)) as u64; } fn observe_first_client_yield(&mut self) { let core = self.core_mut(); let elapsed_ms = stream_elapsed_ms_since(core.stream_started_at); observe_gateway_stage_trace_ms( &mut core.stage_trace, "direct_passthrough_first_client_send", elapsed_ms, ); observe_gateway_stage_trace_ms( &mut core.stage_trace, "stream_first_client_yield", elapsed_ms, ); let request_diagnostics = core.request_diagnostics.clone(); observe_request_accepted_stage_trace_ms( &mut core.stage_trace, request_diagnostics.as_ref(), "frontdoor_to_stream_first_client_yield", ); } fn release_upstream_target_permit_after_first_yield(&mut self) { let core = self.core_mut(); if core._upstream_target_permit.take().is_some() { observe_gateway_stage_trace_ms( &mut core.stage_trace, "stream_upstream_target_permit_release", stream_elapsed_ms_since(core.stream_started_at), ); } } fn record_client_visible_stream_started_if_needed(&mut self) { let Some(core) = self.core.as_mut() else { return; }; core.record_client_visible_stream_started_if_needed(); } async fn finalize(&mut self, downstream_dropped: bool) { let Some(core) = self.core.take() else { return; }; // Move the owned terminal payload into a task before awaiting it. A // client disconnect or an execution timeout may cancel this body // future while terminal admission is backpressured; the handoff must // continue independently so the usage row cannot remain streaming. let task = tokio::spawn(async move { core.finalize(downstream_dropped).await; }); if let Err(_err) = task.await { warn!( event_name = "direct_passthrough_terminal_handoff_failed", log_type = "ops", error_category = "terminal_handoff_failed", "gateway direct passthrough terminal handoff task failed" ); } } } impl Drop for DirectPassthroughFinalizer { fn drop(&mut self) { let Some(core) = self.core.take() else { return; }; observe_gateway_stage_ms("stream_finalizer_enqueue", 0); if let Ok(handle) = tokio::runtime::Handle::try_current() { handle.spawn(async move { core.finalize(true).await; }); } } } fn enqueue_stream_candidate_status_update( state: &AppState, snapshot: LocalRequestCandidateStatusSnapshot, status_update: SchedulerRequestCandidateStatusUpdate, ) -> Option { let Err(record) = try_enqueue_local_request_candidate_status_snapshot(state, &snapshot, status_update) else { return None; }; if state.request_candidate_queue.is_some() { return Some(record); } // Without an async queue, preserve the first-byte path's existing // fire-and-handoff behavior. Queue saturation uses the bounded deferred // record above and does not create one task per waiter. let state = state.clone(); tokio::spawn(async move { persist_local_request_candidate_status_record(&state, record).await; }); None } impl DirectPassthroughFinalizerCore { fn record_client_visible_stream_started_if_needed(&mut self) { if self.stream_started_recorded || self.client_stream_bytes == 0 { return; } self.stream_started_recorded = true; if !self.pending_recorded { self.pending_recorded = true; self.state.usage_runtime.record_pending( self.state.usage_lifecycle_data_state().as_ref(), self.lifecycle_seed.clone(), ); } self.state.usage_runtime.record_stream_started( self.state.usage_lifecycle_data_state().as_ref(), &self.lifecycle_seed, self.status_code, self.usage_stream_telemetry.as_ref(), ); if let Some(snapshot) = self.request_candidate_status_snapshot.take() { self.deferred_request_candidate_status_record = enqueue_stream_candidate_status_update( &self.state, snapshot, SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Streaming, status_code: Some(self.status_code), error_type: None, error_message: None, latency_ms: None, started_at_unix_ms: Some(self.candidate_started_unix_secs), finished_at_unix_ms: None, }, ); } } async fn finalize(mut self, mut downstream_dropped: bool) { self.record_client_visible_stream_started_if_needed(); observe_gateway_stage_ms( "stream_total", stream_elapsed_ms_since(self.stream_started_at), ); let stream_terminal_summary = finalize_stream_usage_observer( &mut self.stream_usage_observer, self.stream_usage_report_context.as_ref(), &mut self.stream_usage_observer_buffered, ); let DirectPassthroughFinalizerCore { state, plan, trace_id, report_kind, report_context, lifecycle_seed: _, direct_stream_finalize_kind, stream_started_at, stage_trace, request_diagnostics, request_id_for_log, candidate_id, request_candidate_status_snapshot: _, deferred_request_candidate_status_record, candidate_started_unix_secs, status_code, headers, stream_usage_report_context, stream_usage_observer: _, stream_usage_observer_buffered: _, provider_error_inspection: _, max_stream_body_buffer_bytes: _, provider_buffered_body, buffered_body, provider_body_truncated, client_body_truncated, client_stream_completion_tracker: _, requires_anthropic_message_stop: _, client_visible_stream_completed, usage_stream_telemetry, telemetry, provider_stream_bytes, client_stream_bytes: _, last_client_chunk_elapsed_ms: _, pending_recorded: _, stream_started_recorded: _, terminal_failure, _provider_pool_in_flight_guard, _upstream_target_permit, } = self; // Queue backpressure must not keep scarce upstream/provider permits // occupied after the client-visible stream has already ended. drop(_provider_pool_in_flight_guard); drop(_upstream_target_permit); if let Some(record) = deferred_request_candidate_status_record { persist_local_request_candidate_status_record(&state, record).await; } 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, request_id = %request_id_for_log, candidate_id = ?candidate_id.as_deref(), "gateway treats direct passthrough downstream close after terminal SSE event as completed" ); downstream_dropped = false; } if let Some(failure) = terminal_failure { record_manual_proxy_stream_error(&state, &plan).await; let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes, )); let report_context_for_payload = report_context_with_stage_trace( report_context, stage_trace, stream_started_at, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics.as_ref(), stream_started_at, terminal_telemetry.as_ref(), ); submit_midstream_stream_failure( &state, &trace_id, &plan, direct_stream_finalize_kind.as_deref(), report_context_for_payload, headers, terminal_telemetry, &provider_buffered_body, candidate_started_unix_secs, failure, ) .await; return; } if downstream_dropped { let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes, )); let report_context_for_payload = report_context_with_stage_trace( report_context, stage_trace, stream_started_at, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics.as_ref(), stream_started_at, terminal_telemetry.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id, report_kind.unwrap_or_default(), report_context_for_payload, 499, headers, &provider_buffered_body, provider_body_truncated, &buffered_body, client_body_truncated, stream_terminal_summary, terminal_telemetry, ); record_stream_terminal_usage( &state, &plan, usage_payload.report_context.as_ref(), &usage_payload, true, ) .await; record_local_request_candidate_status( &state, &plan, 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; } let mut stream_terminal_summary = stream_terminal_summary; maybe_apply_kiro_prompt_cache_usage_to_stream_summary( &state, &plan, report_context.as_ref(), &mut stream_terminal_summary, ) .await; let requires_observed_terminal_event = stream_requires_observed_terminal_event( plan.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.is_some(); let terminal_telemetry = Some(build_terminal_stream_telemetry( stream_started_at, telemetry.as_ref(), usage_stream_telemetry.as_ref(), provider_stream_bytes, )); let report_context_for_payload = report_context_with_stage_trace( report_context, stage_trace, stream_started_at, terminal_telemetry.as_ref(), ); let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics.as_ref(), stream_started_at, terminal_telemetry.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id.clone(), report_kind.unwrap_or_default(), report_context_for_payload, status_code, headers, &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, request_id = %request_id_for_log, candidate_id = ?candidate_id.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, LocalExecutionEffectContext { plan: &plan, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect), ) .await; apply_local_execution_effect( &state, LocalExecutionEffectContext { plan: &plan, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect), ) .await; apply_local_execution_effect( &state, LocalExecutionEffectContext { plan: &plan, report_context: usage_payload.report_context.as_ref(), }, LocalExecutionEffect::PoolSuccessStream { payload: &usage_payload, }, ) .await; } record_stream_terminal_usage( &state, &plan, usage_payload.report_context.as_ref(), &usage_payload, false, ) .await; record_local_request_candidate_status( &state, &plan, 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, usage_payload).await { warn!( event_name = "execution_report_submit_failed", log_type = "ops", trace_id = %trace_id, request_id = %request_id_for_log, candidate_id = ?candidate_id.as_deref(), report_scope = "direct_passthrough_stream", error_category = "stream_report_submit_failed", "gateway failed to submit direct passthrough stream execution report" ); } } } } fn build_direct_passthrough_inline_body_stream( finalizer: DirectPassthroughFinalizer, prefetched_body: VecDeque>, response: DirectUpstreamResponse, upstream_started_at: Instant, stream_first_byte_timeout: Option, ) -> impl futures_util::Stream> + Send + 'static { let state = DirectPassthroughInlineBodyState::new( finalizer, prefetched_body, response, upstream_started_at, stream_first_byte_timeout, ); futures_stream::unfold(state, |state| async move { state.next_item().await }) } struct DirectPassthroughInlineBodyState { finalizer: Option, upstream: Option, upstream_control_filter: Option, upstream_started_at: Instant, stream_first_byte_timeout: Option, observed_first_body_poll: bool, observed_first_client_yield: bool, upstream_done: bool, control_filter_flushed: bool, terminal_error_sent: bool, finalized: bool, } impl DirectPassthroughInlineBodyState { fn new( finalizer: DirectPassthroughFinalizer, prefetched_body: VecDeque>, response: DirectUpstreamResponse, upstream_started_at: Instant, stream_first_byte_timeout: Option, ) -> Self { Self { finalizer: Some(finalizer), upstream: Some(direct_upstream_response_byte_stream( prefetched_body, response, )), upstream_control_filter: Some(SseControlBlockFilter::default()), upstream_started_at, stream_first_byte_timeout, observed_first_body_poll: false, observed_first_client_yield: false, upstream_done: false, control_filter_flushed: false, terminal_error_sent: false, finalized: false, } } async fn next_item(mut self) -> Option<(Result, Self)> { if self.finalized { return None; } if self .finalizer .as_ref() .is_some_and(DirectPassthroughFinalizer::completed_native_anthropic_stream) { self.upstream.take(); self.finalized = true; drop(self.finalizer.take()); return None; } if !self.observed_first_body_poll { self.observed_first_body_poll = true; if let Some(finalizer) = self.finalizer.as_mut() { finalizer.observe_first_body_poll(); } } loop { if self.upstream_done || self .finalizer .as_ref() .and_then(DirectPassthroughFinalizer::terminal_failure) .is_some() { break; } let item = self.next_upstream_item().await; let Some(item) = item else { self.upstream_done = true; break; }; let chunk = match item { Ok(chunk) => chunk, Err(message) => { self.log_upstream_read_error_and_fail(message); self.upstream_done = true; break; } }; if chunk.is_empty() { continue; } let observed_at = Instant::now(); let Some(chunk) = self .finalizer .as_mut() .and_then(|finalizer| finalizer.prepare_upstream_chunk(chunk)) else { continue; }; let provider_error_detected = if let Some(finalizer) = self.finalizer.as_mut() { finalizer.observe_upstream_chunk(&chunk, observed_at); finalizer.terminal_failure().is_some() } else { false }; if let Some(client_chunk) = filter_upstream_sse_control_chunk(&mut self.upstream_control_filter, chunk) { self.prepare_client_chunk_yield(&client_chunk); self.terminal_error_sent |= provider_error_detected; if self .finalizer .as_ref() .is_some_and(DirectPassthroughFinalizer::completed_native_anthropic_stream) { self.upstream.take(); self.upstream_done = true; } return Some((Ok(client_chunk), self)); } } if let Some(finalizer) = self.finalizer.as_mut() { finalizer.fail_if_anthropic_message_stop_missing(); } if !self.control_filter_flushed && self .finalizer .as_ref() .and_then(DirectPassthroughFinalizer::terminal_failure) .is_none() { self.control_filter_flushed = true; if let Some(client_chunk) = flush_upstream_sse_control_filter(&mut self.upstream_control_filter) { self.prepare_client_chunk_yield(&client_chunk); return Some((Ok(client_chunk), self)); } } if !self.terminal_error_sent { if let Some(finalizer) = self.finalizer.as_mut() { if let Some(failure) = finalizer.terminal_failure() { self.terminal_error_sent = true; match encode_terminal_sse_error_event_for_plan(&finalizer.core().plan, failure) { Ok(error_event) => { self.prepare_client_chunk_yield(&error_event); return Some((Ok(error_event), self)); } Err(err) => finalizer.log_terminal_error_event_encode_failed(err), } } } } self.finalize(false).await; None } async fn next_upstream_item(&mut self) -> Option> { let needs_first_byte_timeout = self .finalizer .as_ref() .is_some_and(|finalizer| !finalizer.ttfb_observed()); let upstream = self.upstream.as_mut()?; if needs_first_byte_timeout { match await_direct_passthrough_first_item( upstream.next(), self.upstream_started_at, self.stream_first_byte_timeout, ) .await { Ok(item) => item, Err(timeout) => { if let Some(finalizer) = self.finalizer.as_mut() { finalizer.set_terminal_failure(build_stream_transport_failure_report( "first_byte_timeout", stream_first_byte_timeout_message(timeout), 504, )); } None } } } else { upstream.next().await } } fn prepare_client_chunk_yield(&mut self, chunk: &Bytes) { let Some(finalizer) = self.finalizer.as_mut() else { return; }; finalizer.observe_client_chunk(chunk); if !self.observed_first_client_yield { self.observed_first_client_yield = true; finalizer.release_upstream_target_permit_after_first_yield(); finalizer.record_client_visible_stream_started_if_needed(); finalizer.observe_first_client_yield(); } } fn log_upstream_read_error_and_fail(&mut self, message: String) { let Some(finalizer) = self.finalizer.as_mut() else { return; }; if finalizer.completed_native_anthropic_stream() { let core = finalizer.core(); debug!( event_name = "direct_passthrough_read_error_ignored_after_anthropic_stop", log_type = "debug", trace_id = %core.trace_id, request_id = %core.request_id_for_log, candidate_id = ?core.candidate_id.as_deref(), error_category = "upstream_body_read_failed", "gateway ignored direct passthrough teardown error after Anthropic message_stop" ); return; } let core = finalizer.core(); warn!( event_name = "direct_passthrough_body_read_error", log_type = "ops", trace_id = %core.trace_id, request_id = %core.request_id_for_log, candidate_id = ?core.candidate_id.as_deref(), upstream_bytes = core.provider_stream_bytes, error_category = "upstream_body_read_failed", "gateway direct passthrough upstream body read failed" ); finalizer.set_terminal_failure(build_stream_transport_failure_report( "execution_runtime_stream_read_error", message, 502, )); } async fn finalize(&mut self, downstream_dropped: bool) { if self.finalized { return; } self.finalized = true; self.upstream.take(); if let Some(finalizer) = self.finalizer.as_mut() { finalizer.finalize(downstream_dropped).await; } self.finalizer.take(); } } impl Drop for DirectPassthroughInlineBodyState { fn drop(&mut self) { // `finalized` only prevents another poll from entering finalization; // the finalizer may still be waiting for its handoff task to finish. // Keep the fallback armed while the finalizer is present. if self.finalizer.is_none() { return; } self.upstream.take(); if let Some(finalizer) = self.finalizer.take() { observe_gateway_stage_ms("stream_finalizer_enqueue", 0); if let Ok(handle) = tokio::runtime::Handle::try_current() { handle.spawn(async move { let mut finalizer = finalizer; finalizer.finalize(true).await; }); } } } } async fn record_stream_pending_lifecycle( state: &AppState, lifecycle_seed: &LifecycleUsageSeed, stage_trace: &mut RequestStageTrace, ) { let usage_pending_started_at = Instant::now(); let usage_data = state.usage_lifecycle_data_state().as_ref().clone(); state .usage_runtime .record_pending_direct(&usage_data, lifecycle_seed.clone()) .await; observe_gateway_stage_trace_ms( stage_trace, "stream_usage_pending", usage_pending_started_at.elapsed().as_millis() as u64, ); } fn should_defer_stream_pending_for_direct_inline( state: &AppState, plan: &ExecutionPlan, plan_kind: &str, report_context: Option<&Value>, ) -> bool { if direct_passthrough_mode() != DirectPassthroughMode::Inline { return false; } #[cfg(test)] if state .execution_runtime_override_base_url() .is_some_and(|value| !value.trim().is_empty()) { return false; } if plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND { return false; } if is_openai_responses_family_format(plan.provider_api_format.as_str()) || is_openai_responses_family_format(plan.client_api_format.as_str()) { 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; } 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, pending_recorded: bool, ) -> Result>, GatewayError> { let DirectUpstreamStreamExecution { request_id: _, candidate_id: _, status_code, mut headers, upstream_content_length: _, provider_api_format: _, stream_summary_report_context: _, prefetched_body, stream_precommit_committed: _, response, started_at: upstream_started_at, response_observation, stream_first_byte_timeout, upstream_target_permit, } = execution; let requires_anthropic_message_stop = status_code == 200 && response_headers_indicate_sse(&headers) && plan .provider_api_format .eq_ignore_ascii_case("claude:messages") && plan .client_api_format .eq_ignore_ascii_case("claude:messages"); 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, response_observation.request_started_at_unix_ms, response_observation.response_headers_observed_at_unix_ms, &response_observation.request_order_id, ); spawn_local_oauth_success_effect( state.clone(), &plan, report_context.as_ref(), LocalOAuthSuccessEffect { status_code, request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms), request_order_id: Some(&response_observation.request_order_id), }, ); 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()); let max_stream_body_buffer_bytes = resolve_stream_body_buffer_limit(state).await; let request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); let passthrough_mode = direct_passthrough_mode(); if passthrough_mode == DirectPassthroughMode::Legacy { state.usage_runtime.record_stream_started( state.usage_lifecycle_data_state().as_ref(), &lifecycle_seed, status_code, None, ); if let Some(snapshot) = request_candidate_status_snapshot.as_ref() { record_local_request_candidate_status_snapshot( state, 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; } } let response_header_rules_started_at = Instant::now(); apply_endpoint_response_header_rules(state, &plan, &mut headers, None).await?; observe_gateway_stage_ms( "stream_response_header_rules", response_header_rules_started_at.elapsed().as_millis() as u64, ); 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"); if passthrough_mode == DirectPassthroughMode::Inline { 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 stream_usage_report_context = normalized_stream_report_context.clone().or_else(|| { Some(serde_json::json!({ "provider_api_format": plan.provider_api_format.as_str(), "client_api_format": plan.client_api_format.as_str(), })) }); let stream_usage_observer = stream_usage_report_context .as_ref() .map(|_| StreamingStandardTerminalObserver::default()); observe_gateway_stage_trace_ms( &mut stage_trace, "stream_response_ready", stream_elapsed_ms_since(stream_started_at), ); let request_diagnostics = current_request_diagnostics(); observe_request_accepted_stage_trace_ms( &mut stage_trace, request_diagnostics.as_ref(), "frontdoor_to_stream_response_ready", ); let finalizer = DirectPassthroughFinalizer::new(DirectPassthroughFinalizerCore { state: state.clone(), plan, trace_id: trace_id.to_string(), report_kind, report_context, lifecycle_seed, direct_stream_finalize_kind, stream_started_at, stage_trace, request_diagnostics, request_id_for_log, candidate_id, request_candidate_status_snapshot, deferred_request_candidate_status_record: None, candidate_started_unix_secs, status_code, headers: headers_for_report, stream_usage_report_context, stream_usage_observer, stream_usage_observer_buffered: Vec::new(), provider_error_inspection: ProviderStreamErrorInspection::default(), max_stream_body_buffer_bytes, provider_buffered_body: Vec::new(), buffered_body: Vec::new(), provider_body_truncated: false, client_body_truncated: false, client_stream_completion_tracker: ClientVisibleStreamCompletionTracker::default(), requires_anthropic_message_stop, client_visible_stream_completed: false, usage_stream_telemetry: None, telemetry: None, provider_stream_bytes: 0, client_stream_bytes: 0, last_client_chunk_elapsed_ms: 0, pending_recorded, stream_started_recorded: false, terminal_failure: None, _provider_pool_in_flight_guard: in_flight_guard, _upstream_target_permit: upstream_target_permit, }); let body_stream = build_direct_passthrough_inline_body_stream( finalizer, prefetched_body, response, upstream_started_at, stream_first_byte_timeout, ); return Ok(Some(build_client_response_from_parts( status_code, &headers, Body::from_stream(body_stream), trace_id, Some(decision), )?)); } let (tx, rx) = mpsc::channel::>(direct_passthrough_channel_capacity()); 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 request_diagnostics_for_report = current_request_diagnostics(); observe_request_accepted_stage_trace_ms( &mut stage_trace, request_diagnostics_for_report.as_ref(), "frontdoor_to_stream_response_ready", ); let stage_trace_for_report = stage_trace; 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; record_stream_pre_first_byte_spawn(); 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 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_error_inspection = ProviderStreamErrorInspection::default(); 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 requires_anthropic_message_stop = requires_anthropic_message_stop; 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 provider_error_forwarded_to_client = false; let mut upstream = direct_upstream_response_byte_stream(prefetched_body, 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_transport_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) => { if requires_anthropic_message_stop && client_visible_stream_completed { debug!( event_name = "direct_passthrough_read_error_ignored_after_anthropic_stop", log_type = "debug", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error_category = "upstream_body_read_failed", "gateway ignored direct passthrough teardown error after Anthropic message_stop" ); break; } 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_category = "upstream_body_read_failed", "gateway direct passthrough upstream body read failed" ); terminal_failure = Some(build_stream_transport_failure_report( "execution_runtime_stream_read_error", message, 502, )); break; } }; if chunk.is_empty() { continue; } if requires_anthropic_message_stop && client_visible_stream_completed { continue; } let mut provider_chunk = chunk; if requires_anthropic_message_stop { if let Some(terminal_end) = client_stream_completion_tracker .observe_anthropic_message_stop_terminal_end(provider_chunk.as_ref()) { provider_chunk.truncate(terminal_end); client_visible_stream_completed = true; } } 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 mut sent_client_chunk = false; if let Some(client_chunk) = filter_upstream_sse_control_chunk( &mut upstream_control_filter, provider_chunk.clone(), ) { sent_client_chunk = forward_direct_passthrough_client_chunk( &tx, client_chunk, &mut downstream_dropped, &mut client_visible_stream_completed, &mut client_stream_completion_tracker, !requires_anthropic_message_stop, &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, !observed_first_client_send, 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 = provider_error_inspection.observe( stream_usage_report_context.as_ref(), provider_chunk.as_ref(), ); if let Some(error_body_json) = provider_private_error_body_json { provider_error_forwarded_to_client = sent_client_chunk; let error_status_code = resolve_provider_stream_error_status_code( plan_for_report.provider_api_format.as_str(), status_code, &error_body_json, ); terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, )); break; } if requires_anthropic_message_stop && client_visible_stream_completed { break; } } drop(upstream); drop(_provider_pool_in_flight_guard); drop(_upstream_target_permit); if terminal_failure.is_none() && !downstream_dropped && requires_anthropic_message_stop && !client_visible_stream_completed { terminal_failure = Some(build_anthropic_premature_eof_failure( "upstream Anthropic stream ended before message_stop", )); } 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, !requires_anthropic_message_stop, &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, !observed_first_client_send, 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 && !provider_error_forwarded_to_client) { match encode_terminal_sse_error_event_for_plan(&plan_for_report, 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, true, &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, !observed_first_client_send, 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_category = "terminal_error_event_encode_failed", "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(), stream_started_at_for_report, terminal_telemetry.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, ) .await; 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(), stream_started_at_for_report, terminal_telemetry.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(), stream_started_at_for_report, terminal_telemetry.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, ) .await; 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_category = "stream_report_submit_failed", "gateway failed to submit direct passthrough stream execution report" ); } } }); let body_stream = build_sse_body_stream( Vec::new(), rx, false, false, requires_anthropic_message_stop, 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) fn execute_execution_runtime_stream<'a>( state: &'a AppState, plan: ExecutionPlan, trace_id: &'a str, decision: &'a GatewayControlDecision, plan_kind: &'a str, report_kind: Option, report_context: Option, ) -> Pin>, GatewayError>> + Send + 'a>> { Box::pin(async move { let mut cancellation_guard = AttemptCancellationGuard::disarmed( state, STREAM_ATTEMPT_CANCELLED_ERROR_TYPE, STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE, ); let result = execute_execution_runtime_stream_inner( state, plan, trace_id, decision, plan_kind, report_kind, report_context, None, None, &mut cancellation_guard, ) .await; // The attempt reached its own terminal path, or handed settlement to the // stream finalizer that now lives in the response body. cancellation_guard.disarm(); result }) } #[allow(clippy::too_many_arguments)] pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>( state: &'a AppState, plan: ExecutionPlan, trace_id: &'a str, decision: &'a GatewayControlDecision, plan_kind: &'a str, report_kind: Option, report_context: Option, ) -> Pin< Box< dyn Future>, GatewayError>> + Send + 'a, >, > { Box::pin(async move { let mut retry_scope = AiAttemptRetryScope::Candidate; let mut fallback_response = None; let mut cancellation_guard = AttemptCancellationGuard::disarmed( state, STREAM_ATTEMPT_CANCELLED_ERROR_TYPE, STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE, ); let result = execute_execution_runtime_stream_inner( state, plan, trace_id, decision, plan_kind, report_kind, report_context, Some(&mut retry_scope), Some(&mut fallback_response), &mut cancellation_guard, ) .await; // The attempt reached its own terminal path, or handed settlement to the // stream finalizer that now lives in the response body. cancellation_guard.disarm(); let response = result?; Ok(match response { Some(response) => AiAttemptExecutionOutcome::Responded(response), None => AiAttemptExecutionOutcome::Retry { scope: retry_scope, fallback_response, }, }) }) } async fn maybe_build_stream_transport_error_stop_response( state: &AppState, plan: &ExecutionPlan, report_context: Option<&Value>, trace_id: &str, decision: &GatewayControlDecision, error_type: &str, error_message: &str, elapsed_ms: u64, ) -> Result>, GatewayError> { let analysis = crate::orchestration::resolve_local_transport_failover_analysis_for_attempt( state, plan, report_context, ) .await; if !matches!(analysis.decision, LocalFailoverDecision::StopLocalFailover) { return Ok(None); } crate::execution_runtime::build_transport_error_stop_response( state, plan, report_context, trace_id, decision, http::StatusCode::BAD_GATEWAY.as_u16(), error_type, error_message, elapsed_ms, ) .await .map(Some) } #[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection async fn execute_execution_runtime_stream_inner( state: &AppState, mut plan: ExecutionPlan, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, report_kind: Option, mut report_context: Option, mut retry_scope_out: Option<&mut AiAttemptRetryScope>, mut retry_fallback_out: Option<&mut Option>>, cancellation_guard: &mut AttemptCancellationGuard, ) -> 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 request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); let defer_stream_pending_for_direct_inline = should_defer_stream_pending_for_direct_inline( state, &plan, plan_kind, report_context.as_ref(), ); let candidate_started_unix_secs = current_request_candidate_unix_ms(); let provider_in_flight_started_at = Instant::now(); let mut provider_pool_in_flight_guard = match acquire_provider_pool_execution_guard(state, &plan).await? { ProviderPoolInFlightAdmission::Acquired(guard) => guard, ProviderPoolInFlightAdmission::Saturated { limit } => { record_local_runtime_candidate_skip_reason( state, trace_id, "provider_key_concurrency_limit_reached", ); if let Some(retry_scope) = retry_scope_out.as_deref_mut() { *retry_scope = AiAttemptRetryScope::Candidate; } if let Some(snapshot) = request_candidate_status_snapshot.as_ref() { record_local_request_candidate_status_snapshot( state, snapshot, SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Skipped, status_code: Some(http::StatusCode::TOO_MANY_REQUESTS.as_u16()), error_type: Some("provider_key_concurrency_limit_reached".to_string()), error_message: Some(format!( "provider key concurrency limit reached: {limit}" )), latency_ms: Some(0), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(candidate_started_unix_secs), }, ) .await; } return Ok(None); } }; observe_gateway_stage_trace_ms( &mut stage_trace, "stream_provider_in_flight", provider_in_flight_started_at.elapsed().as_millis() as u64, ); // Inline passthrough records its lifecycle seed after upstream headers are // available. Avoid constructing a throwaway seed on the common path. let mut lifecycle_seed = (!defer_stream_pending_for_direct_inline) .then(|| build_lifecycle_usage_seed(&plan, report_context.as_ref())); let mut lifecycle_pending_recorded = false; if let Some(seed) = lifecycle_seed.as_ref() { record_stream_pending_lifecycle(state, seed, &mut stage_trace).await; lifecycle_pending_recorded = true; } if let Some(snapshot) = request_candidate_status_snapshot.clone() { record_local_request_candidate_status_snapshot( state, &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; } // From here the attempt owns non-terminal rows, and everything that could // settle them runs inside the downstream request future. Arm the guard so a // client disconnect before the stream finalizer exists still settles them. cancellation_guard.arm( &plan, report_context.as_ref(), request_candidate_status_snapshot.as_ref(), candidate_started_unix_secs, stream_started_at, ); 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()); match maybe_execute_grok_stream(&plan, report_context.as_ref()).await { Ok(Some(grok_stream)) => { return execute_stream_from_frame_stream_with_retry_scope( 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, lifecycle_pending_recorded, grok_stream.frame_stream, false, provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), None, ) .await; } Ok(None) => {} Err(_err) => { let transport_error_message = "Grok stream execution unavailable".to_string(); 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_category = "grok_execution_unavailable", "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(transport_error_message.clone()), latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; if let Some(response) = maybe_build_stream_transport_error_stop_response( state, &plan, report_context.as_ref(), trace_id, decision, "grok_execution_unavailable", transport_error_message.as_str(), stream_elapsed_ms_since(stream_started_at), ) .await? { return Ok(Some(response)); } 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_with_retry_scope( 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, lifecycle_pending_recorded, windsurf_stream.frame_stream, false, provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), None, ) .await; } Ok(None) => {} Err(_err) => { let transport_error_message = "Windsurf stream execution unavailable".to_string(); 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_category = "windsurf_execution_unavailable", "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(transport_error_message.clone()), latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; if let Some(response) = maybe_build_stream_transport_error_stop_response( state, &plan, report_context.as_ref(), trace_id, decision, "windsurf_native_execution_unavailable", transport_error_message.as_str(), stream_elapsed_ms_since(stream_started_at), ) .await? { return Ok(Some(response)); } 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_with_retry_scope( 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, lifecycle_pending_recorded, kiro_web_search.frame_stream, false, provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), None, ) .await; } Ok(None) => {} Err(_err) => { let transport_error_message = "Kiro web search execution unavailable".to_string(); 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_category = "kiro_web_search_unavailable", "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(transport_error_message.clone()), latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; if let Some(response) = maybe_build_stream_transport_error_stop_response( state, &plan, report_context.as_ref(), trace_id, decision, "kiro_web_search_mcp_unavailable", transport_error_message.as_str(), stream_elapsed_ms_since(stream_started_at), ) .await? { return Ok(Some(response)); } 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_with_retry_scope( 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, lifecycle_pending_recorded, chatgpt_web_image.frame_stream, false, provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), None, ) .await; } Ok(None) => {} Err(_err) => { let transport_error_message = "ChatGPT-Web image execution unavailable".to_string(); 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_category = "chatgpt_web_image_execution_unavailable", "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(transport_error_message.clone()), latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; if let Some(response) = maybe_build_stream_transport_error_stop_response( state, &plan, report_context.as_ref(), trace_id, decision, "chatgpt_web_image_execution_unavailable", transport_error_message.as_str(), stream_elapsed_ms_since(stream_started_at), ) .await? { return Ok(Some(response)); } 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_candidate_failure( state, &plan, report_context.as_ref(), candidate_started_unix_secs, &err, ) .await; } return Err(err); } Err(InProcessStreamExecutionError::Transport(_err)) => { let transport_error_message = "Execution runtime unavailable".to_string(); 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_category = "execution_runtime_unavailable", "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(transport_error_message.clone()), latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; if let Some(response) = maybe_build_stream_transport_error_stop_response( state, &plan, report_context.as_ref(), trace_id, decision, "execution_runtime_unavailable", transport_error_message.as_str(), stream_elapsed_ms_since(stream_started_at), ) .await? { return Ok(Some(response)); } 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 Box::pin(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(), lifecycle_pending_recorded, )) .await; } if !lifecycle_pending_recorded { let seed = lifecycle_seed .get_or_insert_with(|| build_lifecycle_usage_seed(&plan, report_context.as_ref())); record_stream_pending_lifecycle(state, seed, &mut stage_trace).await; lifecycle_pending_recorded = true; } let report_context = attach_provider_response_headers_to_report_context( report_context, &execution.headers, execution.response_observation.request_started_at_unix_ms, execution .response_observation .response_headers_observed_at_unix_ms, &execution.response_observation.request_order_id, ); let stream_precommit_committed = execution.stream_precommit_committed; let frame_stream = build_direct_execution_frame_stream(execution).boxed(); return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, lifecycle_pending_recorded, frame_stream, stream_precommit_committed, provider_pool_in_flight_guard.take(), retry_scope_out, retry_fallback_out, None, ) .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_candidate_failure( state, &plan, report_context.as_ref(), candidate_started_unix_secs, &err, ) .await; } return Err(err); } Err(InProcessStreamExecutionError::Transport(_err)) => { let transport_error_message = "Execution runtime unavailable".to_string(); 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_category = "execution_runtime_unavailable", "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(transport_error_message.clone()), latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; if let Some(response) = maybe_build_stream_transport_error_stop_response( state, &plan, report_context.as_ref(), trace_id, decision, "execution_runtime_unavailable", transport_error_message.as_str(), stream_elapsed_ms_since(stream_started_at), ) .await? { return Ok(Some(response)); } 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 Box::pin(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(), lifecycle_pending_recorded, )) .await; } if !lifecycle_pending_recorded { let seed = lifecycle_seed.get_or_insert_with(|| { build_lifecycle_usage_seed(&plan, report_context.as_ref()) }); record_stream_pending_lifecycle(state, seed, &mut stage_trace).await; lifecycle_pending_recorded = true; } let report_context = attach_provider_response_headers_to_report_context( report_context, &execution.headers, execution.response_observation.request_started_at_unix_ms, execution .response_observation .response_headers_observed_at_unix_ms, &execution.response_observation.request_order_id, ); let stream_precommit_committed = execution.stream_precommit_committed; let frame_stream = build_direct_execution_frame_stream(execution).boxed(); return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, lifecycle_pending_recorded, frame_stream, stream_precommit_committed, provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), None, ) .await; } let remote_request_started_at_unix_ms = current_request_candidate_unix_ms(); let remote_request_order_id = uuid::Uuid::now_v7().to_string(); let response = match post_stream_plan_to_remote_execution_runtime( state, remote_execution_runtime_base_url, Some(trace_id), &plan, ) .await { Ok(response) => response, Err(_err) => { let transport_error_message = "Remote execution runtime unavailable".to_string(); 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_category = "execution_runtime_unavailable", "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(transport_error_message.clone()), latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; if let Some(response) = maybe_build_stream_transport_error_stop_response( state, &plan, report_context.as_ref(), trace_id, decision, "execution_runtime_unavailable", transport_error_message.as_str(), stream_elapsed_ms_since(stream_started_at), ) .await? { return Ok(Some(response)); } 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 remote_response_observed_at_unix_ms = current_request_candidate_unix_ms(); let remote_fallback_observation = ExecutionResponseObservation { request_started_at_unix_ms: remote_request_started_at_unix_ms, response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms, request_order_id: remote_request_order_id, }; let frame_stream = response .bytes_stream() .map_err(|err| IoError::other(err.to_string())) .boxed(); return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, lifecycle_pending_recorded, frame_stream, false, provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), Some(remote_fallback_observation), ) .await; } } fn decode_stream_data_chunk( chunk_b64: Option<&str>, text: Option<&str>, ) -> Result, GatewayError> { decode_stream_data_chunk_with_limit(chunk_b64, text, MAX_EXECUTION_STREAM_DATA_CHUNK_BYTES) } fn decode_stream_data_chunk_with_limit( chunk_b64: Option<&str>, text: Option<&str>, max_bytes: usize, ) -> Result, GatewayError> { if let Some(chunk_b64) = chunk_b64 { return decode_base64_body_with_limit(chunk_b64, max_bytes) .map_err(|err| GatewayError::Internal(err.to_string())); } let text = text.unwrap_or_default().as_bytes(); if text.len() > max_bytes { return Err(GatewayError::Internal(format!( "execution runtime stream data chunk exceeds {max_bytes} bytes" ))); } Ok(text.to_vec()) } fn response_headers_indicate_sse(headers: &BTreeMap) -> bool { headers .iter() .find(|(name, _)| name.eq_ignore_ascii_case("content-type")) .map(|(_, value)| value.as_str()) .map(str::trim) .filter(|value| !value.is_empty()) .is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream")) } fn report_context_upstream_is_stream(report_context: Option<&Value>) -> bool { report_context .and_then(|value| value.get(UPSTREAM_IS_STREAM_KEY)) .and_then(Value::as_bool) .unwrap_or(false) } fn response_headers_have_octet_stream_content_type(headers: &BTreeMap) -> bool { headers .iter() .find(|(name, _)| name.eq_ignore_ascii_case("content-type")) .map(|(_, value)| value.as_str()) .and_then(|value| value.split(';').next()) .map(str::trim) .is_some_and(|value| value.eq_ignore_ascii_case("application/octet-stream")) } fn response_headers_have_only_identity_content_encoding( headers: &BTreeMap, ) -> bool { headers .iter() .find(|(name, _)| name.eq_ignore_ascii_case("content-encoding")) .map(|(_, value)| value.as_str()) .is_none_or(|value| { value .split(',') .map(str::trim) .all(|coding| coding.is_empty() || coding.eq_ignore_ascii_case("identity")) }) } fn plan_kind_uses_text_event_stream(plan_kind: &str) -> bool { matches!( plan_kind, OPENAI_CHAT_STREAM_PLAN_KIND | OPENAI_RESPONSES_STREAM_PLAN_KIND | OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND | OPENAI_IMAGE_STREAM_PLAN_KIND | CLAUDE_CHAT_STREAM_PLAN_KIND | CLAUDE_CLI_STREAM_PLAN_KIND | GEMINI_CHAT_STREAM_PLAN_KIND | GEMINI_CLI_STREAM_PLAN_KIND | GEMINI_INTERACTIONS_STREAM_PLAN_KIND ) } fn should_normalize_declared_stream_response_headers( plan_kind: &str, status_code: u16, headers: &BTreeMap, report_context: Option<&Value>, ) -> bool { plan_kind_uses_text_event_stream(plan_kind) && (200..300).contains(&status_code) && report_context_upstream_is_stream(report_context) && response_headers_have_octet_stream_content_type(headers) && response_headers_have_only_identity_content_encoding(headers) && !headers .keys() .any(|name| name.eq_ignore_ascii_case("content-length")) } fn normalize_declared_stream_response_headers(headers: &mut BTreeMap) { headers.retain(|name, _| { !name.eq_ignore_ascii_case("content-encoding") && !name.eq_ignore_ascii_case("content-length") && !name.eq_ignore_ascii_case("content-type") }); headers.insert("content-type".to_string(), "text/event-stream".to_string()); } fn parse_prefetched_sync_json_body(body: &[u8]) -> Option { let stripped = strip_utf8_bom_and_ws(body); serde_json::from_slice::(stripped).ok() } fn resolve_provider_stream_error_status_code( provider_api_format: &str, upstream_status_code: u16, body_json: &Value, ) -> u16 { if (200..300).contains(&upstream_status_code) && provider_api_format .trim() .eq_ignore_ascii_case("claude:messages") { anthropic_error_status_code(body_json) } else { resolve_local_sync_error_status_code(upstream_status_code, body_json) } } fn anthropic_premature_eof_error_body(message: &str) -> Value { serde_json::json!({ "type": "error", "error": { "type": "api_error", "message": message, } }) } fn build_anthropic_premature_eof_failure(message: &str) -> StreamFailureReport { let body_json = anthropic_premature_eof_error_body(message); build_stream_failure_from_provider_error_body( anthropic_error_status_code(&body_json), &body_json, ) } 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 encode_anthropic_terminal_sse_error_event( failure: &StreamFailureReport, ) -> Result { let payload = serde_json::to_string(&serde_json::json!({ "type": "error", "error": { "type": "api_error", "message": failure.error_message, } })) .map_err(|err| IoError::other(err.to_string()))?; Ok(Bytes::from(format!("event: error\ndata: {payload}\n\n"))) } fn encode_terminal_sse_error_event_for_plan( plan: &ExecutionPlan, failure: &StreamFailureReport, ) -> Result { if plan .client_api_format .trim() .eq_ignore_ascii_case("claude:messages") && plan .provider_api_format .trim() .eq_ignore_ascii_case("claude:messages") { encode_anthropic_terminal_sse_error_event(failure) } else { encode_terminal_sse_error_event(failure) } } 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, anthropic_message_stop_terminates_body: bool, keepalive_interval: Duration, ) -> impl futures_util::Stream> + Send + 'static { stream! { let mut upstream_control_filter = filter_control_blocks.then(SseControlBlockFilter::default); let mut anthropic_completion_tracker = anthropic_message_stop_terminates_body.then(ClientVisibleStreamCompletionTracker::default); let mut sent_prefetched_chunk = false; for chunk in prefetched_chunks_for_body { if let Some(mut chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { let completed = truncate_at_anthropic_message_stop( anthropic_completion_tracker.as_mut(), &mut chunk, ); sent_prefetched_chunk = true; yield Ok(chunk); if completed { return; } } } 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(mut chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { let completed = truncate_at_anthropic_message_stop( anthropic_completion_tracker.as_mut(), &mut chunk, ); yield Ok(chunk); if completed { break; } } } Err(err) => yield Err(err), } } _ = keepalive.tick() => { yield Ok(Bytes::from_static(SSE_KEEPALIVE_BYTES)); } } } if !anthropic_completion_tracker .as_ref() .is_some_and(|tracker| tracker.completed) { 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(mut chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { let completed = truncate_at_anthropic_message_stop( anthropic_completion_tracker.as_mut(), &mut chunk, ); yield Ok(chunk); if completed { break; } } } Err(err) => yield Err(err), } } if !anthropic_completion_tracker .as_ref() .is_some_and(|tracker| tracker.completed) { if let Some(chunk) = flush_upstream_sse_control_filter(&mut upstream_control_filter) { yield Ok(chunk); } } } } } fn truncate_at_anthropic_message_stop( tracker: Option<&mut ClientVisibleStreamCompletionTracker>, chunk: &mut Bytes, ) -> bool { let Some(tracker) = tracker else { return false; }; if let Some(terminal_end) = tracker.observe_anthropic_message_stop_terminal_end(chunk.as_ref()) { chunk.truncate(terminal_end); return true; } false } #[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)> { find_sse_record_boundary(buffer) } fn sse_block_has_data_line(block: &[u8]) -> bool { let Ok(text) = std::str::from_utf8(block) else { return true; }; text.split(['\r', '\n']) .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.split(['\r', '\n']) .any(|line| line.trim_start().starts_with("data:")) } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum SseTerminalPolicy { AnyKnown, AnthropicMessageStop, } #[derive(Default)] struct ClientVisibleStreamCompletionTracker { line_buffer: Vec, event_type: Option, data_payload: String, has_data_payload: bool, record_bytes: usize, dropping_oversized_record: bool, discarded_line_nonempty: bool, skip_next_lf: bool, completed: bool, } impl ClientVisibleStreamCompletionTracker { fn observe_chunk(&mut self, chunk: &[u8]) -> bool { self.observe_chunk_terminal_end(chunk); self.completed } fn observe_chunk_terminal_end(&mut self, chunk: &[u8]) -> Option { self.observe_chunk_terminal_end_with_policy(chunk, SseTerminalPolicy::AnyKnown) } fn observe_anthropic_message_stop(&mut self, chunk: &[u8]) -> bool { self.observe_anthropic_message_stop_terminal_end(chunk); self.completed } fn observe_anthropic_message_stop_terminal_end(&mut self, chunk: &[u8]) -> Option { self.observe_chunk_terminal_end_with_policy(chunk, SseTerminalPolicy::AnthropicMessageStop) } fn observe_chunk_terminal_end_with_policy( &mut self, chunk: &[u8], policy: SseTerminalPolicy, ) -> Option { if self.completed { return None; } if chunk.is_empty() { return None; } for (index, byte) in chunk.iter().enumerate() { self.record_bytes = self.record_bytes.saturating_add(1); if self.record_bytes > SSE_TERMINAL_DETECTOR_MAX_RECORD_BYTES && !self.dropping_oversized_record { self.dropping_oversized_record = true; self.discarded_line_nonempty = !self.line_buffer.is_empty(); self.line_buffer.clear(); self.reset_current_event(); } if self.skip_next_lf { self.skip_next_lf = false; if *byte == b'\n' { continue; } } if self.dropping_oversized_record { match *byte { b'\n' => self.finish_discarded_line(), b'\r' => { self.finish_discarded_line(); self.skip_next_lf = true; } _ => self.discarded_line_nonempty = true, } continue; } match *byte { b'\n' => self.finish_line(policy), b'\r' => { self.finish_line(policy); self.skip_next_lf = true; } _ => self.line_buffer.push(*byte), } if self.completed { let terminal_end = if *byte == b'\r' && chunk.get(index + 1) == Some(&b'\n') { index + 2 } else { index + 1 }; return Some(terminal_end); } } None } fn finish_line(&mut self, policy: SseTerminalPolicy) { 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(policy); self.reset_current_event(); self.record_bytes = 0; 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 finish_discarded_line(&mut self) { if !self.discarded_line_nonempty { self.dropping_oversized_record = false; self.record_bytes = 0; self.reset_current_event(); } self.discarded_line_nonempty = false; } fn current_event_is_terminal(&self, policy: SseTerminalPolicy) -> bool { match policy { SseTerminalPolicy::AnyKnown => { 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)) } SseTerminalPolicy::AnthropicMessageStop => { let payload_type = self .has_data_payload .then(|| serde_json::from_str::(&self.data_payload).ok()) .flatten() .and_then(|value| { value .get("type") .and_then(serde_json::Value::as_str) .map(ToOwned::to_owned) }); payload_type.as_deref() == Some("message_stop") && self .event_type .as_deref() .is_none_or(|event_type| event_type == "message_stop") } } } 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, } #[derive(Clone)] struct PostStopFrameReadBudget { remaining: Arc, } impl PostStopFrameReadBudget { fn new() -> Self { Self { remaining: Arc::new(AtomicUsize::new(POST_STOP_FRAME_READ_BUDGET_INACTIVE)), } } fn activate(&self, already_buffered: usize) -> bool { let remaining = ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES.saturating_sub(already_buffered); let activated = self .remaining .compare_exchange( POST_STOP_FRAME_READ_BUDGET_INACTIVE, remaining, Ordering::AcqRel, Ordering::Acquire, ) .is_ok(); activated && already_buffered > ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES } } struct PostStopLimitedStreamReader { stream: S, current: Option, budget: PostStopFrameReadBudget, } impl PostStopLimitedStreamReader { fn new(stream: S, budget: PostStopFrameReadBudget) -> Self { Self { stream, current: None, budget, } } fn activate_post_stop_budget(&mut self, already_buffered: usize) -> bool { let over_limit = self.budget.activate(already_buffered); let remaining = self.budget.remaining.load(Ordering::Acquire); self.trim_current_to_budget(remaining, true); over_limit } fn trim_current_to_budget(&mut self, remaining: usize, detach_backing: bool) { if remaining == POST_STOP_FRAME_READ_BUDGET_INACTIVE { return; } if remaining == 0 { self.current = None; return; } if let Some(current) = self.current.as_mut() { if detach_backing || current.len() > remaining { let retained = current.len().min(remaining); // Detach even a small slice because it can retain a giant // producer allocation across post-stop backpressure. *current = Bytes::copy_from_slice(¤t[..retained]); } } } } impl AsyncRead for PostStopLimitedStreamReader where S: Stream> + Unpin, { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { let this = self.get_mut(); if buf.remaining() == 0 { return Poll::Ready(Ok(())); } let mut empty_chunks = 0usize; loop { let remaining = this.budget.remaining.load(Ordering::Acquire); if remaining == 0 { this.current = None; return Poll::Ready(Ok(())); } this.trim_current_to_budget(remaining, false); if let Some(current) = this.current.as_mut() { let read = current.len().min(buf.remaining()); if read > 0 { buf.put_slice(¤t.split_to(read)); if remaining != POST_STOP_FRAME_READ_BUDGET_INACTIVE { let previous = this.budget.remaining.fetch_sub(read, Ordering::AcqRel); debug_assert!(previous != POST_STOP_FRAME_READ_BUDGET_INACTIVE); debug_assert!(previous >= read); } } if current.is_empty() { this.current = None; } if read > 0 { return Poll::Ready(Ok(())); } } match Pin::new(&mut this.stream).poll_next(cx) { Poll::Pending => return Poll::Pending, Poll::Ready(Some(Ok(chunk))) if chunk.is_empty() => { empty_chunks += 1; if empty_chunks >= POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL { cx.waker().wake_by_ref(); return Poll::Pending; } } Poll::Ready(Some(Ok(chunk))) => { this.current = Some(chunk); if remaining != POST_STOP_FRAME_READ_BUDGET_INACTIVE { this.trim_current_to_budget(remaining, true); } } Poll::Ready(Some(Err(err))) => return Poll::Ready(Err(err)), Poll::Ready(None) => return Poll::Ready(Ok(())), } } } } fn activate_post_stop_frame_read_budget( lines: &mut FramedRead, LinesCodec>, ) -> bool { let already_buffered = lines.read_buffer().len(); let over_limit = already_buffered > ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES; let reader_over_limit = lines.get_mut().activate_post_stop_budget(already_buffered); let retained = if over_limit { 0 } else { already_buffered }; let mut bounded = bytes::BytesMut::with_capacity(retained); bounded.extend_from_slice(&lines.read_buffer()[..retained]); *lines.read_buffer_mut() = bounded; reader_over_limit || over_limit } 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 serialized_stream_frame_len(frame: &StreamFrame) -> usize { serde_json::to_vec(frame).map_or(usize::MAX, |encoded| encoded.len()) } fn execution_stream_frame_codec() -> LinesCodec { LinesCodec::new_with_max_length(crate::execution_runtime::MAX_EXECUTION_STREAM_FRAME_LINE_BYTES) } 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.usage_lifecycle_data_state().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, force_prefetch: bool, ) -> bool { StreamCommitPolicy::for_response( direct_stream_finalize_kind.is_some(), content_type, provider_api_format, client_api_format, has_private_stream_normalizer, has_local_stream_rewriter, force_prefetch, ) .commits_on_response_headers() } fn prefetched_openai_responses_body_has_output_boundary(body: &[u8]) -> bool { let Ok(text) = std::str::from_utf8(body) else { return true; }; for line in text.lines() { let Some(data) = line.trim().strip_prefix("data:").map(str::trim) else { continue; }; if data.is_empty() { continue; } if data == "[DONE]" { return true; } let Ok(event) = serde_json::from_str::(data) else { continue; }; let event_type = event.get("type").and_then(Value::as_str).map(str::trim); if !event_type.is_some_and(|event_type| { matches!( event_type, "response.created" | "response.in_progress" | "response.queued" ) }) { return true; } } false } 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 record_prefetch_success_failover( state: &AppState, plan: &ExecutionPlan, report_context: Option<&Value>, elapsed_ms: u64, ) { let finished_at = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, plan, report_context, SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, status_code: Some(200), error_type: Some("success_failover_pattern".to_string()), error_message: Some("HTTP 200 response matched a precommit failover rule".to_string()), latency_ms: Some(elapsed_ms), started_at_unix_ms: None, finished_at_unix_ms: Some(finished_at), }, ) .await; } 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, stage_trace: RequestStageTrace, lifecycle_pending_recorded: bool, frame_stream: BoxStream<'static, Result>, in_flight_guard: Option, ) -> Result>, GatewayError> { execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, decision, plan_kind, report_kind, report_context, candidate_started_unix_secs, stream_started_at, stage_trace, lifecycle_pending_recorded, frame_stream, false, in_flight_guard, None, None, None, ) .await } #[allow(clippy::too_many_arguments)] async fn execute_stream_from_frame_stream_with_retry_scope( 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, lifecycle_pending_recorded: bool, frame_stream: BoxStream<'static, Result>, stream_precommit_committed: bool, in_flight_guard: Option, mut retry_scope_out: Option<&mut AiAttemptRetryScope>, mut retry_fallback_out: Option<&mut Option>>, fallback_response_observation: 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()); if !lifecycle_pending_recorded { record_stream_pending_lifecycle(state, &lifecycle_seed, &mut stage_trace).await; } let max_stream_body_buffer_bytes = resolve_stream_body_buffer_limit(state).await; 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 = PostStopLimitedStreamReader::new(frame_stream, PostStopFrameReadBudget::new()); let mut lines = FramedRead::new(reader, execution_stream_frame_codec()); 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, response_observation, } = first_frame.payload else { return Err(GatewayError::Internal( "execution runtime stream must start with headers frame".to_string(), )); }; let response_observation = response_observation .or(fallback_response_observation) .unwrap_or(ExecutionResponseObservation { request_started_at_unix_ms: candidate_started_unix_secs, response_headers_observed_at_unix_ms: current_request_candidate_unix_ms(), request_order_id: uuid::Uuid::now_v7().to_string(), }); let mut report_context = attach_provider_response_headers_to_report_context( report_context, &headers, response_observation.request_started_at_unix_ms, response_observation.response_headers_observed_at_unix_ms, &response_observation.request_order_id, ); spawn_local_oauth_success_effect( state.clone(), &plan, report_context.as_ref(), LocalOAuthSuccessEffect { status_code, request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms), request_order_id: Some(&response_observation.request_order_id), }, ); 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 failure_disposition = classify_failure_disposition( &plan.provider_api_format, failover_analysis.classification, status_code, ); if let Some(retry_scope) = retry_scope_out.as_deref_mut() { *retry_scope = ai_attempt_retry_scope_from_failure_disposition(failure_disposition); } if failure_disposition.preserve_upstream_error { if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() { let mut fallback_headers = headers.clone(); apply_endpoint_response_header_rules( state, &plan, &mut fallback_headers, provider_body_json.as_ref(), ) .await?; *retry_fallback = Some(attach_control_metadata_headers( build_client_response_from_parts( status_code, &fallback_headers, Body::from(provider_error_body.clone()), trace_id, Some(decision), )?, Some(request_id), candidate_id, )?); } } 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_with_handoff( state, &plan, payload.report_context.as_ref(), &payload, ) .await; 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, )?)); } let response = if (300..400).contains(&status_code) { build_client_response_from_parts_with_mutator( client_status_code, &client_response_headers, Body::from(client_error_body), trace_id, Some(decision), |headers| { headers.insert( http::HeaderName::from_static("x-aether-upstream-status"), http::HeaderValue::from_str(&status_code.to_string()) .map_err(|error| GatewayError::Internal(error.to_string()))?, ); Ok(()) }, )? } else { build_client_response_from_parts( client_status_code, &client_response_headers, Body::from(client_error_body), trace_id, Some(decision), )? }; return Ok(Some(attach_control_metadata_headers( response, 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 normalized_declared_stream_headers = private_stream_normalizer.is_none() && local_stream_rewriter.is_none() && should_normalize_declared_stream_response_headers( plan_kind, status_code, &upstream_headers, report_context.as_ref(), ); if normalized_declared_stream_headers { normalize_declared_stream_response_headers(&mut headers); debug!( event_name = "execution_runtime_stream_content_type_corrected", 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(), upstream_content_type = upstream_content_type.unwrap_or("-"), "gateway normalized declared upstream stream response headers for the client" ); } let prefetch_for_cyber_failover = is_openai_responses_family_format(plan.provider_api_format.as_str()) && crate::orchestration::routing_execution_policy_from_report_context( report_context.as_ref(), ) .is_some_and(|policy| policy.cyber_continue_failover); let prefetch_failover_policy = crate::orchestration::resolve_local_failover_policy(state, &plan, report_context.as_ref()) .await; let prefetch_success_patterns = prefetch_failover_policy .routing_rules .success_failover_patterns .iter() .map(|rule| (&rule.pattern, &rule.status_codes)) .chain( prefetch_failover_policy .success_failover_patterns .iter() .map(|rule| (&rule.pattern, &rule.status_codes)), ) .filter(|(_, status_codes)| status_codes.is_empty() || status_codes.contains(&200)) .filter_map(|(pattern, _)| regex::Regex::new(pattern.trim()).ok()) .collect::>(); let stream_commit_policy = StreamCommitPolicy::for_response( direct_stream_finalize_kind.is_some(), 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(), prefetch_for_cyber_failover || !prefetch_success_patterns.is_empty(), ) .with_precommit_wait(Duration::from_millis( plan.timeouts .as_ref() .and_then(|timeouts| timeouts.first_byte_ms) .unwrap_or(30_000) .max(1), )); let reuse_committed_precommit = stream_precommit_committed && stream_commit_policy.is_native_anthropic() && prefetch_success_patterns.is_empty(); let skip_direct_finalize_prefetch = stream_commit_policy.commits_on_response_headers() || reuse_committed_precommit; let limit_direct_finalize_prefetch = should_limit_direct_finalize_prefetch(plan_kind, local_stream_rewriter.is_some()) || stream_commit_policy.requires_bounded_frame_wait() || !prefetch_success_patterns.is_empty(); let mut stream_commit_gate = StreamCommitGate::new(stream_commit_policy); let mut prefetch_client_completion_tracker = ClientVisibleStreamCompletionTracker::default(); let mut prefetched_client_visible_stream_completed = false; let mut prefetched_anthropic_message_stop_observed_at = None; let mut prefetched_anthropic_post_stop_buffer_over_limit = false; if reuse_committed_precommit { stream_commit_gate.commit(); } let mut prefetched_chunks: Vec = Vec::new(); let mut provider_prefetched_body = Vec::new(); let mut provider_prefetched_body_truncated = false; let mut prefetched_body = Vec::new(); let mut prefetched_inspection_body = Vec::new(); let mut prefetched_inspection_body_truncated = false; 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; let precommit_started_at = Instant::now(); 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 (stream_commit_policy.requires_bounded_frame_wait() || prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES) && prefetched_inspection_body.len() < MAX_STREAM_PREFETCH_BYTES { let next_frame_result = if limit_direct_finalize_prefetch { let prefetch_timeout = stream_commit_policy .max_precommit_wait() .map(|max_wait| max_wait.saturating_sub(precommit_started_at.elapsed())) .unwrap_or(REWRITTEN_STREAM_PREFETCH_TIMEOUT); if prefetch_timeout.is_zero() && !stream_commit_policy.requires_bounded_frame_wait() { stream_commit_gate.commit(); 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 = stream_commit_policy .max_precommit_wait() .unwrap_or(REWRITTEN_STREAM_PREFETCH_TIMEOUT) .as_millis() as u64, "gateway reached bounded stream precommit deadline" ); break; } match tokio::time::timeout( prefetch_timeout, next_stream_frame(&mut buffered_frames, &mut lines), ) .await { Ok(result) => result, Err(_) => { if stream_commit_policy.requires_bounded_frame_wait() { let failure = build_stream_transport_failure_report( "first_byte_timeout", "Upstream did not produce a semantic event before the first byte deadline", 504, ); 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, candidate_started_unix_secs, stream_elapsed_ms_since(stream_started_at), failure, retry_scope_out.as_deref_mut(), ) .await; } stream_commit_gate.commit(); 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 = stream_commit_policy .max_precommit_wait() .unwrap_or(REWRITTEN_STREAM_PREFETCH_TIMEOUT) .as_millis() as u64, "gateway stopped bounded 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, candidate_started_unix_secs, stream_elapsed_ms_since(stream_started_at), failure, None, ) .await; } }) else { if stream_commit_policy.requires_bounded_frame_wait() && stream_commit_gate.is_uncommitted() { let error_body_json = anthropic_premature_eof_error_body( "upstream stream ended before the first semantic event", ); let error_status_code = anthropic_error_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, status_code, error_status_code, error_body_json, retry_scope_out.as_deref_mut(), None, ) .await; } 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), ); state.usage_runtime.record_stream_started( state.usage_lifecycle_data_state().as_ref(), &lifecycle_seed, status_code, prefetched_usage_telemetry.as_ref(), ); } let mut 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, candidate_started_unix_secs, stream_elapsed_ms_since(stream_started_at), failure, None, ) .await; } }; if chunk.is_empty() { continue; } if stream_commit_policy.is_native_anthropic() { if prefetched_client_visible_stream_completed { continue; } if let Some(terminal_end) = prefetch_client_completion_tracker .observe_anthropic_message_stop_terminal_end(&chunk) { chunk.truncate(terminal_end); prefetched_client_visible_stream_completed = true; prefetched_anthropic_message_stop_observed_at .get_or_insert_with(Instant::now); prefetched_anthropic_post_stop_buffer_over_limit |= activate_post_stop_frame_read_budget(&mut lines); } } append_stream_capture_bytes( &mut provider_prefetched_body, &chunk, MAX_STREAM_PREFETCH_BYTES, &mut provider_prefetched_body_truncated, ); append_stream_capture_bytes( &mut prefetched_inspection_body, &chunk, MAX_STREAM_PREFETCH_BYTES, &mut prefetched_inspection_body_truncated, ); if !prefetch_success_patterns.is_empty() && crate::orchestration::attempt_identity_from_report_context( report_context.as_ref(), ) .is_some() { let response_text = String::from_utf8_lossy(&prefetched_inspection_body); if prefetch_success_patterns.iter().any(|pattern| pattern.is_match(&response_text)) && crate::orchestration::classify_local_failover( &prefetch_failover_policy, crate::orchestration::LocalFailoverInput::new(status_code, Some(&response_text)), ) == crate::orchestration::LocalFailoverClassification::RetrySuccessPattern { record_prefetch_success_failover(state, &plan, report_context.as_ref(), stream_elapsed_ms_since(stream_started_at)).await; if let Some(retry_scope) = retry_scope_out.as_deref_mut() { *retry_scope = AiAttemptRetryScope::Candidate; } warn!(event_name = "local_stream_candidate_retry_scheduled", log_type = "event", trace_id, request_id, status_code, "gateway retrying after a precommit success pattern match"); return Ok(None); } } let semantic_commit_ready = match stream_commit_gate.observe_provider_bytes(&chunk) { StreamPrecommitObservation::Pending => false, StreamPrecommitObservation::Commit => true, StreamPrecommitObservation::UpstreamError { status_code: error_status_code, body_json: error_body_json, } => { let error_status_code = if plan .provider_api_format .eq_ignore_ascii_case("claude:messages") { anthropic_error_status_code(&error_body_json) } else { error_status_code }; 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, status_code, error_status_code, error_body_json, retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), ) .await; } }; if !semantic_commit_ready || private_stream_normalizer.is_some() { if let Some(error_body_json) = extract_provider_private_stream_error_body( report_context.as_ref(), &prefetched_inspection_body, ) { let error_status_code = resolve_provider_stream_error_status_code( plan.provider_api_format.as_str(), 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, status_code, error_status_code, error_body_json, retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), ) .await; } } let inspection = if stream_commit_policy.requires_bounded_frame_wait() { StreamPrefetchInspection::NeedMore } else { inspect_prefetched_stream_body( &upstream_headers, &prefetched_inspection_body, ) }; match inspection { StreamPrefetchInspection::EmbeddedError(body_json) => { let error_status_code = resolve_provider_stream_error_status_code( plan.provider_api_format.as_str(), status_code, &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, status_code, error_status_code, body_json, retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), ) .await; } 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)) => { if let Some(record) = outcome.response_history_record { crate::ai_serving::persist_response_history_record( state, record, ) .await; } 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, candidate_started_unix_secs, stream_elapsed_ms_since(stream_started_at), failure, None, ) .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, candidate_started_unix_secs, stream_elapsed_ms_since(stream_started_at), failure, None, ) .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, candidate_started_unix_secs, stream_elapsed_ms_since(stream_started_at), failure, None, ) .await; } } } else { normalized_chunk }; if !rewritten_chunk.is_empty() { prefetched_body.extend_from_slice(&rewritten_chunk); prefetched_chunks.push(Bytes::from(rewritten_chunk)); } if semantic_commit_ready || (matches!(inspection, StreamPrefetchInspection::NonError) && (prefetch_success_patterns.is_empty() || response_headers_indicate_sse(&upstream_headers) || parse_prefetched_sync_json_body(&prefetched_inspection_body) .is_some()) && (!prefetch_for_cyber_failover || prefetched_openai_responses_body_has_output_boundary( &prefetched_inspection_body, ))) { break; } } StreamFramePayload::Telemetry { telemetry: frame_telemetry, } => { prefetched_telemetry = Some(frame_telemetry); } StreamFramePayload::Eof { summary } => { if stream_commit_policy.requires_bounded_frame_wait() && stream_commit_gate.is_uncommitted() { let error_body_json = anthropic_premature_eof_error_body( "upstream stream ended before the first semantic event", ); let error_status_code = anthropic_error_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, status_code, error_status_code, error_body_json, retry_scope_out.as_deref_mut(), None, ) .await; } 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_kind = ?error.kind, error_phase = ?error.phase, upstream_status = ?error.upstream_status, "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, candidate_started_unix_secs, stream_elapsed_ms_since(stream_started_at), build_stream_failure_from_execution_error(&error), retry_scope_out.as_deref_mut(), ) .await; } StreamFramePayload::Headers { .. } => {} } } } if stream_commit_gate.is_uncommitted() { stream_commit_gate.commit(); } let prefetched_response_history_persisted = if let Some(record) = local_stream_rewriter .as_mut() .and_then(|rewriter| rewriter.take_response_history_record()) { crate::ai_serving::persist_response_history_record(state, record).await; true } else { false }; 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.usage_lifecycle_data_state().as_ref(), &lifecycle_seed, status_code, initial_usage_telemetry.as_ref(), ); if let Some(snapshot) = request_candidate_status_snapshot { let latency_ms = prefetched_telemetry .as_ref() .and_then(|telemetry| telemetry.elapsed_ms); record_local_request_candidate_status_snapshot( state, &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 native_anthropic_stream_for_report = stream_commit_policy.is_native_anthropic(); let plan_for_report = plan; let emit_passthrough_sse_terminal_error = (skip_direct_finalize_prefetch || stream_commit_policy.requires_bounded_frame_wait() || normalized_declared_stream_headers) && (response_headers_indicate_sse(&upstream_headers) || normalized_declared_stream_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 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(); let mut provider_error_inspection = ProviderStreamErrorInspection::default(); 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 = if native_anthropic_stream_for_report { client_stream_completion_tracker .observe_anthropic_message_stop(&prefetched_body_for_report) } else { client_stream_completion_tracker.observe_chunk(&prefetched_body_for_report) }; let mut anthropic_post_stop_drain_started_at = (native_anthropic_stream_for_report && client_visible_stream_completed) .then(|| prefetched_anthropic_message_stop_observed_at.unwrap_or_else(Instant::now)); let mut anthropic_post_stop_buffer_over_limit = prefetched_anthropic_post_stop_buffer_over_limit; if anthropic_post_stop_drain_started_at.is_some() && prefetched_anthropic_message_stop_observed_at.is_none() { anthropic_post_stop_buffer_over_limit |= activate_post_stop_frame_read_budget(&mut lines); } let mut anthropic_post_stop_drain_frames = 0usize; let mut anthropic_post_stop_drain_bytes = 0usize; 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 mut provider_error_forwarded_to_client = false; 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_category = "stream_normalization_restore_failed", "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(error_body_json) = provider_error_inspection .observe(stream_usage_report_context.as_ref(), replay_chunk) { provider_error_forwarded_to_client = !prefetched_body_for_report.is_empty(); let error_status_code = resolve_provider_stream_error_status_code( plan_for_report.provider_api_format.as_str(), status_code, &error_body_json, ); terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, )); } 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_category = "stream_rewrite_restore_failed", "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 prefetched_response_history_persisted { if let Some(rewriter) = local_stream_rewriter.as_mut() { let _ = rewriter.take_response_history_record(); } } } if terminal_failure.is_none() && !reached_eof { loop { let draining_after_anthropic_stop = anthropic_post_stop_drain_started_at.is_some(); let next_frame_result = if let Some(drain_started_at) = anthropic_post_stop_drain_started_at { if anthropic_post_stop_drain_frames >= ANTHROPIC_POST_STOP_DRAIN_MAX_FRAMES || anthropic_post_stop_drain_bytes >= ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES || anthropic_post_stop_buffer_over_limit { break; } let remaining = ANTHROPIC_POST_STOP_DRAIN_MAX_WAIT .saturating_sub(drain_started_at.elapsed()); if remaining.is_zero() { break; } match tokio::time::timeout( remaining, next_stream_frame(&mut buffered_frames, &mut lines), ) .await { Ok(result) => result, Err(_) => break, } } else { tokio::select! { biased; _ = tx.closed(), if !downstream_dropped => { 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) => { if native_anthropic_stream_for_report && client_visible_stream_completed { debug!( event_name = "stream_execution_frame_decode_ignored_after_anthropic_stop", log_type = "debug", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error_category = "stream_frame_decode_failed", "gateway ignored execution runtime teardown error after Anthropic message_stop" ); break; } 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_category = "stream_frame_decode_failed", "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; } else if native_anthropic_stream_for_report && !client_visible_stream_completed { terminal_failure = Some(build_anthropic_premature_eof_failure( "upstream Anthropic stream ended before message_stop", )); } break; }; if draining_after_anthropic_stop { anthropic_post_stop_drain_frames = anthropic_post_stop_drain_frames.saturating_add(1); let frame_bytes = serialized_stream_frame_len(&observed_frame.frame); if frame_bytes > ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES .saturating_sub(anthropic_post_stop_drain_bytes) { break; } anthropic_post_stop_drain_bytes = anthropic_post_stop_drain_bytes.saturating_add(frame_bytes); } 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 } => { if native_anthropic_stream_for_report && client_visible_stream_completed { continue; } let first_data_before = usage_stream_telemetry .as_ref() .and_then(|telemetry| telemetry.ttfb_ms) .is_some(); if maybe_capture_first_stream_event_telemetry( stream_started_at_for_report, frame_observed_at, telemetry.as_ref(), &mut usage_stream_telemetry, ) { state_for_report.usage_runtime.record_stream_started( state_for_report.usage_lifecycle_data_state().as_ref(), &lifecycle_seed_for_report, status_code, usage_stream_telemetry.as_ref(), ); } 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 mut 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_category = "stream_chunk_decode_failed", "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; } if native_anthropic_stream_for_report { if let Some(terminal_end) = client_stream_completion_tracker .observe_anthropic_message_stop_terminal_end(&chunk) { chunk.truncate(terminal_end); client_visible_stream_completed = true; if anthropic_post_stop_drain_started_at.is_none() { anthropic_post_stop_drain_started_at = Some(Instant::now()); anthropic_post_stop_buffer_over_limit |= activate_post_stop_frame_read_budget(&mut lines); } } } 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_category = "stream_chunk_normalize_failed", "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 = provider_error_inspection .observe(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_category = "stream_chunk_rewrite_failed", "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 provider_private_error_body_json.is_none() { if let Some(record) = local_stream_rewriter .as_mut() .and_then(|rewriter| rewriter.take_response_history_record()) { crate::ai_serving::persist_response_history_record( &state_for_report, record, ) .await; } } if rewritten_chunk.is_empty() { if let Some(error_body_json) = provider_private_error_body_json { let error_status_code = resolve_provider_stream_error_status_code( plan_for_report.provider_api_format.as_str(), 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; cancelling execution runtime stream" ); downstream_dropped = true; break; } else { if !native_anthropic_stream_for_report { 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, ); provider_error_forwarded_to_client = provider_private_error_body_json.is_some(); } if let Some(error_body_json) = provider_private_error_body_json { let error_status_code = resolve_provider_stream_error_status_code( plan_for_report.provider_api_format.as_str(), 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 { // The first Data frame records the live streaming transition. Later // telemetry frames only refine the terminal accumulator; persisting // every elapsed-time update would create one usage task per frame. usage_stream_telemetry = Some(usage_frame_telemetry); } telemetry = Some(frame_telemetry); } StreamFramePayload::Eof { summary } => { stream_terminal_summary = merge_stream_terminal_summary(stream_terminal_summary.take(), summary); if native_anthropic_stream_for_report && !client_visible_stream_completed { terminal_failure = Some(build_anthropic_premature_eof_failure( "upstream Anthropic stream ended before message_stop", )); } break; } StreamFramePayload::Error { error } => { if native_anthropic_stream_for_report && client_visible_stream_completed { debug!( event_name = "stream_execution_error_frame_ignored_after_anthropic_stop", log_type = "debug", trace_id = %trace_id_owned, request_id = %request_id_for_report_log, candidate_id = ?candidate_id_for_report.as_deref(), error_kind = ?error.kind, error_phase = ?error.phase, upstream_status = ?error.upstream_status, "gateway ignored execution runtime error frame after Anthropic message_stop" ); continue; } 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_kind = ?error.kind, error_phase = ?error.phase, upstream_status = ?error.upstream_status, "execution runtime stream emitted error frame" ); terminal_failure = Some(build_stream_failure_from_execution_error(&error)); break; } StreamFramePayload::Headers { .. } => {} } } } drop(lines); drop(buffered_frames); drop(_provider_pool_in_flight_guard); 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 = provider_error_inspection .observe(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_category = "stream_flush_rewrite_failed", "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 provider_private_error_body_json.is_none() { if let Some(record) = local_stream_rewriter .as_mut() .and_then(|rewriter| rewriter.take_response_history_record()) { crate::ai_serving::persist_response_history_record( &state_for_report, record, ) .await; } } 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_provider_stream_error_status_code( plan_for_report.provider_api_format.as_str(), 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_category = "stream_normalization_flush_failed", "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() { let finish_result = rewriter.finish(); if let Some(record) = rewriter.take_response_history_record() { crate::ai_serving::persist_response_history_record(&state_for_report, record) .await; } match finish_result { 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_category = "stream_rewrite_flush_failed", "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 terminal_failure.is_none() { if let Some(record) = local_stream_rewriter .as_mut() .and_then(|rewriter| rewriter.take_response_history_record()) { crate::ai_serving::persist_response_history_record(&state_for_report, record).await; } } 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 && !provider_error_forwarded_to_client { Some(encode_terminal_sse_error_event_for_plan( &plan_for_report, 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_category = "terminal_error_event_encode_failed", "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(), stream_started_at_for_report, terminal_telemetry.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, ) .await; 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(), stream_started_at_for_report, terminal_telemetry.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(), stream_started_at_for_report, terminal_telemetry.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, ) .await; 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_category = "stream_report_submit_failed", "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, native_anthropic_stream_for_report, 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::{ atomic::{AtomicBool, AtomicUsize, Ordering}, Arc, Mutex, }; use std::time::{Duration, Instant}; use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope}; use aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED; use aether_contracts::{ ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, ExecutionTimeouts, RequestBody, StandardizedUsage, StreamFrame, StreamFramePayload, StreamFrameType, }; use aether_data::repository::candidates::InMemoryRequestCandidateRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::repository::proxy_nodes::{InMemoryProxyNodeRepository, StoredProxyNode}; use aether_data::repository::usage::InMemoryUsageReadRepository; use aether_data_contracts::repository::candidates::{ PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate, UpsertRequestCandidateRecord, }; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use aether_data_contracts::repository::settlement::{ StoredUsageSettlement, UsageSettlementInput, }; use aether_data_contracts::repository::usage::UsageReadRepository; use aether_data_contracts::repository::usage::{ StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState, }; use aether_data_contracts::DataLayerError; use aether_usage_runtime::{ apply_usage_body_capture_policy_to_event, UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageEvent, UsageEventData, UsageEventType, UsageRecordWriter, UsageRequestRecordLevel, UsageRuntimeAccess, UsageRuntimeConfig, UsageSettlementWriter, }; use async_stream::stream; use async_trait::async_trait; 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, http::StatusCode, response::IntoResponse, Json, Router, }; use base64::Engine as _; use futures_util::StreamExt as _; use serde_json::{json, Value}; use tokio::sync::{mpsc, watch, Notify}; use super::{ activate_post_stop_frame_read_budget, build_direct_execution_frame_stream, build_sse_body_stream, build_stream_failure_report, build_stream_sync_payload, client_format_allows_proxy_generated_sse_control_blocks, direct_upstream_response_byte_stream, encode_terminal_sse_error_event_for_plan, ensure_stream_terminal_summary_for_missing_observed_finish, execute_execution_runtime_stream, execute_in_process_stream_with_oauth_retry, execute_stream_from_frame_stream, execute_stream_from_frame_stream_with_retry_scope, execution_stream_frame_codec, maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary, normalize_declared_stream_response_headers, parse_direct_passthrough_mode, prefetch_direct_stream_error_body, prefetched_openai_responses_body_has_output_boundary, record_sync_terminal_usage_with_handoff, record_sync_terminal_usage_with_handoff_after_spawn, resolve_provider_stream_error_status_code, select_direct_anthropic_prefetch_wait, should_limit_direct_finalize_prefetch, should_normalize_declared_stream_response_headers, 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, wrap_non_json_binary_stream_error_for_client, ClientVisibleStreamCompletionTracker, DirectPassthroughFinalizer, DirectPassthroughFinalizerCore, DirectPassthroughInlineBodyState, DirectPassthroughMode, PostStopFrameReadBudget, PostStopLimitedStreamReader, ProviderStreamErrorInspection, ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL, PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES, }; use crate::control::GatewayControlDecision; use crate::stage_metrics::RequestStageTrace; use crate::tunnel::{tunnel_protocol, TunnelProxyConn}; use crate::AppState; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; fn provider_catalog_stop_429_for_plan( plan: &ExecutionPlan, ) -> InMemoryProviderCatalogReadRepository { provider_catalog_for_plan( plan, Some(json!({ "failover_rules": { "stop_status_codes": [429] } })), ) } fn provider_catalog_for_plan( plan: &ExecutionPlan, provider_config: Option, ) -> InMemoryProviderCatalogReadRepository { let credential_state = AppState::new() .expect("credential state should build") .with_data_state_for_tests( crate::data::GatewayDataState::disabled() .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ); let encrypted_api_key = credential_state .seal_provider_catalog_key_api_key( &plan.provider_id, &plan.key_id, "plain-upstream-key", ) .expect("api key should encrypt"); 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, provider_config, ); 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()])), encrypted_api_key, None, None, Some(json!({ "openai:chat": 1 })), None, None, None, None, ) .expect("key transport should build"); InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![endpoint], vec![key]) } #[test] fn non_json_upstream_error_body_is_not_projected_to_clients() { let secret = b"Bearer upstream-secret https://user:password@example.test/private"; let body = wrap_non_json_binary_stream_error_for_client( "openai_chat_stream", &BTreeMap::from([("content-type".to_string(), "text/plain".to_string())]), secret, ) .expect("error body projection should succeed") .expect("non-JSON errors should receive a client projection"); assert_eq!(body["error"]["message"], "Upstream request failed"); assert!(!body.to_string().contains("upstream-secret")); assert!(!body.to_string().contains("password")); } fn provider_catalog_for_stream_auth_plan( plan: &ExecutionPlan, provider_type: &str, auth_type: &str, auth_config: Option, ) -> InMemoryProviderCatalogReadRepository { 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"); 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( plan.url.clone(), None, None, Some(2), None, None, None, None, ) .expect("endpoint transport should build"); let encrypted_auth_config = auth_config.map(|config| { encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, &config.to_string()) .expect("auth config should encrypt") }); let key = StoredProviderCatalogKey::new( plan.key_id.clone(), plan.provider_id.clone(), plan.key_id.clone(), auth_type.to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(json!([plan.provider_api_format.clone()])), None, encrypted_auth_config, None, None, None, None, None, None, ) .expect("key transport should build"); InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![endpoint], vec![key]) } fn direct_stream_test_plan(request_id: &str, url: String) -> ExecutionPlan { ExecutionPlan { request_id: request_id.to_string(), candidate_id: Some(format!("candidate-{request_id}")), provider_name: Some("codex".to_string()), provider_id: format!("provider-{request_id}"), endpoint_id: format!("endpoint-{request_id}"), key_id: format!("key-{request_id}"), method: "POST".to_string(), url, headers: BTreeMap::from([ ("content-type".to_string(), "application/json".to_string()), ( "authorization".to_string(), "AgentAssertion stale-task".to_string(), ), ]), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({"stream": true})), stream: true, client_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5".to_string()), proxy: None, transport_profile: None, timeouts: Some(ExecutionTimeouts { connect_ms: Some(30_000), first_byte_ms: Some(30_000), ..ExecutionTimeouts::default() }), } } fn agent_identity_test_auth_config(task_id: &str) -> Value { json!({ "provider_type": "codex", "auth_mode": "agentIdentity", "agent_runtime_id": "runtime-test", "agent_private_key": "MC4CAQAwBQYDK2VwBCIEIAcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcH", "task_id": task_id }) } fn generic_oauth_test_auth_config(provider_type: &str) -> Value { json!({ "provider_type": provider_type, "access_token": "stale-access-token", "refresh_token": "refresh-token", "expires_at": 4_102_444_800_u64 }) } async fn collect_direct_execution_body( mut execution: crate::execution_runtime::DirectUpstreamStreamExecution, ) -> Result, String> { let prefetched_body = std::mem::take(&mut execution.prefetched_body); let mut stream = direct_upstream_response_byte_stream(prefetched_body, execution.response); let mut body = Vec::new(); while let Some(item) = stream.next().await { body.extend_from_slice(&item?); } Ok(body) } fn codex_cyber_policy_plan(request_id: &str) -> ExecutionPlan { ExecutionPlan { request_id: request_id.to_string(), candidate_id: Some(format!("candidate-{request_id}")), provider_name: Some("codex".to_string()), provider_id: format!("provider-{request_id}"), endpoint_id: format!("endpoint-{request_id}"), key_id: format!("key-{request_id}"), method: "POST".to_string(), url: "https://chatgpt.com/backend-api/codex/responses".to_string(), headers: BTreeMap::from([ ("content-type".to_string(), "application/json".to_string()), ("accept".to_string(), "text/event-stream".to_string()), ]), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.5", "input": [], "stream": true })), stream: true, client_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.5".to_string()), proxy: None, transport_profile: None, timeouts: None, } } async fn execute_prefetched_codex_cyber_policy_failure( continue_failover: bool, ) -> Option> { let request_id = if continue_failover { "req-cyber-policy-retry" } else { "req-cyber-policy-stop" }; let plan = codex_cyber_policy_plan(request_id); let provider_catalog = provider_catalog_for_plan(&plan, None); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( Arc::new(provider_catalog), DEVELOPMENT_ENCRYPTION_KEY, ); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests(data_state); let upstream_setup = "event: response.created\ndata: {\"type\":\"response.created\"}\n\n"; let upstream_error = "event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"type\":\"invalid_request\",\"message\":\"cyber policy rejected the request\",\"code\":\"cyber_policy_violation\",\"param\":\"input\"}}}\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(), )]), response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(upstream_setup.to_string()), }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(upstream_error.to_string()), }, })); yield Ok::(ndjson_frame(StreamFrame::eof())); } .boxed(); execute_stream_from_frame_stream( &state, plan, &format!("trace-{request_id}"), &test_decision(), "openai_responses_stream", Some("openai_responses_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:responses", "client_api_format": "openai:responses", "routing_execution_policy": { "cyber_continue_failover": continue_failover } })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), true, frame_stream, None, ) .await .expect("execution should succeed") } async fn execute_prefetched_transport_failure( stop_on_transport_errors: bool, ) -> AiAttemptExecutionOutcome> { let request_id = if stop_on_transport_errors { "req-prefetch-transport-stop" } else { "req-prefetch-transport-retry" }; let plan = native_anthropic_stream_plan(request_id); let provider_config = stop_on_transport_errors.then(|| { json!({ "failover_rules": { "stop_on_transport_errors": true, } }) }); let provider_catalog = provider_catalog_for_plan(&plan, provider_config); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( Arc::new(provider_catalog), DEVELOPMENT_ENCRYPTION_KEY, ); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests(data_state); 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(), )]), response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Error, payload: StreamFramePayload::Error { error: ExecutionError { kind: ExecutionErrorKind::Internal, phase: ExecutionPhase::StreamRead, message: "connection reset before first body byte".to_string(), upstream_status: None, retryable: true, failover_recommended: true, }, }, })); } .boxed(); let mut retry_scope = AiAttemptRetryScope::Provider; let response = execute_stream_from_frame_stream_with_retry_scope( &state, plan, &format!("trace-{request_id}"), &test_decision(), "claude_chat_stream", Some("claude_chat_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "claude:messages", "client_api_format": "claude:messages" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), true, frame_stream, false, None, Some(&mut retry_scope), None, None, ) .await .expect("prefetch transport execution should resolve"); match response { Some(response) => AiAttemptExecutionOutcome::Responded(response), None => AiAttemptExecutionOutcome::Retry { scope: retry_scope, fallback_response: None, }, } } async fn execute_prefetched_http_status_failure( continue_failover: bool, ) -> AiAttemptExecutionOutcome> { let request_id = if continue_failover { "req-prefetch-http-continue" } else { "req-prefetch-http-stop" }; let plan = native_anthropic_stream_plan(request_id); let failover_rules = if continue_failover { json!({"continue_status_codes": [500]}) } else { json!({"stop_status_codes": [500]}) }; let provider_catalog = provider_catalog_for_plan( &plan, Some(json!({ "failover_rules": failover_rules, })), ); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( Arc::new(provider_catalog), DEVELOPMENT_ENCRYPTION_KEY, ); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests(data_state); 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(), )]), response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Error, payload: StreamFramePayload::Error { error: ExecutionError { kind: ExecutionErrorKind::Internal, phase: ExecutionPhase::StreamRead, message: "upstream returned 500 before the first body byte".to_string(), upstream_status: Some(500), retryable: true, failover_recommended: true, }, }, })); } .boxed(); let mut retry_scope = AiAttemptRetryScope::Provider; let response = execute_stream_from_frame_stream_with_retry_scope( &state, plan, &format!("trace-{request_id}"), &test_decision(), "claude_chat_stream", Some("claude_chat_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "claude:messages", "client_api_format": "claude:messages" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), true, frame_stream, false, None, Some(&mut retry_scope), None, None, ) .await .expect("prefetch HTTP status execution should resolve"); match response { Some(response) => AiAttemptExecutionOutcome::Responded(response), None => AiAttemptExecutionOutcome::Retry { scope: retry_scope, fallback_response: None, }, } } async fn execute_generic_sse_precommit( chunks: Vec<&str>, routing_policy: Value, provider_config: Option, stall: bool, ) -> Option> { execute_generic_stream_precommit( chunks, routing_policy, provider_config, stall, "text/event-stream", ) .await } async fn execute_generic_stream_precommit( chunks: Vec<&str>, routing_policy: Value, provider_config: Option, stall: bool, content_type: &str, ) -> Option> { let request_id = format!("generic-precommit-{}", uuid::Uuid::new_v4()); let mut plan = native_anthropic_stream_plan(&request_id); plan.provider_api_format = "openai:responses".to_string(); plan.client_api_format = "openai:responses".to_string(); plan.timeouts = Some(ExecutionTimeouts { first_byte_ms: Some(20), ..Default::default() }); let provider_catalog = provider_catalog_for_plan(&plan, provider_config); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( Arc::new(provider_catalog), DEVELOPMENT_ENCRYPTION_KEY, ); let state = AppState::new() .unwrap() .with_data_state_for_tests(data_state); let chunks = chunks.into_iter().map(str::to_string).collect::>(); let content_type = content_type.to_string(); let frames = stream! { yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Headers, payload: StreamFramePayload::Headers { status_code: 200, headers: BTreeMap::from([("content-type".to_string(), content_type)]), response_observation: None, }, })); for chunk in chunks { yield Ok(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { text: Some(chunk), chunk_b64: None }, })); } if stall { std::future::pending::<()>().await; } yield Ok(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Eof, payload: StreamFramePayload::Eof { summary: None }, })); } .boxed(); let mut scope = AiAttemptRetryScope::Provider; execute_stream_from_frame_stream_with_retry_scope( &state, plan, "trace-generic-precommit", &test_decision(), "openai_responses_stream", Some("openai_responses_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "openai:responses", "client_api_format": "openai:responses", "routing_execution_policy": routing_policy, })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), true, frames, false, None, Some(&mut scope), None, None, ) .await .unwrap() } #[tokio::test] async fn generic_stream_success_regex_matches_fragmented_plain_body() { assert!(execute_generic_stream_precommit( vec!["upstream CAPACITY ", "exhausted"], json!({"failover_rules": {"success_failover_patterns": [{"pattern": "(?i)capacity.*exhausted"}]}}), None, false, "text/plain", ).await.is_none()); } #[tokio::test] async fn generic_stream_200_json_error_obeys_global_stop_rules() { for stop in [false, true] { let response = execute_generic_stream_precommit( vec![r#"{"error":{"type":"server_error","message":"do not retry"}}"#], if stop { json!({"failover_rules": {"error_stop_patterns": [{"pattern": "do not retry"}]}}) } else { json!({}) }, None, false, "application/json", ).await; assert_eq!(response.is_some(), stop); if let Some(response) = response { assert!(response.status().is_server_error()); } } } #[tokio::test] async fn generic_sse_200_setup_then_error_retries_before_client_output() { let response = execute_generic_sse_precommit(vec![ "event: response.created\ndata: {\"type\":\"response.created\"}\n\n", "event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"error\":{\"type\":\"server_error\",\"message\":\"capacity exhausted\"}}}\n\n", ], json!({}), None, false).await; assert!(response.is_none()); } #[tokio::test] async fn generic_sse_global_stop_rule_overrides_retryable_embedded_error() { let response = execute_generic_sse_precommit(vec![ "event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"error\":{\"type\":\"server_error\",\"message\":\"capacity exhausted\"}}}\n\n", ], json!({ "failover_rules": { "error_stop_patterns": [{ "pattern": "capacity" }] } }), None, false).await; let response = response.expect("global stop must return a terminal response"); assert!(response.status().is_server_error()); to_bytes(response.into_body(), usize::MAX).await.unwrap(); } #[tokio::test] async fn generic_sse_success_regex_applies_to_global_and_provider_rules() { let rule = json!({ "success_failover_patterns": [{ "pattern": "(?i)CAPACITY" }] }); for (global, provider) in [ (json!({ "failover_rules": rule.clone() }), None), (json!({}), Some(json!({ "failover_rules": rule }))), ] { let response = execute_generic_sse_precommit(vec![ "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"capacity exhausted\"}\n\n", ], global, provider, false).await; assert!(response.is_none()); } } #[tokio::test] async fn generic_sse_late_error_does_not_replay_committed_content() { let response = execute_generic_sse_precommit(vec![ "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n", "event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"error\":{\"type\":\"server_error\",\"message\":\"late failure\"}}}\n\n", ], json!({}), None, false).await.expect("committed stream must not retry"); assert_eq!(response.status(), axum::http::StatusCode::OK); let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); assert!(String::from_utf8_lossy(&body).contains("hello")); } #[tokio::test] async fn generic_sse_setup_timeout_always_retries() { let response = execute_generic_sse_precommit( vec!["event: response.created\ndata: {\"type\":\"response.created\"}\n\n"], json!({}), None, true, ) .await; assert!(response.is_none()); } fn native_anthropic_stream_plan(request_id: &str) -> ExecutionPlan { ExecutionPlan { request_id: request_id.to_string(), candidate_id: Some(format!("candidate-{request_id}")), provider_name: Some("custom".to_string()), provider_id: format!("provider-{request_id}"), endpoint_id: format!("endpoint-{request_id}"), key_id: format!("key-{request_id}"), method: "POST".to_string(), url: "https://api.anthropic.com/v1/messages".to_string(), headers: BTreeMap::from([ ("content-type".to_string(), "application/json".to_string()), ("accept".to_string(), "text/event-stream".to_string()), ]), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "claude-sonnet-4-6", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 32, "stream": true })), stream: true, client_api_format: "claude:messages".to_string(), provider_api_format: "claude:messages".to_string(), model_name: Some("claude-sonnet-4-6".to_string()), proxy: None, transport_profile: None, timeouts: None, } } fn antigravity_gemini_stream_plan(request_id: &str) -> ExecutionPlan { ExecutionPlan { request_id: request_id.to_string(), candidate_id: Some(format!("candidate-{request_id}")), provider_name: Some("antigravity".to_string()), provider_id: format!("provider-{request_id}"), endpoint_id: format!("endpoint-{request_id}"), key_id: format!("key-{request_id}"), method: "POST".to_string(), url: "https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent".to_string(), headers: BTreeMap::from([ ("content-type".to_string(), "application/json".to_string()), ("accept".to_string(), "text/event-stream".to_string()), ]), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gemini-3.7-flash-tiered", "contents": [{"role": "user", "parts": [{"text": "validate"}]}] })), stream: true, client_api_format: "openai:responses".to_string(), provider_api_format: "gemini:generate_content".to_string(), model_name: Some("gemini-3.7-flash-tiered".to_string()), proxy: None, transport_profile: None, timeouts: None, } } struct StreamDropFlag(Arc); impl Drop for StreamDropFlag { fn drop(&mut self) { self.0.store(true, Ordering::SeqCst); } } fn direct_anthropic_test_finalizer(request_id: &str) -> DirectPassthroughFinalizer { let state = AppState::new().expect("app state should build"); let plan = native_anthropic_stream_plan(request_id); let lifecycle_seed = aether_usage_runtime::build_lifecycle_usage_seed(&plan, None); DirectPassthroughFinalizer::new(DirectPassthroughFinalizerCore { state, trace_id: format!("trace-{request_id}"), report_kind: None, report_context: None, lifecycle_seed, direct_stream_finalize_kind: None, stream_started_at: Instant::now(), stage_trace: RequestStageTrace::from_env(), request_diagnostics: None, request_id_for_log: request_id.to_string(), candidate_id: plan.candidate_id.clone(), request_candidate_status_snapshot: None, deferred_request_candidate_status_record: None, candidate_started_unix_secs: crate::clock::current_unix_ms(), status_code: 200, headers: BTreeMap::from([( "content-type".to_string(), "text/event-stream".to_string(), )]), stream_usage_report_context: None, stream_usage_observer: None, stream_usage_observer_buffered: Vec::new(), provider_error_inspection: ProviderStreamErrorInspection::default(), max_stream_body_buffer_bytes: super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, provider_buffered_body: Vec::new(), buffered_body: Vec::new(), provider_body_truncated: false, client_body_truncated: false, client_stream_completion_tracker: ClientVisibleStreamCompletionTracker::default(), requires_anthropic_message_stop: true, client_visible_stream_completed: false, usage_stream_telemetry: None, telemetry: None, provider_stream_bytes: 0, client_stream_bytes: 0, last_client_chunk_elapsed_ms: 0, pending_recorded: false, stream_started_recorded: false, terminal_failure: None, _provider_pool_in_flight_guard: None, _upstream_target_permit: None, plan, }) } fn discard_direct_test_finalizer(state: &mut DirectPassthroughInlineBodyState) { if let Some(mut finalizer) = state.finalizer.take() { finalizer.core.take(); } } fn direct_anthropic_inline_state( request_id: &str, items: Vec>, ) -> DirectPassthroughInlineBodyState { DirectPassthroughInlineBodyState { finalizer: Some(direct_anthropic_test_finalizer(request_id)), upstream: Some(futures_util::stream::iter(items).boxed()), upstream_control_filter: Some(super::SseControlBlockFilter::default()), upstream_started_at: Instant::now(), stream_first_byte_timeout: None, observed_first_body_poll: false, observed_first_client_yield: false, upstream_done: false, control_filter_flushed: false, terminal_error_sent: false, finalized: false, } } async fn execute_native_anthropic_prefetch_stream( request_id: &str, chunks: Vec, ) -> AiAttemptExecutionOutcome> { execute_native_anthropic_prefetch_stream_with_terminal_error(request_id, chunks, None).await } async fn execute_native_anthropic_prefetch_stream_with_terminal_error( request_id: &str, chunks: Vec, terminal_error: Option, ) -> AiAttemptExecutionOutcome> { let plan = native_anthropic_stream_plan(request_id); let provider_catalog = provider_catalog_for_plan(&plan, None); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( Arc::new(provider_catalog), DEVELOPMENT_ENCRYPTION_KEY, ); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests(data_state); 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(), )]), response_observation: None, }, })); for chunk in chunks { yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(chunk), }, })); } if let Some(error) = terminal_error { yield Err::(std::io::Error::other(error)); } else { yield Ok::(ndjson_frame(StreamFrame::eof())); } } .boxed(); let mut retry_scope = AiAttemptRetryScope::Candidate; let mut fallback_response = None; let response = execute_stream_from_frame_stream_with_retry_scope( &state, plan, &format!("trace-{request_id}"), &test_decision(), "claude_chat_stream", Some("claude_chat_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "claude:messages", "client_api_format": "claude:messages" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), true, frame_stream, false, None, Some(&mut retry_scope), Some(&mut fallback_response), None, ) .await .expect("native Anthropic stream execution should succeed"); match response { Some(response) => AiAttemptExecutionOutcome::Responded(response), None => AiAttemptExecutionOutcome::Retry { scope: retry_scope, fallback_response, }, } } 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") } #[tokio::test] async fn agent_identity_stream_error_prefetch_is_bounded_and_replayed() { let upstream_body = format!( "{}{}", "x".repeat(crate::execution_runtime::MAX_ERROR_BODY_BYTES), "body-after-inspection-limit" ); let expected_body = upstream_body.clone().into_bytes(); let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("address should resolve"); let server = tokio::spawn(async move { let app = Router::new().route( "/responses", any(move || { let body = upstream_body.clone(); async move { (StatusCode::UNAUTHORIZED, body) } }), ); axum::serve(listener, app) .await .expect("server should start"); }); let plan = direct_stream_test_plan("bounded-agent-error", format!("http://{addr}/responses")); let mut execution = crate::execution_runtime::DirectSyncExecutionRuntime::new() .execute_stream(&plan) .await .expect("stream headers should execute"); let inspected = prefetch_direct_stream_error_body(&mut execution) .await .expect("error body should be inspected"); assert_eq!( inspected.len(), crate::execution_runtime::MAX_ERROR_BODY_BYTES ); let replayed = collect_direct_execution_body(execution) .await .expect("prefetched response should replay"); assert_eq!(replayed, expected_body); server.abort(); } #[tokio::test] async fn non_agent_stream_401_is_not_prefetched_and_body_passes_through() { let upstream_body = br#"{"error":{"code":"ordinary_unauthorized","message":"sign in"}}"#; let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("address should resolve"); let server = tokio::spawn(async move { let app = Router::new().route( "/responses", any(|| async { ( StatusCode::UNAUTHORIZED, [(header::CONTENT_TYPE, "application/json")], upstream_body.as_slice(), ) }), ); axum::serve(listener, app) .await .expect("server should start"); }); let mut plan = direct_stream_test_plan("non-agent-401", format!("http://{addr}/responses")); plan.provider_name = Some("openai".to_string()); let repository = provider_catalog_for_stream_auth_plan(&plan, "openai", "api_key", None); let state = AppState::new() .expect("state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_provider_transport_reader_for_tests( Arc::new(repository), DEVELOPMENT_ENCRYPTION_KEY, ), ); let execution = execute_in_process_stream_with_oauth_retry( &state, &mut plan, "trace-non-agent-401", None, ) .await .expect("stream request should execute"); assert!(execution.prefetched_body.is_empty()); let replayed = collect_direct_execution_body(execution) .await .expect("response body should pass through"); assert_eq!(replayed, upstream_body); server.abort(); } #[tokio::test] async fn agent_identity_stream_non_task_401_replays_original_body_without_refresh() { let upstream_body = br#"{"error":{"code":"account_disabled","message":"account unavailable"}}"#; let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("address should resolve"); let task_registration_hits = Arc::new(AtomicUsize::new(0)); let task_registration_hits_for_server = Arc::clone(&task_registration_hits); let server = tokio::spawn(async move { let app = Router::new() .route( "/responses", any(|| async { ( StatusCode::UNAUTHORIZED, [(header::CONTENT_TYPE, "application/json")], upstream_body.as_slice(), ) }), ) .route( "/api/accounts/v1/agent/runtime-test/task/register", any(move || { let hits = Arc::clone(&task_registration_hits_for_server); async move { hits.fetch_add(1, Ordering::SeqCst); Json(json!({"task_id": "unexpected-task"})) } }), ); axum::serve(listener, app) .await .expect("server should start"); }); let mut plan = direct_stream_test_plan("agent-non-task-401", format!("http://{addr}/responses")); let repository = Arc::new(provider_catalog_for_stream_auth_plan( &plan, "codex", "oauth", Some(agent_identity_test_auth_config("task-old")), )); let oauth_refresh = aether_provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ Arc::new( aether_provider_transport::CodexAgentIdentityRefreshAdapter::default() .with_auth_api_base_url_for_tests(format!("http://{addr}/api/accounts")), ), ]); let state = AppState::new() .expect("state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( repository, ) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); let execution = execute_in_process_stream_with_oauth_retry( &state, &mut plan, "trace-agent-non-task-401", None, ) .await .expect("stream request should execute"); assert!(!execution.prefetched_body.is_empty()); assert_eq!(task_registration_hits.load(Ordering::SeqCst), 0); let replayed = collect_direct_execution_body(execution) .await .expect("response body should replay"); assert_eq!(replayed, upstream_body); server.abort(); } #[tokio::test] async fn agent_identity_stream_invalid_task_refreshes_and_retries_once() { let upstream_hits = Arc::new(AtomicUsize::new(0)); let upstream_hits_for_server = Arc::clone(&upstream_hits); let task_registration_hits = Arc::new(AtomicUsize::new(0)); let task_registration_hits_for_server = Arc::clone(&task_registration_hits); let observed_authorization = Arc::new(Mutex::new(Vec::::new())); let observed_authorization_for_server = Arc::clone(&observed_authorization); let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("address should resolve"); let server = tokio::spawn(async move { let app = Router::new() .route( "/responses", any(move |request: Request| { let hits = Arc::clone(&upstream_hits_for_server); let authorizations = Arc::clone(&observed_authorization_for_server); async move { let authorization = request .headers() .get(header::AUTHORIZATION) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(); authorizations .lock() .expect("authorization mutex should lock") .push(authorization); if hits.fetch_add(1, Ordering::SeqCst) == 0 { ( StatusCode::UNAUTHORIZED, Json(json!({ "error": { "code": "invalid_task_id", "message": "registered task expired" } })), ) .into_response() } else { (StatusCode::OK, Json(json!({"ok": true}))).into_response() } } }), ) .route( "/api/accounts/v1/agent/runtime-test/task/register", any(move || { let hits = Arc::clone(&task_registration_hits_for_server); async move { hits.fetch_add(1, Ordering::SeqCst); Json(json!({"task_id": "task-new"})) } }), ); axum::serve(listener, app) .await .expect("server should start"); }); let mut plan = direct_stream_test_plan("agent-invalid-task", format!("http://{addr}/responses")); let repository = Arc::new(provider_catalog_for_stream_auth_plan( &plan, "codex", "oauth", Some(agent_identity_test_auth_config("task-old")), )); let oauth_refresh = aether_provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ Arc::new( aether_provider_transport::CodexAgentIdentityRefreshAdapter::default() .with_auth_api_base_url_for_tests(format!("http://{addr}/api/accounts")), ), ]); let state = AppState::new() .expect("state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( repository, ) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); let transport = state .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id) .await .expect("transport should load") .expect("transport should exist"); let initial_authorization = match state .resolve_local_oauth_request_auth(&transport) .await .expect("initial Agent Identity auth should resolve") .expect("initial Agent Identity auth should exist") { aether_provider_transport::LocalResolvedOAuthRequestAuth::Header { name, value } => { assert_eq!(name, "authorization"); value } aether_provider_transport::LocalResolvedOAuthRequestAuth::Kiro(_) => { panic!("Agent Identity should resolve to header auth") } }; assert!( aether_provider_transport::codex_agent_identity_authorization_matches_transport( &transport, &initial_authorization, ) ); plan.headers .insert("authorization".to_string(), initial_authorization.clone()); let execution = execute_in_process_stream_with_oauth_retry( &state, &mut plan, "trace-agent-invalid-task", None, ) .await .expect("stream request should recover"); assert_eq!(execution.status_code, 200); assert!(execution.prefetched_body.is_empty()); assert_eq!(upstream_hits.load(Ordering::SeqCst), 2); assert_eq!(task_registration_hits.load(Ordering::SeqCst), 1); let authorizations = observed_authorization .lock() .expect("authorization mutex should lock"); assert_eq!(authorizations.len(), 2); assert_eq!(authorizations[0], initial_authorization); assert!(authorizations[1].starts_with("AgentAssertion ")); assert_ne!(authorizations[1], authorizations[0]); drop(authorizations); let replayed = collect_direct_execution_body(execution) .await .expect("retried response body should read"); assert_eq!( serde_json::from_slice::(&replayed).expect("response should be JSON"), json!({"ok": true}) ); server.abort(); } #[tokio::test] async fn native_anthropic_embedded_auth_error_refreshes_oauth_and_retries_once() { let upstream_hits = Arc::new(AtomicUsize::new(0)); let upstream_hits_for_server = Arc::clone(&upstream_hits); let refresh_hits = Arc::new(AtomicUsize::new(0)); let refresh_hits_for_server = Arc::clone(&refresh_hits); let observed_authorization = Arc::new(Mutex::new(Vec::::new())); let observed_authorization_for_server = Arc::clone(&observed_authorization); let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("address should resolve"); let server = tokio::spawn(async move { let app = Router::new() .route( "/v1/messages", any(move |request: Request| { let hits = Arc::clone(&upstream_hits_for_server); let authorizations = Arc::clone(&observed_authorization_for_server); async move { authorizations .lock() .expect("authorization mutex should lock") .push( request .headers() .get(header::AUTHORIZATION) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), ); let body = if hits.fetch_add(1, Ordering::SeqCst) == 0 { concat!( "event: error\n", "data: {\"type\":\"error\",\"error\":{\"type\":\"authentication_error\",\"message\":\"expired token\"}}\n\n", ) } else { concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", "event: message_stop\n", "data: {\"type\":\"message_stop\"}\n\n", ) }; ( StatusCode::OK, [(header::CONTENT_TYPE, "text/event-stream")], body, ) .into_response() } }), ) .route( "/oauth/token", any(move || { let hits = Arc::clone(&refresh_hits_for_server); async move { hits.fetch_add(1, Ordering::SeqCst); Json(json!({ "access_token": "fresh-access-token", "refresh_token": "fresh-refresh-token", "expires_in": 3600, "token_type": "Bearer" })) } }), ); axum::serve(listener, app) .await .expect("server should start"); }); let mut plan = native_anthropic_stream_plan("anthropic-embedded-oauth-refresh"); plan.url = format!("http://{addr}/v1/messages"); plan.provider_name = Some("claude_code".to_string()); plan.headers.insert( "authorization".to_string(), "Bearer stale-access-token".to_string(), ); let repository = Arc::new(provider_catalog_for_stream_auth_plan( &plan, "claude_code", "oauth", Some(generic_oauth_test_auth_config("claude_code")), )); let oauth_refresh = aether_provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ Arc::new( aether_provider_transport::GenericOAuthRefreshAdapter::default() .with_token_url_for_tests( "claude_code", format!("http://{addr}/oauth/token"), ), ), ]); let state = AppState::new() .expect("state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( repository, ) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); let execution = execute_in_process_stream_with_oauth_retry( &state, &mut plan, "trace-anthropic-embedded-oauth-refresh", None, ) .await .expect("embedded authentication error should recover"); let replayed = collect_direct_execution_body(execution) .await .expect("retried response should read"); assert_eq!(upstream_hits.load(Ordering::SeqCst), 2); assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); assert!(String::from_utf8_lossy(&replayed).contains("event: message_start")); assert_eq!( observed_authorization .lock() .expect("authorization mutex should lock") .as_slice(), [ "Bearer stale-access-token".to_string(), "Bearer fresh-access-token".to_string(), ] ); server.abort(); } #[tokio::test] async fn native_anthropic_http_permission_error_does_not_refresh_oauth() { let upstream_hits = Arc::new(AtomicUsize::new(0)); let upstream_hits_for_server = Arc::clone(&upstream_hits); let refresh_hits = Arc::new(AtomicUsize::new(0)); let refresh_hits_for_server = Arc::clone(&refresh_hits); let permission_body = concat!( "{\"type\":\"error\",\"error\":{", "\"type\":\"permission_error\",", "\"message\":\"this token is not authorized for the workspace\"}}", ); let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("address should resolve"); let server = tokio::spawn(async move { let app = Router::new() .route( "/v1/messages", any(move || { let hits = Arc::clone(&upstream_hits_for_server); async move { hits.fetch_add(1, Ordering::SeqCst); ( StatusCode::FORBIDDEN, [(header::CONTENT_TYPE, "application/json")], permission_body, ) } }), ) .route( "/oauth/token", any(move || { let hits = Arc::clone(&refresh_hits_for_server); async move { hits.fetch_add(1, Ordering::SeqCst); Json(json!({ "access_token": "unexpected-access-token", "refresh_token": "unexpected-refresh-token", "expires_in": 3600, "token_type": "Bearer" })) } }), ); axum::serve(listener, app) .await .expect("server should start"); }); let mut plan = native_anthropic_stream_plan("anthropic-http-oauth-permission"); plan.url = format!("http://{addr}/v1/messages"); plan.provider_name = Some("claude_code".to_string()); plan.headers.insert( "authorization".to_string(), "Bearer stale-access-token".to_string(), ); let repository = Arc::new(provider_catalog_for_stream_auth_plan( &plan, "claude_code", "oauth", Some(generic_oauth_test_auth_config("claude_code")), )); let oauth_refresh = aether_provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ Arc::new( aether_provider_transport::GenericOAuthRefreshAdapter::default() .with_token_url_for_tests( "claude_code", format!("http://{addr}/oauth/token"), ), ), ]); let state = AppState::new() .expect("state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( repository, ) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); let execution = execute_in_process_stream_with_oauth_retry( &state, &mut plan, "trace-anthropic-http-oauth-permission", None, ) .await .expect("permission response should remain available"); assert_eq!(execution.status_code, StatusCode::FORBIDDEN.as_u16()); let replayed = collect_direct_execution_body(execution) .await .expect("permission response should replay"); assert_eq!(upstream_hits.load(Ordering::SeqCst), 1); assert_eq!(refresh_hits.load(Ordering::SeqCst), 0); assert_eq!(replayed, permission_body.as_bytes()); server.abort(); } #[test] fn native_anthropic_oauth_prefetch_respects_short_first_byte_timeout() { let now = Instant::now(); let precommit_started_at = now .checked_sub(Duration::from_millis(10)) .expect("precommit start should be representable"); let upstream_started_at = now .checked_sub(Duration::from_millis(90)) .expect("upstream start should be representable"); let first_byte_wait = select_direct_anthropic_prefetch_wait( precommit_started_at, Duration::from_millis(750), upstream_started_at, Some(Duration::from_millis(100)), false, now, ); assert_eq!(first_byte_wait.remaining, Duration::from_millis(10)); assert!(!first_byte_wait.commit_on_timeout); let precommit_wait = select_direct_anthropic_prefetch_wait( now.checked_sub(Duration::from_millis(750)) .expect("precommit start should be representable"), Duration::from_millis(750), upstream_started_at, Some(Duration::from_secs(5)), false, now, ); assert!(precommit_wait.remaining.is_zero()); assert!(precommit_wait.commit_on_timeout); } #[tokio::test] async fn native_anthropic_oauth_pending_events_do_not_start_a_second_precommit_wait() { let request_id = "anthropic-oauth-single-precommit"; let plan = native_anthropic_stream_plan(request_id); let provider_catalog = provider_catalog_for_plan(&plan, None); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( Arc::new(provider_catalog), DEVELOPMENT_ENCRYPTION_KEY, ); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests(data_state); 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(), )]), response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(": ping\n\n".to_string()), }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some( "event: future_event\ndata: {\"type\":\"future_event\",\"value\":1}\n\n" .to_string(), ), }, })); tokio::time::sleep(Duration::from_secs(2)).await; yield Ok::(ndjson_frame(StreamFrame::eof())); } .boxed(); let response = tokio::time::timeout( Duration::from_millis(300), execute_stream_from_frame_stream_with_retry_scope( &state, plan, "trace-anthropic-oauth-single-precommit", &test_decision(), "claude_chat_stream", Some("claude_chat_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "claude:messages", "client_api_format": "claude:messages" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), true, frame_stream, true, None, None, None, None, ), ) .await .expect("the committed OAuth prefetch must not be followed by another 750 ms wait") .expect("frame stream execution should resolve"); assert!(response.is_some()); } #[tokio::test] async fn native_anthropic_embedded_auth_error_does_not_refresh_api_key() { let upstream_hits = Arc::new(AtomicUsize::new(0)); let upstream_hits_for_server = Arc::clone(&upstream_hits); let listener = crate::test_support::bind_loopback_listener() .await .expect("listener should bind"); let addr = listener.local_addr().expect("address should resolve"); let upstream_body = concat!( "event: error\n", "data: {\"type\":\"error\",\"error\":{\"type\":\"authentication_error\",\"message\":\"invalid key\"}}\n\n", ); let server = tokio::spawn(async move { let app = Router::new().route( "/v1/messages", any(move || { let hits = Arc::clone(&upstream_hits_for_server); async move { hits.fetch_add(1, Ordering::SeqCst); ( StatusCode::OK, [(header::CONTENT_TYPE, "text/event-stream")], upstream_body, ) } }), ); axum::serve(listener, app) .await .expect("server should start"); }); let mut plan = native_anthropic_stream_plan("anthropic-embedded-api-key"); plan.url = format!("http://{addr}/v1/messages"); plan.provider_name = Some("claude_code".to_string()); plan.headers .insert("x-api-key".to_string(), "invalid-api-key".to_string()); let repository = Arc::new(provider_catalog_for_stream_auth_plan( &plan, "claude_code", "api_key", None, )); let state = AppState::new() .expect("state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( repository, ) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ); let execution = execute_in_process_stream_with_oauth_retry( &state, &mut plan, "trace-anthropic-embedded-api-key", None, ) .await .expect("API-key error response should remain available"); let replayed = collect_direct_execution_body(execution) .await .expect("prefetched API-key error should replay"); assert_eq!(upstream_hits.load(Ordering::SeqCst), 1); assert_eq!(replayed, upstream_body.as_bytes()); server.abort(); } struct BlockingStreamingRequestCandidateRepository { inner: InMemoryRequestCandidateRepository, block_streaming: AtomicBool, streaming_started: Notify, release_streaming: Notify, } impl Default for BlockingStreamingRequestCandidateRepository { fn default() -> Self { Self { inner: InMemoryRequestCandidateRepository::default(), block_streaming: AtomicBool::new(true), streaming_started: Notify::new(), release_streaming: Notify::new(), } } } #[async_trait] impl RequestCandidateReadRepository for BlockingStreamingRequestCandidateRepository { async fn list_by_request_id( &self, request_id: &str, ) -> Result, DataLayerError> { self.inner.list_by_request_id(request_id).await } async fn list_recent( &self, limit: usize, ) -> Result, DataLayerError> { self.inner.list_recent(limit).await } async fn list_by_provider_id( &self, provider_id: &str, limit: usize, ) -> Result, DataLayerError> { self.inner.list_by_provider_id(provider_id, limit).await } async fn list_finalized_by_endpoint_ids_since( &self, endpoint_ids: &[String], since_unix_secs: u64, limit: usize, ) -> Result, DataLayerError> { self.inner .list_finalized_by_endpoint_ids_since(endpoint_ids, since_unix_secs, limit) .await } async fn count_finalized_statuses_by_endpoint_ids_since( &self, endpoint_ids: &[String], since_unix_secs: u64, ) -> Result, DataLayerError> { self.inner .count_finalized_statuses_by_endpoint_ids_since(endpoint_ids, since_unix_secs) .await } async fn aggregate_finalized_timeline_by_endpoint_ids_since( &self, endpoint_ids: &[String], since_unix_secs: u64, until_unix_secs: u64, segments: u32, ) -> Result, DataLayerError> { self.inner .aggregate_finalized_timeline_by_endpoint_ids_since( endpoint_ids, since_unix_secs, until_unix_secs, segments, ) .await } } #[async_trait] impl RequestCandidateWriteRepository for BlockingStreamingRequestCandidateRepository { async fn upsert( &self, candidate: UpsertRequestCandidateRecord, ) -> Result { if candidate.status == RequestCandidateStatus::Streaming && self.block_streaming.swap(false, Ordering::AcqRel) { self.streaming_started.notify_one(); self.release_streaming.notified().await; } self.inner.upsert(candidate).await } async fn delete_created_before( &self, created_before_unix_secs: u64, limit: usize, ) -> Result { self.inner .delete_created_before(created_before_unix_secs, limit) .await } } fn stage_metric_count(stage: &str) -> u64 { crate::stage_metrics::gateway_stage_metric_samples() .into_iter() .find(|sample| { sample.name == "gateway_stage_latency_count" && sample .labels .iter() .any(|label| label.key == "stage" && label.value == stage) }) .map(|sample| sample.value) .unwrap_or_default() } #[derive(Clone)] struct BlockingUsageAccess { policy_started: Arc, release_policy: Arc, } #[async_trait] impl UsageRecordWriter for BlockingUsageAccess { async fn upsert_usage_record( &self, _record: UpsertUsageRecord, ) -> Result, DataLayerError> { Ok(None) } } #[async_trait] impl UsageSettlementWriter for BlockingUsageAccess { fn has_usage_settlement_writer(&self) -> bool { false } async fn settle_usage( &self, _input: UsageSettlementInput, ) -> Result, DataLayerError> { Ok(None) } } #[async_trait] impl UsageBillingEventEnricher for BlockingUsageAccess { async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> { Ok(()) } } #[async_trait] impl UsageRuntimeAccess for BlockingUsageAccess { fn has_usage_writer(&self) -> bool { true } fn has_usage_worker_queue(&self) -> bool { false } fn usage_worker_queue(&self) -> Option> { None } fn supports_first_byte_usage_fast_path(&self) -> bool { false } async fn body_capture_policy(&self) -> Result { self.policy_started.notify_one(); self.release_policy.notified().await; Ok(UsageBodyCapturePolicy::default()) } } #[tokio::test] async fn inline_first_chunk_does_not_wait_for_candidate_streaming_persistence() { let request_id = "req-inline-first-chunk-candidate-handoff"; let request_candidate_repository = Arc::new(BlockingStreamingRequestCandidateRepository::default()); let usage_repository = Arc::new(InMemoryUsageReadRepository::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 = codex_cyber_policy_plan(request_id); let lifecycle_seed = aether_usage_runtime::build_lifecycle_usage_seed(&plan, None); let request_candidate_status_snapshot = crate::request_candidate_runtime::snapshot_local_request_candidate_status(&plan, None); let stream_started_at = Instant::now(); let finalizer = DirectPassthroughFinalizer::new(DirectPassthroughFinalizerCore { state, trace_id: "trace-inline-first-chunk-candidate-handoff".to_string(), report_kind: None, report_context: None, lifecycle_seed, direct_stream_finalize_kind: None, stream_started_at, stage_trace: RequestStageTrace::from_env(), request_diagnostics: None, request_id_for_log: request_id.to_string(), candidate_id: plan.candidate_id.clone(), request_candidate_status_snapshot, deferred_request_candidate_status_record: None, candidate_started_unix_secs: crate::clock::current_unix_ms(), status_code: 200, headers: BTreeMap::new(), stream_usage_report_context: None, stream_usage_observer: None, stream_usage_observer_buffered: Vec::new(), provider_error_inspection: ProviderStreamErrorInspection::default(), max_stream_body_buffer_bytes: super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, provider_buffered_body: Vec::new(), buffered_body: Vec::new(), provider_body_truncated: false, client_body_truncated: false, client_stream_completion_tracker: ClientVisibleStreamCompletionTracker::default(), requires_anthropic_message_stop: false, client_visible_stream_completed: false, usage_stream_telemetry: Some(ExecutionTelemetry { ttfb_ms: Some(7), elapsed_ms: Some(7), upstream_bytes: Some(2), }), telemetry: None, provider_stream_bytes: 2, client_stream_bytes: 0, last_client_chunk_elapsed_ms: 0, pending_recorded: false, stream_started_recorded: false, terminal_failure: None, _provider_pool_in_flight_guard: None, _upstream_target_permit: None, plan, }); let mut body_state = DirectPassthroughInlineBodyState { finalizer: Some(finalizer), upstream: None, upstream_control_filter: None, upstream_started_at: stream_started_at, stream_first_byte_timeout: None, observed_first_body_poll: true, observed_first_client_yield: false, upstream_done: false, control_filter_flushed: false, terminal_error_sent: false, finalized: false, }; let first_yield_count = stage_metric_count("stream_first_client_yield"); body_state.prepare_client_chunk_yield(&Bytes::from_static(b"hi")); assert!(body_state.observed_first_client_yield); assert!( stage_metric_count("stream_first_client_yield") > first_yield_count, "the first-yield metric must be recorded before candidate persistence completes" ); tokio::time::timeout( Duration::from_secs(1), request_candidate_repository.streaming_started.notified(), ) .await .expect("candidate streaming persistence should be handed off"); assert!( request_candidate_repository .list_by_request_id(request_id) .await .expect("candidate read should succeed") .is_empty(), "the first chunk must return while candidate persistence is still blocked" ); let live_usage = tokio::time::timeout(Duration::from_secs(2), async { loop { if let Some(usage) = usage_repository .find_by_request_id(request_id) .await .expect("usage read should succeed") { if usage.status == "streaming" { break usage; } } tokio::task::yield_now().await; } }) .await .expect("ordered usage lifecycle should reach streaming independently"); assert_eq!(live_usage.billing_status, "pending"); assert_eq!(live_usage.first_byte_time_ms, Some(7)); request_candidate_repository.release_streaming.notify_one(); let candidate = tokio::time::timeout(Duration::from_secs(1), async { loop { if let Some(candidate) = request_candidate_repository .list_by_request_id(request_id) .await .expect("candidate read should succeed") .into_iter() .next() { break candidate; } tokio::task::yield_now().await; } }) .await .expect("candidate handoff should finish after release"); assert_eq!(candidate.status, RequestCandidateStatus::Streaming); if let Some(mut finalizer) = body_state.finalizer.take() { finalizer.core.take(); } } #[tokio::test] async fn frame_stream_records_deferred_pending_before_waiting_for_headers() { let request_id = "req-frame-stream-deferred-pending"; let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_usage_repository_for_tests(Arc::clone( &usage_repository, )), ) .with_usage_runtime_for_tests(UsageRuntimeConfig { enabled: true, ..UsageRuntimeConfig::default() }); let plan = codex_cyber_policy_plan(request_id); let release_headers = Arc::new(Notify::new()); let release_headers_for_stream = Arc::clone(&release_headers); let frame_stream = stream! { release_headers_for_stream.notified().await; 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(), )]), response_observation: None, }, })); } .boxed(); let state_for_execution = state.clone(); let execution = tokio::spawn(async move { execute_stream_from_frame_stream( &state_for_execution, plan, "trace-frame-stream-deferred-pending", &test_decision(), "openai_responses_stream", None, Some(json!({ "provider_api_format": "openai:responses", "client_api_format": "openai:responses", })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), false, frame_stream, None, ) .await }); let pending = tokio::time::timeout(Duration::from_secs(1), async { loop { if let Some(usage) = usage_repository .find_by_request_id(request_id) .await .expect("usage should read") { break usage; } tokio::task::yield_now().await; } }) .await .expect("deferred pending usage should be recorded before headers"); assert_eq!(pending.status, "pending"); assert_eq!(pending.billing_status, "pending"); assert!( !execution.is_finished(), "frame execution should still be waiting for upstream headers" ); execution.abort(); let _ = execution.await; } #[tokio::test] async fn sync_terminal_handoff_survives_cancellation_during_admission_backpressure() { let blocker_request_id = "req-sync-terminal-admission-blocker"; let target_request_id = "req-sync-terminal-handoff-cancelled"; let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_usage_repository_for_tests(Arc::clone( &usage_repository, )), ) .with_usage_runtime_for_tests(UsageRuntimeConfig { enabled: true, terminal_submission_max_in_flight: 1, ..UsageRuntimeConfig::default() }); let policy_started = Arc::new(Notify::new()); let release_policy = Arc::new(Notify::new()); let blocker = BlockingUsageAccess { policy_started: Arc::clone(&policy_started), release_policy: Arc::clone(&release_policy), }; state .usage_runtime .submit_terminal_event( &blocker, UsageEvent::new( UsageEventType::Completed, blocker_request_id, UsageEventData { provider_name: "test".to_string(), model: "test-model".to_string(), status_code: Some(200), ..UsageEventData::default() }, ), ) .await; tokio::time::timeout(Duration::from_secs(1), policy_started.notified()) .await .expect("first terminal submission should be blocked in body policy"); let plan = codex_cyber_policy_plan(target_request_id); let payload = build_stream_sync_payload( "trace-sync-terminal-handoff-cancelled", "openai_responses_stream".to_string(), Some(json!({ "request_id": target_request_id, "provider_api_format": "openai:responses", "client_api_format": "openai:responses" })), 500, BTreeMap::new(), Some(json!({"error": "synthetic terminal failure"})), None, None, ); let state_for_handoff = state.clone(); let child_started = Arc::new(Notify::new()); let release_child = Arc::new(Notify::new()); let child_started_for_handoff = Arc::clone(&child_started); let release_child_for_handoff = Arc::clone(&release_child); let handoff = tokio::spawn(async move { record_sync_terminal_usage_with_handoff_after_spawn( &state_for_handoff, &plan, payload.report_context.as_ref(), &payload, async move { child_started_for_handoff.notify_one(); release_child_for_handoff.notified().await; }, ) .await; }); tokio::time::timeout(Duration::from_secs(1), child_started.notified()) .await .expect("detached terminal child should start before cancellation"); handoff.abort(); if let Err(err) = handoff.await { assert!(err.is_cancelled(), "terminal handoff task should not panic"); } release_child.notify_one(); tokio::time::timeout(Duration::from_secs(1), async { loop { if state .usage_runtime .metrics_snapshot() .terminal_submission_pending >= 2 { break; } tokio::task::yield_now().await; } }) .await .expect("second terminal submission should reach the ordered backlog"); release_policy.notify_one(); tokio::time::timeout(Duration::from_secs(2), async { loop { if state .usage_runtime .metrics_snapshot() .terminal_submission_pending == 0 && state .usage_runtime .metrics_snapshot() .lifecycle_submission_pending == 0 { break; } tokio::task::yield_now().await; } }) .await .expect("detached terminal handoff should release admission"); let record = tokio::time::timeout(Duration::from_secs(2), async { loop { if let Some(record) = usage_repository .find_by_request_id(target_request_id) .await .expect("usage repository read should succeed") { break record; } tokio::task::yield_now().await; } }) .await .expect("detached terminal handoff should persist the target row"); assert_eq!(record.status, "failed"); assert_eq!(record.billing_status, "void"); } #[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")); } #[test] fn client_visible_terminal_tracker_reports_the_exact_record_boundary() { let message_stop = b"event: message_stop\r\ndata: {\"type\":\"message_stop\"}\r\n\r\n"; let trailing_error = b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"api_error\"}}\n\n"; let chunk = [message_stop.as_slice(), trailing_error.as_slice()].concat(); let mut tracker = ClientVisibleStreamCompletionTracker::default(); assert_eq!( tracker.observe_chunk_terminal_end(&chunk), Some(message_stop.len()) ); assert!(tracker.completed); } #[test] fn anthropic_terminal_tracker_ignores_non_message_stop_terminals() { let mut tracker = ClientVisibleStreamCompletionTracker::default(); assert!(!tracker.observe_anthropic_message_stop(b"data: [DONE]\n\n")); assert!(!tracker.observe_anthropic_message_stop( b"event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n" )); assert!(tracker.observe_anthropic_message_stop( b"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" )); } #[test] fn terminal_tracker_caps_multiline_record_and_resumes_after_boundary() { let line = b"data: short-payload\r\n"; let repeated = super::SSE_TERMINAL_DETECTOR_MAX_RECORD_BYTES / line.len() + 2; let oversized_record = line.repeat(repeated); let mut tracker = ClientVisibleStreamCompletionTracker::default(); assert!(!tracker.observe_anthropic_message_stop(&oversized_record)); assert!(tracker.dropping_oversized_record); assert!(tracker.line_buffer.is_empty()); assert!(tracker.data_payload.is_empty()); assert!(!tracker.observe_anthropic_message_stop(b"\r\n")); assert!(!tracker.dropping_oversized_record); assert!(tracker.observe_anthropic_message_stop( b"event: message_stop\r\ndata: {\"type\":\"message_stop\"}\r\n\r\n" )); } #[test] fn stream_capture_hard_caps_a_single_oversized_chunk() { let mut buffer = vec![1, 2]; let mut truncated = false; super::append_stream_capture_bytes(&mut buffer, &[3, 4, 5, 6], 4, &mut truncated); assert_eq!(buffer, vec![1, 2, 3, 4]); assert!(truncated); } #[test] fn stream_capture_encoding_defensively_caps_an_oversized_slice() { let (body, state) = super::build_stream_body_capture_with_limit(b"abcdef", false, 3); let decoded = base64::engine::general_purpose::STANDARD .decode(body.expect("bounded capture should be encoded")) .expect("capture should be valid base64"); assert_eq!(decoded, b"abc"); assert_eq!(state, Some(UsageBodyCaptureState::Truncated)); } #[test] fn execution_stream_data_chunk_decode_is_bounded_before_allocation() { assert_eq!( super::decode_stream_data_chunk_with_limit(Some("YWJj"), None, 3) .expect("three decoded bytes"), b"abc" ); assert!(super::decode_stream_data_chunk_with_limit(Some("YWJjZA=="), None, 3).is_err()); assert!(super::decode_stream_data_chunk_with_limit(None, Some("abcd"), 3).is_err()); } #[test] fn stream_capture_policy_hard_caps_full_and_basic_analysis_buffers() { let oversized_chunk = vec![b'x'; super::BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES + 1]; let full_limit = super::stream_body_buffer_limit_for_record_level(UsageRequestRecordLevel::Full); assert_eq!( full_limit, crate::execution_runtime::MAX_STREAM_BODY_CAPTURE_BYTES ); let mut full_buffer = Vec::new(); let mut full_truncated = false; super::append_stream_capture_bytes( &mut full_buffer, &oversized_chunk, full_limit, &mut full_truncated, ); assert_eq!(full_buffer, oversized_chunk); assert!(!full_truncated); let (full_body, full_state) = super::build_stream_body_capture(&full_buffer, full_truncated); assert!(full_body.is_some()); assert_eq!(full_state, Some(UsageBodyCaptureState::Inline)); drop(full_body); let basic_limit = super::stream_body_buffer_limit_for_record_level(UsageRequestRecordLevel::Basic); assert_eq!(basic_limit, super::BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES); let mut basic_buffer = Vec::new(); let mut basic_truncated = false; super::append_stream_capture_bytes( &mut basic_buffer, &oversized_chunk, basic_limit, &mut basic_truncated, ); assert_eq!(basic_buffer.len(), basic_limit); assert!(basic_truncated); let (basic_body, basic_state) = super::build_stream_body_capture(&basic_buffer, basic_truncated); assert!(basic_body.is_some()); assert_eq!(basic_state, Some(UsageBodyCaptureState::Truncated)); let mut event = UsageEvent::new( UsageEventType::Completed, "req-basic-stream-capture", UsageEventData { provider_name: "provider".to_string(), model: "model".to_string(), response_body: basic_body.map(Value::String), response_body_state: basic_state, client_response_body: Some(json!("captured client body")), client_response_body_state: Some(UsageBodyCaptureState::Truncated), ..UsageEventData::default() }, ); apply_usage_body_capture_policy_to_event( UsageBodyCapturePolicy { record_level: UsageRequestRecordLevel::Basic, }, &mut event, ); assert_eq!(event.data.response_body, None); assert_eq!( event.data.response_body_state, Some(UsageBodyCaptureState::Disabled) ); assert_eq!(event.data.client_response_body, None); assert_eq!( event.data.client_response_body_state, Some(UsageBodyCaptureState::Disabled) ); } #[test] fn provider_error_inspection_detects_response_failed_at_every_chunk_boundary() { let body = concat!( "event: response.created\n", "data: {\"type\":\"response.created\",\"response\":{\"status\":\"in_progress\"}}\n\n", "event: response.failed\n", "data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"type\":\"invalid_request\",\"message\":\"cyber policy rejected the request\",\"code\":\"cyber_policy_violation\",\"param\":\"input\"}}}\n\n", ) .as_bytes(); for split in 1..body.len() { let mut inspection = ProviderStreamErrorInspection::default(); let detected = inspection .observe(None, &body[..split]) .or_else(|| inspection.observe(None, &body[split..])) .unwrap_or_else(|| panic!("response.failed was missed at byte split {split}")); assert_eq!( detected.pointer("/error/code"), Some(&json!("cyber_policy_violation")), "string provider code changed at byte split {split}" ); assert_eq!( detected.pointer("/error/param"), Some(&json!("input")), "provider error fields changed at byte split {split}" ); } let mut inspection = ProviderStreamErrorInspection::default(); let mut detected = None; for byte in body.chunks(1) { if let Some(error_body) = inspection.observe(None, byte) { detected = Some(error_body); break; } } let detected = detected.expect("byte-wise response.failed stream should be detected"); assert_eq!( detected.pointer("/error/code"), Some(&json!("cyber_policy_violation")) ); assert_eq!(detected.pointer("/error/param"), Some(&json!("input"))); } #[test] fn provider_error_inspection_bounds_oversized_chunks_and_keeps_boundary_detection() { let error_event = concat!( "event: response.failed\n", "data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"cyber_policy_violation\"}}}\n\n", ) .as_bytes(); let mut prefix_chunk = vec![b'x'; PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES + 64]; prefix_chunk[..error_event.len()].copy_from_slice(error_event); let mut inspection = ProviderStreamErrorInspection::default(); let detected = inspection .observe(None, &prefix_chunk) .expect("error at the bounded chunk prefix should be detected"); assert_eq!( detected.pointer("/error/code"), Some(&json!("cyber_policy_violation")) ); let mut suffix_chunk = vec![b'x'; PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES + 64]; let suffix_start = suffix_chunk.len() - error_event.len(); suffix_chunk[suffix_start..].copy_from_slice(error_event); let mut inspection = ProviderStreamErrorInspection::default(); let detected = inspection .observe(None, &suffix_chunk) .expect("error at the bounded chunk suffix should be detected"); assert_eq!( detected.pointer("/error/code"), Some(&json!("cyber_policy_violation")) ); // The JSON payload is split across chunks. The previous rolling tail // must still be combined with the prefix of the oversized chunk. let split = b"event: response.failed\ndata: {".len(); let mut inspection = ProviderStreamErrorInspection::default(); assert!(inspection.observe(None, &error_event[..split]).is_none()); let mut continuation = vec![b'x'; PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES + 64]; let continuation_len = error_event.len() - split; continuation[..continuation_len].copy_from_slice(&error_event[split..]); let detected = inspection .observe(None, &continuation) .expect("error split across an oversized chunk boundary should be detected"); assert_eq!( detected.pointer("/error/code"), Some(&json!("cyber_policy_violation")) ); } #[tokio::test] async fn prefetched_codex_cyber_policy_violation_stops_failover_by_default() { let response = execute_prefetched_codex_cyber_policy_failure(false) .await .expect("default Codex cyber policy handling should return the provider error"); assert_eq!(response.status(), axum::http::StatusCode::BAD_REQUEST); } #[tokio::test] async fn prefetched_codex_cyber_policy_violation_retries_when_routing_strategy_is_enabled() { assert!( execute_prefetched_codex_cyber_policy_failure(true) .await .is_none(), "enabling cyber failover should retry the next candidate" ); } #[tokio::test] async fn prefetched_transport_failure_retries_by_default() { assert!(matches!( execute_prefetched_transport_failure(false).await, AiAttemptExecutionOutcome::Retry { scope: AiAttemptRetryScope::Candidate, fallback_response: None, } )); } #[tokio::test] async fn prefetched_transport_failure_can_stop_without_matching_http_status_rules() { let AiAttemptExecutionOutcome::Responded(response) = execute_prefetched_transport_failure(true).await else { panic!("transport stop policy should return a local response"); }; assert_eq!(response.status(), StatusCode::BAD_GATEWAY); } #[tokio::test] async fn prefetched_http_error_frame_honors_continue_status_codes() { assert!(matches!( execute_prefetched_http_status_failure(true).await, AiAttemptExecutionOutcome::Retry { .. } )); } #[tokio::test] async fn prefetched_http_error_frame_honors_stop_status_codes() { let AiAttemptExecutionOutcome::Responded(response) = execute_prefetched_http_status_failure(false).await else { panic!("HTTP stop policy should return the upstream error"); }; assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); } #[tokio::test] async fn malformed_antigravity_function_call_streams_thought_then_fails_in_band() { let request_id = "req-antigravity-malformed-function-call"; let plan = antigravity_gemini_stream_plan(request_id); let provider_catalog = provider_catalog_for_plan( &plan, Some(json!({ "failover_rules": { "continue_status_codes": [502] } })), ); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( Arc::new(provider_catalog), "development-key", ); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests(data_state); 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(), )]), response_observation: None, }, })); for chunk in [ r#"data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thought":true,"text":"Validating the document."}]} }],"modelVersion":"gemini-3.7-flash-tiered"}} "#, r#"data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"signature","text":""}]},"finishReason":"MALFORMED_FUNCTION_CALL","finishMessage":"Malformed function call: Function call is empty - no input to parse."}],"modelVersion":"gemini-3.7-flash-tiered"}} "#, ] { yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(chunk.to_string()), }, })); } yield Ok::(ndjson_frame(StreamFrame::eof())); } .boxed(); let mut retry_scope = AiAttemptRetryScope::Provider; let response = execute_stream_from_frame_stream_with_retry_scope( &state, plan, "trace-antigravity-malformed-function-call", &test_decision(), OPENAI_RESPONSES_STREAM_PLAN_KIND, Some("openai_responses_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "gemini:generate_content", "client_api_format": "openai:responses", "needs_conversion": true })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), true, frame_stream, false, None, Some(&mut retry_scope), None, None, ) .await .expect("malformed Antigravity stream should return a client stream") .expect("the first reasoning delta should commit the selected candidate"); assert_eq!(response.status(), StatusCode::OK); let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body = String::from_utf8(body.to_vec()).expect("response body should be utf8"); assert!( body.contains("event: response.reasoning_summary_text.delta\n"), "{body}" ); assert!( body.contains("\"delta\":\"Validating the document.\""), "{body}" ); assert!(body.contains("event: response.failed\n"), "{body}"); assert!( body.contains("\"code\":\"MALFORMED_FUNCTION_CALL\""), "{body}" ); assert!( body.contains( "\"message\":\"Malformed function call: Function call is empty - no input to parse.\"" ), "{body}" ); assert!(!body.contains("unsupported_finish_reason"), "{body}"); assert_eq!(retry_scope, AiAttemptRetryScope::Provider); } 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})), } } const LOCAL_TUNNEL_TEST_PSK: &str = "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc="; const LOCAL_TUNNEL_TEST_GENERATION: &str = "stream-test-generation-1"; fn authenticated_local_tunnel_test_state() -> AppState { let node = StoredProxyNode::new( "node-1".to_string(), "Node 1".to_string(), "127.0.0.1".to_string(), 0, false, "online".to_string(), 30, 1, 0, 0, 0, 0, true, true, 1, ) .expect("tunnel node should build") .with_runtime_fields( None, None, None, None, Some(json!({ "tunnel_security": { "mode": TUNNEL_SECURITY_NON_TLS_REQUIRED, "encryption_key": LOCAL_TUNNEL_TEST_PSK, } })), None, None, None, None, None, None, ) .with_tunnel_generation(LOCAL_TUNNEL_TEST_GENERATION.to_string()); let data = crate::data::GatewayDataState::with_proxy_node_repository_for_tests(Arc::new( InMemoryProxyNodeRepository::seed([node]), )) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY); AppState::new() .expect("app state should build") .with_data_state_for_tests(data) } async fn recv_tunnel_test_frame( proxy_rx: &mut aether_runtime::BoundedQueueReceiver, description: &str, ) -> Message { tokio::time::timeout(Duration::from_secs(5), proxy_rx.recv()) .await .unwrap_or_else(|_| panic!("timed out waiting for {description}")) .unwrap_or_else(|| panic!("proxy channel closed before {description}")) } 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 execution_stream_frame_codec_has_a_bounded_line_length() { assert_eq!( execution_stream_frame_codec().max_length(), crate::execution_runtime::MAX_EXECUTION_STREAM_FRAME_LINE_BYTES ); } #[test] fn post_stop_reader_yields_after_bounded_empty_chunks() { let polls = Arc::new(AtomicUsize::new(0)); let polls_for_stream = Arc::clone(&polls); let stream = futures_util::stream::poll_fn(move |_| { let poll = polls_for_stream.fetch_add(1, Ordering::SeqCst); if poll < POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL * 2 { std::task::Poll::Ready(Some(Ok::(Bytes::new()))) } else { std::task::Poll::Ready(Some(Ok::(Bytes::from_static(b"x")))) } }); let mut reader = Box::pin(PostStopLimitedStreamReader::new( stream, PostStopFrameReadBudget::new(), )); let waker = futures_util::task::noop_waker(); let mut context = std::task::Context::from_waker(&waker); let mut storage = [0u8; 1]; let mut read_buf = tokio::io::ReadBuf::new(&mut storage); let result = tokio::io::AsyncRead::poll_read(reader.as_mut(), &mut context, &mut read_buf); assert!(result.is_pending()); assert_eq!( polls.load(Ordering::SeqCst), POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL ); } #[tokio::test] async fn post_stop_activation_trims_prefetched_current_item_immediately() { const GIANT_TAIL_BYTES: usize = 4 * 1024 * 1024; let mut combined = ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some( "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n".to_string(), ), }, }) .to_vec(); combined.resize(combined.len() + GIANT_TAIL_BYTES, b'x'); let frame_stream = futures_util::stream::iter([Ok::(Bytes::from(combined))]); let reader = PostStopLimitedStreamReader::new(frame_stream, PostStopFrameReadBudget::new()); let mut lines = tokio_util::codec::FramedRead::new(reader, execution_stream_frame_codec()); super::read_next_frame(&mut lines) .await .expect("frame should decode") .expect("data frame should exist"); assert!(lines .get_ref() .current .as_ref() .is_some_and(|current| current.len() > ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES)); let already_buffered = lines.read_buffer().len(); lines.read_buffer_mut().reserve(GIANT_TAIL_BYTES); assert!(lines.read_buffer().capacity() >= GIANT_TAIL_BYTES); assert!(!activate_post_stop_frame_read_budget(&mut lines)); let retained = lines .get_ref() .current .as_ref() .map(Bytes::len) .unwrap_or_default(); assert_eq!( retained, ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES.saturating_sub(already_buffered) ); assert_eq!(lines.read_buffer().len(), already_buffered); assert!(lines.read_buffer().capacity() <= ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES); } #[test] fn post_stop_activation_releases_over_limit_framed_buffer() { let frame_stream = futures_util::stream::empty::>(); let reader = PostStopLimitedStreamReader::new(frame_stream, PostStopFrameReadBudget::new()); let mut lines = tokio_util::codec::FramedRead::new(reader, execution_stream_frame_codec()); lines .read_buffer_mut() .resize(ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES + 1, b'x'); assert!(activate_post_stop_frame_read_budget(&mut lines)); assert!(lines.read_buffer().is_empty()); assert_eq!(lines.read_buffer().capacity(), 0); } #[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()), provider_actual_service_tier: Some("priority".to_string()), unknown_event_count: 1, ..ExecutionStreamTerminalSummary::default() }), Some(ExecutionStreamTerminalSummary { standardized_usage: Some(observed_usage), response_id: Some("resp_123".to_string()), provider_actual_service_tier: Some("default".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_eq!( merged.provider_actual_service_tier.as_deref(), Some("default") ); 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 native_anthropic_event_stream_uses_bounded_precommit() { assert!(!should_skip_direct_finalize_prefetch( Some("claude_cli_sync_finalize"), Some("text/event-stream"), "claude:messages", "claude:messages", false, false, false, )); } #[test] fn native_anthropic_terminal_error_uses_anthropic_sse_shape() { let plan = native_anthropic_stream_plan("anthropic-terminal-error-shape"); let failure = build_stream_failure_report( "execution_runtime_stream_read_error", "upstream disconnected", 502, ); let event = encode_terminal_sse_error_event_for_plan(&plan, &failure) .expect("terminal event should encode"); let event = String::from_utf8(event.to_vec()).expect("event should be utf8"); assert!(event.starts_with("event: error\ndata: ")); assert!(event.contains("\"type\":\"error\"")); assert!(event.contains("\"type\":\"api_error\"")); assert!(event.contains("Upstream response stream failed")); assert!(!event.contains("upstream disconnected")); assert!(!event.contains("[DONE]")); } #[tokio::test] async fn native_anthropic_error_before_semantic_event_allows_failover() { let unknown = "event: future_event\ndata: {\"type\":\"future_event\",\"value\":1}\n\n"; let upstream_error = "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n"; let outcome = execute_native_anthropic_prefetch_stream( "req-anthropic-precommit-error", vec![unknown.to_string(), upstream_error.to_string()], ) .await; let AiAttemptExecutionOutcome::Retry { scope, fallback_response: Some(fallback_response), } = outcome else { panic!("precommit 529 should retry with the upstream response preserved") }; assert_eq!(scope, AiAttemptRetryScope::Provider); assert_eq!(fallback_response.status(), StatusCode::OK); let fallback_body = to_bytes(fallback_response.into_body(), usize::MAX) .await .expect("fallback response body should read"); assert_eq!( fallback_body.as_ref(), format!("{unknown}{upstream_error}").as_bytes() ); } #[tokio::test] async fn native_anthropic_eof_before_semantic_event_allows_failover() { let outcome = execute_native_anthropic_prefetch_stream( "req-anthropic-precommit-eof", vec![": ping\n\n".to_string()], ) .await; let AiAttemptExecutionOutcome::Retry { scope, fallback_response, } = outcome else { panic!("EOF before the first semantic event should retry another endpoint") }; assert_eq!(scope, AiAttemptRetryScope::Endpoint); assert!(fallback_response.is_none()); } #[tokio::test] async fn native_anthropic_auth_error_moves_to_the_next_credential() { let upstream_error = concat!( "event: error\n", "data: {\"type\":\"error\",\"error\":{\"type\":\"authentication_error\",\"message\":\"invalid credential\"}}\n\n", ); let outcome = execute_native_anthropic_prefetch_stream( "req-anthropic-precommit-auth-error", vec![upstream_error.to_string()], ) .await; let AiAttemptExecutionOutcome::Retry { scope, fallback_response: Some(fallback_response), } = outcome else { panic!("precommit authentication error should retry another credential") }; assert_eq!(scope, AiAttemptRetryScope::Credential); let fallback_body = to_bytes(fallback_response.into_body(), usize::MAX) .await .expect("fallback response body should read"); assert_eq!(fallback_body.as_ref(), upstream_error.as_bytes()); } #[tokio::test] async fn native_anthropic_semantic_event_commits_before_later_error() { let raw = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", "event: content_block_delta\n", "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n", "event: error\n", "data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"late\"}}\n\n", ); let outcome = execute_native_anthropic_prefetch_stream( "req-anthropic-postcommit-error", vec![raw.to_string()], ) .await; let AiAttemptExecutionOutcome::Responded(response) = outcome else { panic!("a semantic event should commit the selected candidate") }; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); assert_eq!(body.as_ref(), raw.as_bytes()); } #[tokio::test] async fn native_anthropic_frame_error_after_commit_emits_anthropic_terminal_event() { let message_start = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", "event: content_block_delta\n", "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n", ); let original_error = "upstream disconnected after message_start"; let outcome = execute_native_anthropic_prefetch_stream_with_terminal_error( "req-anthropic-postcommit-frame-error", vec![message_start.to_string()], Some(original_error.to_string()), ) .await; let AiAttemptExecutionOutcome::Responded(response) = outcome else { panic!("a semantic event should commit the selected candidate") }; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body = String::from_utf8(body.to_vec()).expect("body should be utf8"); assert!(body.starts_with(message_start)); assert!(body.contains("event: error\ndata: {\"type\":\"error\"")); assert!(body.contains("\"type\":\"api_error\"")); assert!(body.contains("Execution runtime stream protocol failed")); assert!(!body.contains(original_error)); assert!(!body.contains("[DONE]")); } #[tokio::test] async fn native_anthropic_eof_after_commit_emits_anthropic_terminal_event() { let message_start = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", "event: content_block_delta\n", "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n", ); let outcome = execute_native_anthropic_prefetch_stream( "req-anthropic-postcommit-eof", vec![message_start.to_string()], ) .await; let AiAttemptExecutionOutcome::Responded(response) = outcome else { panic!("a semantic event should commit the selected candidate") }; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body = String::from_utf8(body.to_vec()).expect("body should be utf8"); assert!(body.starts_with(message_start)); assert!(body.contains("event: error\ndata: {\"type\":\"error\"")); assert!(body.contains("ended before message_stop")); assert!(!body.contains("[DONE]")); } #[tokio::test] async fn native_anthropic_done_marker_does_not_replace_message_stop() { let message_start = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", "event: content_block_delta\n", "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n", ); let done = "data: [DONE]\n\n"; let outcome = execute_native_anthropic_prefetch_stream( "req-anthropic-done-without-message-stop", vec![message_start.to_string(), done.to_string()], ) .await; let AiAttemptExecutionOutcome::Responded(response) = outcome else { panic!("text output should commit the selected candidate") }; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); let body = String::from_utf8(body.to_vec()).expect("body should be utf8"); assert!(body.starts_with(message_start)); assert!(body.contains(done)); assert!(body.contains("event: error\ndata: {\"type\":\"error\"")); assert!(body.contains("ended before message_stop")); } #[tokio::test] async fn native_anthropic_hanging_tail_does_not_delay_body_eof_and_is_bounded() { let request_id = "req-anthropic-hanging-tail"; let plan = native_anthropic_stream_plan(request_id); let provider_catalog = provider_catalog_for_plan(&plan, None); let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_usage_repository_for_tests(Arc::clone( &usage_repository, )) .with_provider_catalog_reader(Arc::new(provider_catalog)) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ) .with_usage_runtime_for_tests(UsageRuntimeConfig { enabled: true, ..UsageRuntimeConfig::default() }); let message_start = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", ); let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; let stream_dropped = Arc::new(AtomicBool::new(false)); let drop_flag = StreamDropFlag(Arc::clone(&stream_dropped)); let frame_stream = stream! { let _drop_flag = drop_flag; 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(), )]), response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(message_start.to_string()), }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(message_stop.to_string()), }, })); std::future::pending::<()>().await; } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-anthropic-hanging-tail", &test_decision(), "claude_chat_stream", Some("claude_chat_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "claude:messages", "client_api_format": "claude:messages" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), false, frame_stream, None, ) .await .expect("stream execution should succeed") .expect("stream execution should return a response"); let body = tokio::time::timeout( Duration::from_secs(1), to_bytes(response.into_body(), usize::MAX), ) .await .expect("client body EOF must not wait for the hanging producer tail") .expect("client body should read"); assert_eq!( body.as_ref(), format!("{message_start}{message_stop}").as_bytes() ); tokio::time::timeout(Duration::from_secs(1), async { while !stream_dropped.load(Ordering::SeqCst) { tokio::task::yield_now().await; } }) .await .expect("producer tail should be dropped after the bounded drain window"); let stored_usage = tokio::time::timeout(Duration::from_secs(2), async { loop { let usage = usage_repository .find_by_request_id(request_id) .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("bounded drain timeout should still settle usage successfully"); assert_eq!(stored_usage.status_code, Some(200)); } #[tokio::test] async fn native_anthropic_unterminated_oversized_tail_respects_read_budget() { const TAIL_CHUNK_BYTES: usize = 4 * 1024 * 1024; let request_id = "req-anthropic-oversized-tail"; let plan = native_anthropic_stream_plan(request_id); let provider_catalog = provider_catalog_for_plan(&plan, None); let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let state = AppState::new() .expect("app state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_usage_repository_for_tests(Arc::clone( &usage_repository, )) .with_provider_catalog_reader(Arc::new(provider_catalog)) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ) .with_usage_runtime_for_tests(UsageRuntimeConfig { enabled: true, ..UsageRuntimeConfig::default() }); let message_start = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", ); let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; let stream_dropped = Arc::new(AtomicBool::new(false)); let drop_flag = StreamDropFlag(Arc::clone(&stream_dropped)); let tail_chunks_polled = Arc::new(AtomicUsize::new(0)); let tail_chunks_polled_for_stream = Arc::clone(&tail_chunks_polled); let tail_chunk = Bytes::from(vec![b'x'; TAIL_CHUNK_BYTES]); let tail_chunk_count = 32; let frame_stream = stream! { let _drop_flag = drop_flag; 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(), )]), response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(message_start.to_string()), }, })); yield Ok::(ndjson_frame(StreamFrame { frame_type: StreamFrameType::Data, payload: StreamFramePayload::Data { chunk_b64: None, text: Some(message_stop.to_string()), }, })); for _ in 0..tail_chunk_count { tail_chunks_polled_for_stream.fetch_add(1, Ordering::SeqCst); yield Ok::(tail_chunk.clone()); } std::future::pending::<()>().await; } .boxed(); let response = execute_stream_from_frame_stream( &state, plan, "trace-anthropic-oversized-tail", &test_decision(), "claude_chat_stream", Some("claude_chat_stream_success".to_string()), Some(json!({ "request_id": request_id, "candidate_id": format!("candidate-{request_id}"), "candidate_index": 0, "retry_index": 0, "provider_api_format": "claude:messages", "client_api_format": "claude:messages" })), crate::clock::current_unix_ms(), Instant::now(), RequestStageTrace::from_env(), false, frame_stream, None, ) .await .expect("stream execution should succeed") .expect("stream execution should return a response"); let body = tokio::time::timeout( Duration::from_secs(1), to_bytes(response.into_body(), usize::MAX), ) .await .expect("client body EOF must not wait for the oversized unterminated tail") .expect("client body should read"); assert_eq!( body.as_ref(), format!("{message_start}{message_stop}").as_bytes() ); tokio::time::timeout(Duration::from_secs(1), async { while !stream_dropped.load(Ordering::SeqCst) { tokio::task::yield_now().await; } }) .await .expect("oversized tail producer should be released by the read budget"); assert!( tail_chunks_polled.load(Ordering::SeqCst) <= 1, "post-stop drain must retain at most one atomic upstream stream item" ); let stored_usage = tokio::time::timeout(Duration::from_secs(2), async { loop { let usage = usage_repository .find_by_request_id(request_id) .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("bounded oversized tail should still settle usage successfully"); assert_eq!(stored_usage.status_code, Some(200)); } #[tokio::test] async fn native_anthropic_frame_error_after_message_stop_is_ignored() { let message_start = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", ); let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; let outcome = execute_native_anthropic_prefetch_stream_with_terminal_error( "req-anthropic-error-after-message-stop", vec![message_start.to_string(), message_stop.to_string()], Some("connection reset after message_stop".to_string()), ) .await; let AiAttemptExecutionOutcome::Responded(response) = outcome else { panic!("message_stop should keep the selected candidate committed") }; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); assert_eq!( body.as_ref(), format!("{message_start}{message_stop}").as_bytes() ); } #[tokio::test] async fn native_anthropic_same_chunk_stops_at_message_stop_record() { let message_start = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", ); let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; let trailing_error = concat!( "event: error\n", "data: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"after stop\"}}\n\n", ); let outcome = execute_native_anthropic_prefetch_stream( "req-anthropic-same-chunk-message-stop", vec![format!("{message_start}{message_stop}{trailing_error}")], ) .await; let AiAttemptExecutionOutcome::Responded(response) = outcome else { panic!("message_stop should complete the selected candidate") }; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); assert_eq!( body.as_ref(), format!("{message_start}{message_stop}").as_bytes() ); } #[tokio::test] async fn direct_anthropic_stops_at_message_stop_and_ignores_teardown_error() { let message_start = Bytes::from_static( b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{}}\n\n", ); let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; let trailing_error = concat!( "event: error\n", "data: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"after stop\"}}\n\n", ); let state = direct_anthropic_inline_state( "req-direct-anthropic-message-stop", vec![ Ok(message_start.clone()), Ok(Bytes::from(format!("{message_stop}{trailing_error}"))), Err("connection reset after message_stop".to_string()), ], ); let (first, state) = state .next_item() .await .expect("message_start should stream"); assert_eq!(first.expect("message_start should succeed"), message_start); let (second, state) = state.next_item().await.expect("message_stop should stream"); assert_eq!( second.expect("message_stop should succeed").as_ref(), message_stop.as_bytes() ); assert!(state.next_item().await.is_none()); } #[tokio::test] async fn direct_anthropic_clean_eof_after_message_start_emits_one_error() { let message_start = Bytes::from_static( b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{}}\n\n", ); let state = direct_anthropic_inline_state( "req-direct-anthropic-premature-eof", vec![Ok(message_start.clone())], ); let (first, state) = state .next_item() .await .expect("message_start should stream"); assert_eq!(first.expect("message_start should succeed"), message_start); let (error, mut state) = state .next_item() .await .expect("premature EOF should emit an Anthropic error event"); let error = String::from_utf8(error.expect("error event should succeed").to_vec()) .expect("error event should be utf8"); assert!(error.starts_with("event: error\ndata: ")); assert!(error.contains("ended before message_stop")); discard_direct_test_finalizer(&mut state); } #[tokio::test] async fn direct_anthropic_provider_error_is_not_duplicated() { let message_start = Bytes::from_static( b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{}}\n\n", ); let provider_error = Bytes::from_static( b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}\n\n", ); let state = direct_anthropic_inline_state( "req-direct-anthropic-provider-error", vec![Ok(message_start.clone()), Ok(provider_error.clone())], ); let (first, state) = state .next_item() .await .expect("message_start should stream"); assert_eq!(first.expect("message_start should succeed"), message_start); let (second, state) = state .next_item() .await .expect("provider error should stream"); assert_eq!( second.expect("provider error should succeed"), provider_error ); assert!(state.terminal_error_sent); assert!(state.next_item().await.is_none()); } #[test] fn postcommit_anthropic_errors_use_the_precommit_status_taxonomy() { for (error_type, expected_status) in [ ("request_too_large", 413), ("overloaded_error", 529), ("api_error", 500), ] { let body = json!({ "type": "error", "error": { "type": error_type, "message": "upstream failure" } }); assert_eq!( resolve_provider_stream_error_status_code("claude:messages", 200, &body), expected_status, ); } } #[tokio::test] async fn native_anthropic_unknown_event_is_replayed_byte_for_byte() { let unknown = "event: future_event\ndata: {\"type\":\"future_event\",\"value\":1}\n\n"; let message_start = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{}}\n\n", ); let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; let expected = format!("{unknown}{message_start}{message_stop}"); let outcome = execute_native_anthropic_prefetch_stream( "req-anthropic-unknown-replay", vec![ unknown.to_string(), message_start.to_string(), message_stop.to_string(), ], ) .await; let AiAttemptExecutionOutcome::Responded(response) = outcome else { panic!("unknown event should not terminate the stream") }; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("response body should read"); assert_eq!(body.as_ref(), expected.as_bytes()); } #[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, 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, false, )); } #[test] fn declared_stream_response_headers_are_normalized_without_body_inspection() { let mut headers = BTreeMap::from([ ( "Content-Type".to_string(), "Application/Octet-Stream; charset=binary".to_string(), ), ("Content-Encoding".to_string(), "identity".to_string()), ("x-upstream-header".to_string(), "preserved".to_string()), ]); assert!(should_normalize_declared_stream_response_headers( OPENAI_CHAT_STREAM_PLAN_KIND, 200, &headers, Some(&json!({"upstream_is_stream": true})), )); headers.insert("Content-Length".to_string(), "4096".to_string()); normalize_declared_stream_response_headers(&mut headers); assert_eq!( headers.get("content-type").map(String::as_str), Some("text/event-stream") ); assert!(!headers .keys() .any(|name| name.eq_ignore_ascii_case("content-encoding"))); assert!(!headers .keys() .any(|name| name.eq_ignore_ascii_case("content-length"))); assert_eq!( headers.get("x-upstream-header").map(String::as_str), Some("preserved") ); } #[test] fn declared_stream_header_normalization_requires_success_and_stream_context() { let headers = BTreeMap::from([( "content-type".to_string(), "application/octet-stream".to_string(), )]); assert!(!should_normalize_declared_stream_response_headers( OPENAI_CHAT_STREAM_PLAN_KIND, 500, &headers, Some(&json!({"upstream_is_stream": true})), )); assert!(!should_normalize_declared_stream_response_headers( OPENAI_CHAT_STREAM_PLAN_KIND, 200, &headers, Some(&json!({"upstream_is_stream": false})), )); assert!(!should_normalize_declared_stream_response_headers( OPENAI_CHAT_STREAM_PLAN_KIND, 200, &BTreeMap::from([("content-type".to_string(), "text/event-stream".to_string(),)]), Some(&json!({"upstream_is_stream": true})), )); assert!(!should_normalize_declared_stream_response_headers( OPENAI_CHAT_STREAM_PLAN_KIND, 200, &BTreeMap::from([("content-type".to_string(), "application/json".to_string(),)]), Some(&json!({"upstream_is_stream": true})), )); assert!(!should_normalize_declared_stream_response_headers( OPENAI_CHAT_STREAM_PLAN_KIND, 200, &BTreeMap::from([("content-type".to_string(), "text/plain".to_string(),)]), Some(&json!({"upstream_is_stream": true})), )); assert!(!should_normalize_declared_stream_response_headers( OPENAI_CHAT_STREAM_PLAN_KIND, 200, &BTreeMap::from([ ( "content-type".to_string(), "application/octet-stream".to_string(), ), ("content-encoding".to_string(), "gzip".to_string()), ]), Some(&json!({"upstream_is_stream": true})), )); assert!(!should_normalize_declared_stream_response_headers( OPENAI_CHAT_STREAM_PLAN_KIND, 200, &BTreeMap::from([ ( "content-type".to_string(), "application/octet-stream".to_string(), ), ("content-length".to_string(), "128".to_string()), ]), Some(&json!({"upstream_is_stream": true})), )); assert!(!should_normalize_declared_stream_response_headers( GEMINI_FILES_DOWNLOAD_PLAN_KIND, 200, &headers, Some(&json!({"upstream_is_stream": true})), )); } #[test] fn keeps_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, false, )); } #[test] fn cyber_failover_setting_forces_prefetch_for_event_streams() { assert!(!should_skip_direct_finalize_prefetch( Some("openai_responses_sync_finalize"), Some("text/event-stream"), "openai:responses", "openai:responses", false, false, true, )); } #[test] fn cyber_prefetch_waits_through_response_setup_until_output() { assert!(!prefetched_openai_responses_body_has_output_boundary( b"event: response.created\ndata: {\"type\":\"response.created\"}\n\n" )); assert!(prefetched_openai_responses_body_has_output_boundary( b"event: response.created\ndata: {\"type\":\"response.created\"}\n\nevent: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\n" )); } #[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 direct_passthrough_mode_defaults_inline_and_accepts_legacy() { assert_eq!( parse_direct_passthrough_mode(""), DirectPassthroughMode::Inline ); assert_eq!( parse_direct_passthrough_mode("inline"), DirectPassthroughMode::Inline ); assert_eq!( parse_direct_passthrough_mode("legacy"), DirectPassthroughMode::Legacy ); assert_eq!( parse_direct_passthrough_mode("mpsc"), DirectPassthroughMode::Legacy ); } #[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 native_anthropic_sse_body_ends_at_message_stop_while_sender_is_alive() { let (tx, rx) = mpsc::channel::>(1); let message_stop = Bytes::from_static(b"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"); tx.send(Ok(message_stop.clone())) .await .expect("message_stop should send"); let mut body_stream = Box::pin(build_sse_body_stream( Vec::new(), rx, true, false, true, Duration::from_secs(60), )); let chunk = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("message_stop should arrive immediately") .expect("stream should yield message_stop") .expect("message_stop should be successful"); assert_eq!(chunk, message_stop); assert!( tokio::time::timeout(Duration::from_millis(50), body_stream.next()) .await .expect("body EOF must not wait for the producer") .is_none() ); assert!(tx.is_closed(), "body EOF should drop the receiver"); } #[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, false, 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, 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, false, 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, false, 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, false, 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, false, 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(), )]), response_observation: None, }, })); 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(), true, 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("Execution runtime stream failed")); 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(), true, 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_ENCRYPTION_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(), true, 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_ENCRYPTION_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(), true, 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!(record.response_body.is_none()); assert!(record.response_body_ref.is_none()); assert!(record.client_response_body.is_none()); assert!(record.client_response_body_ref.is_none()); } #[tokio::test] async fn execute_stream_from_frame_stream_honors_client_disconnect_policy() { for cancel_on_client_disconnect in [false, true] { 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 = crate::request_lifecycle::run_request(async move { crate::request_lifecycle::configure_client_disconnect( aether_routing_core::RoutingExecutionPolicy { cancel_on_client_disconnect, ..Default::default() }, ); 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(), true, frame_stream, None, ) .await .map(|response| response.expect("execution should return a client response")) }) .await .expect("execution should succeed"); 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); if !cancel_on_client_disconnect { release_terminal.notify_one(); } let expected_candidate_status = if cancel_on_client_disconnect { RequestCandidateStatus::Cancelled } else { RequestCandidateStatus::Success }; let expected_usage_status = if cancel_on_client_disconnect { "cancelled" } else { "completed" }; 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 == expected_candidate_status) { break candidates; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("candidate should be marked cancelled"); assert_eq!( candidates[0].status_code, Some(if cancel_on_client_disconnect { 499 } else { 200 }) ); assert_eq!( candidates[0].error_type.as_deref(), cancel_on_client_disconnect.then_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 == expected_usage_status) { break usage.expect("cancelled usage should exist"); } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .expect("usage should be marked cancelled"); if !cancel_on_client_disconnect { assert_ne!(stored_usage.billing_status, "void"); 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); continue; } assert_eq!(stored_usage.billing_status, "void"); assert_eq!(stored_usage.status_code, Some(499)); assert_eq!(stored_usage.input_tokens, 0); assert_eq!(stored_usage.output_tokens, 0); assert_eq!(stored_usage.total_tokens, 0); 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" ); release_terminal.notify_one(); assert!( tokio::time::timeout( Duration::from_millis(100), terminal_frame_drained.notified() ) .await .is_err(), "upstream frame stream should stop when the client disconnects" ); } } #[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(), true, 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(), true, 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 execution_task = tokio::spawn(async move { 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 }); 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()); assert!(!execution_task.is_finished()); release_text.notify_one(); text_seen.notified().await; let response = tokio::time::timeout(Duration::from_secs(1), execution_task) .await .expect("semantic text should commit the response") .expect("execution task should complete") .expect("execution should succeed") .expect("execution should return a client response"); 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", "upstream_is_stream": true, })), ) .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_system_config_values_for_tests([( "request_record_level".to_string(), json!("full"), )]), ) .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.is_none()); assert_eq!( usage.client_response_headers.as_ref().unwrap()["content-type"], json!("application/json") ); assert_eq!( usage.response_headers.as_ref().unwrap()["content-type"], json!("text/html") ); assert!( usage.response_body.is_none(), "upstream redirect did not include a body" ); assert_eq!(usage.client_response_body.as_ref(), Some(&body_json)); 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"], "/" ); 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_sanitizes_local_tunnel_error_before_first_data() { let state = authenticated_local_tunnel_test_state(); 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, ) .with_tunnel_generation(LOCAL_TUNNEL_TEST_GENERATION.to_string()) .with_authenticated_key(LOCAL_TUNNEL_TEST_PSK.to_string()), )); 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 recv_tunnel_test_frame(&mut proxy_rx, "headers frame").await { 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 recv_tunnel_test_frame(&mut proxy_rx, "body frame").await { 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, "Upstream response stream failed"); assert!(!error_message.contains(original_error)); } #[tokio::test] async fn execute_execution_runtime_stream_emits_terminal_sse_error_event_after_body_started() { let state = authenticated_local_tunnel_test_state(); 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, ) .with_tunnel_generation(LOCAL_TUNNEL_TEST_GENERATION.to_string()) .with_authenticated_key(LOCAL_TUNNEL_TEST_PSK.to_string()), )); 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 recv_tunnel_test_frame(&mut proxy_rx, "headers frame").await { 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 recv_tunnel_test_frame(&mut proxy_rx, "body frame").await { 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("Upstream response stream failed")); assert!(!body.contains(original_error)); assert!(body.contains("data: [DONE]\n\n")); } }