first pass of merging in warp (doesn't build)

This commit is contained in:
Ryan Ward
2026-07-01 16:08:58 -05:00
parent 2f64909469
commit 4770ac06b5
3662 changed files with 414574 additions and 89772 deletions
+176 -36
View File
@@ -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()
}
+170 -6
View File
@@ -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));
}
+145 -16
View File
@@ -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,
}
}
+129 -11
View File
@@ -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"
);
}
+2 -2
View File
@@ -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;