merge(main): sync latest main into security branch

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<String>,
+    candidate: Option<LocalRequestCandidateStatusSnapshot>,
+    // Boxed: the guard lives inside the stream request future, which is already
+    // very large, and `UsageEventData` is a wide struct.
+    usage_seed: Option<Box<UsageEventData>>,
+    request_diagnostics: Option<Arc<RequestDiagnostics>>,
+    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<Arc<StreamCandidateWatchdogProgress>>,
+    armed: Option<ArmedAttempt>,
+}
+
+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<Value> {
+        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<InMemoryUsageReadRepository>,
+        request_candidate_repository: &Arc<InMemoryRequestCandidateRepository>,
+    ) -> 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<Value>,
+        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<StoredRequestUsageAudit> {
+        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());
+    }
+}

apps/aether-gateway/src/execution_runtime/fallback.rs

@@ -156,7 +156,22 @@ pub(crate) fn should_fallback_to_control_sync(
         return true;
     };
 
-    body_json.get("error").is_some()
+    sync_body_has_embedded_error(Some(body_json))
+}
+
+/// Mirrors the error-like body markers used by the formats layer. Successful OpenAI Responses
+/// bodies contain `"error": null`, which must not route them through error finalization.
+fn sync_body_has_embedded_error(body_json: Option<&serde_json::Value>) -> bool {
+    let Some(object) = body_json.and_then(serde_json::Value::as_object) else {
+        return false;
+    };
+
+    object.get("error").is_some_and(|error| !error.is_null())
+        || object.get("status").and_then(serde_json::Value::as_str) == Some("failed")
+        || object
+            .get("type")
+            .and_then(serde_json::Value::as_str)
+            .is_some_and(|value| value == "error")
 }
 
 pub(crate) fn should_finalize_sync_response(report_kind: Option<&str>) -> bool {
@@ -168,7 +183,7 @@ pub(crate) fn resolve_core_sync_error_finalize_report_kind(
     result: &ExecutionResult,
     body_json: Option<&serde_json::Value>,
 ) -> Option<String> {
-    let has_embedded_error = body_json.is_some_and(|value| value.get("error").is_some());
+    let has_embedded_error = sync_body_has_embedded_error(body_json);
     if result.status_code < 400 && !has_embedded_error {
         return None;
     }
@@ -510,6 +525,74 @@ mod tests {
         );
     }
 
+    #[test]
+    fn successful_responses_body_with_null_error_stays_on_success_path() {
+        let result = ExecutionResult {
+            request_id: "req-1".to_string(),
+            candidate_id: None,
+            status_code: 200,
+            headers: Default::default(),
+            response_observation: None,
+            body: None,
+            telemetry: None,
+            error: None,
+        };
+        let body_json = serde_json::json!({
+            "id": "resp_1",
+            "object": "response",
+            "status": "completed",
+            "error": null,
+            "output": [],
+        });
+
+        assert_eq!(
+            resolve_core_sync_error_finalize_report_kind(
+                "openai_responses_sync",
+                &result,
+                Some(&body_json)
+            ),
+            None
+        );
+        assert!(!should_fallback_to_control_sync(
+            "openai_responses_sync",
+            &result,
+            Some(&body_json),
+            true,
+            false,
+            false,
+        ));
+    }
+
+    #[test]
+    fn error_like_success_status_bodies_still_map_to_error_finalize() {
+        let result = ExecutionResult {
+            request_id: "req-1".to_string(),
+            candidate_id: None,
+            status_code: 200,
+            headers: Default::default(),
+            response_observation: None,
+            body: None,
+            telemetry: None,
+            error: None,
+        };
+
+        for body_json in [
+            serde_json::json!({"status": "failed", "error": null}),
+            serde_json::json!({"type": "error"}),
+            serde_json::json!({"error": {"message": "boom"}}),
+        ] {
+            assert_eq!(
+                resolve_core_sync_error_finalize_report_kind(
+                    "openai_responses_sync",
+                    &result,
+                    Some(&body_json)
+                ),
+                Some("openai_responses_sync_finalize".to_string()),
+                "error-like body must not escape through the success path: {body_json}"
+            );
+        }
+    }
+
     #[test]
     fn stream_failover_marks_chat_errors() {
         assert!(should_fallback_to_control_stream(

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;

apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs

@@ -455,7 +455,10 @@ fn gemini_part_is_client_semantic(part: &Value) -> bool {
         return true;
     }
     if part.get("thought").and_then(Value::as_bool) == Some(true) {
-        return false;
+        return part
+            .get("text")
+            .and_then(Value::as_str)
+            .is_some_and(|text| !text.is_empty());
     }
     if part.keys().all(|key| key == "thoughtSignature") {
         return false;
@@ -584,17 +587,12 @@ mod tests {
     }
 
     #[test]
-    fn gemini_gate_waits_through_thought_and_commits_on_text() {
+    fn gemini_gate_commits_on_first_nonempty_thought() {
         let mut gate = StreamCommitGate::new(gemini_policy());
         let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"checking\"}]}}]}}\n\n";
-        let text = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"answer\"}]}}]}}\n\n";
 
         assert_eq!(
             gate.observe_provider_bytes(thought),
-            StreamPrecommitObservation::Pending
-        );
-        assert_eq!(
-            gate.observe_provider_bytes(text),
             StreamPrecommitObservation::Commit
         );
         assert_eq!(gate.state(), StreamCommitState::Committed);
@@ -615,7 +613,7 @@ mod tests {
     #[test]
     fn gemini_gate_rejects_malformed_function_call_before_commit() {
         let mut gate = StreamCommitGate::new(gemini_policy());
-        let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"calling\"}]}}]}}\n\n";
+        let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thoughtSignature\":\"signature\",\"text\":\"\"}]}}]}}\n\n";
         let malformed = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thoughtSignature\":\"signature\",\"text\":\"\"}]},\"finishReason\":\"MALFORMED_FUNCTION_CALL\",\"finishMessage\":\"Malformed function call: Function call is empty - no input to parse.\"}]}}\n\n";
 
         assert_eq!(

apps/aether-gateway/src/execution_runtime/stream/execution.rs

@@ -79,6 +79,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;
@@ -151,6 +152,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;
@@ -3718,17 +3724,30 @@ pub(crate) fn execute_execution_runtime_stream<'a>(
     report_kind: Option<String>,
     report_context: Option<serde_json::Value>,
 ) -> Pin<Box<dyn Future<Output = Result<Option<Response<Body>>, 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)]
@@ -3750,7 +3769,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,
@@ -3760,8 +3784,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 {
@@ -3807,6 +3836,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,
@@ -3817,6 +3847,7 @@ async fn execute_execution_runtime_stream_inner(
     mut report_context: Option<serde_json::Value>,
     mut retry_scope_out: Option<&mut AiAttemptRetryScope>,
     mut retry_fallback_out: Option<&mut Option<Response<Body>>>,
+    cancellation_guard: &mut AttemptCancellationGuard,
 ) -> Result<Option<Response<Body>>, GatewayError> {
     let stream_started_at = Instant::now();
     let mut stage_trace = RequestStageTrace::from_env();
@@ -3900,6 +3931,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
@@ -10764,7 +10805,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn malformed_antigravity_function_call_retries_before_stream_commit() {
+    async fn malformed_antigravity_function_call_streams_thought_then_fails_in_band() {
         let request_id = "req-antigravity-malformed-function-call";
         let plan = antigravity_gemini_stream_plan(request_id);
         let provider_catalog = provider_catalog_for_plan(
@@ -10843,10 +10884,35 @@ mod tests {
             None,
         )
         .await
-        .expect("malformed Antigravity stream should resolve through failover");
+        .expect("malformed Antigravity stream should return a client stream")
+        .expect("the first reasoning delta should commit the selected candidate");
 
-        assert!(response.is_none());
-        assert_eq!(retry_scope, AiAttemptRetryScope::Candidate);
+        assert_eq!(response.status(), StatusCode::OK);
+        let body = to_bytes(response.into_body(), usize::MAX)
+            .await
+            .expect("response body should read");
+        let body = String::from_utf8(body.to_vec()).expect("response body should be utf8");
+        assert!(
+            body.contains("event: response.reasoning_summary_text.delta\n"),
+            "{body}"
+        );
+        assert!(
+            body.contains("\"delta\":\"Validating the document.\""),
+            "{body}"
+        );
+        assert!(body.contains("event: response.failed\n"), "{body}");
+        assert!(
+            body.contains("\"code\":\"MALFORMED_FUNCTION_CALL\""),
+            "{body}"
+        );
+        assert!(
+            body.contains(
+                "\"message\":\"Malformed function call: Function call is empty - no input to parse.\""
+            ),
+            "{body}"
+        );
+        assert!(!body.contains("unsupported_finish_reason"), "{body}");
+        assert_eq!(retry_scope, AiAttemptRetryScope::Provider);
     }
 
     fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {

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<Arc<Self>> {
+        STREAM_CANDIDATE_WATCHDOG_PROGRESS.try_with(Arc::clone).ok()
+    }
+
     pub(crate) async fn scope<F>(self: Arc<Self>, future: F) -> F::Output
     where
         F: Future,