From 9282cce1d643f8dbbf04d679b0fbbfb35066d920 Mon Sep 17 00:00:00 2001 From: stabey <36232531+stabey@users.noreply.github.com> Date: Wed, 2 Sep 2026 02:07:26 +0800 Subject: [PATCH] fix(gateway): settle stream attempts dropped before first byte A local stream attempt writes its `usage` row and its `request_candidates` slot as `pending` in `execute_execution_runtime_stream_inner`, then awaits the provider's response headers. Everything after that point runs inside the downstream request future, so a client disconnect drops it: the dispatch `.await` never resumes and nothing settles either row. The stream finalizer that already covers this only exists once upstream headers have arrived, so the pre-first-byte window has no owner at all. Both rows stay `pending` until the maintenance sweeper rewrites them as a 504 timeout ten minutes later, losing the real outcome, the real latency, and the 499. `AttemptCancellationGuard` takes that window. It is created disarmed, so an attempt dropped before it owns any row does not grow a settlement row it never had; it is armed as soon as the attempt owns its non-terminal rows, and the stream wrappers disarm it the moment the attempt returns, from where settlement belongs to the transport. On a cancelling drop it settles the candidate slot through the same snapshot writer the `pending` write above it uses, and the usage row through a terminal `Cancelled` event. The guard outlives the request future, so what it captures is retained for the whole attempt. It therefore holds no request body: the plan carries the provider request body and the report context carries the client request body, and keeping both would double the request-body residency of every in-flight stream attempt to serve a path that almost never runs. Simply omitting them is not safe either, because a terminal write is body-capture-authoritative: with both absent the seed carries the typed `none` marker, which clears the stored capture rather than leaving it alone. `build_usage_event_data_seed_describing_request_bodies` is the third option -- it derives every capture state, body reference, request type and derived request fact from the real plan and report context, and leaves out only the two body values -- so the guard's snapshot is small and its terminal write preserves the capture the `pending` write recorded. The stream candidate first-byte watchdog also drops the attempt future, but it settles the attempt itself through `build_transport_error_stop_response`. It now marks the attempt abandoned before returning so the guard stands down instead of racing a 499 against the watchdog's 504. Co-Authored-By: Claude Opus 5 --- .../execution_runtime/attempt_cancellation.rs | 568 ++++++++++++++++++ .../src/execution_runtime/mod.rs | 1 + .../src/execution_runtime/stream/execution.rs | 67 ++- .../execution_runtime/transport_failure.rs | 19 + .../src/executor/candidate_loop.rs | 4 + .../src/tests/ai_execute/lifecycle.rs | 109 ++++ crates/aether-usage/runtime/src/lib.rs | 3 +- crates/aether-usage/runtime/src/write.rs | 198 +++++- 8 files changed, 943 insertions(+), 26 deletions(-) create mode 100644 apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs diff --git a/apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs b/apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs new file mode 100644 index 000000000..f538810e6 --- /dev/null +++ b/apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs @@ -0,0 +1,568 @@ +//! Terminal settlement for a local stream attempt whose future is dropped +//! mid-flight. +//! +//! A local stream attempt writes its `usage` row and its `request_candidates` +//! slot as `pending` before it dispatches to the provider, then keeps running +//! inside the downstream request future. When the client disconnects, axum drops +//! that future: the remaining `.await`s never resume and nothing settles either +//! row. They stay `pending` until the maintenance sweeper rewrites them as a 504 +//! timeout roughly ten minutes later, which loses the real outcome and the real +//! latency. +//! +//! The stream transport therefore keeps a guard alive across the window between +//! the `pending` write and terminal settlement, and settles the attempt from +//! `Drop` when that window is left by cancellation instead of by a terminal +//! state. + +use std::sync::Arc; +use std::time::Instant; + +use aether_contracts::ExecutionPlan; +use aether_data_contracts::repository::candidates::RequestCandidateStatus; +use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate; +use aether_usage_runtime::{ + build_usage_event_data_seed_describing_request_bodies, UsageEvent, UsageEventData, + UsageEventType, +}; +use serde_json::{json, Value}; +use tracing::warn; + +use crate::clock::current_unix_ms as current_request_candidate_unix_ms; +use crate::execution_runtime::attempt_lifecycle::CLIENT_CANCELLED_STATUS_CODE; +use crate::execution_runtime::transport_failure::StreamCandidateWatchdogProgress; +use crate::log_ids::short_request_id; +use crate::request_candidate_runtime::{ + record_local_request_candidate_status_snapshot, LocalRequestCandidateStatusSnapshot, +}; +use crate::request_diagnostics::{ + attach_request_diagnostics_to_report_context, current_request_diagnostics, RequestDiagnostics, +}; +use crate::AppState; + +fn elapsed_ms_since(started_at: Instant) -> u64 { + started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 +} + +/// The facts the guard needs to settle the attempt it is watching. +/// +/// This is held for the whole attempt, so it is deliberately free of request +/// bodies. A request body can be megabytes, and holding one per in-flight +/// attempt would cost far more than the row it settles: the usage seed is built +/// with [`build_usage_event_data_seed_describing_request_bodies`], which derives +/// every capture state, body reference and derived request fact from the real +/// plan and report context but keeps neither body. The terminal write it +/// produces therefore preserves the capture the `pending` write recorded instead +/// of clearing it. +struct ArmedAttempt { + request_id: String, + candidate_id: Option, + candidate: Option, + // Boxed: the guard lives inside the stream request future, which is already + // very large, and `UsageEventData` is a wide struct. + usage_seed: Option>, + request_diagnostics: Option>, + candidate_started_unix_ms: u64, + candidate_started_at: Instant, +} + +/// Settles an attempt as cancelled when its future is dropped before the +/// transport reaches a terminal state. +/// +/// The guard is created disarmed and stays inert until [`Self::arm`] is called, +/// so an attempt that is dropped before it owns any `pending` row does not grow +/// a settlement row it never had. The owner disarms it as soon as the attempt +/// completes, whichever way it completes: from that point terminal settlement +/// belongs to the transport (for streams, to the stream finalizer that lives in +/// the response body), and the guard must not write a second terminal state. +/// +/// A stream candidate also runs under a first-byte watchdog that drops the +/// attempt future when it gives up. That drop is not a client disconnect and the +/// watchdog settles the attempt itself, so the guard stands down for it. +pub(crate) struct AttemptCancellationGuard { + state: AppState, + error_type: &'static str, + error_message: &'static str, + watchdog: Option>, + armed: Option, +} + +impl AttemptCancellationGuard { + pub(crate) fn disarmed( + state: &AppState, + error_type: &'static str, + error_message: &'static str, + ) -> Self { + Self { + state: state.clone(), + error_type, + error_message, + watchdog: StreamCandidateWatchdogProgress::current(), + armed: None, + } + } + + /// Takes ownership of the attempt's settlement until it is disarmed. + pub(crate) fn arm( + &mut self, + plan: &ExecutionPlan, + report_context: Option<&Value>, + candidate: Option<&LocalRequestCandidateStatusSnapshot>, + candidate_started_unix_ms: u64, + candidate_started_at: Instant, + ) { + let usage_seed = self.state.usage_runtime.is_enabled().then(|| { + Box::new(build_usage_event_data_seed_describing_request_bodies( + plan, + report_context, + )) + }); + self.armed = Some(ArmedAttempt { + request_id: plan.request_id.clone(), + candidate_id: plan.candidate_id.clone(), + candidate: candidate.cloned(), + usage_seed, + request_diagnostics: current_request_diagnostics(), + candidate_started_unix_ms, + candidate_started_at, + }); + } + + pub(crate) fn disarm(&mut self) { + self.armed = None; + } +} + +/// Writes the candidate terminal row and the terminal usage event for an attempt +/// that never reached its own terminal path. +async fn settle_cancelled_attempt( + state: AppState, + armed: ArmedAttempt, + error_type: &'static str, + error_message: &'static str, +) { + let ArmedAttempt { + request_id, + candidate_id: _, + candidate, + usage_seed, + request_diagnostics, + candidate_started_unix_ms, + candidate_started_at, + } = armed; + let terminal_unix_ms = current_request_candidate_unix_ms(); + let latency_ms = elapsed_ms_since(candidate_started_at); + + if let Some(candidate) = candidate.as_ref() { + record_local_request_candidate_status_snapshot( + &state, + candidate, + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Cancelled, + status_code: Some(CLIENT_CANCELLED_STATUS_CODE), + error_type: Some(error_type.to_string()), + error_message: Some(error_message.to_string()), + latency_ms: Some(latency_ms), + started_at_unix_ms: Some(candidate_started_unix_ms), + finished_at_unix_ms: Some(terminal_unix_ms), + }, + ) + .await; + } + + let Some(usage_data) = usage_seed else { + return; + }; + let mut usage_data = *usage_data; + // The seed was built when the attempt was armed, so it predates the + // diagnostics it should carry. Attaching them to the seed's metadata is the + // same write the report context would have carried into a seed built here: + // both land the same keys in the same object. + usage_data.request_metadata = attach_request_diagnostics_to_report_context( + usage_data.request_metadata.take(), + request_diagnostics.as_ref(), + ); + usage_data.status_code = Some(CLIENT_CANCELLED_STATUS_CODE); + usage_data.error_message = Some(error_message.to_string()); + usage_data.error_category = Some("cancelled".to_string()); + usage_data.response_time_ms = Some(latency_ms); + let error_body = json!({ + "error": { + "type": error_type, + "message": error_message, + "code": CLIENT_CANCELLED_STATUS_CODE + } + }); + usage_data.response_headers = Some(json!({"content-type": "application/json"})); + usage_data.response_body = Some(error_body.clone()); + usage_data.client_response_headers = Some(json!({"content-type": "application/json"})); + usage_data.client_response_body = Some(error_body); + + state + .usage_runtime + .record_terminal_event_direct( + state.usage_lifecycle_data_state().as_ref(), + UsageEvent::new(UsageEventType::Cancelled, request_id, usage_data), + ) + .await; +} + +impl Drop for AttemptCancellationGuard { + fn drop(&mut self) { + let Some(armed) = self.armed.take() else { + return; + }; + if self + .watchdog + .as_ref() + .is_some_and(|watchdog| watchdog.abandoned()) + { + return; + } + let state = self.state.clone(); + let error_type = self.error_type; + let error_message = self.error_message; + // `Drop` cannot await, and the settlement writes touch the database. + // Hand them to the runtime so they survive the dropped request future. + let Ok(handle) = tokio::runtime::Handle::try_current() else { + warn!( + event_name = "local_attempt_cancellation_guard_no_runtime", + log_type = "ops", + request_id = %short_request_id(armed.request_id.as_str()), + candidate_id = ?armed.candidate_id, + error_type, + "gateway could not settle dropped local attempt because no Tokio runtime is available" + ); + return; + }; + handle.spawn(async move { + settle_cancelled_attempt(state, armed, error_type, error_message).await; + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aether_contracts::RequestBody; + use aether_data::repository::candidates::InMemoryRequestCandidateRepository; + use aether_data::repository::usage::InMemoryUsageReadRepository; + use aether_data_contracts::repository::candidates::RequestCandidateReadRepository; + use aether_data_contracts::repository::usage::{ + StoredRequestUsageAudit, UsageBodyCaptureState, UsageReadRepository, UsageWriteRepository, + }; + use aether_usage_runtime::{ + build_lifecycle_usage_seed, build_pending_usage_record, UsageRuntimeConfig, + }; + use std::collections::BTreeMap; + use std::time::Duration; + + use crate::request_candidate_runtime::{ + ensure_execution_request_candidate_slot, snapshot_local_request_candidate_status, + }; + + const TEST_ERROR_TYPE: &str = "local_stream_attempt_cancelled"; + const TEST_ERROR_MESSAGE: &str = + "Local stream attempt was dropped before terminal finalization."; + + fn test_stream_plan(request_id: &str) -> ExecutionPlan { + ExecutionPlan { + request_id: request_id.to_string(), + candidate_id: None, + provider_name: Some("Anthropic".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.test/v1/messages".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({"stream": true, "service_tier": "priority"})), + stream: true, + client_api_format: "claude:messages".to_string(), + provider_api_format: "claude:messages".to_string(), + model_name: Some("claude-sonnet-4-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + } + } + + fn test_report_context() -> Option { + Some(json!({ + "candidate_index": 0, + "retry_index": 0, + "user_id": "user-cancel", + "api_key_id": "api-key-cancel", + "client_api_format": "claude:messages", + "provider_api_format": "claude:messages", + "request_path": "/v1/messages", + "request_path_and_query": "/v1/messages?beta=true", + "upstream_url": "https://example.test/v1/messages", + "mapped_model": "claude-sonnet-4-5", + "original_request_body": {"stream": true, "messages": []}, + })) + } + + fn test_state( + usage_repository: &Arc, + request_candidate_repository: &Arc, + ) -> AppState { + AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests( + Arc::clone(request_candidate_repository), + Arc::clone(usage_repository), + ), + ) + .with_usage_runtime_for_tests(UsageRuntimeConfig { + enabled: true, + ..UsageRuntimeConfig::default() + }) + } + + /// Writes the `pending` rows the same way a stream attempt does before it + /// dispatches to the provider, and returns the candidate slot snapshot the + /// attempt owns from that point on. + async fn record_pending_attempt( + state: &AppState, + plan: &mut ExecutionPlan, + report_context: &mut Option, + candidate_started_unix_ms: u64, + ) -> LocalRequestCandidateStatusSnapshot { + ensure_execution_request_candidate_slot(state, plan, report_context).await; + state.usage_runtime.record_pending( + state.usage_lifecycle_data_state().as_ref(), + build_lifecycle_usage_seed(plan, report_context.as_ref()), + ); + let snapshot = snapshot_local_request_candidate_status(plan, report_context.as_ref()) + .expect("attempt should own a candidate slot"); + 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_ms), + finished_at_unix_ms: None, + }, + ) + .await; + snapshot + } + + async fn wait_for_usage_status( + usage_repository: &InMemoryUsageReadRepository, + request_id: &str, + status: &str, + ) -> Option { + for _ in 0..50 { + if let Some(usage) = usage_repository + .find_by_request_id(request_id) + .await + .expect("usage should read") + { + if usage.status == status { + return Some(usage); + } + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + None + } + + #[tokio::test] + async fn armed_guard_settles_a_dropped_attempt_as_cancelled() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = test_state(&usage_repository, &request_candidate_repository); + let mut plan = test_stream_plan("stream-cancel-guard-request"); + let mut report_context = test_report_context(); + let candidate_started_unix_ms = current_request_candidate_unix_ms(); + let snapshot = record_pending_attempt( + &state, + &mut plan, + &mut report_context, + candidate_started_unix_ms, + ) + .await; + + { + let mut guard = + AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE); + guard.arm( + &plan, + report_context.as_ref(), + Some(&snapshot), + candidate_started_unix_ms, + Instant::now(), + ); + } + + let usage = wait_for_usage_status( + usage_repository.as_ref(), + "stream-cancel-guard-request", + "cancelled", + ) + .await + .expect("cancelled usage should be recorded"); + assert_eq!(usage.billing_status, "void"); + assert_eq!(usage.status_code, Some(CLIENT_CANCELLED_STATUS_CODE)); + assert_eq!(usage.error_category.as_deref(), Some("cancelled")); + assert!(usage.response_time_ms.is_some()); + + let candidates = request_candidate_repository + .list_by_request_id("stream-cancel-guard-request") + .await + .expect("candidates should read"); + let candidate = candidates.first().expect("candidate row should exist"); + assert_eq!(candidate.status, RequestCandidateStatus::Cancelled); + assert_eq!(candidate.status_code, Some(CLIENT_CANCELLED_STATUS_CODE)); + assert_eq!(candidate.error_type.as_deref(), Some(TEST_ERROR_TYPE)); + assert!(candidate.finished_at_unix_ms.is_some()); + } + + /// The guard holds no request body, so its settlement write must describe the + /// capture rather than deny it: a typed `none` capture state would clear the + /// stored request body instead of leaving it alone. + #[tokio::test] + async fn settling_a_dropped_attempt_leaves_the_captured_request_body_alone() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = test_state(&usage_repository, &request_candidate_repository); + let mut plan = test_stream_plan("stream-cancel-guard-capture"); + let mut report_context = test_report_context(); + let candidate_started_unix_ms = current_request_candidate_unix_ms(); + let snapshot = record_pending_attempt( + &state, + &mut plan, + &mut report_context, + candidate_started_unix_ms, + ) + .await; + // Stand in for a write that already captured this request's body. + let captured_body = json!({"stream": true, "service_tier": "priority"}); + let mut capture = build_pending_usage_record( + &plan, + report_context.as_ref(), + current_request_candidate_unix_ms() / 1_000, + ) + .expect("pending usage record should build"); + capture.provider_request_body = Some(captured_body.clone()); + capture.provider_request_body_state = Some(UsageBodyCaptureState::Inline); + usage_repository + .upsert(capture) + .await + .expect("captured request body should upsert"); + + { + let mut guard = + AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE); + guard.arm( + &plan, + report_context.as_ref(), + Some(&snapshot), + candidate_started_unix_ms, + Instant::now(), + ); + } + + let usage = wait_for_usage_status( + usage_repository.as_ref(), + "stream-cancel-guard-capture", + "cancelled", + ) + .await + .expect("cancelled usage should be recorded"); + assert_eq!(usage.provider_request_body, Some(captured_body)); + assert_ne!( + usage.provider_request_body_state, + Some(UsageBodyCaptureState::None) + ); + } + + #[tokio::test] + async fn guard_stands_down_when_the_watchdog_abandons_the_attempt() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = test_state(&usage_repository, &request_candidate_repository); + let mut plan = test_stream_plan("stream-watchdog-guard-request"); + let mut report_context = test_report_context(); + let candidate_started_unix_ms = current_request_candidate_unix_ms(); + let snapshot = record_pending_attempt( + &state, + &mut plan, + &mut report_context, + candidate_started_unix_ms, + ) + .await; + + let watchdog = StreamCandidateWatchdogProgress::shared(); + Arc::clone(&watchdog) + .scope(async { + let mut guard = + AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE); + guard.arm( + &plan, + report_context.as_ref(), + Some(&snapshot), + candidate_started_unix_ms, + Instant::now(), + ); + // The watchdog gives up and takes over settlement before the + // abandoned attempt is dropped. + watchdog.mark_abandoned(); + }) + .await; + + assert!(wait_for_usage_status( + usage_repository.as_ref(), + "stream-watchdog-guard-request", + "cancelled", + ) + .await + .is_none()); + } + + #[tokio::test] + async fn disarmed_guard_leaves_the_attempt_pending() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = test_state(&usage_repository, &request_candidate_repository); + let mut plan = test_stream_plan("stream-disarmed-guard-request"); + let mut report_context = test_report_context(); + let candidate_started_unix_ms = current_request_candidate_unix_ms(); + let snapshot = record_pending_attempt( + &state, + &mut plan, + &mut report_context, + candidate_started_unix_ms, + ) + .await; + + { + let mut guard = + AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE); + guard.arm( + &plan, + report_context.as_ref(), + Some(&snapshot), + candidate_started_unix_ms, + Instant::now(), + ); + guard.disarm(); + } + + assert!(wait_for_usage_status( + usage_repository.as_ref(), + "stream-disarmed-guard-request", + "cancelled", + ) + .await + .is_none()); + } +} diff --git a/apps/aether-gateway/src/execution_runtime/mod.rs b/apps/aether-gateway/src/execution_runtime/mod.rs index 89fbf1653..c8596db2a 100644 --- a/apps/aether-gateway/src/execution_runtime/mod.rs +++ b/apps/aether-gateway/src/execution_runtime/mod.rs @@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; pub(crate) mod admission; +pub(crate) mod attempt_cancellation; pub(crate) mod attempt_lifecycle; mod chatgpt_web_image; mod constants; diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 536326765..ba3369701 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -78,6 +78,7 @@ use crate::api::response::{ 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; @@ -149,6 +150,11 @@ 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; @@ -3655,17 +3661,30 @@ pub(crate) fn execute_execution_runtime_stream<'a>( report_kind: Option, report_context: Option, ) -> Pin>, GatewayError>> + Send + 'a>> { - Box::pin(execute_execution_runtime_stream_inner( - state, - plan, - trace_id, - decision, - plan_kind, - report_kind, - report_context, - None, - None, - )) + 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)] @@ -3687,7 +3706,12 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>( Box::pin(async move { let mut retry_scope = AiAttemptRetryScope::Candidate; let mut fallback_response = None; - let response = execute_execution_runtime_stream_inner( + 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, @@ -3697,8 +3721,13 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>( report_context, Some(&mut retry_scope), Some(&mut fallback_response), + &mut cancellation_guard, ) - .await?; + .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 { @@ -3744,6 +3773,7 @@ async fn maybe_build_stream_transport_error_stop_response( .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, @@ -3754,6 +3784,7 @@ async fn execute_execution_runtime_stream_inner( 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(); @@ -3837,6 +3868,16 @@ async fn execute_execution_runtime_stream_inner( ) .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 diff --git a/apps/aether-gateway/src/execution_runtime/transport_failure.rs b/apps/aether-gateway/src/execution_runtime/transport_failure.rs index ccd532adf..c3f1885fb 100644 --- a/apps/aether-gateway/src/execution_runtime/transport_failure.rs +++ b/apps/aether-gateway/src/execution_runtime/transport_failure.rs @@ -22,6 +22,7 @@ const TRANSPORT_ERROR_CLIENT_MESSAGE: &str = #[derive(Debug, Default)] pub(crate) struct StreamCandidateWatchdogProgress { terminal_started: AtomicBool, + abandoned: AtomicBool, } tokio::task_local! { @@ -37,6 +38,24 @@ impl StreamCandidateWatchdogProgress { self.terminal_started.load(Ordering::Acquire) } + /// The watchdog gave up waiting and settles this attempt itself. + /// + /// The attempt future is dropped once the watchdog returns, so its own + /// cancellation guard must stay out of the way instead of racing the + /// watchdog's terminal rows with a cancellation. + pub(crate) fn mark_abandoned(&self) { + self.abandoned.store(true, Ordering::Release); + } + + pub(crate) fn abandoned(&self) -> bool { + self.abandoned.load(Ordering::Acquire) + } + + /// The watchdog watching the attempt on this task, if it runs under one. + pub(crate) fn current() -> Option> { + STREAM_CANDIDATE_WATCHDOG_PROGRESS.try_with(Arc::clone).ok() + } + pub(crate) async fn scope(self: Arc, future: F) -> F::Output where F: Future, diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 9a689089d..42f262a78 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -1418,6 +1418,10 @@ where let outcome = match execution_result { Some(result) => result.map(StreamCandidateWatchdogOutcome::Executed), None => { + // The abandoned attempt is dropped when this function returns. + // Claim its settlement before that so its cancellation guard does + // not race the watchdog rows written just below. + watchdog_progress.mark_abandoned(); let finished_at_unix_ms = current_unix_ms(); let request_id = short_request_id(plan.request_id.as_str()); let provider_name = plan.provider_name.as_deref().unwrap_or("-"); diff --git a/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs b/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs index c3cae80b1..6e26a8976 100644 --- a/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs +++ b/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs @@ -14,6 +14,9 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi use aether_data_contracts::repository::candidate_selection::{ StoredMinimalCandidateSelectionRow, StoredProviderModelMapping, }; +use aether_data_contracts::repository::candidates::{ + RequestCandidateReadRepository, RequestCandidateStatus, +}; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; @@ -427,6 +430,112 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects_impl() { upstream_handle.abort(); } +#[test] +fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byte() { + run_lifecycle_test( + "gateway_settles_stream_attempt_when_client_disconnects_before_first_byte", + gateway_settles_stream_attempt_when_client_disconnects_before_first_byte_impl, + ); +} + +async fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byte_impl() { + // The execution runtime accepts the plan and then goes quiet, so the attempt + // is parked between its `pending` rows and the first upstream byte. + let execution_runtime = Router::new().route( + "/v1/execute/stream", + any(|_request: Request| async move { + tokio::time::sleep(std::time::Duration::from_secs(30)).await; + StatusCode::OK + }), + ); + + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-client-openai-stream-precommit-disconnect")), + sample_local_openai_auth_snapshot( + "api-key-openai-lifecycle-local-1", + "user-openai-lifecycle-local-1", + ), + )])); + let candidate_selection_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + sample_local_openai_candidate_row(), + ])); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![sample_local_openai_provider()], + vec![sample_local_openai_endpoint()], + vec![sample_local_openai_key()], + )); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let gateway = build_router_with_state( + build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests( + GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests( + auth_repository, + candidate_selection_repository, + provider_catalog_repository, + Arc::clone(&request_candidate_repository), + DEVELOPMENT_ENCRYPTION_KEY, + ), + ), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let request = reqwest::Client::new() + .post(format!("{gateway_url}/v1/chat/completions")) + .header(http::header::CONTENT_TYPE, "application/json") + .header( + http::header::AUTHORIZATION, + "Bearer sk-client-openai-stream-precommit-disconnect", + ) + .header( + TRACE_ID_HEADER, + "trace-openai-chat-stream-precommit-disconnect-123", + ) + .body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}") + .send(); + + // Drop the in-flight request the way a downstream client does when its own + // first-byte timeout fires, before any response header exists. + assert!( + tokio::time::timeout(std::time::Duration::from_millis(750), request) + .await + .is_err(), + "the execution runtime should not have answered before the client gave up" + ); + + let mut stored_candidates = Vec::new(); + for _ in 0..200 { + stored_candidates = request_candidate_repository + .list_by_request_id("trace-openai-chat-stream-precommit-disconnect-123") + .await + .expect("request candidate trace should read"); + if stored_candidates + .iter() + .any(|candidate| candidate.status == RequestCandidateStatus::Cancelled) + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + + let cancelled = stored_candidates + .iter() + .find(|candidate| candidate.status == RequestCandidateStatus::Cancelled) + .unwrap_or_else(|| { + panic!("dropped stream attempt should settle as cancelled: {stored_candidates:?}") + }); + assert_eq!(cancelled.status_code, Some(499)); + assert_eq!( + cancelled.error_type.as_deref(), + Some("local_stream_attempt_cancelled") + ); + assert!(cancelled.finished_at_unix_ms.is_some()); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + #[test] fn gateway_returns_error_body_when_prefetch_detects_embedded_stream_error() { run_lifecycle_test( diff --git a/crates/aether-usage/runtime/src/lib.rs b/crates/aether-usage/runtime/src/lib.rs index 0be0b5b97..e37d38eef 100644 --- a/crates/aether-usage/runtime/src/lib.rs +++ b/crates/aether-usage/runtime/src/lib.rs @@ -61,7 +61,8 @@ pub use write::{ build_sync_terminal_usage_event, build_sync_terminal_usage_outcome, build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed, build_terminal_usage_context_seed, build_terminal_usage_event_from_outcome, - build_terminal_usage_event_from_seed, build_usage_event_data_seed, LifecycleUsageSeed, + build_terminal_usage_event_from_seed, build_usage_event_data_seed, + build_usage_event_data_seed_describing_request_bodies, LifecycleUsageSeed, StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed, TerminalUsageOutcome, TerminalUsageSeed, UsageTerminalState, }; diff --git a/crates/aether-usage/runtime/src/write.rs b/crates/aether-usage/runtime/src/write.rs index ba332fdab..cfd61fab2 100644 --- a/crates/aether-usage/runtime/src/write.rs +++ b/crates/aether-usage/runtime/src/write.rs @@ -72,6 +72,20 @@ struct RuntimeRequestCaptureSeed { provider_request: Option, provider_request_body_ref: Option, body_states: UsageBodyStatesSeed, + request_has_inline_body: bool, + provider_request_has_inline_body: bool, +} + +/// Whether a seed keeps the request bodies it describes, or only describes them. +/// +/// A holder that has to outlive the request itself pays for every byte it keeps, +/// and a request body can be megabytes. [`RequestBodyCapture::Describe`] computes +/// every capture state, reference and derived fact from the real plan and report +/// context, and leaves out only the body values themselves. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RequestBodyCapture { + Keep, + Describe, } #[derive(Debug, Clone, PartialEq)] @@ -804,7 +818,8 @@ pub fn build_terminal_usage_context_seed( report_context: Option<&Value>, ) -> TerminalUsageContextSeed { let context = report_context.and_then(Value::as_object); - let request_capture = build_runtime_request_capture_seed(plan, context); + let request_capture = + build_runtime_request_capture_seed(plan, context, RequestBodyCapture::Keep); let client_contract = context_string(context, "client_contract") .or_else(|| context_string(context, "client_api_format")) .or_else(|| non_empty_str(Some(plan.client_api_format.as_str()))) @@ -1691,16 +1706,34 @@ pub fn build_usage_event_data_seed( plan: &ExecutionPlan, report_context: Option<&Value>, ) -> UsageEventData { - build_usage_event_data_seed_with_detail(plan, report_context) + build_usage_event_data_seed_with_detail(plan, report_context, RequestBodyCapture::Keep) +} + +/// Builds the same seed as [`build_usage_event_data_seed`] without keeping the +/// request bodies. +/// +/// This is for a caller that has to hold a seed for the whole life of an attempt +/// so it can still write a terminal row if the attempt is dropped: a request body +/// can be megabytes, and holding one per in-flight attempt is far more expensive +/// than the row it would eventually capture. Every capture state, body reference +/// and derived request fact is still computed from the real plan and report +/// context, so the resulting terminal write preserves the capture an earlier +/// non-terminal write recorded rather than clearing it. +pub fn build_usage_event_data_seed_describing_request_bodies( + plan: &ExecutionPlan, + report_context: Option<&Value>, +) -> UsageEventData { + build_usage_event_data_seed_with_detail(plan, report_context, RequestBodyCapture::Describe) } fn build_usage_event_data_seed_with_detail( plan: &ExecutionPlan, report_context: Option<&Value>, + capture: RequestBodyCapture, ) -> UsageEventData { let context = report_context.and_then(Value::as_object); let routing = build_runtime_routing_seed(plan, context); - let request_capture = build_runtime_request_capture_seed(plan, context); + let request_capture = build_runtime_request_capture_seed(plan, context, capture); let api_format = context_string(context, "client_api_format") .or_else(|| non_empty_str(Some(plan.client_api_format.as_str()))); let endpoint_api_format = context_string(context, "provider_api_format") @@ -1714,7 +1747,10 @@ fn build_usage_event_data_seed_with_detail( let request_type = Some(infer_request_type_from_contracts( api_format.as_deref(), endpoint_api_format.as_deref(), - request_capture.provider_request.as_ref(), + request_capture + .provider_request + .as_ref() + .or_else(|| provider_request_body_ref_for_inference(plan, context)), )); let api_family = api_format .as_deref() @@ -1737,9 +1773,9 @@ fn build_usage_event_data_seed_with_detail( build_runtime_request_metadata_seed_from_parts( plan, context, - request_capture.request_body.is_some(), + request_capture.request_has_inline_body, request_capture.request_body_ref.as_deref(), - request_capture.provider_request.is_some(), + request_capture.provider_request_has_inline_body, request_capture.provider_request_body_ref.as_deref(), plan.body.body_bytes_b64.as_deref(), ), @@ -1999,17 +2035,29 @@ fn plan_has_inline_json_body_for_usage(plan: &ExecutionPlan) -> bool { fn build_runtime_request_capture_seed( plan: &ExecutionPlan, context: Option<&Map>, + capture: RequestBodyCapture, ) -> RuntimeRequestCaptureSeed { - let request_body = context_body_value(context, "original_request_body"); + // Presence, not the value, is what every capture state and derived fact is + // built from, so both capture modes agree on all of them. + let request_has_inline_body = context_has_inline_body(context, "original_request_body"); + let provider_request_has_inline_body = + context_has_inline_body(context, "provider_request_body") + || plan_has_inline_json_body_for_usage(plan); + let (request_body, provider_request) = match capture { + RequestBodyCapture::Keep => ( + context_body_value(context, "original_request_body"), + context_body_value(context, "provider_request_body") + .or_else(|| plan_json_body_capture_for_usage(plan)), + ), + RequestBodyCapture::Describe => (None, None), + }; let request_body_ref = context_string(context, "request_body_ref"); - let provider_request = context_body_value(context, "provider_request_body") - .or_else(|| plan_json_body_capture_for_usage(plan)); let provider_request_body_ref = context_string(context, "provider_request_body_ref") .or_else(|| non_empty_str(plan.body.body_ref.as_deref())); let body_states = build_runtime_body_states_seed_from_parts( - request_body.is_some(), + request_has_inline_body, request_body_ref.as_deref(), - provider_request.is_some(), + provider_request_has_inline_body, provider_request_body_ref.as_deref(), plan.body.body_bytes_b64.is_some(), ); @@ -2020,9 +2068,26 @@ fn build_runtime_request_capture_seed( provider_request, provider_request_body_ref, body_states, + request_has_inline_body, + provider_request_has_inline_body, } } +/// Borrows the provider request body that request-type inference reads, without +/// cloning it. +fn provider_request_body_ref_for_inference<'a>( + plan: &'a ExecutionPlan, + context: Option<&'a Map>, +) -> Option<&'a Value> { + context_value_ref(context, "provider_request_body") + .filter(|value| !value.is_null()) + .or_else(|| { + plan_has_inline_json_body_for_usage(plan) + .then_some(plan.body.json_body.as_ref()) + .flatten() + }) +} + fn build_runtime_request_metadata_seed( plan: &ExecutionPlan, context: Option<&Map>, @@ -3493,7 +3558,8 @@ mod tests { build_streaming_usage_event_from_owned_seed, build_streaming_usage_record, build_sync_terminal_usage_event, build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed, build_terminal_usage_context_seed, - build_terminal_usage_event_from_seed, build_usage_event_data_seed, decode_body_for_storage, + build_terminal_usage_event_from_seed, build_usage_event_data_seed, + build_usage_event_data_seed_describing_request_bodies, decode_body_for_storage, extract_token_counts_from_json, extract_token_counts_from_value, headers_to_json, mask_header_value, mask_sensitive_body_fields, mask_sensitive_headers_in_json_value, parse_sse_body_for_storage, resolve_error_message, trim_owned_non_empty_string, @@ -6842,6 +6908,114 @@ mod tests { assert!(body_size.get("provider_request_body").is_some()); } + #[test] + fn describing_request_bodies_matches_the_capturing_seed_apart_from_the_bodies() { + let plan = ExecutionPlan { + request_id: "req-seed-describe-1".to_string(), + candidate_id: Some("cand-seed-describe-1".to_string()), + provider_name: Some("OpenAI".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.com/v1/chat/completions".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({ + "model": "gpt-5", + "service_tier": "priority", + "reasoning": {"effort": "high"} + })), + stream: false, + client_api_format: "openai:chat".to_string(), + provider_api_format: "openai:chat".to_string(), + model_name: Some("gpt-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + let report_context = json!({ + "client_api_format": "openai:chat", + "provider_api_format": "openai:chat", + "original_request_body": {"model": "gpt-5", "messages": []}, + "original_headers": {"accept": "application/json"} + }); + + let captured = build_usage_event_data_seed(&plan, Some(&report_context)); + let described = + build_usage_event_data_seed_describing_request_bodies(&plan, Some(&report_context)); + + // Only the two heavy values differ. + assert!(captured.request_body.is_some()); + assert!(captured.provider_request_body.is_some()); + assert_eq!(described.request_body, None); + assert_eq!(described.provider_request_body, None); + + // Everything a terminal write reads to decide what to do with the stored + // capture is identical, so the described seed preserves it rather than + // clearing it. + assert_eq!( + described.request_body_state, + Some(UsageBodyCaptureState::Inline) + ); + assert_eq!( + described.provider_request_body_state, + Some(UsageBodyCaptureState::Inline) + ); + assert_eq!(described.request_body_state, captured.request_body_state); + assert_eq!( + described.provider_request_body_state, + captured.provider_request_body_state + ); + assert_eq!(described.request_body_ref, captured.request_body_ref); + assert_eq!( + described.provider_request_body_ref, + captured.provider_request_body_ref + ); + assert_eq!(described.request_type, captured.request_type); + assert_eq!(described.request_metadata, captured.request_metadata); + assert_eq!( + described.provider_request_headers, + captured.provider_request_headers + ); + assert_eq!(described.request_headers, captured.request_headers); + } + + #[test] + fn describing_request_bodies_keeps_the_unavailable_marker_for_raw_bodies() { + let mut plan = ExecutionPlan { + request_id: "req-seed-describe-2".to_string(), + candidate_id: None, + provider_name: Some("OpenAI".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.com/v1/chat/completions".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({"model": "gpt-5"})), + stream: false, + client_api_format: "openai:chat".to_string(), + provider_api_format: "openai:chat".to_string(), + model_name: Some("gpt-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + plan.body.json_body = None; + plan.body.body_bytes_b64 = Some("eyJtb2RlbCI6ICJncHQtNSJ9".to_string()); + + let described = build_usage_event_data_seed_describing_request_bodies(&plan, None); + + assert_eq!( + described.provider_request_body_state, + Some(UsageBodyCaptureState::Unavailable) + ); + } + #[test] fn masks_known_sensitive_header_values() { let token = "Bearer eyJhbGciOiJSUzI1NiJ9.payload-here.signature-tail";