Initial public release of Warp.
Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b
This commit is contained in:
@@ -0,0 +1,366 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use futures::future::Either;
|
||||
use futures::StreamExt;
|
||||
use instant::Instant;
|
||||
use warpui::r#async::Timer;
|
||||
|
||||
use crate::server::server_api::ai::AgentRunEvent;
|
||||
use crate::server::server_api::ServerApi;
|
||||
|
||||
pub(crate) const DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10];
|
||||
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;
|
||||
|
||||
/// 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>,
|
||||
/// 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.
|
||||
pub reconnect_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.
|
||||
/// 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 {
|
||||
Self {
|
||||
run_ids,
|
||||
since_sequence,
|
||||
reconnect_backoff_steps: DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: Some(DEFAULT_AGENT_EVENT_PROACTIVE_RECONNECT),
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tells the shared driver whether to continue or stop after a handled event.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum AgentEventConsumerControlFlow {
|
||||
Continue,
|
||||
#[cfg_attr(not(test), allow(dead_code))]
|
||||
Stop,
|
||||
}
|
||||
|
||||
/// High-level connection state updates emitted by the shared driver.
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum AgentEventDriverState {
|
||||
Connected,
|
||||
RetryScheduled {
|
||||
/// Number of consecutive failed reconnect cycles since the last
|
||||
/// successful stream open or event delivery.
|
||||
consecutive_failures: usize,
|
||||
/// Delay before the next reconnect attempt.
|
||||
backoff: Duration,
|
||||
/// Whether the retry happened before the stream ever connected.
|
||||
is_initial_connect: bool,
|
||||
},
|
||||
/// A healthy stream was intentionally recycled after
|
||||
/// `proactive_reconnect_after`.
|
||||
ProactiveReconnect,
|
||||
}
|
||||
|
||||
/// Parsed items emitted by an [`AgentEventSource`].
|
||||
pub(crate) enum AgentEventSourceItem {
|
||||
Open,
|
||||
Event(AgentRunEvent),
|
||||
}
|
||||
|
||||
cfg_if::cfg_if! {
|
||||
if #[cfg(target_family = "wasm")] {
|
||||
type AgentEventSourceStream =
|
||||
futures::stream::LocalBoxStream<'static, Result<AgentEventSourceItem>>;
|
||||
} else {
|
||||
type AgentEventSourceStream =
|
||||
futures::stream::BoxStream<'static, Result<AgentEventSourceItem>>;
|
||||
}
|
||||
}
|
||||
|
||||
/// Opens a stream of parsed agent events for one or more run IDs.
|
||||
#[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],
|
||||
since_sequence: i64,
|
||||
) -> Result<AgentEventSourceStream>;
|
||||
}
|
||||
|
||||
/// [`AgentEventSource`] backed by [`ServerApi::stream_agent_events`].
|
||||
pub(crate) struct ServerApiAgentEventSource {
|
||||
server_api: Arc<ServerApi>,
|
||||
}
|
||||
|
||||
impl ServerApiAgentEventSource {
|
||||
pub(crate) fn new(server_api: Arc<ServerApi>) -> Self {
|
||||
Self { server_api }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
impl AgentEventSource for ServerApiAgentEventSource {
|
||||
async fn open_stream(
|
||||
&self,
|
||||
run_ids: &[String],
|
||||
since_sequence: i64,
|
||||
) -> Result<AgentEventSourceStream> {
|
||||
let stream = self
|
||||
.server_api
|
||||
.stream_agent_events(run_ids, since_sequence)
|
||||
.await?;
|
||||
|
||||
let stream = stream.filter_map(|event_result| async move {
|
||||
match event_result {
|
||||
Ok(reqwest_eventsource::Event::Open) => Some(Ok(AgentEventSourceItem::Open)),
|
||||
Ok(reqwest_eventsource::Event::Message(message)) => {
|
||||
match serde_json::from_str::<AgentRunEvent>(&message.data) {
|
||||
Ok(event) => Some(Ok(AgentEventSourceItem::Event(event))),
|
||||
Err(err) => {
|
||||
log::warn!("Skipping malformed agent event from SSE stream: {err}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => Some(Err(anyhow!("SSE stream error: {err:?}"))),
|
||||
}
|
||||
});
|
||||
|
||||
cfg_if::cfg_if! {
|
||||
if #[cfg(target_family = "wasm")] {
|
||||
Ok(stream.boxed_local())
|
||||
} else {
|
||||
Ok(stream.boxed())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Consumes events produced by [`run_agent_event_driver`].
|
||||
///
|
||||
/// Errors from [`on_event`](Self::on_event) are fatal because the event could not
|
||||
/// be safely processed. [`persist_cursor`](Self::persist_cursor) and
|
||||
/// [`on_driver_state`](Self::on_driver_state) are treated as best-effort hooks:
|
||||
/// their errors are logged and the driver continues.
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
pub(crate) trait AgentEventConsumer: Send {
|
||||
async fn on_event(&mut self, event: AgentRunEvent) -> Result<AgentEventConsumerControlFlow>;
|
||||
|
||||
async fn persist_cursor(&mut self, _sequence: i64) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn on_driver_state(&mut self, _state: AgentEventDriverState) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs a reconnecting agent-event stream until the consumer stops it or a
|
||||
/// fatal event-processing error occurs.
|
||||
pub(crate) async fn run_agent_event_driver<S, C>(
|
||||
source: S,
|
||||
config: AgentEventDriverConfig,
|
||||
consumer: &mut C,
|
||||
) -> Result<()>
|
||||
where
|
||||
S: AgentEventSource,
|
||||
C: AgentEventConsumer,
|
||||
{
|
||||
let mut since_sequence = config.since_sequence;
|
||||
let mut failures = 0usize;
|
||||
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
|
||||
}
|
||||
Err(err) => {
|
||||
failures += 1;
|
||||
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
|
||||
log_stream_failure(
|
||||
&config.run_ids,
|
||||
failures,
|
||||
backoff,
|
||||
&err,
|
||||
config.failures_before_error_log,
|
||||
);
|
||||
notify_driver_state(
|
||||
consumer,
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures: failures,
|
||||
backoff,
|
||||
is_initial_connect: !has_connected_once,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
Timer::after(backoff).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let proactive_reconnect_deadline = config
|
||||
.proactive_reconnect_after
|
||||
.map(|duration| Instant::now() + duration);
|
||||
|
||||
loop {
|
||||
let next_item = if let Some(deadline) = proactive_reconnect_deadline {
|
||||
let remaining = deadline.saturating_duration_since(Instant::now());
|
||||
if remaining.is_zero() {
|
||||
NextDriverItem::ProactiveReconnect
|
||||
} else {
|
||||
let next_stream_item = stream.next();
|
||||
let reconnect_timer = Timer::after(remaining);
|
||||
futures::pin_mut!(next_stream_item);
|
||||
futures::pin_mut!(reconnect_timer);
|
||||
match futures::future::select(next_stream_item, reconnect_timer).await {
|
||||
Either::Left((stream_item, _)) => NextDriverItem::StreamItem(stream_item),
|
||||
Either::Right(_) => NextDriverItem::ProactiveReconnect,
|
||||
}
|
||||
}
|
||||
} else {
|
||||
NextDriverItem::StreamItem(stream.next().await)
|
||||
};
|
||||
|
||||
match next_item {
|
||||
NextDriverItem::ProactiveReconnect => {
|
||||
notify_driver_state(consumer, AgentEventDriverState::ProactiveReconnect).await;
|
||||
break;
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Ok(AgentEventSourceItem::Open))) => {
|
||||
failures = 0;
|
||||
log::info!("Agent event stream opened for {:?}", config.run_ids);
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Ok(AgentEventSourceItem::Event(event)))) => {
|
||||
failures = 0;
|
||||
if event.sequence <= since_sequence {
|
||||
continue;
|
||||
}
|
||||
|
||||
let event_sequence = event.sequence;
|
||||
let control_flow = consumer.on_event(event).await?;
|
||||
since_sequence = event_sequence;
|
||||
|
||||
if let Err(err) = consumer.persist_cursor(since_sequence).await {
|
||||
log::warn!(
|
||||
"Ignoring agent event cursor persistence failure at sequence {since_sequence}: {err:#}"
|
||||
);
|
||||
}
|
||||
|
||||
if matches!(control_flow, AgentEventConsumerControlFlow::Stop) {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Err(err))) => {
|
||||
failures += 1;
|
||||
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
|
||||
log_stream_failure(
|
||||
&config.run_ids,
|
||||
failures,
|
||||
backoff,
|
||||
&err,
|
||||
config.failures_before_error_log,
|
||||
);
|
||||
notify_driver_state(
|
||||
consumer,
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures: failures,
|
||||
backoff,
|
||||
is_initial_connect: false,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
Timer::after(backoff).await;
|
||||
break;
|
||||
}
|
||||
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
|
||||
);
|
||||
notify_driver_state(
|
||||
consumer,
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures: failures,
|
||||
backoff,
|
||||
is_initial_connect: false,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
Timer::after(backoff).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum NextDriverItem {
|
||||
StreamItem(Option<Result<AgentEventSourceItem>>),
|
||||
ProactiveReconnect,
|
||||
}
|
||||
|
||||
async fn notify_driver_state<C: AgentEventConsumer>(
|
||||
consumer: &mut C,
|
||||
state: AgentEventDriverState,
|
||||
) {
|
||||
if let Err(err) = consumer.on_driver_state(state.clone()).await {
|
||||
log::warn!("Ignoring agent event driver state callback error for {state:?}: {err:#}");
|
||||
}
|
||||
}
|
||||
|
||||
fn log_stream_failure(
|
||||
run_ids: &[String],
|
||||
failures: usize,
|
||||
backoff: Duration,
|
||||
err: &anyhow::Error,
|
||||
failures_before_error_log: usize,
|
||||
) {
|
||||
if agent_event_failures_exceeded_threshold(failures, failures_before_error_log) {
|
||||
log::error!(
|
||||
"Agent event stream failed {failures} consecutive times for {:?}, retrying in {backoff:?}: {err:#}",
|
||||
run_ids
|
||||
);
|
||||
} else {
|
||||
log::warn!(
|
||||
"Agent event stream failed for {:?}, retrying in {backoff:?}: {err:#}",
|
||||
run_ids
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn agent_event_backoff(failures: usize, backoff_steps: &[u64]) -> Duration {
|
||||
let safe_steps = if backoff_steps.is_empty() {
|
||||
DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS
|
||||
} else {
|
||||
backoff_steps
|
||||
};
|
||||
let index = failures.saturating_sub(1).min(safe_steps.len() - 1);
|
||||
Duration::from_secs(safe_steps[index])
|
||||
}
|
||||
|
||||
pub(crate) fn agent_event_failures_exceeded_threshold(failures: usize, threshold: usize) -> bool {
|
||||
failures >= threshold
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::anyhow;
|
||||
use async_trait::async_trait;
|
||||
use futures::stream::{self, BoxStream};
|
||||
use futures::StreamExt;
|
||||
|
||||
use super::*;
|
||||
use crate::server::server_api::ai::AgentRunEvent;
|
||||
|
||||
const ZERO_BACKOFF_STEPS: &[u64] = &[0];
|
||||
|
||||
struct FakeAgentEventSource {
|
||||
responses: Mutex<VecDeque<anyhow::Result<Vec<anyhow::Result<AgentEventSourceItem>>>>>,
|
||||
}
|
||||
|
||||
impl FakeAgentEventSource {
|
||||
fn new(responses: Vec<anyhow::Result<Vec<anyhow::Result<AgentEventSourceItem>>>>) -> Self {
|
||||
Self {
|
||||
responses: Mutex::new(VecDeque::from(responses)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentEventSource for FakeAgentEventSource {
|
||||
async fn open_stream(
|
||||
&self,
|
||||
_run_ids: &[String],
|
||||
_since_sequence: i64,
|
||||
) -> anyhow::Result<BoxStream<'static, anyhow::Result<AgentEventSourceItem>>> {
|
||||
let response = self
|
||||
.responses
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.expect("fake response missing");
|
||||
let stream = response?;
|
||||
Ok(stream::iter(stream).boxed())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct RecordingConsumer {
|
||||
handled_sequences: Vec<i64>,
|
||||
persisted_sequences: Vec<i64>,
|
||||
driver_states: Vec<AgentEventDriverState>,
|
||||
stop_after: usize,
|
||||
fail_persist_cursor: bool,
|
||||
fail_driver_state: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentEventConsumer for RecordingConsumer {
|
||||
async fn on_event(
|
||||
&mut self,
|
||||
event: AgentRunEvent,
|
||||
) -> anyhow::Result<AgentEventConsumerControlFlow> {
|
||||
self.handled_sequences.push(event.sequence);
|
||||
if self.handled_sequences.len() >= self.stop_after {
|
||||
Ok(AgentEventConsumerControlFlow::Stop)
|
||||
} else {
|
||||
Ok(AgentEventConsumerControlFlow::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_cursor(&mut self, sequence: i64) -> anyhow::Result<()> {
|
||||
self.persisted_sequences.push(sequence);
|
||||
if self.fail_persist_cursor {
|
||||
Err(anyhow!("persist failed"))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_driver_state(&mut self, state: AgentEventDriverState) -> anyhow::Result<()> {
|
||||
self.driver_states.push(state);
|
||||
if self.fail_driver_state {
|
||||
Err(anyhow!("state callback failed"))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn make_run_event(
|
||||
sequence: i64,
|
||||
event_type: &str,
|
||||
run_id: &str,
|
||||
ref_id: Option<&str>,
|
||||
) -> AgentRunEvent {
|
||||
AgentRunEvent {
|
||||
event_type: event_type.to_string(),
|
||||
run_id: run_id.to_string(),
|
||||
ref_id: ref_id.map(|value| value.to_string()),
|
||||
execution_id: None,
|
||||
occurred_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
sequence,
|
||||
}
|
||||
}
|
||||
|
||||
fn ok_stream(
|
||||
items: Vec<anyhow::Result<AgentEventSourceItem>>,
|
||||
) -> anyhow::Result<Vec<anyhow::Result<AgentEventSourceItem>>> {
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_skips_duplicate_sequences_and_persists_new_cursor() {
|
||||
let source = FakeAgentEventSource::new(vec![ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
2,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-2"),
|
||||
))),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
3,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-3"),
|
||||
))),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
4,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-4"),
|
||||
))),
|
||||
])]);
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 2,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 2,
|
||||
reconnect_backoff_steps: DEFAULT_AGENT_EVENT_RECONNECT_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![3, 4]);
|
||||
assert_eq!(consumer.persisted_sequences, vec![3, 4]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_resets_failures_after_successful_event_delivery() {
|
||||
let source = FakeAgentEventSource::new(vec![
|
||||
ok_stream(vec![Ok(AgentEventSourceItem::Open), Err(anyhow!("boom-1"))]),
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
1,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-1"),
|
||||
))),
|
||||
Err(anyhow!("boom-2")),
|
||||
]),
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
2,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-2"),
|
||||
))),
|
||||
]),
|
||||
]);
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 2,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 0,
|
||||
reconnect_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();
|
||||
|
||||
let retry_failures = consumer
|
||||
.driver_states
|
||||
.into_iter()
|
||||
.filter_map(|state| match state {
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures,
|
||||
..
|
||||
} => Some(consecutive_failures),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(retry_failures, vec![1, 1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_ignores_persist_cursor_errors() {
|
||||
let source = FakeAgentEventSource::new(vec![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,
|
||||
fail_persist_cursor: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 0,
|
||||
reconnect_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]);
|
||||
assert_eq!(consumer.persisted_sequences, vec![1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_ignores_driver_state_errors() {
|
||||
let source = FakeAgentEventSource::new(vec![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,
|
||||
fail_driver_state: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 0,
|
||||
reconnect_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]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_retries_initial_connection_until_stream_opens() {
|
||||
let source = FakeAgentEventSource::new(vec![
|
||||
Err(anyhow!("boom-1")),
|
||||
Err(anyhow!("boom-2")),
|
||||
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()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 0,
|
||||
reconnect_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]);
|
||||
let retry_failures = consumer
|
||||
.driver_states
|
||||
.into_iter()
|
||||
.filter_map(|state| match state {
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures,
|
||||
is_initial_connect,
|
||||
..
|
||||
} if is_initial_connect => Some(consecutive_failures),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(retry_failures, vec![1, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backoff_escalates_then_caps() {
|
||||
assert_eq!(
|
||||
agent_event_backoff(1, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(1)
|
||||
);
|
||||
assert_eq!(
|
||||
agent_event_backoff(2, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(2)
|
||||
);
|
||||
assert_eq!(
|
||||
agent_event_backoff(3, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(5)
|
||||
);
|
||||
assert_eq!(
|
||||
agent_event_backoff(4, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(10)
|
||||
);
|
||||
assert_eq!(
|
||||
agent_event_backoff(100, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(10)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failure_threshold_is_reached_at_and_above_limit() {
|
||||
assert!(!agent_event_failures_exceeded_threshold(4, 5));
|
||||
assert!(agent_event_failures_exceeded_threshold(5, 5));
|
||||
assert!(agent_event_failures_exceeded_threshold(6, 5));
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use futures::future::Either;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use warpui::r#async::Timer;
|
||||
|
||||
use crate::ai::agent::ReceivedMessageInput;
|
||||
use crate::server::server_api::ai::{AIClient, AgentRunEvent, ReadAgentMessageResponse};
|
||||
|
||||
pub(crate) const DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
/// Hydrates `new_message` agent events into full message payloads and delivery
|
||||
/// acknowledgements.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MessageHydrator {
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
fetch_timeout: Duration,
|
||||
}
|
||||
|
||||
impl MessageHydrator {
|
||||
pub(crate) fn new(ai_client: Arc<dyn AIClient>) -> Self {
|
||||
Self::with_fetch_timeout(ai_client, DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT)
|
||||
}
|
||||
|
||||
pub(crate) fn with_fetch_timeout(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
fetch_timeout: Duration,
|
||||
) -> Self {
|
||||
Self {
|
||||
ai_client,
|
||||
fetch_timeout,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn hydrate_event_for_recipient(
|
||||
&self,
|
||||
event: &AgentRunEvent,
|
||||
recipient_run_id: &str,
|
||||
) -> Option<ReceivedMessageInput> {
|
||||
if event.event_type != "new_message" || event.run_id != recipient_run_id {
|
||||
return None;
|
||||
}
|
||||
|
||||
let message = match self.read_message_from_event_with_timeout(event).await {
|
||||
Ok(message) => message,
|
||||
Err(err) => {
|
||||
log::warn!(
|
||||
"Failed to hydrate agent message for event ref_id={:?}: {err:#}",
|
||||
event.ref_id
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(ReceivedMessageInput {
|
||||
message_id: message.message_id,
|
||||
sender_agent_id: message.sender_run_id,
|
||||
addresses: vec![recipient_run_id.to_string()],
|
||||
subject: message.subject,
|
||||
message_body: message.body,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub(crate) async fn read_message_with_timeout(
|
||||
&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);
|
||||
|
||||
match futures::future::select(read_message, timeout).await {
|
||||
Either::Left((result, _)) => {
|
||||
result.with_context(|| format!("Failed to read agent message {message_id}"))
|
||||
}
|
||||
Either::Right(_) => Err(anyhow!("Timed out reading agent message {message_id}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_family = "wasm")]
|
||||
pub(crate) async fn read_message_with_timeout(
|
||||
&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}"))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_message_from_event_with_timeout(
|
||||
&self,
|
||||
event: &AgentRunEvent,
|
||||
) -> Result<ReadAgentMessageResponse> {
|
||||
let Some(message_id) = event.ref_id.as_deref() else {
|
||||
return Err(anyhow!("Agent event is missing ref_id"));
|
||||
};
|
||||
self.read_message_with_timeout(message_id).await
|
||||
}
|
||||
|
||||
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"))
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_messages_delivered_best_effort<'a, I>(
|
||||
&self,
|
||||
message_ids: I,
|
||||
) -> Vec<(String, anyhow::Error)>
|
||||
where
|
||||
I: IntoIterator<Item = &'a str>,
|
||||
{
|
||||
let mut failures = Vec::new();
|
||||
// TODO(REMOTE-1266): Parallelize delivery acknowledgements for bursty
|
||||
// batches once the parent-bridge restore path is hardened enough to
|
||||
// tolerate a concurrent FuturesUnordered/join_all flow here.
|
||||
for message_id in message_ids {
|
||||
if let Err(err) = self.mark_message_delivered(message_id).await {
|
||||
failures.push((message_id.to_string(), err));
|
||||
}
|
||||
}
|
||||
failures
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use mockall::predicate::eq;
|
||||
|
||||
use super::*;
|
||||
use crate::server::server_api::ai::{
|
||||
AIClient, AgentRunEvent, MockAIClient, ReadAgentMessageResponse,
|
||||
};
|
||||
|
||||
fn make_run_event(
|
||||
sequence: i64,
|
||||
event_type: &str,
|
||||
run_id: &str,
|
||||
ref_id: Option<&str>,
|
||||
) -> AgentRunEvent {
|
||||
AgentRunEvent {
|
||||
event_type: event_type.to_string(),
|
||||
run_id: run_id.to_string(),
|
||||
ref_id: ref_id.map(|value| value.to_string()),
|
||||
execution_id: None,
|
||||
occurred_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
sequence,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hydrator_reads_new_message_for_matching_run() {
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client
|
||||
.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()),
|
||||
})
|
||||
});
|
||||
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
|
||||
let hydrator = MessageHydrator::new(ai_client);
|
||||
let event = make_run_event(7, "new_message", "child-run", Some("msg-123"));
|
||||
|
||||
let hydrated = hydrator
|
||||
.hydrate_event_for_recipient(&event, "child-run")
|
||||
.await
|
||||
.expect("expected hydrated message");
|
||||
|
||||
assert_eq!(hydrated.message_id, "msg-123");
|
||||
assert_eq!(hydrated.sender_agent_id, "parent-run");
|
||||
assert_eq!(hydrated.subject, "Need a redirect");
|
||||
assert_eq!(hydrated.message_body, "Switch to the failing test first.");
|
||||
assert_eq!(hydrated.addresses, vec!["child-run".to_string()]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hydrator_ignores_events_for_other_runs() {
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(MockAIClient::new());
|
||||
let hydrator = MessageHydrator::new(ai_client);
|
||||
let event = make_run_event(7, "new_message", "other-run", Some("msg-123"));
|
||||
|
||||
assert!(hydrator
|
||||
.hydrate_event_for_recipient(&event, "child-run")
|
||||
.await
|
||||
.is_none());
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
//! Shared agent-event stream utilities used by orchestration consumers and
|
||||
//! third-party harness bridges.
|
||||
|
||||
mod driver;
|
||||
mod message_hydrator;
|
||||
|
||||
#[cfg(test)]
|
||||
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,
|
||||
};
|
||||
pub(crate) use driver::{
|
||||
run_agent_event_driver, AgentEventConsumer, AgentEventConsumerControlFlow,
|
||||
AgentEventDriverConfig, ServerApiAgentEventSource,
|
||||
};
|
||||
pub(crate) use message_hydrator::MessageHydrator;
|
||||
|
||||
#[cfg(test)]
|
||||
mod driver_tests;
|
||||
#[cfg(test)]
|
||||
mod message_hydrator_tests;
|
||||
Reference in New Issue
Block a user