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]>
This commit is contained in:
stabey
2026-09-04 17:40:32 +08:00
committed by ZheFox
co-authored by Claude Opus 5
parent 14744abd57
commit 9282cce1d6
8 changed files with 943 additions and 26 deletions
@@ -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,