diff --git a/Cargo.lock b/Cargo.lock index 3823839e9..a2171f728 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -305,6 +305,7 @@ dependencies = [ "futures-util", "hmac", "http", + "http-body", "http-body-util", "hyper", "hyper-util", diff --git a/apps/aether-gateway/Cargo.toml b/apps/aether-gateway/Cargo.toml index f664afcff..8337ac465 100644 --- a/apps/aether-gateway/Cargo.toml +++ b/apps/aether-gateway/Cargo.toml @@ -62,6 +62,7 @@ flate2.workspace = true futures-util.workspace = true hmac.workspace = true http.workspace = true +http-body = "1" http-body-util = "0.1" hyper = { version = "1", features = ["client", "server", "http1", "http2"] } hyper-util = { version = "0.1", features = ["client-legacy", "client-pool", "server-auto", "service", "tokio"] } diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 97b9d0f91..da0d04f08 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -13607,10 +13607,12 @@ mod tests { } #[tokio::test] - async fn execute_stream_from_frame_stream_cancels_upstream_when_client_drops_body() { - let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); - let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); - let state = AppState::new() + 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( @@ -13622,39 +13624,39 @@ mod tests { 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! { + 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", )); @@ -13669,119 +13671,157 @@ mod tests { } .boxed(); - let response = execute_stream_from_frame_stream( - &state, - plan, - "trace-client-drop-cancels-upstream", - &test_decision(), - "openai_chat_stream", - None, - Some(json!({ - "request_id": "req-client-drop-cancels-upstream", - "candidate_id": "cand-client-drop-cancels-upstream", - "candidate_index": 0, - "retry_index": 0, - "provider_api_format": "openai:chat", - "client_api_format": "openai:chat" - })), - crate::clock::current_unix_ms(), - Instant::now(), - RequestStageTrace::from_env(), - true, - frame_stream, - None, - ) - .await - .expect("execution should succeed") - .expect("execution should return a client response"); + 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; + 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!( + }) + .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); - let candidates = tokio::time::timeout(Duration::from_secs(1), async { - loop { - let candidates = request_candidate_repository - .list_by_request_id("req-client-drop-cancels-upstream") - .await - .expect("request candidates should read"); - if candidates - .first() - .is_some_and(|candidate| candidate.status == RequestCandidateStatus::Cancelled) - { - break candidates; - } - tokio::time::sleep(Duration::from_millis(10)).await; + tokio::time::sleep(Duration::from_millis(30)).await; + drop(body_stream); + if !cancel_on_client_disconnect { + release_terminal.notify_one(); } - }) - .await - .expect("candidate should be marked cancelled"); - assert_eq!(candidates[0].status_code, Some(499)); - assert_eq!( - candidates[0].error_type.as_deref(), - Some("downstream_disconnect") - ); - - let stored_usage = tokio::time::timeout(Duration::from_secs(1), async { - loop { - let usage = usage_repository - .find_by_request_id("req-client-drop-cancels-upstream") - .await - .expect("usage should read"); - if usage - .as_ref() - .is_some_and(|usage| usage.status == "cancelled") - { - break usage.expect("cancelled usage should exist"); + 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; } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("usage should be marked cancelled"); - 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" - ); + .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] diff --git a/apps/aether-gateway/src/executor/orchestration.rs b/apps/aether-gateway/src/executor/orchestration.rs index 7fd6ebca6..e79f96038 100644 --- a/apps/aether-gateway/src/executor/orchestration.rs +++ b/apps/aether-gateway/src/executor/orchestration.rs @@ -822,15 +822,20 @@ where let started_at = Instant::now(); let (tx, rx) = mpsc::channel::>(1); let request_diagnostics = current_request_diagnostics(); + let cancel_on_disconnect = crate::request_lifecycle::cancel_on_client_disconnect(); tokio::spawn(async move { scope_request_diagnostics_with(request_diagnostics, async move { - let bytes = standard_text_sync_heartbeat_final_bytes( + let completion = standard_text_sync_heartbeat_final_bytes( client_api_format.as_str(), redaction_slot.as_ref(), - execute(state, parts, trace_id, decision, plan_kind, started_at).await, - ) - .await; + tokio::select! { + biased; + _ = tx.closed(), if cancel_on_disconnect => return, + result = execute(state, parts, trace_id, decision, plan_kind, started_at) => result, + }, + ); + let bytes = completion.await; let _ = tx.send(Ok(Bytes::from(bytes))).await; }) .await; @@ -1097,23 +1102,26 @@ fn build_openai_image_sync_heartbeat_shell_response( let started_at = Instant::now(); let (tx, rx) = mpsc::channel::>(1); let request_diagnostics = current_request_diagnostics(); + let cancel_on_disconnect = crate::request_lifecycle::cancel_on_client_disconnect(); tokio::spawn(async move { scope_request_diagnostics_with(request_diagnostics, async move { - let bytes = openai_image_sync_heartbeat_final_bytes( - execute_openai_image_sync_heartbeat_attempts( - state, - request_path, - trace_id, - decision, - plan_kind, - attempts, - transfer_tracker, - started_at, - ) - .await, - ) - .await; + let execution = execute_openai_image_sync_heartbeat_attempts( + state, + request_path, + trace_id, + decision, + plan_kind, + attempts, + transfer_tracker, + started_at, + ); + let outcome = tokio::select! { + biased; + _ = tx.closed(), if cancel_on_disconnect => return, + result = execution => result, + }; + let bytes = openai_image_sync_heartbeat_final_bytes(outcome).await; let _ = tx.send(Ok(Bytes::from(bytes))).await; }) .await; @@ -2331,6 +2339,45 @@ mod tests { .expect("background completion should release admission"); } + #[tokio::test] + async fn standard_text_sync_heartbeat_cancels_when_routing_policy_enables_it() { + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (mut release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + let response = crate::request_lifecycle::run_request(async move { + crate::request_lifecycle::configure_client_disconnect( + aether_routing_core::RoutingExecutionPolicy { + cancel_on_client_disconnect: true, + ..Default::default() + }, + ); + let (parts, _) = http::Request::builder() + .method("POST") + .uri("/v1/responses") + .body(()) + .unwrap() + .into_parts(); + build_standard_text_sync_heartbeat_shell_response( + AppState::new().unwrap(), + parts, + "trace-heartbeat-disconnect".to_string(), + test_standard_text_heartbeat_decision(), + TEST_STANDARD_TEXT_SYNC_PLAN_KIND.to_string(), + move |_, _, _, _, _, _| async move { + started_tx.send(()).unwrap(); + release_rx.await.unwrap(); + Ok(LocalExecutionRequestOutcome::NoPath) + }, + ) + }) + .await + .unwrap(); + started_rx.await.unwrap(); + drop(response); + tokio::time::timeout(Duration::from_secs(1), release_tx.closed()) + .await + .expect("heartbeat must drop upstream execution immediately"); + } + #[tokio::test] async fn standard_text_sync_heartbeat_propagates_request_diagnostics_to_terminal_usage() { let (state, usage_repository) = heartbeat_usage_test_state(json!({ diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 5d8886ce4..0f8f77fd8 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -1035,7 +1035,7 @@ pub(crate) async fn proxy_request( ConnectInfo(remote_addr): ConnectInfo, request: Request, ) -> Result, GatewayError> { - crate::request_diagnostics::scope_request_diagnostics(Box::pin(proxy_request_inner( + crate::request_lifecycle::run_request(Box::pin(proxy_request_inner( state, remote_addr, request, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs index 5749b714c..23774c5f5 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs @@ -68,7 +68,12 @@ pub(super) async fn relay_bound_connection( state: &AppState, context: &WebSocketRequestContext, ) { + let mut client_connected = true; loop { + if !client_connected && !bound.turn_state.response_in_flight() { + close_bound_upstream(bound).await; + break; + } let active_turn_deadline = bound.turn_state.attempt().map(|turn| turn.deadline()); tokio::select! { _ = wait_for_optional_deadline(active_turn_deadline.map(|deadline| deadline.deadline)) => { @@ -100,8 +105,12 @@ pub(super) async fn relay_bound_connection( ).await; break; } - client_message = client_socket.next() => { + client_message = client_socket.next(), if client_connected => { let Some(client_message) = client_message else { + if retain_disconnected_turn(bound) { + client_connected = false; + continue; + } finalize_active_turn( bound, state, @@ -111,6 +120,10 @@ pub(super) async fn relay_bound_connection( break; }; let Ok(client_message) = client_message else { + if retain_disconnected_turn(bound) { + client_connected = false; + continue; + } warn!( event_name = "responses_websocket_client_receive_failed", log_type = "ops", @@ -127,6 +140,12 @@ pub(super) async fn relay_bound_connection( close_bound_upstream(bound).await; break; }; + if matches!(client_message, AxumWsMessage::Close(_)) + && retain_disconnected_turn(bound) + { + client_connected = false; + continue; + } match Box::pin(forward_client_message( client_message, bound, @@ -559,6 +578,7 @@ pub(super) async fn relay_bound_connection( let mut relay_send_error = None; let mut relay_serialization_failed = false; match relay_directive { + _ if !client_connected => {} Some(ResponsesWebSocketRelayDirective::ForwardOriginal) => { let client_frame = match parsed_upstream_frame.as_ref().map(|frame| { bound @@ -673,6 +693,10 @@ pub(super) async fn relay_bound_connection( break; } if let Some(error) = relay_send_error { + if terminal_outcome.is_none() && retain_disconnected_turn(bound) { + client_connected = false; + continue; + } warn!( event_name = "responses_websocket_client_send_failed", log_type = "ops", @@ -737,6 +761,20 @@ pub(super) async fn relay_bound_connection( } } +fn retain_disconnected_turn(bound: &mut BoundResponsesConnection) -> bool { + if bound + .turn_state + .attempt() + .is_none_or(|attempt| attempt.cancel_on_client_disconnect()) + { + return false; + } + bound + .turn_state + .record_client_delivery_aborted(CLIENT_DELIVERY_FAILED_REASON); + true +} + struct PendingContinuationRegistration { user_id: String, api_key_id: String, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs index 4fa9d7b07..34b73d12b 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs @@ -845,6 +845,13 @@ fn websocket_auth_rejection_error(rejection: GatewayLocalAuthRejection) -> Gatew } impl ResponsesProviderAttempt { + pub(super) fn cancel_on_client_disconnect(&self) -> bool { + crate::orchestration::routing_execution_policy_from_report_context( + self.lifecycle.report_context(), + ) + .is_some_and(|policy| policy.cancel_on_client_disconnect) + } + /// Releases all per-turn capacity before terminal persistence starts. /// Provider-pool runtime tokens normally use an awaited removal. The /// bounded wait prevents a broken runtime backend from stalling the relay; diff --git a/apps/aether-gateway/src/lib.rs b/apps/aether-gateway/src/lib.rs index 630717169..57bd703a2 100644 --- a/apps/aether-gateway/src/lib.rs +++ b/apps/aether-gateway/src/lib.rs @@ -71,6 +71,7 @@ mod rate_limit; mod request_candidate_queue; mod request_candidate_runtime; mod request_diagnostics; +mod request_lifecycle; mod roles; mod router; mod routing; diff --git a/apps/aether-gateway/src/request_lifecycle.rs b/apps/aether-gateway/src/request_lifecycle.rs new file mode 100644 index 000000000..107b1f7ea --- /dev/null +++ b/apps/aether-gateway/src/request_lifecycle.rs @@ -0,0 +1,336 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use aether_routing_core::RoutingExecutionPolicy; +use axum::body::{Body, Bytes, HttpBody}; +use http::Response; +use http_body::{Frame, SizeHint}; +use http_body_util::BodyExt; + +use crate::request_diagnostics::{scope_request_diagnostics_with, RequestDiagnostics}; +use crate::GatewayError; + +tokio::task_local! { + static CANCEL_ON_CLIENT_DISCONNECT: Arc; +} + +pub(crate) fn configure_client_disconnect(policy: RoutingExecutionPolicy) { + let _ = CANCEL_ON_CLIENT_DISCONNECT.try_with(|cancel| { + cancel.store(policy.cancel_on_client_disconnect, Ordering::Release); + }); +} + +pub(crate) fn cancel_on_client_disconnect() -> bool { + CANCEL_ON_CLIENT_DISCONNECT + .try_with(|cancel| cancel.load(Ordering::Acquire)) + .unwrap_or(false) +} + +pub(crate) async fn run_request(future: F) -> Result, GatewayError> +where + F: Future, GatewayError>> + Send + 'static, +{ + let cancel = Arc::new(AtomicBool::new(true)); + let diagnostics = Arc::new(RequestDiagnostics::default()); + let cancel_for_response = Arc::clone(&cancel); + let future = CANCEL_ON_CLIENT_DISCONNECT.scope( + Arc::clone(&cancel), + scope_request_diagnostics_with(Some(Arc::clone(&diagnostics)), async move { + let response = future.await?; + if cancel_for_response.load(Ordering::Acquire) { + return Ok(response); + } + Ok(response.map(|body| { + Body::new(CompleteOnDisconnectBody { + body: Some(body), + diagnostics, + }) + })) + }), + ); + CompleteOnDisconnectRequest { + future: Some(Box::pin(future)), + cancel, + } + .await +} + +struct CompleteOnDisconnectRequest +where + F: Future, GatewayError>> + Send + 'static, +{ + future: Option>>, + cancel: Arc, +} + +impl Future for CompleteOnDisconnectRequest +where + F: Future, GatewayError>> + Send + 'static, +{ + type Output = Result, GatewayError>; + + fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { + let result = self + .future + .as_mut() + .expect("request future") + .as_mut() + .poll(context); + if result.is_ready() { + self.future.take(); + } + result + } +} + +impl Drop for CompleteOnDisconnectRequest +where + F: Future, GatewayError>> + Send + 'static, +{ + fn drop(&mut self) { + if self.cancel.load(Ordering::Acquire) { + return; + } + if let (Some(future), Ok(runtime)) = + (self.future.take(), tokio::runtime::Handle::try_current()) + { + runtime.spawn(async move { + if let Ok(response) = future.await { + drain_body(response.into_body()).await; + } + }); + } + } +} + +struct CompleteOnDisconnectBody { + body: Option, + diagnostics: Arc, +} + +impl HttpBody for CompleteOnDisconnectBody { + type Data = Bytes; + type Error = axum::Error; + + fn poll_frame( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let Some(body) = self.body.as_mut() else { + return Poll::Ready(None); + }; + let result = Pin::new(body).poll_frame(context); + if matches!(result, Poll::Ready(None | Some(Err(_)))) { + self.body.take(); + } + result + } + + fn is_end_stream(&self) -> bool { + self.body.as_ref().is_none_or(HttpBody::is_end_stream) + } + + fn size_hint(&self) -> SizeHint { + self.body + .as_ref() + .map(HttpBody::size_hint) + .unwrap_or_else(|| SizeHint::with_exact(0)) + } +} + +impl Drop for CompleteOnDisconnectBody { + fn drop(&mut self) { + let Some(body) = self.body.take().filter(|body| !body.is_end_stream()) else { + return; + }; + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(scope_request_diagnostics_with( + Some(Arc::clone(&self.diagnostics)), + drain_body(body), + )); + } + } +} + +async fn drain_body(mut body: Body) { + while let Some(frame) = body.frame().await { + if frame.is_err() { + break; + } + } +} + +#[cfg(test)] +mod tests { + use std::io; + use std::time::Duration; + + use futures_util::stream; + use http::HeaderMap; + use http_body_util::StreamBody; + use tokio::sync::{mpsc, oneshot}; + + use super::*; + + #[tokio::test] + async fn disconnected_request_finishes_and_keeps_admission_and_diagnostics() { + let gate = aether_runtime::ConcurrencyGate::new("disconnect_request", 1); + let permit = gate.try_acquire().unwrap(); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let (finished_tx, finished_rx) = oneshot::channel(); + let request = tokio::spawn(run_request(async move { + let _permit = permit; + configure_client_disconnect(RoutingExecutionPolicy::default()); + started_tx.send(()).unwrap(); + release_rx.await.unwrap(); + assert!(crate::request_diagnostics::current_request_diagnostics().is_some()); + finished_tx.send(()).unwrap(); + Ok(Response::new(Body::empty())) + })); + started_rx.await.unwrap(); + request.abort(); + assert!(request.await.unwrap_err().is_cancelled()); + assert_eq!(gate.snapshot().in_flight, 1); + release_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(1), finished_rx) + .await + .unwrap() + .unwrap(); + assert_eq!(gate.snapshot().in_flight, 0); + } + + #[tokio::test] + async fn enabled_cancellation_and_unresolved_requests_drop_immediately() { + for resolve_policy in [false, true] { + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel::<()>(); + let request = tokio::spawn(run_request(async move { + if resolve_policy { + configure_client_disconnect(RoutingExecutionPolicy { + cancel_on_client_disconnect: true, + ..Default::default() + }); + } + started_tx.send(()).unwrap(); + release_rx.await.unwrap(); + Ok(Response::new(Body::empty())) + })); + started_rx.await.unwrap(); + request.abort(); + assert!(request.await.unwrap_err().is_cancelled()); + assert!(release_tx.send(()).is_err()); + } + } + + #[tokio::test] + async fn disconnected_body_drains_without_buffering_and_holds_admission() { + for consume_first_chunk in [false, true] { + let gate = aether_runtime::ConcurrencyGate::new("disconnect_body", 1); + let permit = gate.try_acquire().unwrap(); + let (sender, receiver) = mpsc::channel(1); + let (finished_tx, finished_rx) = oneshot::channel(); + let response = run_request(async move { + configure_client_disconnect(RoutingExecutionPolicy::default()); + let body = Body::from_stream(stream::unfold( + (receiver, finished_tx, permit), + |(mut receiver, finished_tx, permit)| async move { + match receiver.recv().await { + Some(bytes) => { + Some((Ok::<_, io::Error>(bytes), (receiver, finished_tx, permit))) + } + None => { + assert!(crate::request_diagnostics::current_request_diagnostics() + .is_some()); + finished_tx.send(()).unwrap(); + None + } + } + }, + )); + Ok(Response::new(body)) + }) + .await + .unwrap(); + let mut body = response.into_body(); + if consume_first_chunk { + sender.send(Bytes::from_static(b"first")).await.unwrap(); + assert_eq!( + body.frame().await.unwrap().unwrap().into_data().unwrap(), + "first" + ); + } + drop(body); + assert_eq!(gate.snapshot().in_flight, 1); + tokio::time::timeout(Duration::from_secs(1), async { + for _ in 0..100 { + sender.send(Bytes::from_static(b"remaining")).await.unwrap(); + } + drop(sender); + finished_rx.await.unwrap(); + }) + .await + .unwrap(); + assert_eq!(gate.snapshot().in_flight, 0); + } + } + + #[tokio::test] + async fn enabled_cancellation_drops_stream_receiver() { + let (sender, receiver) = mpsc::channel::>(1); + let response = run_request(async move { + configure_client_disconnect(RoutingExecutionPolicy { + cancel_on_client_disconnect: true, + ..Default::default() + }); + Ok(Response::new(Body::from_stream(stream::unfold( + receiver, + |mut receiver| async { receiver.recv().await.map(|item| (item, receiver)) }, + )))) + }) + .await + .unwrap(); + drop(response); + assert!(sender.is_closed()); + } + + #[tokio::test] + async fn connected_response_preserves_headers_size_hint_and_trailers() { + let response = run_request(async { + configure_client_disconnect(RoutingExecutionPolicy::default()); + Ok(Response::builder() + .status(201) + .header("x-test", "unchanged") + .body(Body::from("hello")) + .unwrap()) + }) + .await + .unwrap(); + assert_eq!(response.status(), 201); + assert_eq!(response.headers()["x-test"], "unchanged"); + assert_eq!(response.body().size_hint().exact(), Some(5)); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "hello" + ); + + let mut trailers = HeaderMap::new(); + trailers.insert("x-finished", "yes".parse().unwrap()); + let response = run_request(async move { + configure_client_disconnect(RoutingExecutionPolicy::default()); + let frames = stream::iter([ + Ok::<_, io::Error>(Frame::data(Bytes::from_static(b"hello"))), + Ok(Frame::trailers(trailers)), + ]); + Ok(Response::new(Body::new(StreamBody::new(frames)))) + }) + .await + .unwrap(); + let collected = response.into_body().collect().await.unwrap(); + assert_eq!(collected.trailers().unwrap()["x-finished"], "yes"); + assert_eq!(collected.to_bytes(), "hello"); + } +} diff --git a/apps/aether-gateway/src/routing/resolver.rs b/apps/aether-gateway/src/routing/resolver.rs index b636abd02..0c0904579 100644 --- a/apps/aether-gateway/src/routing/resolver.rs +++ b/apps/aether-gateway/src/routing/resolver.rs @@ -57,7 +57,7 @@ pub(crate) fn resolve_gateway_routing_policy( let config = serde_json::from_value::(input.group_config_json.clone()) .map_err(|_| invalid_routing_group_config())?; - resolve_routing_policy( + let policy = resolve_routing_policy( &config, RoutingPolicyInput { group_id: input.group_id, @@ -73,7 +73,9 @@ pub(crate) fn resolve_gateway_routing_policy( phase: input.phase, }, ) - .map_err(routing_policy_error) + .map_err(routing_policy_error)?; + crate::request_lifecycle::configure_client_disconnect(policy.execution_policy); + Ok(policy) } pub(crate) fn resolve_gateway_static_default_routing_policy( @@ -82,6 +84,7 @@ pub(crate) fn resolve_gateway_static_default_routing_policy( let Some(default_policy) = static_default_policy_fields(input.group_config_json)? else { return Ok(None); }; + crate::request_lifecycle::configure_client_disconnect(default_policy.execution_policy); Ok(Some(ResolvedRoutingPolicy { group_id: input.group_id.map(str::to_string), @@ -163,6 +166,10 @@ fn static_default_policy_fields( default_policy.get("cyber_continue_failover"), "cyber_continue_failover", )?, + cancel_on_client_disconnect: routing_bool_field( + default_policy.get("cancel_on_client_disconnect"), + "cancel_on_client_disconnect", + )?, }; Ok(Some(RoutingDefaultPolicy { @@ -238,7 +245,8 @@ mod tests { "default_policy": { "priority_mode": "global_key", "scheduling_mode": "load_balance", - "keep_priority_on_conversion": true + "keep_priority_on_conversion": true, + "cancel_on_client_disconnect": true }, "allowed_models": ["legacy-model"], "model_policies": [], diff --git a/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs b/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs index 6e26a8976..109748e3e 100644 --- a/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs +++ b/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs @@ -55,6 +55,24 @@ fn hash_api_key(value: &str) -> String { format!("{:x}", hasher.finalize()) } +async fn build_cancelling_gateway(state: crate::AppState) -> Router { + state + .data + .update_routing_group( + "system-default", + aether_data_contracts::repository::routing_profiles::UpdateRoutingGroupRecord { + config_json: Some(json!({"default_policy": {"cancel_on_client_disconnect": true}})), + version: Some(2), + updated_at: 2, + ..Default::default() + }, + ) + .await + .expect("routing policy should update") + .expect("default strategy should exist"); + build_router_with_state(state) +} + fn sample_local_openai_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot { StoredAuthApiKeySnapshot::new( user_id.to_string(), @@ -384,7 +402,7 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects_impl() { vec![sample_local_openai_key()], )); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); - let gateway = build_router_with_state( + let gateway = build_cancelling_gateway( 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( @@ -395,7 +413,7 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects_impl() { DEVELOPMENT_ENCRYPTION_KEY, ), ), - ); + ).await; let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() @@ -467,7 +485,7 @@ async fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byt vec![sample_local_openai_key()], )); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); - let gateway = build_router_with_state( + let gateway = build_cancelling_gateway( 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( @@ -478,7 +496,7 @@ async fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byt DEVELOPMENT_ENCRYPTION_KEY, ), ), - ); + ).await; let (gateway_url, gateway_handle) = start_server(gateway).await; let request = reqwest::Client::new() diff --git a/crates/aether-billing/src/event_enrichment.rs b/crates/aether-billing/src/event_enrichment.rs index a6dd1ccd1..d2267a0f9 100644 --- a/crates/aether-billing/src/event_enrichment.rs +++ b/crates/aether-billing/src/event_enrichment.rs @@ -1,8 +1,8 @@ use aether_data_contracts::repository::billing::StoredBillingModelContext; use aether_data_contracts::repository::usage::{ extract_provider_cache_ttl_minutes_from_metadata, resolve_provider_cache_ttl_minutes, - resolve_provider_service_tier_from_request_capture, USAGE_AVAILABLE_METADATA_KEY, - USAGE_PRICING_AVAILABLE_METADATA_KEY, + resolve_provider_service_tier_from_request_capture, CANCELLED_REQUEST_FEE_METADATA_KEY, + USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY, }; use aether_data_contracts::DataLayerError; use aether_usage_runtime::{UsageEvent, UsageEventType}; @@ -40,6 +40,18 @@ pub async fn enrich_usage_event_with_billing( data: &dyn BillingModelContextLookup, event: &mut UsageEvent, ) -> Result<(), DataLayerError> { + if matches!(event.event_type, UsageEventType::Cancelled) { + event.data.total_cost_usd = Some(0.0); + event.data.actual_total_cost_usd = Some(0.0); + if let Some(metadata) = event + .data + .request_metadata + .as_mut() + .and_then(Value::as_object_mut) + { + metadata.remove(CANCELLED_REQUEST_FEE_METADATA_KEY); + } + } // Session transports such as Codex Live expose lifecycle telemetry but no // authoritative token/cost object. Do not run request-based pricing with // zero default tokens: that would turn "unknown" into a fabricated charge. @@ -65,7 +77,10 @@ pub async fn enrich_usage_event_with_billing( clear_usage_costs(event); return Ok(()); } - if !matches!(event.event_type, UsageEventType::Completed) { + if !matches!( + event.event_type, + UsageEventType::Completed | UsageEventType::Cancelled + ) { event.data.total_cost_usd = Some(0.0); event.data.actual_total_cost_usd = Some(0.0); return Ok(()); @@ -189,7 +204,10 @@ fn calculate_billing_computation( } else { usage_event_image_count(&event.data).unwrap_or(0) }; - let request_count = if failed { + let cancelled = matches!(event.event_type, UsageEventType::Cancelled); + let request_count = if cancelled { + 1 + } else if failed { 0 } else if is_image_usage && image_count > 0 { image_count @@ -197,7 +215,7 @@ fn calculate_billing_computation( 1 }; let processing_tiers = usage_event_processing_tiers(&event.data); - let input = BillingUsageInput { + let mut input = BillingUsageInput { task_type: if is_image_usage { "image".to_string() } else { @@ -237,6 +255,16 @@ fn calculate_billing_computation( .or(pricing.provider_api_key_cache_ttl_minutes), }; + if cancelled { + input.input_tokens = 0; + input.output_tokens = 0; + input.cache_creation_tokens = 0; + input.cache_creation_ephemeral_5m_tokens = 0; + input.cache_creation_ephemeral_1h_tokens = 0; + input.cache_read_tokens = 0; + input.image_count = 0; + } + BillingService::new() .calculate(pricing, &input) .map_err(|err| { @@ -356,9 +384,32 @@ fn apply_billing_computation( pricing: &BillingModelPricingSnapshot, computation: BillingComputation, ) -> Result<(), DataLayerError> { + let cancelled = matches!(event.event_type, UsageEventType::Cancelled); + if cancelled + && !computation + .pricing_resolution + .price_per_request + .is_some_and(|price| price > 0.0) + { + return Ok(()); + } event.data.total_cost_usd = Some(computation.cost_result.cost); event.data.actual_total_cost_usd = Some(computation.actual_total_cost); - merge_billing_snapshot_metadata(&mut event.data.request_metadata, pricing, &computation) + merge_billing_snapshot_metadata(&mut event.data.request_metadata, pricing, &computation)?; + if cancelled { + if let Some(metadata) = event + .data + .request_metadata + .as_mut() + .and_then(Value::as_object_mut) + { + metadata.insert( + CANCELLED_REQUEST_FEE_METADATA_KEY.to_string(), + Value::Bool(true), + ); + } + } + Ok(()) } fn map_pricing_context(context: StoredBillingModelContext) -> BillingModelPricingSnapshot { @@ -1272,8 +1323,11 @@ mod tests { } #[tokio::test] - async fn cancelled_usage_event_remains_unbilled() { - let lookup = TestLookup { + async fn cancelled_usage_bills_only_configured_request_fee() { + for (request_type, request_price) in + [("chat", None), ("chat", Some(0.02)), ("image", Some(0.02))] + { + let lookup = TestLookup { name_context: Some( StoredBillingModelContext::new( "provider-1".to_string(), @@ -1284,7 +1338,7 @@ mod tests { "global-model-1".to_string(), "gpt-5".to_string(), None, - Some(0.02), + request_price, Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0,"cache_creation_price_per_1m":3.75,"cache_read_price_per_1m":0.30}]})), Some("model-1".to_string()), Some("gpt-5-upstream".to_string()), @@ -1296,61 +1350,69 @@ mod tests { ), model_id_context: None, }; - let mut event = UsageEvent::new( - UsageEventType::Cancelled, - "req-billing-cancelled-1", - UsageEventData { - provider_name: "OpenAI".to_string(), - model: "gpt-5".to_string(), - provider_id: Some("provider-1".to_string()), - provider_api_key_id: Some("key-1".to_string()), - request_type: Some("chat".to_string()), - api_format: Some("openai:responses".to_string()), - endpoint_api_format: Some("openai:responses".to_string()), - input_tokens: Some(1_000), - output_tokens: Some(500), - cache_read_input_tokens: Some(100), - status_code: Some(499), - ..UsageEventData::default() - }, - ); + let mut event = UsageEvent::new( + UsageEventType::Cancelled, + "req-billing-cancelled-1", + UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5".to_string(), + provider_id: Some("provider-1".to_string()), + provider_api_key_id: Some("key-1".to_string()), + request_type: Some(request_type.to_string()), + api_format: Some("openai:responses".to_string()), + endpoint_api_format: Some("openai:responses".to_string()), + input_tokens: Some(1_000), + output_tokens: Some(500), + cache_read_input_tokens: Some(100), + status_code: Some(499), + request_metadata: Some( + json!({"cancelled_request_fee": true, "image_count": 3}), + ), + ..UsageEventData::default() + }, + ); - enrich_usage_event_with_billing(&lookup, &mut event) - .await - .expect("billing should succeed"); + enrich_usage_event_with_billing(&lookup, &mut event) + .await + .expect("billing should succeed"); - assert_eq!(event.data.total_cost_usd, Some(0.0)); - assert_eq!(event.data.actual_total_cost_usd, Some(0.0)); - assert_eq!( - event - .data - .request_metadata - .as_ref() - .and_then(|value| value.get("billing_snapshot")) - .and_then(|value| value.get("status")) - .and_then(Value::as_str), - None - ); - assert_eq!( - event - .data - .request_metadata - .as_ref() - .and_then(|value| value.get("billing_dimensions")) - .and_then(|value| value.get("input_tokens")) - .and_then(Value::as_i64), - None - ); - assert_eq!( - event - .data - .request_metadata - .as_ref() - .and_then(|value| value.get("billing_dimensions")) - .and_then(|value| value.get("cache_read_tokens")) - .and_then(Value::as_i64), - None - ); + let expected_cost = request_price.unwrap_or(0.0); + assert_eq!(event.data.total_cost_usd, Some(expected_cost)); + assert_eq!(event.data.actual_total_cost_usd, Some(expected_cost * 0.5)); + assert_eq!(event.data.input_tokens, Some(1_000)); + assert_eq!(event.data.output_tokens, Some(500)); + let metadata = event.data.request_metadata.as_ref().unwrap(); + assert_eq!( + aether_data_contracts::repository::usage::cancelled_request_fee_is_billable(Some( + metadata + )), + request_price.is_some() + ); + if request_price.is_some() { + assert_eq!( + metadata.pointer("/billing_snapshot/cost_breakdown/request_cost"), + Some(&json!(expected_cost)) + ); + assert_eq!( + metadata.pointer("/billing_dimensions/input_tokens"), + Some(&json!(0)) + ); + assert_eq!( + metadata.pointer("/billing_dimensions/output_tokens"), + Some(&json!(0)) + ); + assert_eq!( + metadata.pointer("/billing_dimensions/cache_read_tokens"), + Some(&json!(0)) + ); + assert_eq!( + metadata.pointer("/billing_dimensions/request_count"), + Some(&json!(1)) + ); + } else { + assert!(metadata.get("billing_snapshot").is_none()); + } + } } #[tokio::test] diff --git a/crates/aether-data/contracts/src/repository/usage/metadata_policy.rs b/crates/aether-data/contracts/src/repository/usage/metadata_policy.rs index d9463ca9d..70b2fdca2 100644 --- a/crates/aether-data/contracts/src/repository/usage/metadata_policy.rs +++ b/crates/aether-data/contracts/src/repository/usage/metadata_policy.rs @@ -20,6 +20,14 @@ use super::{ const UPSTREAM_IS_STREAM_KEY: &str = "upstream_is_stream"; const PLAN_USAGE_RESERVATION_TOKEN_KEY: &str = "plan_usage_reservation_token"; const BODY_SIZE_BASIS: &str = "serialized gateway request bodies after normalization"; +pub const CANCELLED_REQUEST_FEE_METADATA_KEY: &str = "cancelled_request_fee"; + +pub fn cancelled_request_fee_is_billable(metadata: Option<&Value>) -> bool { + metadata + .and_then(|metadata| metadata.get(CANCELLED_REQUEST_FEE_METADATA_KEY)) + .and_then(Value::as_bool) + .unwrap_or(false) +} /// Projects request metadata onto the persistence contract. Unknown fields and malformed values /// are discarded instead of being recursively copied into an audit row. @@ -48,6 +56,7 @@ pub fn sanitize_usage_request_metadata_object(source: &Map) -> Op PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY, "transport_error", "is_free_tier", + CANCELLED_REQUEST_FEE_METADATA_KEY, USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY, ] { diff --git a/crates/aether-routing-core/src/model.rs b/crates/aether-routing-core/src/model.rs index 823c3e314..03b9412bf 100644 --- a/crates/aether-routing-core/src/model.rs +++ b/crates/aether-routing-core/src/model.rs @@ -39,6 +39,8 @@ pub struct RoutingExecutionPolicy { pub enable_cf_heartbeat: bool, #[serde(default, skip_serializing_if = "is_false")] pub cyber_continue_failover: bool, + #[serde(default, skip_serializing_if = "is_false")] + pub cancel_on_client_disconnect: bool, } impl<'de> Deserialize<'de> for RoutingExecutionPolicy { @@ -56,6 +58,8 @@ impl<'de> Deserialize<'de> for RoutingExecutionPolicy { enable_standard_text_sync_heartbeat: bool, #[serde(default)] cyber_continue_failover: bool, + #[serde(default)] + cancel_on_client_disconnect: bool, } let value = LegacyCompatibleExecutionPolicy::deserialize(deserializer)?; @@ -64,6 +68,7 @@ impl<'de> Deserialize<'de> for RoutingExecutionPolicy { || value.enable_openai_image_sync_heartbeat || value.enable_standard_text_sync_heartbeat, cyber_continue_failover: value.cyber_continue_failover, + cancel_on_client_disconnect: value.cancel_on_client_disconnect, }) } } @@ -107,6 +112,36 @@ fn is_false(value: &bool) -> bool { !*value } +#[cfg(test)] +mod execution_policy_tests { + use super::*; + + #[test] + fn cancellation_defaults_off_and_round_trips_with_legacy_heartbeat() { + let default: RoutingDefaultPolicy = serde_json::from_str("{}").unwrap(); + assert!(!default.execution_policy.cancel_on_client_disconnect); + let policy: RoutingDefaultPolicy = serde_json::from_value(serde_json::json!({ + "cancel_on_client_disconnect": true, + "enable_standard_text_sync_heartbeat": true + })) + .unwrap(); + assert!(policy.execution_policy.cancel_on_client_disconnect); + assert!(policy.execution_policy.enable_cf_heartbeat); + let encoded = serde_json::to_value(&policy).unwrap(); + assert_eq!(encoded["cancel_on_client_disconnect"], true); + assert_eq!( + serde_json::from_value::(encoded).unwrap(), + policy + ); + assert!( + serde_json::from_value::(serde_json::json!({ + "cancel_on_client_disconnect": "true" + })) + .is_err() + ); + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] pub struct RoutingModelPolicy { pub model: String, diff --git a/crates/aether-testing/integration/tests/responses_websocket_e2e.rs b/crates/aether-testing/integration/tests/responses_websocket_e2e.rs index afc501151..35215cfff 100644 --- a/crates/aether-testing/integration/tests/responses_websocket_e2e.rs +++ b/crates/aether-testing/integration/tests/responses_websocket_e2e.rs @@ -25,6 +25,9 @@ use aether_data_contracts::repository::global_models::{ use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; +use aether_data_contracts::repository::routing_profiles::{ + RoutingGroupLookupKey, UpdateRoutingGroupRecord, +}; use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageAuditListQuery}; use aether_gateway::{build_router_with_state, AppState, GatewayDataConfig, UsageRuntimeConfig}; use aether_testkit::{ManagedPostgresServer, SpawnedServer}; @@ -491,8 +494,7 @@ async fn disabling_the_downstream_key_is_enforced_on_the_next_turn_of_the_same_s Ok(()) } -/// A client that walks away before the provider produced anything must settle -/// as a void row: nothing was produced, so nothing is billed. +/// Immediate cancellation voids token billing before the provider completes. /// /// This is the path with no protocol event to announce it: the relay loop owns /// the turn, and losing the client is an exit the upstream never reports. @@ -506,7 +508,14 @@ async fn disabling_the_downstream_key_is_enforced_on_the_next_turn_of_the_same_s /// `a_closed_client_socket_before_any_terminal_still_voids_the_bill`. #[tokio::test] async fn client_disconnect_before_any_provider_output_settles_a_void_row() -> Result<(), BoxError> { - let harness = Harness::start(UpstreamBehavior::StallAfterCreated).await?; + let harness = Harness::start_configured( + UpstreamBehavior::StallAfterCreated, + ProviderFixture::SingleOpenAiKey, + PiiRedaction::Disabled, + true, + None, + ) + .await?; let mut client = harness.connect().await?; client @@ -542,6 +551,73 @@ async fn client_disconnect_before_any_provider_output_settles_a_void_row() -> Re Ok(()) } +#[tokio::test] +async fn client_disconnect_defaults_to_completing_and_billing_the_turn() -> Result<(), BoxError> { + let harness = Harness::start(UpstreamBehavior::CompleteAfterRelease).await?; + let mut client = harness.connect().await?; + client + .send(response_create(json!({"input": "finish without client"}))) + .await?; + receive_event(&mut client, "response.created").await?; + drop(client); + tokio::time::sleep(Duration::from_millis(50)).await; + harness.upstream.release_completion.notify_one(); + let audits = harness + .usage_audits_where(1, "completed disconnected turn", |audit| { + audit.status == "completed" && audit.billing_status == "settled" + }) + .await?; + assert_eq!(audits.len(), 1); + assert_eq!(audits[0].input_tokens, INPUT_TOKENS); + assert_eq!(audits[0].output_tokens, OUTPUT_TOKENS); + assert_eq!(audits[0].status_code, Some(200)); + Ok(()) +} + +#[tokio::test] +async fn client_disconnect_still_settles_the_per_request_fee_when_aborted() -> Result<(), BoxError> +{ + let harness = Harness::start_configured( + UpstreamBehavior::StallAfterCreated, + ProviderFixture::SingleOpenAiKey, + PiiRedaction::Disabled, + true, + Some(0.02), + ) + .await?; + let mut client = harness.connect().await?; + client + .send(response_create(json!({"input": "cancel with request fee"}))) + .await?; + receive_event(&mut client, "response.created").await?; + drop(client); + let audits = harness + .usage_audits_where(1, "cancelled request fee settlement", |audit| { + audit.status == "cancelled" && audit.billing_status == "settled" + }) + .await?; + assert_eq!(audits.len(), 1); + assert_eq!(audits[0].status_code, Some(499)); + assert_eq!(audits[0].total_tokens, 0); + assert_eq!(audits[0].total_cost_usd, 0.02); + assert_eq!(audits[0].actual_total_cost_usd, 0.02); + let backends = DataBackends::from_config(DataLayerConfig::from_database( + harness.database.config.clone(), + ))?; + let detail = backends + .read() + .usage() + .ok_or("usage reader unavailable")? + .find_by_request_id(&audits[0].request_id) + .await? + .ok_or("usage detail unavailable")?; + assert_eq!( + detail.request_metadata.as_ref().unwrap()["cancelled_request_fee"], + true + ); + Ok(()) +} + /// An upstream that dies mid-turn must surface an error and still settle. #[tokio::test] async fn upstream_drop_mid_turn_reports_an_error_and_settles_the_usage_row() -> Result<(), BoxError> @@ -850,6 +926,16 @@ impl Harness { behavior: UpstreamBehavior, fixture: ProviderFixture, redaction: PiiRedaction, + ) -> Result { + Self::start_configured(behavior, fixture, redaction, false, None).await + } + + async fn start_configured( + behavior: UpstreamBehavior, + fixture: ProviderFixture, + redaction: PiiRedaction, + cancel_on_client_disconnect: bool, + request_price: Option, ) -> Result { let upstream = Arc::new(MockUpstreamState::new(behavior)); let upstream_server = @@ -864,6 +950,16 @@ impl Harness { ) .await?; + if let Some(price) = request_price { + let pool = sqlx::PgPool::connect(&database.config.url).await?; + sqlx::query("UPDATE models SET price_per_request = $1 WHERE id = $2") + .bind(price) + .bind(PROVIDER_MODEL_ID) + .execute(&pool) + .await?; + pool.close().await; + } + let data_config = GatewayDataConfig::from_database_config(database.config.clone()) .with_encryption_key(DEVELOPMENT_ENCRYPTION_KEY); let state = AppState::new()? @@ -877,6 +973,32 @@ impl Harness { ..UsageRuntimeConfig::default() })?; state.ensure_system_default_routing_group().await?; + if cancel_on_client_disconnect { + let backends = + DataBackends::from_config(DataLayerConfig::from_database(database.config.clone()))?; + let mut group = backends + .read() + .routing_groups() + .ok_or("routing reader unavailable")? + .find_routing_group(RoutingGroupLookupKey::SystemDefault) + .await? + .ok_or("default routing group unavailable")?; + group.config_json["default_policy"]["cancel_on_client_disconnect"] = json!(true); + backends + .write() + .routing_groups() + .ok_or("routing writer unavailable")? + .update_routing_group( + &group.id, + UpdateRoutingGroupRecord { + config_json: Some(group.config_json), + version: Some(group.version + 1), + updated_at: group.updated_at + 1, + ..Default::default() + }, + ) + .await?; + } let gateway_server = SpawnedServer::start(build_router_with_state(state)).await?; let websocket_url = format!( "{}/v1/responses", @@ -1093,6 +1215,7 @@ where enum UpstreamBehavior { /// Announce, stream one delta, and complete — the ordinary turn. CompleteEveryTurn, + CompleteAfterRelease, /// Announce the response and then go quiet, leaving the turn in flight. StallAfterCreated, /// Announce the response and then hang up mid-turn. @@ -1118,6 +1241,7 @@ struct MockUpstreamState { events: Mutex>, authorization_headers: Mutex>>, handshakes: Mutex>, + release_completion: tokio::sync::Notify, } #[derive(Debug, Clone)] @@ -1134,6 +1258,7 @@ impl MockUpstreamState { events: Mutex::new(Vec::new()), authorization_headers: Mutex::new(Vec::new()), handshakes: Mutex::new(Vec::new()), + release_completion: tokio::sync::Notify::new(), } } @@ -1215,6 +1340,15 @@ async fn run_mock_upstream( break; } } + UpstreamBehavior::CompleteAfterRelease => { + if send_mock_created(&mut socket, &response_id).await.is_err() { + break; + } + state.release_completion.notified().await; + if send_mock_turn(&mut socket, &response_id).await.is_err() { + break; + } + } UpstreamBehavior::StallAfterCreated => { if send_mock_created(&mut socket, &response_id).await.is_err() { break; diff --git a/crates/aether-usage/runtime/src/record.rs b/crates/aether-usage/runtime/src/record.rs index 62a8feefd..5b8ba7a20 100644 --- a/crates/aether-usage/runtime/src/record.rs +++ b/crates/aether-usage/runtime/src/record.rs @@ -221,6 +221,13 @@ fn lifecycle_status_and_billing( } UsageEventType::Completed => ("completed", "pending"), UsageEventType::Failed => ("failed", "void"), + UsageEventType::Cancelled + if aether_data_contracts::repository::usage::cancelled_request_fee_is_billable( + request_metadata, + ) => + { + ("cancelled", "pending") + } UsageEventType::Cancelled => ("cancelled", "void"), } } @@ -491,6 +498,31 @@ mod tests { assert_eq!(record.first_byte_time_ms, Some(50)); } + #[test] + fn cancelled_request_fee_keeps_cancelled_status_and_pending_billing() { + let event = UsageEvent::new( + UsageEventType::Cancelled, + "req-cancelled-fee", + UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5".to_string(), + total_cost_usd: Some(0.02), + actual_total_cost_usd: Some(0.01), + request_metadata: Some(serde_json::json!({"cancelled_request_fee": true})), + ..Default::default() + }, + ); + let record = build_upsert_usage_record_from_event(&event).unwrap(); + assert_eq!(record.status, "cancelled"); + assert_eq!(record.billing_status, "pending"); + assert_eq!(record.total_cost_usd, Some(0.02)); + assert_eq!(record.actual_total_cost_usd, Some(0.01)); + assert_eq!( + record.request_metadata.unwrap()["cancelled_request_fee"], + true + ); + } + #[test] fn completed_unmetered_session_audit_is_void_without_fabricated_usage() { let record = build_upsert_usage_record_from_event(&UsageEvent { diff --git a/crates/aether-usage/runtime/src/settlement.rs b/crates/aether-usage/runtime/src/settlement.rs index 9d57610ea..89669a060 100644 --- a/crates/aether-usage/runtime/src/settlement.rs +++ b/crates/aether-usage/runtime/src/settlement.rs @@ -5,8 +5,10 @@ use aether_data_contracts::repository::settlement::{ ReconcileUsagePolicyCostInput, StoredUsagePolicyCostReservation, StoredUsageSettlement, UsagePolicyCostReservationState, UsageSettlementInput, }; -use aether_data_contracts::repository::usage::StoredRequestUsageAudit; use aether_data_contracts::repository::usage::PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY; +use aether_data_contracts::repository::usage::{ + cancelled_request_fee_is_billable, StoredRequestUsageAudit, +}; use aether_data_contracts::{DataLayerError, DataLayerError::InvalidInput}; use async_trait::async_trait; @@ -39,6 +41,11 @@ pub async fn reconcile_usage_policy_cost_for_event( } let terminal_state = match event.event_type { UsageEventType::Completed => UsagePolicyCostReservationState::Finalized, + UsageEventType::Cancelled + if cancelled_request_fee_is_billable(event.data.request_metadata.as_ref()) => + { + UsagePolicyCostReservationState::Finalized + } UsageEventType::Failed | UsageEventType::Cancelled => { UsagePolicyCostReservationState::Released } @@ -107,7 +114,10 @@ pub async fn settle_usage_if_needed( usage.user_id.as_deref().and_then(non_empty_trimmed), usage_policy_reservation_token(usage), ) { - let (terminal_state, actual_cost_units) = if usage.status == "completed" { + let (terminal_state, actual_cost_units) = if usage.status == "completed" + || (usage.status == "cancelled" + && cancelled_request_fee_is_billable(usage.request_metadata.as_ref())) + { ( UsagePolicyCostReservationState::Finalized, nonnegative_usd_to_usage_policy_cost_units( @@ -136,7 +146,10 @@ pub async fn settle_usage_if_needed( } } - if usage.status == "cancelled" || usage.billing_status != "pending" { + if usage.billing_status != "pending" + || (usage.status == "cancelled" + && !cancelled_request_fee_is_billable(usage.request_metadata.as_ref())) + { return Ok(()); } let input = UsageSettlementInput { @@ -429,6 +442,61 @@ mod tests { ); } + #[tokio::test] + async fn cancelled_request_fee_settles_wallet_and_finalizes_cost_reservation() { + let writer = TestSettlementWriter { + has_writer: true, + ..Default::default() + }; + let mut usage = sample_usage(); + usage.status = "cancelled".to_string(); + usage.status_code = Some(499); + usage.request_metadata.as_mut().unwrap()["cancelled_request_fee"] = json!(true); + settle_usage_if_needed(&writer, &usage).await.unwrap(); + let inputs = writer.inputs.lock().unwrap(); + assert_eq!(inputs.len(), 1); + assert_eq!(inputs[0].status, "cancelled"); + assert_eq!(inputs[0].actual_total_cost_usd, usage.actual_total_cost_usd); + let reconciliations = writer.reconciliations.lock().unwrap(); + assert_eq!(reconciliations.len(), 1); + assert_eq!( + reconciliations[0].terminal_state, + UsagePolicyCostReservationState::Finalized + ); + assert_eq!(reconciliations[0].actual_cost_units, 75_000_000); + } + + #[tokio::test] + async fn cancelled_request_fee_event_finalizes_cost_reservation() { + let writer = TestSettlementWriter { + has_writer: true, + ..Default::default() + }; + let event = UsageEvent::new( + UsageEventType::Cancelled, + "req-cancelled-fee", + UsageEventData { + user_id: Some("user-1".to_string()), + actual_total_cost_usd: Some(0.01), + request_metadata: Some(json!({ + "cancelled_request_fee": true, + "plan_usage_reservation_token": "server-token" + })), + ..Default::default() + }, + ); + reconcile_usage_policy_cost_for_event(&writer, &event) + .await + .unwrap(); + let reconciliations = writer.reconciliations.lock().unwrap(); + assert_eq!(reconciliations.len(), 1); + assert_eq!( + reconciliations[0].terminal_state, + UsagePolicyCostReservationState::Finalized + ); + assert_eq!(reconciliations[0].actual_cost_units, 1_000_000); + } + #[tokio::test] async fn releases_failed_usage_before_void_settlement() { let writer = TestSettlementWriter { diff --git a/docs/operations/client-disconnect-policy.md b/docs/operations/client-disconnect-policy.md new file mode 100644 index 000000000..fabb8f434 --- /dev/null +++ b/docs/operations/client-disconnect-policy.md @@ -0,0 +1,39 @@ +# 调度策略：取消请求立即打断 + +管理入口：**调度策略配置 → 系统配置 → 取消请求立即打断**。 + +配置保存在当前策略的 `config_json.default_policy.cancel_on_client_disconnect`，默认 `false`。 +旧策略缺失此字段也按关闭处理；不需要数据库结构迁移，不读取同名全局系统设置。 +策略选定后，请求沿用该次解析的配置，不因管理员随后修改策略而改变断连处理。 + +| 配置 | 客户端取消／断连时 | 计费 | +| --- | --- | --- | +| 关闭（默认） | 已选定策略的请求继续执行，后台读取响应直到正常完成或原有超时／上游错误 | 按实际完成结果正常结算 | +| 开启 | 中止仍在执行的请求、停止读取上游响应 | Token、缓存 Token 和图片产出费用不收取；配置了 `price_per_request` 的请求保留一次请求费用及对应倍率 | + +已经取得上游终态的请求不会因最后一跳投递失败而撤销已完成的结算。 +取消按次收费的记录仍显示 `cancelled`／499，但计费状态可为 `settled`；不要仅凭请求状态判断是否收费。 + +## 执行边界 + +- HTTP 同步、SSE、同格式直通和 CF 心跳响应共用请求生命周期保护。 +- Responses WebSocket 断连后只完成当前进行中的 turn，不无限维持空闲连接；开启开关则直接终止当前 turn。 +- 鉴权、请求体接收等尚未选定策略的阶段仍可直接取消。 +- Live／Realtime 长连接会话关闭、异步任务显式取消不是有限 HTTP／Responses 请求的断连续跑，不转成后台常驻会话。 +- 原有上游总超时、首字节超时及故障处理仍有效；这里不创建持久化后台作业，进程退出不能保证继续执行。 + +## 代码调整 + +- `request_lifecycle.rs` 将 HTTP 请求 Future 和响应 Body 的所有权与客户端连接解耦。正常连接保留原来的逐帧路径、响应头、长度提示和 trailers；仅断连时启动后台接管，逐帧丢弃待发送内容，不聚合完整响应。 +- 请求准入凭证随原 Future／Body 保留至完成，防止断连提前释放并发额度；诊断上下文同时保留。 +- CF 心跳后台执行显式监听响应接收端关闭，避免打开开关后仍继续执行。 +- Responses WebSocket 将“客户端已断开”和“上游已完成”分别处理，复用既有终态观察、超时和结算逻辑。 +- 计费统一在 billing enrichment 中计算取消请求的单次费用。`cancelled_request_fee` 由服务端计算后标记，贯穿审计持久化、钱包结算及套餐成本预留结算，避免有费用却仍写为 `void` 或释放成本预留。 + +## 回归覆盖 + +- 新旧配置默认关闭、布尔校验、策略保存及模型级调度编辑保留开关。 +- 响应头前断连、首帧前／首帧后断连、立即取消、并发额度、诊断信息、响应头／trailers 透传。 +- 真实流式执行链路的完成／取消 usage 与候选状态，以及心跳取消。 +- 取消时按次／按 Token／混合定价、图片请求单次费、倍率、审计状态、钱包及成本预留结算。 +- Responses WebSocket 使用真实网关和临时 PostgreSQL 验证断连继续完成、开启后不收费，以及按次费用实际结算。 diff --git a/frontend/src/features/routing/__tests__/routingPolicy.spec.ts b/frontend/src/features/routing/__tests__/routingPolicy.spec.ts index 87ce56635..ed3008f98 100644 --- a/frontend/src/features/routing/__tests__/routingPolicy.spec.ts +++ b/frontend/src/features/routing/__tests__/routingPolicy.spec.ts @@ -24,6 +24,20 @@ describe('routingPolicy', () => { expect(config.default_policy.priority_mode).toBe('provider') expect(config.default_policy.scheduling_mode).toBe('cache_affinity') + expect(config.default_policy.cancel_on_client_disconnect).toBe(false) + }) + + it('preserves cancellation policy across model scheduling edits', () => { + const config = createEmptyRoutingGroupConfig() + config.default_policy.cancel_on_client_disconnect = true + const updated = upsertModelSchedulingRule(config, 'gpt-5', { + priority_mode: 'global_key', + scheduling_mode: 'fixed_order', + }) + expect(normalizeRoutingGroupConfig(updated).default_policy.cancel_on_client_disconnect).toBe(true) + expect(getModelScheduling(updated, 'gpt-5').cancel_on_client_disconnect).toBe(true) + expect(getModelScheduling(updated, 'other-model').cancel_on_client_disconnect).toBe(true) + expect(createEmptyRoutingGroupConfig().default_policy.cancel_on_client_disconnect).toBe(false) }) it('drops the legacy group model allowlist while normalizing config', () => { @@ -77,7 +91,7 @@ describe('routingPolicy', () => { expect(createEmptyRoutingGroupConfig().default_policy.sticky_key_attempts).toBe(2) expect(normalizeRoutingGroupConfig({}).default_policy.sticky_key_attempts).toBe(2) expect(normalizeRoutingGroupConfig({ - default_policy: { priority_mode: 'provider', scheduling_mode: 'cache_affinity', keep_priority_on_conversion: false, sticky_key_attempts: 3, enable_cf_heartbeat: false, cyber_continue_failover: false }, + default_policy: { priority_mode: 'provider', scheduling_mode: 'cache_affinity', keep_priority_on_conversion: false, sticky_key_attempts: 3, enable_cf_heartbeat: false, cyber_continue_failover: false, cancel_on_client_disconnect: false }, }).default_policy.sticky_key_attempts).toBe(3) expect(normalizeStickyKeyAttempts('5')).toBe(5) expect(normalizeStickyKeyAttempts(-1)).toBe(2) diff --git a/frontend/src/features/routing/utils/routingPolicy.ts b/frontend/src/features/routing/utils/routingPolicy.ts index 92f3906c4..7b048c0cb 100644 --- a/frontend/src/features/routing/utils/routingPolicy.ts +++ b/frontend/src/features/routing/utils/routingPolicy.ts @@ -12,6 +12,7 @@ export interface RoutingDefaultPolicy { keep_priority_on_conversion: boolean enable_cf_heartbeat: boolean cyber_continue_failover: boolean + cancel_on_client_disconnect: boolean /** 首个候选的总尝试次数；后续候选始终只尝试 1 次。0 或 1 表示不重试 */ sticky_key_attempts: number } @@ -78,6 +79,7 @@ export function createEmptyRoutingGroupConfig(): RoutingGroupConfig { keep_priority_on_conversion: false, enable_cf_heartbeat: false, cyber_continue_failover: false, + cancel_on_client_disconnect: false, sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS, }, model_policies: [], @@ -374,6 +376,7 @@ export function getModelScheduling( keep_priority_on_conversion: normalized.default_policy.keep_priority_on_conversion, enable_cf_heartbeat: normalized.default_policy.enable_cf_heartbeat, cyber_continue_failover: normalized.default_policy.cyber_continue_failover, + cancel_on_client_disconnect: normalized.default_policy.cancel_on_client_disconnect, sticky_key_attempts: action?.sticky_key_attempts ?? normalized.default_policy.sticky_key_attempts, } } diff --git a/frontend/src/i18n/legacy-admin-messages.ts b/frontend/src/i18n/legacy-admin-messages.ts index 55b60b584..bad7721f1 100644 --- a/frontend/src/i18n/legacy-admin-messages.ts +++ b/frontend/src/i18n/legacy-admin-messages.ts @@ -114,6 +114,8 @@ export const legacyAdminEnglishMessages: Record = { '错误重试次数': 'Error retry count', 'CF保持心跳': 'CF keepalive', 'Cyber继续转移': 'Fail over on Cyber errors', + '取消请求立即打断': 'Abort immediately on client cancellation', + '默认关闭：客户端取消或断开后，服务端继续等待请求完成并正常计费。开启后立即打断且不计费，按次计费的请求仍收取单次请求费用。仅作用于当前调度策略。': 'Off by default: after client cancellation or disconnection, the server completes the request and bills normally. When enabled, requests are aborted immediately without usage charges, except for per-request fees. Applies only to this routing policy.', '调度配置': 'Routing configuration', '先选择调度维度，再配置优先级模式、调度策略和提供商排序。': 'Choose the routing scope, then configure the priority mode, routing strategy, and provider order.', '调度维度': 'Routing scope', diff --git a/frontend/src/views/admin/RoutingProfiles.vue b/frontend/src/views/admin/RoutingProfiles.vue index 00e59b2cd..47699728b 100644 --- a/frontend/src/views/admin/RoutingProfiles.vue +++ b/frontend/src/views/admin/RoutingProfiles.vue @@ -402,7 +402,7 @@ 这些选项作用于当前调度策略。

-
+
+
+
+ 取消请求立即打断 + +
+ +
@@ -970,6 +988,9 @@ const cfHeartbeat = computed(() => ( const cyberContinueFailover = computed(() => ( draft.value?.config_json.default_policy.cyber_continue_failover ?? false )) +const cancelOnClientDisconnect = computed(() => ( + draft.value?.config_json.default_policy.cancel_on_client_disconnect ?? false +)) interface ModelRow { name: string displayName: string @@ -1306,7 +1327,7 @@ function updateKeepPriorityOnConversion(value: boolean): void { } function updateExecutionPolicy( - field: 'enable_cf_heartbeat' | 'cyber_continue_failover', + field: 'enable_cf_heartbeat' | 'cyber_continue_failover' | 'cancel_on_client_disconnect', value: boolean, ): void { if (!draft.value) return