use crate::{
cancellation::{AgentCancellation, is_run_canceled},
providers::{
HttpRequest, HttpTransport, ProviderEvent,
error::{ProviderError, provider_stream_trace_from_error, retryable_provider_error},
transport::provider_stream_deadline_after,
},
};
use std::{
sync::atomic::{AtomicU64, Ordering},
time::{Duration, Instant},
};
const PROVIDER_STREAM_MAX_ATTEMPTS: usize = 3;
const RATE_LIMIT_RETRY_BACKOFFS: [Duration; 3] = [
Duration::from_secs(2),
Duration::from_secs(10),
Duration::from_secs(20),
];
const RATE_LIMIT_MAX_ATTEMPTS: usize = RATE_LIMIT_RETRY_BACKOFFS.len() + 1;
pub(super) const PROVIDER_STREAM_NO_SEMANTIC_PROGRESS_TIMEOUT: Duration = Duration::from_secs(60);
#[derive(Debug)]
pub(super) struct StreamAttemptError {
pub(super) error: anyhow::Error,
pub(super) made_semantic_progress: bool,
pub(super) unsafe_recovery_progress: bool,
pub(super) attempts_used: usize,
}
pub(super) trait ProviderStreamParser {
fn push_chunk_outcome(
&mut self,
chunk: &str,
) -> anyhow::Result<crate::providers::stream::StreamParseOutcome>;
fn finish(&mut self) -> anyhow::Result<Vec<ProviderEvent>>;
fn response_model(&self) -> Option<String> {
None
}
}
impl ProviderStreamParser for crate::providers::stream::StreamParser {
fn push_chunk_outcome(
&mut self,
chunk: &str,
) -> anyhow::Result<crate::providers::stream::StreamParseOutcome> {
self.push_chunk_outcome(chunk)
}
fn finish(&mut self) -> anyhow::Result<Vec<ProviderEvent>> {
self.finish()
}
fn response_model(&self) -> Option<String> {
crate::providers::stream::StreamParser::response_model(self)
}
}
pub(super) fn stream_with_transport_parser<T: HttpTransport, P: ProviderStreamParser>(
transport: &T,
http_request: HttpRequest,
cancellation: &AgentCancellation,
semantic_progress_timeout: Duration,
parser_factory: impl FnMut() -> P,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
stream_with_transport_attempt_result_with_parser::<T, P>(
transport,
http_request,
cancellation,
semantic_progress_timeout,
parser_factory,
on_event,
)
.map_err(|attempt_error| attempt_error.error)
}
pub(super) fn stream_with_transport_attempt_result<T: HttpTransport>(
transport: &T,
http_request: HttpRequest,
cancellation: &AgentCancellation,
semantic_progress_timeout: Duration,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
) -> Result<(), StreamAttemptError> {
stream_with_transport_attempt_result_with_attempt_offset(
transport,
http_request,
cancellation,
semantic_progress_timeout,
0,
on_event,
)
}
pub(super) fn stream_with_transport_attempt_result_with_attempt_offset<T: HttpTransport>(
transport: &T,
http_request: HttpRequest,
cancellation: &AgentCancellation,
semantic_progress_timeout: Duration,
attempt_offset: usize,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
) -> Result<(), StreamAttemptError> {
stream_with_transport_attempt_result_with_parser_and_sleep::<
T,
crate::providers::stream::StreamParser,
_,
>(
transport,
http_request,
cancellation,
semantic_progress_timeout,
crate::providers::stream::StreamParser::default,
attempt_offset,
on_event,
&mut sleep_cancellable,
)
}
fn stream_with_transport_attempt_result_with_parser<T: HttpTransport, P: ProviderStreamParser>(
transport: &T,
http_request: HttpRequest,
cancellation: &AgentCancellation,
semantic_progress_timeout: Duration,
parser_factory: impl FnMut() -> P,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
) -> Result<(), StreamAttemptError> {
let mut sleep = sleep_cancellable;
stream_with_transport_attempt_result_with_parser_and_sleep::<T, P, _>(
transport,
http_request,
cancellation,
semantic_progress_timeout,
parser_factory,
0,
on_event,
&mut sleep,
)
}
#[expect(clippy::too_many_arguments)]
fn stream_with_transport_attempt_result_with_parser_and_sleep<
T: HttpTransport,
P: ProviderStreamParser,
S: FnMut(Duration, &AgentCancellation) -> anyhow::Result<()>,
>(
transport: &T,
http_request: HttpRequest,
cancellation: &AgentCancellation,
semantic_progress_timeout: Duration,
mut parser_factory: impl FnMut() -> P,
attempt_offset: usize,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
sleep: &mut S,
) -> Result<(), StreamAttemptError> {
let mut last_retryable_error = None;
for local_attempt in 1..=RATE_LIMIT_MAX_ATTEMPTS {
let attempt = attempt_offset + local_attempt;
if let Err(error) = cancellation.check() {
return Err(StreamAttemptError {
error,
made_semantic_progress: false,
unsafe_recovery_progress: false,
attempts_used: local_attempt,
});
}
let mut parser = parser_factory();
match stream_attempt(
transport,
http_request.clone(),
cancellation,
semantic_progress_timeout,
&mut parser,
on_event,
attempt,
) {
Ok(()) => return Ok(()),
Err(attempt_error)
if !attempt_error.made_semantic_progress
&& retryable_provider_error(&attempt_error.error) =>
{
let Some(backoff) = provider_retry_backoff_for(&attempt_error.error, local_attempt)
else {
let error = attempt_error.error;
let detail = error.to_string();
return Err(StreamAttemptError {
error: error.context(format!(
"provider request failed after {attempt} attempts: {detail}"
)),
made_semantic_progress: false,
unsafe_recovery_progress: false,
attempts_used: local_attempt,
});
};
last_retryable_error = Some(attempt_error.error);
if let Err(error) = sleep(backoff, cancellation) {
return Err(StreamAttemptError {
error,
made_semantic_progress: false,
unsafe_recovery_progress: false,
attempts_used: local_attempt,
});
}
}
Err(mut attempt_error) => {
attempt_error.attempts_used = local_attempt;
return Err(attempt_error);
}
}
}
let error =
last_retryable_error.unwrap_or_else(|| anyhow::anyhow!("unknown provider retry failure"));
let detail = error.to_string();
Err(StreamAttemptError {
error: error.context(format!(
"provider request failed after {attempt_offset} + {RATE_LIMIT_MAX_ATTEMPTS} attempts: {detail}"
)),
made_semantic_progress: false,
unsafe_recovery_progress: false,
attempts_used: RATE_LIMIT_MAX_ATTEMPTS,
})
}
fn stream_attempt<T: HttpTransport, P: ProviderStreamParser>(
transport: &T,
http_request: HttpRequest,
cancellation: &AgentCancellation,
semantic_progress_timeout: Duration,
parser: &mut P,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
attempt: usize,
) -> Result<(), StreamAttemptError> {
let mut made_semantic_progress = false;
let mut last_semantic_progress = Instant::now();
let mut unsafe_recovery_progress = false;
let semantic_deadline =
AtomicU64::new(provider_stream_deadline_after(semantic_progress_timeout));
let is_responses = http_request.url.contains("/responses");
let requested_model = http_request
.body
.get("model")
.and_then(serde_json::Value::as_str)
.unwrap_or("unknown")
.to_string();
let provider = if http_request.url.contains("chatgpt.com") {
"openai-codex"
} else {
"openai-compatible"
};
let mut request_id = None;
let mut attempt_events = AttemptEventBuffer::default();
let stream_result = transport.stream_json_cancellable_with_response_metadata(
http_request,
cancellation,
&semantic_deadline,
&mut |metadata| request_id = metadata.request_id,
&mut |chunk| {
cancellation.check()?;
let outcome = parser.push_chunk_outcome(chunk)?;
unsafe_recovery_progress |= outcome.unsafe_recovery_progress;
if outcome.semantic_progress {
made_semantic_progress = true;
last_semantic_progress = Instant::now();
semantic_deadline.store(
provider_stream_deadline_after(semantic_progress_timeout),
Ordering::SeqCst,
);
} else if last_semantic_progress.elapsed() >= semantic_progress_timeout {
return Err(ProviderError::stream_terminal(
"provider stream no semantic progress before timeout",
)
.into());
}
for event in outcome.events {
attempt_events.forward(event, on_event)?;
cancellation.check()?;
}
cancellation.check()?;
Ok(())
},
);
if let Err(error) = stream_result {
let error =
incomplete_stream_error(error, made_semantic_progress, unsafe_recovery_progress);
attempt_events.discard();
if is_responses {
emit_response_identity(
on_event,
response_identity(
provider,
attempt,
&requested_model,
parser.response_model(),
request_id,
"error",
),
made_semantic_progress,
unsafe_recovery_progress,
)?;
}
return Err(StreamAttemptError {
error,
made_semantic_progress,
unsafe_recovery_progress,
attempts_used: 1,
});
}
let finish_events = match parser.finish() {
Ok(events) => events,
Err(error) => {
let error =
incomplete_stream_error(error, made_semantic_progress, unsafe_recovery_progress);
attempt_events.discard();
if is_responses {
emit_response_identity(
on_event,
response_identity(
provider,
attempt,
&requested_model,
parser.response_model(),
request_id,
"error",
),
made_semantic_progress,
unsafe_recovery_progress,
)?;
}
return Err(StreamAttemptError {
error,
made_semantic_progress,
unsafe_recovery_progress,
attempts_used: 1,
});
}
};
made_semantic_progress |= finish_events
.iter()
.any(crate::providers::stream::is_semantic_progress_event);
for event in finish_events {
attempt_events
.forward(event, on_event)
.map_err(|error| StreamAttemptError {
error,
made_semantic_progress,
unsafe_recovery_progress,
attempts_used: 1,
})?;
}
if let Err(error) = cancellation.check() {
attempt_events.discard();
return Err(StreamAttemptError {
error,
made_semantic_progress,
unsafe_recovery_progress,
attempts_used: 1,
});
}
attempt_events
.emit_success(on_event, cancellation)
.map_err(|error| StreamAttemptError {
error,
made_semantic_progress,
unsafe_recovery_progress,
attempts_used: 1,
})?;
if is_responses {
emit_response_identity(
on_event,
response_identity(
provider,
attempt,
&requested_model,
parser.response_model(),
request_id,
"success",
),
made_semantic_progress,
unsafe_recovery_progress,
)?;
}
Ok(())
}
#[derive(Default)]
struct AttemptEventBuffer {
service_tier: Option<String>,
pending_done_count: usize,
}
impl AttemptEventBuffer {
fn forward(
&mut self,
event: ProviderEvent,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
match event {
ProviderEvent::ServiceTier(tier) => {
self.service_tier = Some(tier);
Ok(())
}
ProviderEvent::Done => {
self.pending_done_count = self.pending_done_count.saturating_add(1);
Ok(())
}
event => on_event(event),
}
}
fn emit_success(
&mut self,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
cancellation: &AgentCancellation,
) -> anyhow::Result<()> {
if let Some(tier) = self.service_tier.take() {
on_event(ProviderEvent::ServiceTier(tier))?;
cancellation.check()?;
}
let pending_done_count = std::mem::take(&mut self.pending_done_count);
for _ in 0..pending_done_count {
on_event(ProviderEvent::Done)?;
cancellation.check()?;
}
Ok(())
}
fn discard(&mut self) {
self.service_tier = None;
self.pending_done_count = 0;
}
}
fn response_identity(
provider: &str,
attempt: usize,
requested_model: &str,
provider_response_model: Option<String>,
request_id: Option<String>,
outcome: &str,
) -> crate::providers::error::ResponseAttemptIdentity {
crate::providers::error::ResponseAttemptIdentity {
schema_version: 1,
provider: provider.to_string(),
attempt,
requested_model: crate::providers::error::bounded_response_identity_string(requested_model),
provider_response_model: provider_response_model
.map(|value| crate::providers::error::bounded_response_identity_string(&value)),
request_id: request_id
.map(|value| crate::providers::error::bounded_response_identity_string(&value)),
outcome: crate::providers::error::bounded_response_identity_string(outcome),
cache: None,
}
}
fn emit_response_identity(
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
identity: crate::providers::error::ResponseAttemptIdentity,
made_semantic_progress: bool,
unsafe_recovery_progress: bool,
) -> Result<(), StreamAttemptError> {
on_event(ProviderEvent::ResponseIdentity(identity)).map_err(|error| StreamAttemptError {
error,
made_semantic_progress,
unsafe_recovery_progress,
attempts_used: 1,
})
}
fn incomplete_stream_error(
error: anyhow::Error,
made_semantic_progress: bool,
unsafe_recovery_progress: bool,
) -> anyhow::Error {
if !made_semantic_progress
|| is_run_canceled(&error)
|| contains_actionable_provider_error(&error)
{
return error;
}
let stream_trace = provider_stream_trace_from_error(&error);
let unsafe_suffix = if unsafe_recovery_progress {
"; unsafe tool-call progress observed"
} else {
""
};
let mut wrapped = ProviderError::stream_failed_incomplete(format!(
"provider stream ended prematurely after partial response; response is incomplete: {error}{unsafe_suffix}"
));
if let Some(trace) = stream_trace {
wrapped = wrapped.with_stream_trace(trace);
}
wrapped.into()
}
fn contains_actionable_provider_error(error: &anyhow::Error) -> bool {
error
.downcast_ref::<ProviderError>()
.and_then(ProviderError::http_status_code)
.is_some()
}
pub(super) fn sleep_cancellable(
duration: Duration,
cancellation: &AgentCancellation,
) -> anyhow::Result<()> {
const CANCEL_POLL_INTERVAL: Duration = Duration::from_millis(25);
let start = Instant::now();
loop {
cancellation.check()?;
let elapsed = start.elapsed();
if elapsed >= duration {
return Ok(());
}
std::thread::sleep(duration.saturating_sub(elapsed).min(CANCEL_POLL_INTERVAL));
}
}
fn provider_retry_backoff_for(error: &anyhow::Error, attempt: usize) -> Option<Duration> {
if !retryable_provider_error(error) {
return None;
}
if is_rate_limited(error) {
return RATE_LIMIT_RETRY_BACKOFFS
.get(attempt.checked_sub(1)?)
.copied();
}
(attempt < PROVIDER_STREAM_MAX_ATTEMPTS).then(|| provider_retry_backoff(attempt))
}
fn is_rate_limited(error: &anyhow::Error) -> bool {
error
.downcast_ref::<ProviderError>()
.and_then(ProviderError::http_status_code)
== Some(429)
}
static RETRY_JITTER_COUNTER: AtomicU64 = AtomicU64::new(0);
fn provider_retry_backoff(attempt: usize) -> Duration {
let base_ms: u64 = match attempt {
1 => 250,
_ => 750,
};
let bound_ms = base_ms / 5;
let span = bound_ms.saturating_mul(2).saturating_add(1);
let offset =
(RETRY_JITTER_COUNTER.fetch_add(1, Ordering::Relaxed) % span) as i64 - bound_ms as i64;
Duration::from_millis((base_ms as i64 + offset).max(1) as u64)
}