first pass of merging in warp (doesn't build)
This commit is contained in:
@@ -7,46 +7,102 @@ use futures::future::Either;
|
||||
use futures::StreamExt;
|
||||
use galaxyui::r#async::Timer;
|
||||
use instant::Instant;
|
||||
use galaxy_core::errors::AnyhowErrorExt as _;
|
||||
|
||||
use crate::server::retry_strategies::is_transient_http_error;
|
||||
use crate::server::server_api::ai::AgentRunEvent;
|
||||
use crate::server::server_api::presigned_upload::HttpStatusError;
|
||||
use crate::server::server_api::ServerApi;
|
||||
|
||||
pub(crate) const DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10];
|
||||
pub(crate) const DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS: &[u64] = &[30];
|
||||
pub(crate) const DEFAULT_AGENT_EVENT_PROACTIVE_RECONNECT: Duration = Duration::from_secs(14 * 60);
|
||||
pub(crate) const DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG: usize = 5;
|
||||
|
||||
/// Selects which server-side filter shape an [`AgentEventSource`] should use
|
||||
/// when opening a stream.
|
||||
///
|
||||
/// `RunIds` maps to the `?run_ids[]=` query parameter on the SSE endpoint
|
||||
/// and is used by child-only per-conversation streams and the dormant
|
||||
/// Claude wake listener. `AncestorRunId` maps to the `?ancestor_run_id=`
|
||||
/// shape: with `include_self=false` it streams events for every direct
|
||||
/// child of the supplied parent run (the shared-session viewer's pill bar),
|
||||
/// and with `include_self=true` it additionally streams the parent run's
|
||||
/// own events so an owner-side orchestrator can receive child lifecycle
|
||||
/// events plus its own inbox on one ordered stream.
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) enum AgentEventFilter {
|
||||
/// One stream per multiplexed set of run IDs. Matches today's
|
||||
/// `?run_ids[]=` endpoint.
|
||||
RunIds(Vec<String>),
|
||||
/// Stream events for every direct child of the supplied parent run, and
|
||||
/// (when `include_self` is true) the parent run itself. Matches the
|
||||
/// `?ancestor_run_id=` endpoint.
|
||||
AncestorRunId {
|
||||
ancestor_run_id: String,
|
||||
include_self: bool,
|
||||
},
|
||||
}
|
||||
|
||||
impl AgentEventFilter {
|
||||
/// Returns a short debug label used in driver log lines so we don't have
|
||||
/// to format the full `Vec<String>` payload on every retry.
|
||||
pub(crate) fn log_label(&self) -> String {
|
||||
match self {
|
||||
AgentEventFilter::RunIds(ids) => format!("run_ids={ids:?}"),
|
||||
AgentEventFilter::AncestorRunId {
|
||||
ancestor_run_id,
|
||||
include_self,
|
||||
} => format!("ancestor_run_id={ancestor_run_id} include_self={include_self}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for the shared agent-event stream driver.
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct AgentEventDriverConfig {
|
||||
/// Run IDs whose events should be multiplexed into a single stream.
|
||||
pub run_ids: Vec<String>,
|
||||
/// Wire-level filter selecting which run IDs the stream serves. Either a
|
||||
/// concrete multiplexed list of run IDs or an ancestor-scoped child set;
|
||||
/// see [`AgentEventFilter`].
|
||||
pub filter: AgentEventFilter,
|
||||
/// Last fully handled event sequence. Events at or below this cursor are
|
||||
/// ignored on reconnect so the consumer only sees new work.
|
||||
pub since_sequence: i64,
|
||||
/// Exponential-ish reconnect delays, in seconds, used after stream open
|
||||
/// failures, stream errors, and clean stream termination.
|
||||
/// Exponential-ish reconnect delays, in seconds, used after transient
|
||||
/// stream failures (5xx, timeouts, connection resets).
|
||||
pub reconnect_backoff_steps: &'static [u64],
|
||||
/// Reconnect delays for permanent HTTP errors (4xx other than 408/429).
|
||||
/// Typically much slower than transient backoff to reduce log spam while
|
||||
/// still allowing recovery if the error was spurious.
|
||||
pub permanent_error_backoff_steps: &'static [u64],
|
||||
/// Optional deadline for proactively recycling an otherwise healthy stream
|
||||
/// before upstream infrastructure times it out (for example, before Cloud
|
||||
/// Run's 20-minute streaming timeout).
|
||||
pub proactive_reconnect_after: Option<Duration>,
|
||||
/// Failure count at which reconnect logging is escalated from debug to warn.
|
||||
/// Failure count at which actionable reconnect failures are reported at Error level.
|
||||
/// This only affects log severity; retry behavior stays the same.
|
||||
pub failures_before_error_log: usize,
|
||||
}
|
||||
|
||||
impl AgentEventDriverConfig {
|
||||
/// Build the production reconnecting configuration used by long-lived
|
||||
/// orchestration and harness listeners.
|
||||
pub(crate) fn retry_forever(run_ids: Vec<String>, since_sequence: i64) -> Self {
|
||||
/// orchestration and harness listeners, parameterised on the wire filter.
|
||||
pub(crate) fn retry_forever(filter: AgentEventFilter, since_sequence: i64) -> Self {
|
||||
Self {
|
||||
run_ids,
|
||||
filter,
|
||||
since_sequence,
|
||||
reconnect_backoff_steps: DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
|
||||
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: Some(DEFAULT_AGENT_EVENT_PROACTIVE_RECONNECT),
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience: build a `retry_forever` config from a concrete list of
|
||||
/// run IDs. Lets existing call sites keep their current ergonomics.
|
||||
pub(crate) fn retry_forever_run_ids(run_ids: Vec<String>, since_sequence: i64) -> Self {
|
||||
Self::retry_forever(AgentEventFilter::RunIds(run_ids), since_sequence)
|
||||
}
|
||||
}
|
||||
|
||||
/// Tells the shared driver whether to continue or stop after a handled event.
|
||||
@@ -75,6 +131,27 @@ pub(crate) enum AgentEventDriverState {
|
||||
ProactiveReconnect,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub(crate) struct AgentMessageEventMetadata {
|
||||
pub sequence: i64,
|
||||
pub message_id: String,
|
||||
pub occurred_at: String,
|
||||
}
|
||||
|
||||
impl AgentMessageEventMetadata {
|
||||
pub(crate) fn from_event(event: &AgentRunEvent) -> Option<Self> {
|
||||
if event.event_type != "new_message" {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(Self {
|
||||
sequence: event.sequence,
|
||||
message_id: event.ref_id.clone()?,
|
||||
occurred_at: event.occurred_at.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Parsed items emitted by an [`AgentEventSource`].
|
||||
pub(crate) enum AgentEventSourceItem {
|
||||
Open,
|
||||
@@ -91,13 +168,13 @@ cfg_if::cfg_if! {
|
||||
}
|
||||
}
|
||||
|
||||
/// Opens a stream of parsed agent events for one or more run IDs.
|
||||
/// Opens a stream of parsed agent events for the supplied filter.
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
pub(crate) trait AgentEventSource: Send + Sync {
|
||||
async fn open_stream(
|
||||
&self,
|
||||
run_ids: &[String],
|
||||
filter: &AgentEventFilter,
|
||||
since_sequence: i64,
|
||||
) -> Result<AgentEventSourceStream>;
|
||||
}
|
||||
@@ -118,13 +195,28 @@ impl ServerApiAgentEventSource {
|
||||
impl AgentEventSource for ServerApiAgentEventSource {
|
||||
async fn open_stream(
|
||||
&self,
|
||||
run_ids: &[String],
|
||||
filter: &AgentEventFilter,
|
||||
since_sequence: i64,
|
||||
) -> Result<AgentEventSourceStream> {
|
||||
let stream = self
|
||||
.server_api
|
||||
.stream_agent_events(run_ids, since_sequence)
|
||||
.await?;
|
||||
let stream = match filter {
|
||||
AgentEventFilter::RunIds(run_ids) => {
|
||||
self.server_api
|
||||
.stream_agent_events(run_ids, since_sequence)
|
||||
.await?
|
||||
}
|
||||
AgentEventFilter::AncestorRunId {
|
||||
ancestor_run_id,
|
||||
include_self,
|
||||
} => {
|
||||
self.server_api
|
||||
.stream_agent_events_for_ancestor(
|
||||
ancestor_run_id,
|
||||
*include_self,
|
||||
since_sequence,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
|
||||
let stream = stream.filter_map(|event_result| async move {
|
||||
match event_result {
|
||||
@@ -138,7 +230,29 @@ impl AgentEventSource for ServerApiAgentEventSource {
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => Some(Err(anyhow!("SSE stream error: {err:?}"))),
|
||||
Err(err) => {
|
||||
let anyhow_err = match err {
|
||||
reqwest_eventsource::Error::InvalidStatusCode(status_code, response) => {
|
||||
let body = response
|
||||
.text()
|
||||
.await
|
||||
.unwrap_or_else(|err| format!("(no response body: {err:#})"));
|
||||
let status_err = HttpStatusError {
|
||||
status: status_code.as_u16(),
|
||||
body: body.clone(),
|
||||
};
|
||||
anyhow::Error::new(status_err).context(format!(
|
||||
"SSE stream error: invalid status code {status_code}: {body}"
|
||||
))
|
||||
}
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
reqwest_eventsource::Error::Transport(err) => {
|
||||
anyhow::Error::new(err).context("SSE stream error")
|
||||
}
|
||||
err => anyhow!("SSE stream error: {err:?}"),
|
||||
};
|
||||
Some(Err(anyhow_err))
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -188,18 +302,23 @@ where
|
||||
let mut has_connected_once = false;
|
||||
|
||||
loop {
|
||||
let mut stream = match source.open_stream(&config.run_ids, since_sequence).await {
|
||||
Ok(stream) => {
|
||||
failures = 0;
|
||||
has_connected_once = true;
|
||||
notify_driver_state(consumer, AgentEventDriverState::Connected).await;
|
||||
stream
|
||||
}
|
||||
// `open_stream` is lazy for the SSE-backed source: the TCP
|
||||
// connect happens when the stream is first polled, not when
|
||||
// this returns Ok. Wait for the `AgentEventSourceItem::Open`
|
||||
// event below before declaring connectivity, so a server
|
||||
// outage doesn't reset `failures` between every retry.
|
||||
let mut stream = match source.open_stream(&config.filter, since_sequence).await {
|
||||
Ok(stream) => stream,
|
||||
Err(err) => {
|
||||
failures += 1;
|
||||
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
|
||||
let backoff_steps = if is_transient_http_error(&err) {
|
||||
config.reconnect_backoff_steps
|
||||
} else {
|
||||
config.permanent_error_backoff_steps
|
||||
};
|
||||
let backoff = agent_event_backoff(failures, backoff_steps);
|
||||
log_stream_failure(
|
||||
&config.run_ids,
|
||||
&config.filter,
|
||||
failures,
|
||||
backoff,
|
||||
&err,
|
||||
@@ -249,7 +368,12 @@ where
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Ok(AgentEventSourceItem::Open))) => {
|
||||
failures = 0;
|
||||
log::info!("Agent event stream opened for {:?}", config.run_ids);
|
||||
has_connected_once = true;
|
||||
notify_driver_state(consumer, AgentEventDriverState::Connected).await;
|
||||
log::info!(
|
||||
"Agent event stream opened for {}",
|
||||
config.filter.log_label()
|
||||
);
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Ok(AgentEventSourceItem::Event(event)))) => {
|
||||
failures = 0;
|
||||
@@ -273,9 +397,14 @@ where
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Err(err))) => {
|
||||
failures += 1;
|
||||
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
|
||||
let backoff_steps = if is_transient_http_error(&err) {
|
||||
config.reconnect_backoff_steps
|
||||
} else {
|
||||
config.permanent_error_backoff_steps
|
||||
};
|
||||
let backoff = agent_event_backoff(failures, backoff_steps);
|
||||
log_stream_failure(
|
||||
&config.run_ids,
|
||||
&config.filter,
|
||||
failures,
|
||||
backoff,
|
||||
&err,
|
||||
@@ -293,12 +422,15 @@ where
|
||||
Timer::after(backoff).await;
|
||||
break;
|
||||
}
|
||||
// Clean stream closure (server-side close, not an HTTP
|
||||
// error) — always use the transient backoff schedule since
|
||||
// there is no HTTP status to classify.
|
||||
NextDriverItem::StreamItem(None) => {
|
||||
failures += 1;
|
||||
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
|
||||
log::warn!(
|
||||
"Agent event stream closed for {:?}, reconnecting in {backoff:?}",
|
||||
config.run_ids
|
||||
"Agent event stream closed for {}, reconnecting in {backoff:?}",
|
||||
config.filter.log_label()
|
||||
);
|
||||
notify_driver_state(
|
||||
consumer,
|
||||
@@ -332,21 +464,20 @@ async fn notify_driver_state<C: AgentEventConsumer>(
|
||||
}
|
||||
|
||||
fn log_stream_failure(
|
||||
run_ids: &[String],
|
||||
filter: &AgentEventFilter,
|
||||
failures: usize,
|
||||
backoff: Duration,
|
||||
err: &anyhow::Error,
|
||||
failures_before_error_log: usize,
|
||||
) {
|
||||
if agent_event_failures_exceeded_threshold(failures, failures_before_error_log) {
|
||||
let label = filter.log_label();
|
||||
if agent_event_failure_should_log_error(err, failures, failures_before_error_log) {
|
||||
log::error!(
|
||||
"Agent event stream failed {failures} consecutive times for {:?}, retrying in {backoff:?}: {err:#}",
|
||||
run_ids
|
||||
"Agent event stream failed {failures} consecutive times for {label}, retrying in {backoff:?}: {err:#}"
|
||||
);
|
||||
} else {
|
||||
log::warn!(
|
||||
"Agent event stream failed for {:?}, retrying in {backoff:?}: {err:#}",
|
||||
run_ids
|
||||
"Agent event stream failed {failures} consecutive times for {label}, retrying in {backoff:?}: {err:#}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -361,6 +492,15 @@ pub(crate) fn agent_event_backoff(failures: usize, backoff_steps: &[u64]) -> Dur
|
||||
Duration::from_secs(safe_steps[index])
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn agent_event_failures_exceeded_threshold(failures: usize, threshold: usize) -> bool {
|
||||
failures >= threshold
|
||||
}
|
||||
|
||||
pub(crate) fn agent_event_failure_should_log_error(
|
||||
err: &anyhow::Error,
|
||||
failures: usize,
|
||||
threshold: usize,
|
||||
) -> bool {
|
||||
threshold > 0 && failures == threshold && err.is_actionable()
|
||||
}
|
||||
|
||||
@@ -6,9 +6,12 @@ use anyhow::anyhow;
|
||||
use async_trait::async_trait;
|
||||
use futures::stream::{self, BoxStream};
|
||||
use futures::StreamExt;
|
||||
use galaxy_core::errors::AnyhowErrorExt as _;
|
||||
|
||||
use super::*;
|
||||
use crate::ai::agent_events::driver::agent_event_failure_should_log_error;
|
||||
use crate::server::server_api::ai::AgentRunEvent;
|
||||
use crate::server::server_api::presigned_upload::HttpStatusError;
|
||||
|
||||
const ZERO_BACKOFF_STEPS: &[u64] = &[0];
|
||||
|
||||
@@ -28,7 +31,7 @@ impl FakeAgentEventSource {
|
||||
impl AgentEventSource for FakeAgentEventSource {
|
||||
async fn open_stream(
|
||||
&self,
|
||||
_run_ids: &[String],
|
||||
_filter: &AgentEventFilter,
|
||||
_since_sequence: i64,
|
||||
) -> anyhow::Result<BoxStream<'static, anyhow::Result<AgentEventSourceItem>>> {
|
||||
let response = self
|
||||
@@ -136,9 +139,10 @@ async fn driver_skips_duplicate_sequences_and_persists_new_cursor() {
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
|
||||
since_sequence: 2,
|
||||
reconnect_backoff_steps: DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
|
||||
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
@@ -181,9 +185,10 @@ async fn driver_resets_failures_after_successful_event_delivery() {
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
@@ -225,9 +230,10 @@ async fn driver_ignores_persist_cursor_errors() {
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
@@ -258,9 +264,10 @@ async fn driver_ignores_driver_state_errors() {
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
@@ -293,9 +300,10 @@ async fn driver_retries_initial_connection_until_stream_opens() {
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
@@ -350,3 +358,159 @@ fn failure_threshold_is_reached_at_and_above_limit() {
|
||||
assert!(agent_event_failures_exceeded_threshold(5, 5));
|
||||
assert!(agent_event_failures_exceeded_threshold(6, 5));
|
||||
}
|
||||
|
||||
fn make_http_status_error(status: u16) -> anyhow::Error {
|
||||
anyhow::Error::new(HttpStatusError {
|
||||
status,
|
||||
body: "not found".to_string(),
|
||||
})
|
||||
.context("SSE stream error")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn actionable_stream_status_reports_only_at_threshold_crossing() {
|
||||
let err = make_http_status_error(400);
|
||||
assert_eq!(
|
||||
[
|
||||
agent_event_failure_should_log_error(&err, 4, 5),
|
||||
agent_event_failure_should_log_error(&err, 5, 5),
|
||||
agent_event_failure_should_log_error(&err, 6, 5),
|
||||
],
|
||||
[false, true, false]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_threshold_disables_stream_error_escalation() {
|
||||
let err = make_http_status_error(400);
|
||||
assert!(!agent_event_failure_should_log_error(&err, 1, 0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_actionable_stream_statuses_do_not_report_at_threshold() {
|
||||
for status in [408, 429] {
|
||||
let err = make_http_status_error(status);
|
||||
assert!(
|
||||
!agent_event_failure_should_log_error(&err, 5, 5),
|
||||
"status {status}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn server_error_status_reports_at_threshold_crossing() {
|
||||
let err = make_http_status_error(500);
|
||||
assert!(agent_event_failure_should_log_error(&err, 5, 5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_status_error_actionability_follows_status_classification() {
|
||||
assert!(make_http_status_error(400).is_actionable());
|
||||
assert!(make_http_status_error(500).is_actionable());
|
||||
assert!(!make_http_status_error(429).is_actionable());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_uses_slow_backoff_on_permanent_http_error() {
|
||||
let source = FakeAgentEventSource::new(vec![
|
||||
// First attempt: stream opens then returns a 404-enriched error.
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Err(make_http_status_error(404)),
|
||||
]),
|
||||
// Second attempt: succeeds with a stoppable event.
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
1,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-1"),
|
||||
))),
|
||||
]),
|
||||
]);
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 1,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Use asymmetric values: transient=9999s (would block if chosen),
|
||||
// permanent=0s (instant). Proving backoff is 0s confirms the
|
||||
// permanent schedule was selected.
|
||||
let config = AgentEventDriverConfig {
|
||||
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: &[9999],
|
||||
permanent_error_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
|
||||
run_agent_event_driver(source, config, &mut consumer)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(consumer.handled_sequences, vec![1]);
|
||||
// The backoff must be from the permanent schedule (0s), not transient (9999s).
|
||||
let retry_backoff = consumer
|
||||
.driver_states
|
||||
.iter()
|
||||
.find_map(|s| match s {
|
||||
AgentEventDriverState::RetryScheduled { backoff, .. } => Some(*backoff),
|
||||
_ => None,
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(retry_backoff, Duration::from_secs(0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_uses_fast_backoff_on_transient_http_error() {
|
||||
let source = FakeAgentEventSource::new(vec![
|
||||
// First attempt: stream opens then returns a 500-enriched error.
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Err(make_http_status_error(500)),
|
||||
]),
|
||||
// Second attempt: succeeds.
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
1,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-1"),
|
||||
))),
|
||||
]),
|
||||
]);
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 1,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Set permanent backoff to something large so we can verify it was NOT used.
|
||||
let config = AgentEventDriverConfig {
|
||||
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
permanent_error_backoff_steps: &[9999],
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
|
||||
run_agent_event_driver(source, config, &mut consumer)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Should have retried with the fast backoff (0s) and succeeded.
|
||||
assert_eq!(consumer.handled_sequences, vec![1]);
|
||||
// The backoff used should be from ZERO_BACKOFF_STEPS (0s), not permanent (9999s).
|
||||
let retry_backoff = consumer
|
||||
.driver_states
|
||||
.iter()
|
||||
.find_map(|s| match s {
|
||||
AgentEventDriverState::RetryScheduled { backoff, .. } => Some(*backoff),
|
||||
_ => None,
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(retry_backoff, Duration::from_secs(0));
|
||||
}
|
||||
|
||||
@@ -5,20 +5,33 @@ use anyhow::{anyhow, Context, Result};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use futures::future::Either;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use instant::Instant;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use reqwest::Error as ReqwestError;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use galaxyui::r#async::Timer;
|
||||
|
||||
use crate::ai::agent::ReceivedMessageInput;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::server::server_api::ai::{AIClient, AgentRunEvent, ReadAgentMessageResponse};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use crate::server::server_api::presigned_upload::HttpStatusError;
|
||||
use crate::server::server_api::ServerApi;
|
||||
|
||||
pub(crate) const DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const DEFAULT_AGENT_MESSAGE_RETRY_DELAY: Duration = Duration::from_millis(50);
|
||||
|
||||
/// Hydrates `new_message` agent events into full message payloads and delivery
|
||||
/// acknowledgements.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MessageHydrator {
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
task_scoped_server_api: Option<Arc<ServerApi>>,
|
||||
task_id: Option<AmbientAgentTaskId>,
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
fetch_timeout: Duration,
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
retry_delay: Duration,
|
||||
}
|
||||
|
||||
impl MessageHydrator {
|
||||
@@ -26,16 +39,57 @@ impl MessageHydrator {
|
||||
Self::with_fetch_timeout(ai_client, DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT)
|
||||
}
|
||||
|
||||
pub(crate) fn for_task(server_api: Arc<ServerApi>, task_id: AmbientAgentTaskId) -> Self {
|
||||
let ai_client: Arc<dyn AIClient> = server_api.clone();
|
||||
Self {
|
||||
ai_client,
|
||||
task_scoped_server_api: Some(server_api),
|
||||
task_id: Some(task_id),
|
||||
fetch_timeout: DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT,
|
||||
retry_delay: DEFAULT_AGENT_MESSAGE_RETRY_DELAY,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_fetch_timing(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
fetch_timeout: Duration,
|
||||
retry_delay: Duration,
|
||||
) -> Self {
|
||||
Self {
|
||||
ai_client,
|
||||
task_scoped_server_api: None,
|
||||
task_id: None,
|
||||
fetch_timeout,
|
||||
retry_delay,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn with_fetch_timeout(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
fetch_timeout: Duration,
|
||||
) -> Self {
|
||||
Self {
|
||||
ai_client,
|
||||
task_scoped_server_api: None,
|
||||
task_id: None,
|
||||
fetch_timeout,
|
||||
retry_delay: DEFAULT_AGENT_MESSAGE_RETRY_DELAY,
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_message(&self, message_id: &str) -> Result<ReadAgentMessageResponse> {
|
||||
match (self.task_scoped_server_api.as_ref(), self.task_id) {
|
||||
(Some(server_api), Some(task_id)) => {
|
||||
server_api
|
||||
.read_agent_message_for_task(&task_id, message_id)
|
||||
.await
|
||||
}
|
||||
_ => self.ai_client.read_agent_message(message_id).await,
|
||||
}
|
||||
.with_context(|| format!("Failed to read agent message {message_id}"))
|
||||
}
|
||||
|
||||
pub(crate) async fn hydrate_event_for_recipient(
|
||||
&self,
|
||||
event: &AgentRunEvent,
|
||||
@@ -55,6 +109,17 @@ impl MessageHydrator {
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if message.body.is_empty() {
|
||||
log::warn!(
|
||||
"Hydrated empty-body agent message: message_id={} event_sequence={} recipient_run_id={} sender_run_id={} subject={:?} task_id={:?}",
|
||||
message.message_id,
|
||||
event.sequence,
|
||||
recipient_run_id,
|
||||
message.sender_run_id,
|
||||
message.subject,
|
||||
self.task_id.map(|task_id| task_id.to_string())
|
||||
);
|
||||
}
|
||||
|
||||
Some(ReceivedMessageInput {
|
||||
message_id: message.message_id,
|
||||
@@ -70,16 +135,35 @@ impl MessageHydrator {
|
||||
&self,
|
||||
message_id: &str,
|
||||
) -> Result<ReadAgentMessageResponse> {
|
||||
let read_message = self.ai_client.read_agent_message(message_id);
|
||||
let timeout = Timer::after(self.fetch_timeout);
|
||||
futures::pin_mut!(read_message);
|
||||
futures::pin_mut!(timeout);
|
||||
let deadline = Instant::now() + self.fetch_timeout;
|
||||
let mut last_error = None;
|
||||
|
||||
match futures::future::select(read_message, timeout).await {
|
||||
Either::Left((result, _)) => {
|
||||
result.with_context(|| format!("Failed to read agent message {message_id}"))
|
||||
loop {
|
||||
let remaining = deadline.saturating_duration_since(Instant::now());
|
||||
if remaining.is_zero() {
|
||||
return Err(message_read_timeout_error(message_id, last_error));
|
||||
}
|
||||
|
||||
let read_message = self.read_message(message_id);
|
||||
let timeout = Timer::after(remaining);
|
||||
futures::pin_mut!(read_message);
|
||||
futures::pin_mut!(timeout);
|
||||
|
||||
match futures::future::select(read_message, timeout).await {
|
||||
Either::Left((Ok(message), _)) => return Ok(message),
|
||||
Either::Left((Err(err), _)) if should_retry_message_read_error(&err) => {
|
||||
last_error = Some(err);
|
||||
let sleep_duration = self
|
||||
.retry_delay
|
||||
.min(deadline.saturating_duration_since(Instant::now()));
|
||||
if sleep_duration.is_zero() {
|
||||
return Err(message_read_timeout_error(message_id, last_error));
|
||||
}
|
||||
Timer::after(sleep_duration).await;
|
||||
}
|
||||
Either::Left((Err(err), _)) => return Err(err),
|
||||
Either::Right(_) => return Err(message_read_timeout_error(message_id, last_error)),
|
||||
}
|
||||
Either::Right(_) => Err(anyhow!("Timed out reading agent message {message_id}")),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -88,10 +172,7 @@ impl MessageHydrator {
|
||||
&self,
|
||||
message_id: &str,
|
||||
) -> Result<ReadAgentMessageResponse> {
|
||||
self.ai_client
|
||||
.read_agent_message(message_id)
|
||||
.await
|
||||
.with_context(|| format!("Failed to read agent message {message_id}"))
|
||||
self.read_message(message_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_message_from_event_with_timeout(
|
||||
@@ -105,10 +186,15 @@ impl MessageHydrator {
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_message_delivered(&self, message_id: &str) -> Result<()> {
|
||||
self.ai_client
|
||||
.mark_message_delivered(message_id)
|
||||
.await
|
||||
.with_context(|| format!("Failed to mark agent message {message_id} as delivered"))
|
||||
match (self.task_scoped_server_api.as_ref(), self.task_id) {
|
||||
(Some(server_api), Some(task_id)) => {
|
||||
server_api
|
||||
.mark_message_delivered_for_task(&task_id, message_id)
|
||||
.await
|
||||
}
|
||||
_ => self.ai_client.mark_message_delivered(message_id).await,
|
||||
}
|
||||
.with_context(|| format!("Failed to mark agent message {message_id} as delivered"))
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_messages_delivered_best_effort<'a, I>(
|
||||
@@ -130,3 +216,46 @@ impl MessageHydrator {
|
||||
failures
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn should_retry_message_read_error(err: &anyhow::Error) -> bool {
|
||||
// Immediate read-after-event lag can surface as a short-lived 404 before
|
||||
// the message row becomes readable. Restrict retries to status-preserving
|
||||
// eventual-consistency/server responses plus clearly transient transport
|
||||
// failures so permanent 4xxs fail fast instead of timing out.
|
||||
if let Some(status) = message_read_error_status(err) {
|
||||
return matches!(status, 404 | 408 | 429 | 500..=599);
|
||||
}
|
||||
|
||||
err.chain().any(|cause| {
|
||||
cause
|
||||
.downcast_ref::<ReqwestError>()
|
||||
.is_some_and(is_transient_message_read_transport_error)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn message_read_error_status(err: &anyhow::Error) -> Option<u16> {
|
||||
err.chain().find_map(|cause| {
|
||||
cause
|
||||
.downcast_ref::<HttpStatusError>()
|
||||
.map(|http_err| http_err.status)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn is_transient_message_read_transport_error(err: &ReqwestError) -> bool {
|
||||
err.is_timeout() || err.is_connect()
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn message_read_timeout_error(
|
||||
message_id: &str,
|
||||
last_error: Option<anyhow::Error>,
|
||||
) -> anyhow::Error {
|
||||
let timeout_error = anyhow!("Timed out reading agent message {message_id}");
|
||||
match last_error {
|
||||
Some(err) => timeout_error.context(err),
|
||||
None => timeout_error,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use mockall::predicate::eq;
|
||||
|
||||
@@ -6,6 +8,7 @@ use super::*;
|
||||
use crate::server::server_api::ai::{
|
||||
AIClient, AgentRunEvent, MockAIClient, ReadAgentMessageResponse,
|
||||
};
|
||||
use crate::server::server_api::presigned_upload::HttpStatusError;
|
||||
|
||||
fn make_run_event(
|
||||
sequence: i64,
|
||||
@@ -23,6 +26,25 @@ fn make_run_event(
|
||||
}
|
||||
}
|
||||
|
||||
fn make_message_response(message_id: &str) -> ReadAgentMessageResponse {
|
||||
ReadAgentMessageResponse {
|
||||
message_id: message_id.to_string(),
|
||||
sender_run_id: "parent-run".to_string(),
|
||||
subject: "Need a redirect".to_string(),
|
||||
body: "Switch to the failing test first.".to_string(),
|
||||
sent_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
delivered_at: None,
|
||||
read_at: Some("2026-01-01T00:00:01Z".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn http_status_read_error(status: u16) -> anyhow::Error {
|
||||
anyhow::Error::new(HttpStatusError {
|
||||
status,
|
||||
body: format!("status {status} body"),
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hydrator_reads_new_message_for_matching_run() {
|
||||
let mut ai_client = MockAIClient::new();
|
||||
@@ -30,17 +52,7 @@ async fn hydrator_reads_new_message_for_matching_run() {
|
||||
.expect_read_agent_message()
|
||||
.with(eq("msg-123"))
|
||||
.times(1)
|
||||
.returning(|_| {
|
||||
Ok(ReadAgentMessageResponse {
|
||||
message_id: "msg-123".to_string(),
|
||||
sender_run_id: "parent-run".to_string(),
|
||||
subject: "Need a redirect".to_string(),
|
||||
body: "Switch to the failing test first.".to_string(),
|
||||
sent_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
delivered_at: None,
|
||||
read_at: Some("2026-01-01T00:00:01Z".to_string()),
|
||||
})
|
||||
});
|
||||
.returning(|_| Ok(make_message_response("msg-123")));
|
||||
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
|
||||
let hydrator = MessageHydrator::new(ai_client);
|
||||
@@ -69,3 +81,109 @@ async fn hydrator_ignores_events_for_other_runs() {
|
||||
.await
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_message_with_timeout_retries_transient_failures_until_success() {
|
||||
let attempts = Arc::new(AtomicUsize::new(0));
|
||||
let attempts_clone = attempts.clone();
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client
|
||||
.expect_read_agent_message()
|
||||
.with(eq("msg-123"))
|
||||
.times(2)
|
||||
.returning(move |_| {
|
||||
let attempt = attempts_clone.fetch_add(1, Ordering::SeqCst);
|
||||
if attempt == 0 {
|
||||
Err(http_status_read_error(404))
|
||||
} else {
|
||||
Ok(make_message_response("msg-123"))
|
||||
}
|
||||
});
|
||||
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
|
||||
let hydrator = MessageHydrator::with_fetch_timing(
|
||||
ai_client,
|
||||
Duration::from_millis(100),
|
||||
Duration::from_millis(5),
|
||||
);
|
||||
|
||||
let message = hydrator.read_message_with_timeout("msg-123").await.unwrap();
|
||||
|
||||
assert_eq!(attempts.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(message.message_id, "msg-123");
|
||||
assert_eq!(message.body, "Switch to the failing test first.");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_message_with_timeout_times_out_after_retrying_transient_failures() {
|
||||
let attempts = Arc::new(AtomicUsize::new(0));
|
||||
let attempts_clone = attempts.clone();
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client
|
||||
.expect_read_agent_message()
|
||||
.with(eq("msg-123"))
|
||||
.returning(move |_| {
|
||||
attempts_clone.fetch_add(1, Ordering::SeqCst);
|
||||
Err(http_status_read_error(404))
|
||||
});
|
||||
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
|
||||
let hydrator = MessageHydrator::with_fetch_timing(
|
||||
ai_client,
|
||||
Duration::from_millis(120),
|
||||
Duration::from_millis(20),
|
||||
);
|
||||
|
||||
let err = hydrator
|
||||
.read_message_with_timeout("msg-123")
|
||||
.await
|
||||
.expect_err("expected timeout after transient retries");
|
||||
let err_chain = format!("{err:#}");
|
||||
|
||||
assert!(
|
||||
err_chain.contains("Timed out reading agent message msg-123"),
|
||||
"{err:#}"
|
||||
);
|
||||
assert!(
|
||||
attempts.load(Ordering::SeqCst) >= 2,
|
||||
"expected at least one retry before timeout"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_message_with_timeout_does_not_retry_permanent_http_failures() {
|
||||
let attempts = Arc::new(AtomicUsize::new(0));
|
||||
let attempts_clone = attempts.clone();
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client
|
||||
.expect_read_agent_message()
|
||||
.with(eq("msg-123"))
|
||||
.times(1)
|
||||
.returning(move |_| {
|
||||
attempts_clone.fetch_add(1, Ordering::SeqCst);
|
||||
Err(http_status_read_error(403))
|
||||
});
|
||||
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
|
||||
let hydrator = MessageHydrator::with_fetch_timing(
|
||||
ai_client,
|
||||
Duration::from_millis(100),
|
||||
Duration::from_millis(5),
|
||||
);
|
||||
|
||||
let err = hydrator
|
||||
.read_message_with_timeout("msg-123")
|
||||
.await
|
||||
.expect_err("expected permanent failure to fail fast");
|
||||
let err_chain = format!("{err:#}");
|
||||
|
||||
assert!(
|
||||
err_chain.contains("HTTP request failed with status 403"),
|
||||
"{err:#}"
|
||||
);
|
||||
assert_eq!(
|
||||
attempts.load(Ordering::SeqCst),
|
||||
1,
|
||||
"permanent 4xx errors should not retry"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -8,11 +8,11 @@ mod message_hydrator;
|
||||
pub(crate) use driver::{
|
||||
agent_event_backoff, agent_event_failures_exceeded_threshold, AgentEventDriverState,
|
||||
AgentEventSource, AgentEventSourceItem, DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
|
||||
DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS, DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
|
||||
};
|
||||
pub(crate) use driver::{
|
||||
run_agent_event_driver, AgentEventConsumer, AgentEventConsumerControlFlow,
|
||||
AgentEventDriverConfig, ServerApiAgentEventSource,
|
||||
AgentEventDriverConfig, AgentEventFilter, AgentMessageEventMetadata, ServerApiAgentEventSource,
|
||||
};
|
||||
pub(crate) use message_hydrator::MessageHydrator;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user