fix(gateway): settle stream attempts dropped before first byte

A local stream attempt writes its `usage` row and its `request_candidates`
slot as `pending` in `execute_execution_runtime_stream_inner`, then awaits
the provider's response headers. Everything after that point runs inside
the downstream request future, so a client disconnect drops it: the
dispatch `.await` never resumes and nothing settles either row. The stream
finalizer that already covers this only exists once upstream headers have
arrived, so the pre-first-byte window has no owner at all. Both rows stay
`pending` until the maintenance sweeper rewrites them as a 504 timeout ten
minutes later, losing the real outcome, the real latency, and the 499.

`AttemptCancellationGuard` takes that window. It is created disarmed, so
an attempt dropped before it owns any row does not grow a settlement row
it never had; it is armed as soon as the attempt owns its non-terminal
rows, and the stream wrappers disarm it the moment the attempt returns,
from where settlement belongs to the transport. On a cancelling drop it
settles the candidate slot through the same snapshot writer the `pending`
write above it uses, and the usage row through a terminal `Cancelled`
event.

The guard outlives the request future, so what it captures is retained for
the whole attempt. It therefore holds no request body: the plan carries the
provider request body and the report context carries the client request
body, and keeping both would double the request-body residency of every
in-flight stream attempt to serve a path that almost never runs. Simply
omitting them is not safe either, because a terminal write is
body-capture-authoritative: with both absent the seed carries the typed
`none` marker, which clears the stored capture rather than leaving it
alone. `build_usage_event_data_seed_describing_request_bodies` is the third
option -- it derives every capture state, body reference, request type and
derived request fact from the real plan and report context, and leaves out
only the two body values -- so the guard's snapshot is small and its
terminal write preserves the capture the `pending` write recorded.

The stream candidate first-byte watchdog also drops the attempt future, but
it settles the attempt itself through `build_transport_error_stop_response`.
It now marks the attempt abandoned before returning so the guard stands down
instead of racing a 499 against the watchdog's 504.

Co-Authored-By: Claude Opus 5 <[email protected]>

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

@@ -78,6 +78,7 @@ use crate::api::response::{
 use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
 use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
 use crate::control::GatewayControlDecision;
+use crate::execution_runtime::attempt_cancellation::AttemptCancellationGuard;
 use crate::execution_runtime::build_direct_execution_frame_stream;
 use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_stream;
 use crate::execution_runtime::grok::maybe_execute_grok_stream;
@@ -149,6 +150,11 @@ use crate::{
     AppState, GatewayError, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
 };
 
+/// Settlement labels for a stream attempt whose future is dropped before the
+/// transport reaches a terminal state.
+const STREAM_ATTEMPT_CANCELLED_ERROR_TYPE: &str = "local_stream_attempt_cancelled";
+const STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE: &str = "Local stream attempt was dropped before terminal finalization, usually because the client disconnected or the request task was cancelled.";
+
 const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
 const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n";
 const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024;
@@ -3655,17 +3661,30 @@ pub(crate) fn execute_execution_runtime_stream<'a>(
     report_kind: Option<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)]
@@ -3687,7 +3706,12 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>(
     Box::pin(async move {
         let mut retry_scope = AiAttemptRetryScope::Candidate;
         let mut fallback_response = None;
-        let response = execute_execution_runtime_stream_inner(
+        let mut cancellation_guard = AttemptCancellationGuard::disarmed(
+            state,
+            STREAM_ATTEMPT_CANCELLED_ERROR_TYPE,
+            STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE,
+        );
+        let result = execute_execution_runtime_stream_inner(
             state,
             plan,
             trace_id,
@@ -3697,8 +3721,13 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>(
             report_context,
             Some(&mut retry_scope),
             Some(&mut fallback_response),
+            &mut cancellation_guard,
         )
-        .await?;
+        .await;
+        // The attempt reached its own terminal path, or handed settlement to the
+        // stream finalizer that now lives in the response body.
+        cancellation_guard.disarm();
+        let response = result?;
         Ok(match response {
             Some(response) => AiAttemptExecutionOutcome::Responded(response),
             None => AiAttemptExecutionOutcome::Retry {
@@ -3744,6 +3773,7 @@ async fn maybe_build_stream_transport_error_stop_response(
     .map(Some)
 }
 
+#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
 async fn execute_execution_runtime_stream_inner(
     state: &AppState,
     mut plan: ExecutionPlan,
@@ -3754,6 +3784,7 @@ async fn execute_execution_runtime_stream_inner(
     mut report_context: Option<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();
@@ -3837,6 +3868,16 @@ async fn execute_execution_runtime_stream_inner(
         )
         .await;
     }
+    // From here the attempt owns non-terminal rows, and everything that could
+    // settle them runs inside the downstream request future. Arm the guard so a
+    // client disconnect before the stream finalizer exists still settles them.
+    cancellation_guard.arm(
+        &plan,
+        report_context.as_ref(),
+        request_candidate_status_snapshot.as_ref(),
+        candidate_started_unix_secs,
+        stream_started_at,
+    );
     let plan_request_id_for_log = short_request_id(plan.request_id.as_str());
     let provider_name = plan
         .provider_name