Remove orchestration_event_push feature flag; rename poller to streamer (#9265)
## Description Removes the `orchestration_event_push` feature flag and the polling fallback in the orchestration event delivery path. SSE-based event push has been on in dogfood and staging long enough that it's now the only path; the dual-mode `OrchestrationEventPoller` is renamed to `OrchestrationEventStreamer` and only opens persistent SSE connections. - Removed `FeatureFlag::OrchestrationEventPush` and the matching Cargo feature. - Renamed `app/src/ai/blocklist/orchestration_event_poller.rs` (and its tests) to `orchestration_event_streamer.rs`. Renamed the public type to `OrchestrationEventStreamer`; renamed the shared event-injection sink `handle_poll_result` → `handle_event_batch`. - Removed polling-only state (`poll_backoff_index`, `poll_in_flight`), methods (`poll_and_inject`, `start_idle_poll_timer`), and constants (`POLL_BACKOFF_STEPS`, `EVENT_POLL_BATCH_LIMIT`). Kept `event_cursor`, `pending_delivery`, and `conversation_statuses` since they are also used by the SSE path. - Removed the now-dead `AIClient::poll_agent_events` trait method and its `ServerApi` implementation in `app/src/server/server_api/ai.rs`. The `OrchestrationV2` flag continues to gate streamer instantiation and watched-run registration; runtime behavior under v2 is unchanged. Pairs with the warp-server PR that drops the flag server-side: https://github.com/warpdotdev/warp-server/pull/10736 — that PR should land first so the polling endpoint stays available for older clients during rollout. ## Testing - `cargo build -p warp` clean. - `cargo clippy -p warp -p warp_features --all-targets -- -D warnings` clean (after `cargo fmt`). - `cargo test -p warp --lib ai::blocklist::orchestration_event_streamer` — all 7 tests pass. ## Agent Mode - [x] Warp Agent Mode - This PR was created via Warp's AI Agent Mode --- [Plan: Remove orchestration_event_push flag (server + client)](https://staging.warp.dev/drive/notebook/yAlAxMAr4EO65A9LEy40cX) [Conversation](https://staging.warp.dev/conversation/ff5e82cb-bba7-4cbf-ae0b-e51f2c542f4a) --------- Co-authored-by: Oz <oz-agent@warp.dev>
This commit is contained in:
Matthew Albright
@@ -6,7 +6,7 @@ use crate::ai::agent::{
|
||||
AIAgentAction, AIAgentActionResultType, AIAgentActionType, LifecycleEventType,
|
||||
StartAgentExecutionMode, StartAgentResult,
|
||||
};
|
||||
use crate::ai::blocklist::orchestration_event_poller::OrchestrationEventPoller;
|
||||
use crate::ai::blocklist::orchestration_event_streamer::OrchestrationEventStreamer;
|
||||
use crate::ai::blocklist::orchestration_events::OrchestrationEventService;
|
||||
use crate::ai::blocklist::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel};
|
||||
use warp_cli::agent::Harness;
|
||||
@@ -110,8 +110,8 @@ impl StartAgentExecutor {
|
||||
agent_id: id.clone(),
|
||||
});
|
||||
if FeatureFlag::OrchestrationV2.is_enabled() {
|
||||
OrchestrationEventPoller::handle(ctx).update(ctx, |poller, ctx| {
|
||||
poller.register_watched_run_id(
|
||||
OrchestrationEventStreamer::handle(ctx).update(ctx, |streamer, ctx| {
|
||||
streamer.register_watched_run_id(
|
||||
pending.parent_conversation_id,
|
||||
id,
|
||||
ctx,
|
||||
|
||||
@@ -5,7 +5,7 @@ pub mod block;
|
||||
pub mod code_block;
|
||||
mod context_model;
|
||||
mod controller;
|
||||
pub(crate) mod orchestration_event_poller;
|
||||
pub(crate) mod orchestration_event_streamer;
|
||||
pub(crate) mod orchestration_events;
|
||||
mod passive_suggestions;
|
||||
pub(crate) mod task_status_sync_model;
|
||||
|
||||
+43
-180
@@ -22,11 +22,9 @@ use warp_multi_agent_api as api;
|
||||
use warpui::r#async::Timer;
|
||||
use warpui::{Entity, ModelContext, SingletonEntity};
|
||||
|
||||
/// Adaptive polling backoff: 1s, 2s, 5s, then 10s max. Resets to 1s when
|
||||
/// events are found.
|
||||
const POLL_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10];
|
||||
/// Keep each catch-up poll bounded so the event poller can drain backlog without overfetching.
|
||||
const EVENT_POLL_BATCH_LIMIT: i32 = 100;
|
||||
/// Backoff schedule (seconds) reused for the post-restore
|
||||
/// `get_ambient_agent_task` retry: 1s, 2s, 5s, then 10s max.
|
||||
const RESTORE_FETCH_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10];
|
||||
/// How often (milliseconds) the drain timer checks for SSE events.
|
||||
const SSE_DRAIN_INTERVAL_MS: u64 = 500;
|
||||
|
||||
@@ -78,23 +76,18 @@ impl AgentEventConsumer for SseForwardingConsumer {
|
||||
}
|
||||
}
|
||||
|
||||
/// Async network coordinator for v2 orchestration event delivery.
|
||||
/// Owns polling, adaptive backoff, event cursors, watched run_ids,
|
||||
/// lifecycle reporting, delivery confirmation, and self-registration.
|
||||
///
|
||||
/// When the `OrchestrationEventPush` feature flag is enabled the poller
|
||||
/// opens a persistent SSE connection to the server instead of short-polling.
|
||||
/// Async network coordinator for v2 orchestration event delivery via SSE.
|
||||
/// Opens persistent SSE connections to the server and forwards events into
|
||||
/// the OrchestrationEventService. Owns watched run_ids, event cursors used
|
||||
/// for deduplication, lifecycle reporting, and delivery confirmation.
|
||||
/// SSE retries with exponential backoff on failure.
|
||||
pub struct OrchestrationEventPoller {
|
||||
pub struct OrchestrationEventStreamer {
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
server_api: Arc<ServerApi>,
|
||||
watched_run_ids: HashMap<AIConversationId, HashSet<String>>,
|
||||
event_cursor: HashMap<AIConversationId, i64>,
|
||||
poll_backoff_index: HashMap<AIConversationId, usize>,
|
||||
pending_delivery: HashMap<AIConversationId, PendingDeliveryConfirmation>,
|
||||
conversation_statuses: HashMap<AIConversationId, ConversationStatus>,
|
||||
poll_in_flight: HashSet<AIConversationId>,
|
||||
// ---- SSE state ----
|
||||
/// Active SSE connections keyed by conversation.
|
||||
sse_connections: HashMap<AIConversationId, SseConnectionState>,
|
||||
/// Monotonic counter for SSE connection generations. Ensures stale
|
||||
@@ -105,11 +98,11 @@ pub struct OrchestrationEventPoller {
|
||||
restore_fetch_failures: HashMap<AIConversationId, usize>,
|
||||
}
|
||||
|
||||
pub enum OrchestrationEventPollerEvent {
|
||||
pub enum OrchestrationEventStreamerEvent {
|
||||
// Reserved for future use (e.g., status signals to the controller).
|
||||
}
|
||||
|
||||
impl OrchestrationEventPoller {
|
||||
impl OrchestrationEventStreamer {
|
||||
pub fn new(ctx: &mut ModelContext<Self>) -> Self {
|
||||
let provider = ServerApiProvider::as_ref(ctx);
|
||||
let ai_client = provider.get_ai_client();
|
||||
@@ -123,17 +116,15 @@ impl OrchestrationEventPoller {
|
||||
server_api,
|
||||
watched_run_ids: HashMap::new(),
|
||||
event_cursor: HashMap::new(),
|
||||
poll_backoff_index: HashMap::new(),
|
||||
pending_delivery: HashMap::new(),
|
||||
conversation_statuses: HashMap::new(),
|
||||
poll_in_flight: HashSet::new(),
|
||||
sse_connections: HashMap::new(),
|
||||
next_sse_generation: 0,
|
||||
restore_fetch_failures: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Constructs a poller wired to the supplied (mock) clients instead of
|
||||
/// Constructs a streamer wired to the supplied (mock) clients instead of
|
||||
/// looking them up via `ServerApiProvider`. Lets unit tests inject a
|
||||
/// `MockAIClient` while still subscribing to `BlocklistAIHistoryModel`.
|
||||
#[cfg(test)]
|
||||
@@ -151,10 +142,8 @@ impl OrchestrationEventPoller {
|
||||
server_api,
|
||||
watched_run_ids: HashMap::new(),
|
||||
event_cursor: HashMap::new(),
|
||||
poll_backoff_index: HashMap::new(),
|
||||
pending_delivery: HashMap::new(),
|
||||
conversation_statuses: HashMap::new(),
|
||||
poll_in_flight: HashSet::new(),
|
||||
sse_connections: HashMap::new(),
|
||||
next_sse_generation: 0,
|
||||
restore_fetch_failures: HashMap::new(),
|
||||
@@ -215,13 +204,11 @@ impl OrchestrationEventPoller {
|
||||
} => {
|
||||
self.watched_run_ids.remove(conversation_id);
|
||||
self.event_cursor.remove(conversation_id);
|
||||
self.poll_backoff_index.remove(conversation_id);
|
||||
self.pending_delivery.remove(conversation_id);
|
||||
self.conversation_statuses.remove(conversation_id);
|
||||
self.poll_in_flight.remove(conversation_id);
|
||||
self.restore_fetch_failures.remove(conversation_id);
|
||||
// SSE cleanup
|
||||
// task's next send to fail, which terminates the task.
|
||||
// Dropping the SSE connection state closes the channel,
|
||||
// causing the task's next send to fail and terminate.
|
||||
self.sse_connections.remove(conversation_id);
|
||||
}
|
||||
BlocklistAIHistoryEvent::StartedNewConversation { .. }
|
||||
@@ -249,7 +236,7 @@ impl OrchestrationEventPoller {
|
||||
///
|
||||
/// Re-establishes orchestration event delivery state that is not persisted
|
||||
/// directly in memory: watched run_ids, the per-conversation event cursor,
|
||||
/// and — for `Success` parents with watched children — the poll/SSE loop.
|
||||
/// and — for `Success` parents with watched children — the SSE event loop.
|
||||
fn on_restored_conversations(
|
||||
&mut self,
|
||||
conversation_ids: Vec<AIConversationId>,
|
||||
@@ -277,7 +264,7 @@ impl OrchestrationEventPoller {
|
||||
};
|
||||
|
||||
// Shared-session viewers receive updates through session sharing;
|
||||
// polling here would re-inject events the session has already
|
||||
// subscribing here would re-inject events the session has already
|
||||
// processed.
|
||||
if is_viewer {
|
||||
continue;
|
||||
@@ -296,7 +283,7 @@ impl OrchestrationEventPoller {
|
||||
self.conversation_statuses.insert(conv_id, status.clone());
|
||||
|
||||
// Register the conversation's own run_id so lifecycle events for
|
||||
// self are correctly filtered and the SSE/poll loop has a set
|
||||
// self are correctly filtered and the SSE loop has a set
|
||||
// of run_ids to open against.
|
||||
if let Some(ref own) = run_id {
|
||||
self.watched_run_ids
|
||||
@@ -363,7 +350,7 @@ impl OrchestrationEventPoller {
|
||||
match run_result {
|
||||
Ok(task) => {
|
||||
// If the conversation was removed while the fetch was in-flight,
|
||||
// the removal handler already cleaned up all poller state. Return
|
||||
// the removal handler already cleaned up all streamer state. Return
|
||||
// early to avoid recreating watched_run_ids for a deleted conversation.
|
||||
if !self.event_cursor.contains_key(&conv_id) {
|
||||
self.restore_fetch_failures.remove(&conv_id);
|
||||
@@ -415,8 +402,8 @@ impl OrchestrationEventPoller {
|
||||
|
||||
/// Schedules a retry of the post-restore `get_ambient_agent_task` fetch
|
||||
/// after an exponential backoff. The backoff schedule reuses
|
||||
/// `POLL_BACKOFF_STEPS` (1s, 2s, 5s, 10s capped) keyed on a per-conversation
|
||||
/// failure counter. The counter resets on success.
|
||||
/// `RESTORE_FETCH_BACKOFF_STEPS` (1s, 2s, 5s, 10s capped) keyed on a
|
||||
/// per-conversation failure counter. The counter resets on success.
|
||||
fn start_restore_fetch_retry_timer(
|
||||
&mut self,
|
||||
conv_id: AIConversationId,
|
||||
@@ -429,8 +416,10 @@ impl OrchestrationEventPoller {
|
||||
.entry(conv_id)
|
||||
.and_modify(|c| *c += 1)
|
||||
.or_insert(1);
|
||||
let step_index = failures.saturating_sub(1).min(POLL_BACKOFF_STEPS.len() - 1);
|
||||
let backoff = Duration::from_secs(POLL_BACKOFF_STEPS[step_index]);
|
||||
let step_index = failures
|
||||
.saturating_sub(1)
|
||||
.min(RESTORE_FETCH_BACKOFF_STEPS.len() - 1);
|
||||
let backoff = Duration::from_secs(RESTORE_FETCH_BACKOFF_STEPS[step_index]);
|
||||
ctx.spawn(
|
||||
async move { Timer::after(backoff).await },
|
||||
move |me, _, ctx| {
|
||||
@@ -488,9 +477,8 @@ impl OrchestrationEventPoller {
|
||||
let became_success = matches!(¤t_status, ConversationStatus::Success)
|
||||
&& !matches!(previous_status.as_ref(), Some(ConversationStatus::Success));
|
||||
|
||||
// Trigger event delivery when a conversation with watched run_ids
|
||||
// becomes idle. With the event-push flag this opens an SSE stream;
|
||||
// otherwise it falls back to the existing polling loop.
|
||||
// Open an SSE stream when a conversation with watched run_ids
|
||||
// becomes idle.
|
||||
if became_success && self.watched_run_ids.contains_key(&conversation_id) {
|
||||
self.start_event_delivery(conversation_id, ctx);
|
||||
}
|
||||
@@ -507,9 +495,9 @@ impl OrchestrationEventPoller {
|
||||
else {
|
||||
return;
|
||||
};
|
||||
// Shared session viewers must not poll for events — the actual
|
||||
// agent handles event delivery. Polling here would re-inject
|
||||
// events the session has already processed.
|
||||
// Shared session viewers must not subscribe to events — the
|
||||
// actual agent handles event delivery. Subscribing here would
|
||||
// re-inject events the session has already processed.
|
||||
if conversation.is_viewing_shared_session() {
|
||||
return;
|
||||
}
|
||||
@@ -586,93 +574,10 @@ impl OrchestrationEventPoller {
|
||||
);
|
||||
}
|
||||
|
||||
/// Polls the server for events and feeds them into the service queue.
|
||||
fn poll_and_inject(&mut self, conversation_id: AIConversationId, ctx: &mut ModelContext<Self>) {
|
||||
if self.poll_in_flight.contains(&conversation_id) {
|
||||
return;
|
||||
}
|
||||
let Some(watched) = self.watched_run_ids.get(&conversation_id) else {
|
||||
return;
|
||||
};
|
||||
if watched.is_empty() {
|
||||
return;
|
||||
}
|
||||
self.poll_in_flight.insert(conversation_id);
|
||||
let watched: Vec<String> = watched.iter().cloned().collect();
|
||||
let cursor = self
|
||||
.event_cursor
|
||||
.get(&conversation_id)
|
||||
.copied()
|
||||
.unwrap_or(0);
|
||||
|
||||
let ai_client = self.ai_client.clone();
|
||||
let hydrator = MessageHydrator::new(ai_client.clone());
|
||||
|
||||
// Capture own run_id to filter out self-originated lifecycle events.
|
||||
let self_run_id = BlocklistAIHistoryModel::as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.and_then(|c| c.run_id())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_default();
|
||||
|
||||
struct PollResult {
|
||||
events: Vec<crate::server::server_api::ai::AgentRunEvent>,
|
||||
fetched_messages: Vec<ReceivedMessageInput>,
|
||||
}
|
||||
|
||||
let self_run_id_clone = self_run_id.clone();
|
||||
ctx.spawn(
|
||||
async move {
|
||||
let events = ai_client
|
||||
.poll_agent_events(&watched, cursor, EVENT_POLL_BATCH_LIMIT)
|
||||
.await?;
|
||||
|
||||
let mut fetched_messages = Vec::new();
|
||||
for event in &events {
|
||||
if let Some(message) = hydrator
|
||||
.hydrate_event_for_recipient(event, &self_run_id)
|
||||
.await
|
||||
{
|
||||
fetched_messages.push(message);
|
||||
}
|
||||
}
|
||||
|
||||
Ok::<_, anyhow::Error>(PollResult {
|
||||
events,
|
||||
fetched_messages,
|
||||
})
|
||||
},
|
||||
move |me, result, ctx| {
|
||||
me.poll_in_flight.remove(&conversation_id);
|
||||
let self_run_id = self_run_id_clone;
|
||||
let poll_result = match result {
|
||||
Ok(r) => r,
|
||||
Err(err) => {
|
||||
log::warn!("V2 event poll failed for {conversation_id:?}: {err:#}");
|
||||
me.start_idle_poll_timer(conversation_id, ctx);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if poll_result.events.is_empty() {
|
||||
me.start_idle_poll_timer(conversation_id, ctx);
|
||||
return;
|
||||
}
|
||||
|
||||
me.handle_poll_result(
|
||||
conversation_id,
|
||||
&self_run_id,
|
||||
cursor,
|
||||
poll_result.events,
|
||||
poll_result.fetched_messages,
|
||||
ctx,
|
||||
);
|
||||
me.start_idle_poll_timer(conversation_id, ctx);
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn handle_poll_result(
|
||||
/// Feeds a batch of fetched events through the OrchestrationEventService,
|
||||
/// updating the in-memory and persisted cursors and tracking message IDs
|
||||
/// awaiting delivery confirmation.
|
||||
fn handle_event_batch(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
self_run_id: &str,
|
||||
@@ -743,65 +648,23 @@ impl OrchestrationEventPoller {
|
||||
return;
|
||||
}
|
||||
|
||||
// Only reset backoff when events actually produce pending items.
|
||||
self.poll_backoff_index.remove(&conversation_id);
|
||||
|
||||
let pending = build_pending_events(messages, lifecycle_events);
|
||||
OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| {
|
||||
svc.enqueue_polled_events(conversation_id, pending, ctx);
|
||||
svc.enqueue_event_batch(conversation_id, pending, ctx);
|
||||
});
|
||||
}
|
||||
|
||||
/// Starts a background poll timer with adaptive backoff.
|
||||
fn start_idle_poll_timer(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if !self.watched_run_ids.contains_key(&conversation_id) {
|
||||
return;
|
||||
}
|
||||
|
||||
let index = self.poll_backoff_index.entry(conversation_id).or_insert(0);
|
||||
let interval_secs = POLL_BACKOFF_STEPS[(*index).min(POLL_BACKOFF_STEPS.len() - 1)];
|
||||
*index = (*index + 1).min(POLL_BACKOFF_STEPS.len() - 1);
|
||||
|
||||
ctx.spawn(
|
||||
async move { Timer::after(Duration::from_secs(interval_secs)).await },
|
||||
move |me, _, ctx| {
|
||||
// Re-check that the conversation is still idle before polling.
|
||||
let is_success = BlocklistAIHistoryModel::as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.is_some_and(|c| matches!(c.status(), ConversationStatus::Success));
|
||||
if is_success && me.watched_run_ids.contains_key(&conversation_id) {
|
||||
me.poll_and_inject(conversation_id, ctx);
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// ---- SSE event-push methods ----
|
||||
|
||||
/// Chooses between SSE and polling based on the feature flag, then starts
|
||||
/// the appropriate event delivery loop for the given conversation.
|
||||
/// Opens an SSE connection for the given conversation if one isn't already active.
|
||||
fn start_event_delivery(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if self.should_use_sse() {
|
||||
if !self.sse_connections.contains_key(&conversation_id) {
|
||||
self.start_sse_connection(conversation_id, ctx);
|
||||
}
|
||||
} else {
|
||||
self.poll_and_inject(conversation_id, ctx);
|
||||
if !self.sse_connections.contains_key(&conversation_id) {
|
||||
self.start_sse_connection(conversation_id, ctx);
|
||||
}
|
||||
}
|
||||
|
||||
fn should_use_sse(&self) -> bool {
|
||||
FeatureFlag::OrchestrationEventPush.is_enabled()
|
||||
}
|
||||
|
||||
/// Opens a long-lived SSE connection for `conversation_id`. Events are
|
||||
/// sent through an mpsc channel and drained by a periodic timer.
|
||||
fn start_sse_connection(
|
||||
@@ -911,8 +774,8 @@ impl OrchestrationEventPoller {
|
||||
);
|
||||
}
|
||||
|
||||
/// Drains all buffered SSE events and feeds them through the normal
|
||||
/// `handle_poll_result` path.
|
||||
/// Drains all buffered SSE events and feeds them through the
|
||||
/// `handle_event_batch` sink.
|
||||
fn drain_sse_events(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
@@ -951,7 +814,7 @@ impl OrchestrationEventPoller {
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_default();
|
||||
|
||||
self.handle_poll_result(conversation_id, &self_run_id, cursor, events, messages, ctx);
|
||||
self.handle_event_batch(conversation_id, &self_run_id, cursor, events, messages, ctx);
|
||||
}
|
||||
|
||||
/// Tears down the current SSE connection and opens a new one with the
|
||||
@@ -968,11 +831,11 @@ impl OrchestrationEventPoller {
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for OrchestrationEventPoller {
|
||||
type Event = OrchestrationEventPollerEvent;
|
||||
impl Entity for OrchestrationEventStreamer {
|
||||
type Event = OrchestrationEventStreamerEvent;
|
||||
}
|
||||
|
||||
impl SingletonEntity for OrchestrationEventPoller {}
|
||||
impl SingletonEntity for OrchestrationEventStreamer {}
|
||||
|
||||
fn parse_occurred_at(s: &str) -> prost_types::Timestamp {
|
||||
chrono::DateTime::parse_from_rfc3339(s)
|
||||
@@ -1071,5 +934,5 @@ fn build_pending_events(
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "orchestration_event_poller_tests.rs"]
|
||||
#[path = "orchestration_event_streamer_tests.rs"]
|
||||
mod tests;
|
||||
+15
-21
@@ -218,7 +218,7 @@ fn finish_restore_fetch_uses_server_cursor_when_sqlite_is_absent() {
|
||||
let server_api = ServerApiProvider::new_for_test().get();
|
||||
|
||||
let poller = app.add_singleton_model(|ctx| {
|
||||
OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
OrchestrationEventStreamer::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
});
|
||||
|
||||
// Seed event_cursor as on_restored_conversations would before spawning
|
||||
@@ -283,24 +283,22 @@ fn restored_inprogress_parent_defers_delivery_until_success() {
|
||||
// a permissive expectation prevents spurious panics either way.
|
||||
mock.expect_get_ambient_agent_task()
|
||||
.returning(|_| Ok(make_ambient_task_with_event_seq(None)));
|
||||
mock.expect_poll_agent_events()
|
||||
.returning(|_, _, _| Ok(vec![]));
|
||||
mock.expect_update_event_sequence_on_server()
|
||||
.returning(|_, _| Ok(()));
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(mock);
|
||||
let server_api = ServerApiProvider::new_for_test().get();
|
||||
|
||||
let poller = app.add_singleton_model(|ctx| {
|
||||
OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
let streamer = app.add_singleton_model(|ctx| {
|
||||
OrchestrationEventStreamer::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
});
|
||||
|
||||
// Synchronous part of `on_restored_conversations`: cursor seeded,
|
||||
// own run_id watched. No event delivery yet because parent is
|
||||
// InProgress.
|
||||
poller.update(&mut app, |me, ctx| {
|
||||
streamer.update(&mut app, |me, ctx| {
|
||||
me.on_restored_conversations(vec![conversation_id], ctx);
|
||||
});
|
||||
poller.read(&app, |me, _| {
|
||||
streamer.read(&app, |me, _| {
|
||||
assert_eq!(me.event_cursor.get(&conversation_id).copied(), Some(0));
|
||||
assert!(
|
||||
me.watched_run_ids
|
||||
@@ -308,18 +306,14 @@ fn restored_inprogress_parent_defers_delivery_until_success() {
|
||||
.is_some_and(|w| !w.is_empty()),
|
||||
"own run_id should have been registered as watched"
|
||||
);
|
||||
assert!(
|
||||
!me.poll_in_flight.contains(&conversation_id),
|
||||
"InProgress parent must not start polling"
|
||||
);
|
||||
assert!(
|
||||
me.sse_connections.is_empty(),
|
||||
"InProgress parent must not open SSE"
|
||||
);
|
||||
});
|
||||
|
||||
// Transitioning the conversation to Success should trigger event
|
||||
// delivery (poll_and_inject in the non-SSE default path).
|
||||
// Transitioning the conversation to Success should open an SSE
|
||||
// connection for event delivery.
|
||||
history_model.update(&mut app, |model, ctx| {
|
||||
model.update_conversation_status(
|
||||
terminal_view_id,
|
||||
@@ -328,17 +322,17 @@ fn restored_inprogress_parent_defers_delivery_until_success() {
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
poller.read(&app, |me, _| {
|
||||
streamer.read(&app, |me, _| {
|
||||
assert!(
|
||||
me.poll_in_flight.contains(&conversation_id),
|
||||
"Success transition with watched run_ids should start delivery"
|
||||
me.sse_connections.contains_key(&conversation_id),
|
||||
"Success transition with watched run_ids should open an SSE connection"
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_poll_result_persists_max_seq_to_history_model() {
|
||||
fn handle_event_batch_persists_max_seq_to_history_model() {
|
||||
use crate::ai::agent::conversation::{AIConversation, AIConversationId};
|
||||
use crate::persistence::ModelEvent;
|
||||
use crate::server::server_api::ai::MockAIClient;
|
||||
@@ -378,7 +372,7 @@ fn handle_poll_result_persists_max_seq_to_history_model() {
|
||||
let server_api = ServerApiProvider::new_for_test().get();
|
||||
|
||||
let poller = app.add_singleton_model(|ctx| {
|
||||
OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
OrchestrationEventStreamer::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
});
|
||||
|
||||
// Build a poll batch with max sequence = 42. Use an unrecognized
|
||||
@@ -405,7 +399,7 @@ fn handle_poll_result_persists_max_seq_to_history_model() {
|
||||
];
|
||||
|
||||
poller.update(&mut app, |me, ctx| {
|
||||
me.handle_poll_result(
|
||||
me.handle_event_batch(
|
||||
conversation_id,
|
||||
/* self_run_id */ "some-other-run",
|
||||
/* previous_cursor */ 0,
|
||||
@@ -463,7 +457,7 @@ fn finish_restore_fetch_no_ops_when_conversation_deleted_mid_flight() {
|
||||
let server_api = ServerApiProvider::new_for_test().get();
|
||||
|
||||
let poller = app.add_singleton_model(|ctx| {
|
||||
OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
OrchestrationEventStreamer::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
});
|
||||
|
||||
// Seed cursor as on_restored_conversations would.
|
||||
@@ -542,7 +536,7 @@ fn finish_restore_fetch_reconnects_sse_when_children_added_to_open_connection()
|
||||
let server_api = ServerApiProvider::new_for_test().get();
|
||||
|
||||
let poller = app.add_singleton_model(|ctx| {
|
||||
OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
OrchestrationEventStreamer::new_with_clients_for_test(ai_client, server_api, ctx)
|
||||
});
|
||||
|
||||
// Seed the state on_restored_conversations would have set up, then
|
||||
@@ -865,10 +865,10 @@ impl OrchestrationEventService {
|
||||
}
|
||||
}
|
||||
|
||||
/// Accepts pre-built events from the v2 poller and enqueues them
|
||||
/// Accepts pre-built events from the v2 streamer and enqueues them
|
||||
/// for drain by the controller via the normal v1 path.
|
||||
/// Lifecycle events go through coalescing and cap enforcement.
|
||||
pub fn enqueue_polled_events(
|
||||
pub fn enqueue_event_batch(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
events: Vec<PendingEvent>,
|
||||
|
||||
Reference in New Issue
Block a user