From 206995645b8a5177a285b1025bc7477d0c2e0322 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:01:19 +0800 Subject: [PATCH] fix(gateway): share stream first-byte deadline across retries --- .../src/executor/candidate_loop.rs | 82 ++++++++++++++++++- 1 file changed, 80 insertions(+), 2 deletions(-) diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 42f262a78..be1f9f876 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -412,6 +412,7 @@ where decision, plan_kind, transfer_tracker, + request_first_byte_started_at: Instant::now(), }; match run_ai_attempt_loop(&port, plan_and_reports).await? { AiAttemptLoopOutcome::Responded(response) => { @@ -482,6 +483,7 @@ where decision, plan_kind, transfer_tracker, + request_first_byte_started_at: Instant::now(), }; run_dynamic_attempt_loop( &port, @@ -946,6 +948,10 @@ struct StreamAttemptLoopPort<'a> { decision: &'a GatewayControlDecision, plan_kind: &'a str, transfer_tracker: &'a ProviderTransferTracker, + /// All candidates in one downstream stream request share this origin. + /// Without it every retry receives a fresh full first-byte timeout and a + /// 30-second provider timeout can accumulate into a 60-120 second stall. + request_first_byte_started_at: Instant, } #[async_trait] @@ -1058,6 +1064,7 @@ where self.plan_kind, plan, watchdog_report_context, + self.request_first_byte_started_at, stop_on_transport_errors, move || async move { execute_execution_runtime_stream_with_retry_scope( @@ -1085,7 +1092,7 @@ where http::StatusCode::GATEWAY_TIMEOUT.as_u16(), "local_stream_candidate_watchdog_timeout", stream_candidate_watchdog_timeout_message(), - watchdog_started_at.elapsed().as_millis() as u64, + self.request_first_byte_started_at.elapsed().as_millis() as u64, ) .await?, ) @@ -1368,6 +1375,7 @@ async fn execute_stream_candidate_with_watchdog( plan_kind: &str, plan: &aether_contracts::ExecutionPlan, report_context: Option<&serde_json::Value>, + request_first_byte_started_at: Instant, stop_on_transport_errors: bool, execute: impl FnOnce() -> Fut, ) -> Result @@ -1377,6 +1385,7 @@ where > + Send, { let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context); + let request_first_byte_deadline = request_first_byte_started_at + timeout_duration; let candidate_started_at = std::time::Instant::now(); let candidate_started_unix_ms = current_unix_ms(); let permit = match acquire_upstream_execution_gate(state, trace_id).await { @@ -1402,7 +1411,14 @@ where let watchdog_progress = StreamCandidateWatchdogProgress::shared(); let execution = watchdog_progress.clone().scope(execute()); tokio::pin!(execution); - let deadline = tokio::time::sleep(timeout_duration); + // This is an absolute request-level deadline, not a new timeout for this + // candidate. Retries therefore consume only the budget left by earlier + // candidates instead of resetting the full provider timeout. + let candidate_budget_ms = request_first_byte_deadline + .saturating_duration_since(Instant::now()) + .as_millis() + .min(u128::from(u64::MAX)) as u64; + let deadline = tokio::time::sleep_until(request_first_byte_deadline); tokio::pin!(deadline); let execution_result = tokio::select! { biased; @@ -1431,6 +1447,10 @@ where .map(|value| value.to_string()) .unwrap_or_else(|| "-".to_string()); let timeout_ms = u64::try_from(timeout_duration.as_millis()).unwrap_or(u64::MAX); + let request_elapsed_ms = request_first_byte_started_at + .elapsed() + .as_millis() + .min(u128::from(u64::MAX)) as u64; record_local_request_candidate_status( state, plan, @@ -1459,6 +1479,8 @@ where model_name, candidate_index = candidate_index.as_str(), timeout_ms, + candidate_budget_ms, + request_elapsed_ms, "gateway local stream candidate watchdog timed out" ); if stop_on_transport_errors { @@ -2389,6 +2411,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), false, || { std::future::pending::< @@ -2427,6 +2450,56 @@ mod tests { assert_eq!(record.candidate_index, 2); } + #[tokio::test] + async fn stream_candidate_retry_does_not_reset_an_expired_request_first_byte_budget() { + let writer = Arc::new(TestRequestCandidateWriter::default()); + let plan = test_plan(Some(ExecutionTimeouts { + first_byte_ms: Some(250), + ..ExecutionTimeouts::default() + })); + let report_context = test_report_context(); + // Stand in for earlier candidates having already consumed the request's + // complete first-byte budget. A per-candidate watchdog would wait a new + // 250 ms here; the shared absolute deadline must settle immediately. + let request_first_byte_started_at = Instant::now() - Duration::from_millis(300); + + let result = tokio::time::timeout( + Duration::from_millis(100), + execute_stream_candidate_with_watchdog( + writer.as_ref(), + "trace_watchdog_shared_budget", + "claude_cli_stream", + &plan, + Some(&report_context), + request_first_byte_started_at, + false, + || { + std::future::pending::< + Result>, GatewayError>, + >() + }, + ), + ) + .await + .expect("an expired request-level first-byte budget must not restart per candidate"); + + assert!(matches!( + result, + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + } + )) + )); + let records = writer.records.lock().await; + assert_eq!(records.len(), 1); + assert_eq!( + records[0].error_type.as_deref(), + Some("local_stream_candidate_watchdog_timeout") + ); + } + #[tokio::test] async fn stream_candidate_watchdog_can_stop_on_transport_error() { let writer = Arc::new(TestRequestCandidateWriter::default()); @@ -2442,6 +2515,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), true, || { std::future::pending::< @@ -2479,6 +2553,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), true, || async { mark_stream_candidate_watchdog_terminal_started(); @@ -2511,6 +2586,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), true, || async { Err(GatewayError::UpstreamUnavailable { @@ -2550,6 +2626,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), false, || async { panic!("execute future should not run while upstream execution gate is saturated") @@ -2597,6 +2674,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), false, || async { Err(GatewayError::AdmissionTimeout {