Fix stream timeout semantics

This commit is contained in:
fawney19
2026-05-25 20:09:37 +08:00
parent 54c5d5803b
commit d3249485fa
7 changed files with 561 additions and 189 deletions
@@ -26,7 +26,8 @@ use crate::request_candidate_runtime::{
};
use crate::{AppState, GatewayError};
const DEFAULT_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS: u64 = 300_000;
const DEFAULT_STREAM_FIRST_BYTE_WATCHDOG_TIMEOUT_MS: u64 = 30_000;
const DEFAULT_NON_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS: u64 = 300_000;
fn attach_redaction_execution_candidate(response: &mut Response<Body>, candidate_id: Option<&str>) {
if let Some(candidate_id) = candidate_id
@@ -534,18 +535,18 @@ fn resolve_stream_candidate_watchdog_timeout(
.and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or(true);
let timeout_ms = plan
.timeouts
.as_ref()
.and_then(|timeouts| {
if upstream_is_stream {
timeouts.first_byte_ms.or(timeouts.total_ms)
} else {
timeouts.total_ms.or(timeouts.first_byte_ms)
}
})
.unwrap_or(DEFAULT_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS)
.max(1);
let timeout_ms = if upstream_is_stream {
plan.timeouts
.as_ref()
.and_then(|timeouts| timeouts.first_byte_ms)
.unwrap_or(DEFAULT_STREAM_FIRST_BYTE_WATCHDOG_TIMEOUT_MS)
} else {
plan.timeouts
.as_ref()
.and_then(|timeouts| timeouts.total_ms)
.unwrap_or(DEFAULT_NON_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS)
}
.max(1);
Duration::from_millis(timeout_ms)
}
@@ -783,7 +784,24 @@ mod tests {
assert_eq!(
timeout,
Duration::from_millis(DEFAULT_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS)
Duration::from_millis(DEFAULT_STREAM_FIRST_BYTE_WATCHDOG_TIMEOUT_MS)
);
}
#[test]
fn stream_candidate_watchdog_ignores_total_timeout_for_stream_upstream() {
let report_context = json!({"upstream_is_stream": true});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
total_ms: Some(90_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(
timeout,
Duration::from_millis(DEFAULT_STREAM_FIRST_BYTE_WATCHDOG_TIMEOUT_MS)
);
}
@@ -803,17 +821,20 @@ mod tests {
}
#[test]
fn stream_candidate_watchdog_falls_back_to_first_byte_when_upstream_non_stream_lacks_total() {
fn stream_candidate_watchdog_uses_non_stream_default_when_upstream_non_stream_lacks_total() {
let report_context = json!({"upstream_is_stream": false});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(300_000),
first_byte_ms: Some(12_345),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(300_000));
assert_eq!(
timeout,
Duration::from_millis(DEFAULT_NON_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS)
);
}
#[test]