fix(gateway): share stream first-byte deadline across retries

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<Fut>(
     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<StreamCandidateWatchdogOutcome, GatewayError>
@@ -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<AiAttemptExecutionOutcome<Response<Body>>, 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 {