Replay agent events on restore (#9251)
## Description Restore orchestration event delivery on the client after a Warp restart so that a parent conversation continues to receive lifecycle events and inbox messages from its children — including terminal events that arrived while Warp was not running. See the full design in `specs/replay-agent-events-on-restore/PRODUCT.md` and `specs/replay-agent-events-on-restore/TECH.md`. This is a re-land of warpdotdev/warp-internal#24999, which was reverted in warpdotdev/warp-internal#25055 due to a CI race. This version fixes the test compilation issue (missing fields in `AmbientAgentTask` struct literal in `conversation_ended_tombstone_view_tests.rs`). ### What - Persists the per-conversation event cursor across restarts. - Adds `last_event_sequence: Option<i64>` to `AgentConversationData` (SQLite) and `AIConversation`. - New `BlocklistAIHistoryModel::update_event_sequence` helper writes the cursor through `write_updated_conversation_state` after each event batch. - Also persists the cursor to the server (fire-and-forget) so driver / cloud restarts can resume without local SQLite state. - Restores `OrchestrationEventPoller.watched_run_ids` and re-establishes event delivery on `BlocklistAIHistoryEvent::RestoredConversations`. - New `on_restored_conversations` handler issues `GET /agent/runs/{run_id}` for each restored parent and uses the response inline `children` and `last_event_sequence` to populate watched run ids and merge the cursor (`max(SQLite, server)`). - Fetch failures retry with exponential backoff (1s, 2s, 5s, 10s capped) keyed off a per-conversation `restore_fetch_failures` counter; reset on success and on conversation removal. - Gated on `OrchestrationV2`. Shared-session viewers and conversations without children are skipped. - `Success` parents resume delivery immediately; `InProgress` parents defer to the existing `on_conversation_status_updated` path. - Restores V1 lifecycle subscriptions on restart by extending the existing `RestoredConversations` handler in `OrchestrationEventService` to re-register `lifecycle_subscription_routes` for restored child conversations whose parents are present locally. ### Why After a Warp restart, parent conversations were silently receiving no further events from children. The `event_cursor` in `OrchestrationEventPoller` was always initialized to 0, so even if delivery had resumed, every event since the start of the conversation would have replayed and produced duplicate messages. V1 lifecycle subscription routes were also not restored, so V1 parents missed child status transitions. ## Testing - Added unit tests in `app/src/ai/blocklist/orchestration_event_poller_tests.rs` covering: cursor merge from server vs SQLite, retry on `get_ambient_agent_task` failure, V2 gating, shared-session-viewer exclusion, cleanup on delete, and `last_event_sequence` round-trip through `AIConversation::new_restored`. - Added unit test coverage in `app/src/ai/blocklist/orchestration_events_tests.rs` for V1 lifecycle re-registration on restore. - Manual verification per `specs/replay-agent-events-on-restore/TECH.md`. ## Server API dependencies - [x] Does this change rely on a [new server API](https://www.notion.so/warpdev/How-to-add-a-new-full-stack-feature-8412cede405a4ec194b32bdd4b951035?pvs=4#04da1e6a493542d68b3e998c7d339640)? - [x] If so, is the use of this API restricted to client channels that rely on the staging server (e.g. WarpDev)? The companion server change adds: - `last_event_sequence` on `Task` (`ai_tasks` column), surfaced inline on `GET /agent/runs/:run_id`. - `PATCH /agent/runs/:run_id/event-sequence` for client cursor writes. - `children` inline on the `GET /agent/runs/:run_id` response. ## Agent Mode - [x] Warp Agent Mode - This PR was created via Warp's AI Agent Mode Co-Authored-By: Oz <oz-agent@warp.dev> --------- Co-authored-by: Oz <oz-agent@warp.dev>
This commit is contained in:
Matthew Albright
@@ -84,6 +84,7 @@ pub fn convert_conversation_data_to_ai_conversation(
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
RestorationMode::Continue => AgentConversationData {
|
||||
server_conversation_token: Some(
|
||||
@@ -102,6 +103,7 @@ pub fn convert_conversation_data_to_ai_conversation(
|
||||
// dispatch time; adding it here would avoid a round-trip to StreamInit.
|
||||
run_id: None,
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
};
|
||||
|
||||
|
||||
@@ -223,6 +223,11 @@ pub struct AIConversation {
|
||||
/// these conversations — the remote worker's own client handles status
|
||||
/// reporting. TaskStatusSyncModel skips status updates for these.
|
||||
is_remote_child: bool,
|
||||
|
||||
/// The last event sequence number observed from the v2 orchestration
|
||||
/// event log. Used on restore to resume event delivery without
|
||||
/// re-delivering already-processed events.
|
||||
last_event_sequence: Option<i64>,
|
||||
}
|
||||
|
||||
pub(crate) fn artifact_from_fork_proto(
|
||||
@@ -272,6 +277,7 @@ impl AIConversation {
|
||||
agent_name: None,
|
||||
parent_conversation_id: None,
|
||||
is_remote_child: false,
|
||||
last_event_sequence: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -351,6 +357,7 @@ impl AIConversation {
|
||||
parent_conversation_id,
|
||||
run_id,
|
||||
autoexecute_override,
|
||||
last_event_sequence,
|
||||
) = if let Some(data) = conversation_data {
|
||||
let server_conversation_token = data
|
||||
.server_conversation_token
|
||||
@@ -381,6 +388,7 @@ impl AIConversation {
|
||||
} else {
|
||||
AIConversationAutoexecuteMode::default()
|
||||
};
|
||||
let last_event_sequence = data.last_event_sequence;
|
||||
|
||||
(
|
||||
server_conversation_token,
|
||||
@@ -393,6 +401,7 @@ impl AIConversation {
|
||||
parent_conversation_id,
|
||||
run_id,
|
||||
autoexecute_override,
|
||||
last_event_sequence,
|
||||
)
|
||||
} else {
|
||||
(
|
||||
@@ -406,6 +415,7 @@ impl AIConversation {
|
||||
None,
|
||||
None,
|
||||
AIConversationAutoexecuteMode::default(),
|
||||
None,
|
||||
)
|
||||
};
|
||||
|
||||
@@ -448,6 +458,7 @@ impl AIConversation {
|
||||
agent_name,
|
||||
parent_conversation_id,
|
||||
is_remote_child: false,
|
||||
last_event_sequence,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -786,6 +797,16 @@ impl AIConversation {
|
||||
self.parent_conversation_id = Some(id);
|
||||
}
|
||||
|
||||
/// Returns the last observed v2 orchestration event sequence number, if any.
|
||||
pub fn last_event_sequence(&self) -> Option<i64> {
|
||||
self.last_event_sequence
|
||||
}
|
||||
|
||||
/// Updates the last observed v2 orchestration event sequence number.
|
||||
pub fn set_last_event_sequence(&mut self, sequence: i64) {
|
||||
self.last_event_sequence = Some(sequence);
|
||||
}
|
||||
|
||||
/// Returns true if this conversation was spawned by a parent orchestrator agent.
|
||||
pub fn is_child_agent_conversation(&self) -> bool {
|
||||
self.parent_conversation_id.is_some()
|
||||
@@ -2811,6 +2832,7 @@ impl AIConversation {
|
||||
parent_conversation_id: self.parent_conversation_id.map(|id| id.to_string()),
|
||||
run_id: self.task_id.map(|id| id.to_string()),
|
||||
autoexecute_override: Some(self.autoexecute_override.into()),
|
||||
last_event_sequence: self.last_event_sequence,
|
||||
},
|
||||
};
|
||||
ctx.spawn(
|
||||
|
||||
@@ -60,6 +60,8 @@ fn create_test_task(
|
||||
agent_config_snapshot: None,
|
||||
artifacts: vec![],
|
||||
is_sandbox_running: false,
|
||||
last_event_sequence: None,
|
||||
children: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -152,6 +154,7 @@ fn test_display_status_uses_matching_conversation_for_in_progress_task() {
|
||||
parent_conversation_id: None,
|
||||
run_id: Some(task_id.clone()),
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -203,6 +206,7 @@ fn test_display_status_updates_when_blocked_conversation_resumes() {
|
||||
parent_conversation_id: None,
|
||||
run_id: Some(task_id.clone()),
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -278,6 +282,7 @@ fn test_display_status_terminal_task_state_overrides_matching_conversation() {
|
||||
parent_conversation_id: None,
|
||||
run_id: Some(task_id.clone()),
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -329,6 +334,7 @@ fn test_status_filter_uses_display_status_for_task_backed_conversations() {
|
||||
parent_conversation_id: None,
|
||||
run_id: Some(task_id.clone()),
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -763,6 +769,7 @@ fn test_get_tasks_and_conversations_prefers_task_when_task_id_matches_conversati
|
||||
parent_conversation_id: None,
|
||||
run_id: Some(task_id.clone()),
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -819,6 +826,7 @@ fn test_get_tasks_and_conversations_prefers_task_when_server_token_matches() {
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -874,6 +882,7 @@ fn test_get_tasks_and_conversations_keeps_unrelated_tasks_and_conversations() {
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
|
||||
@@ -34,6 +34,8 @@ fn task_with(
|
||||
agent_config_snapshot: None,
|
||||
artifacts: vec![],
|
||||
is_sandbox_running: true,
|
||||
last_event_sequence: None,
|
||||
children: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -262,6 +262,20 @@ pub struct AmbientAgentTask {
|
||||
pub agent_config_snapshot: Option<AgentConfigSnapshot>,
|
||||
#[serde(default, deserialize_with = "deserialize_artifacts")]
|
||||
pub artifacts: Vec<Artifact>,
|
||||
|
||||
/// The last event sequence number recorded for this run by the server.
|
||||
/// Used by orchestration event delivery to resume from the correct
|
||||
/// cursor on restart. Populated by `GET /agent/runs/{run_id}` when the
|
||||
/// server supports it; `None` on older servers.
|
||||
#[serde(default)]
|
||||
pub last_event_sequence: Option<i64>,
|
||||
|
||||
/// The server-recorded `run_id`s of direct children of this run. Used
|
||||
/// by orchestration event-delivery restore to discover children whose
|
||||
/// records may not exist locally (e.g. remote-worker children in the
|
||||
/// driver case). Empty on older servers.
|
||||
#[serde(default)]
|
||||
pub children: Vec<String>,
|
||||
}
|
||||
|
||||
/// Represents a single attachment input from the client (e.g., file upload)
|
||||
|
||||
@@ -452,6 +452,23 @@ impl BlocklistAIHistoryModel {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Updates the persisted `last_event_sequence` for a conversation and
|
||||
/// writes the updated conversation state to SQLite. Used by the
|
||||
/// orchestration event poller after draining an event batch to keep the
|
||||
/// cursor durable across restarts.
|
||||
pub fn update_event_sequence(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
sequence: i64,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) else {
|
||||
return;
|
||||
};
|
||||
conversation.set_last_event_sequence(sequence);
|
||||
conversation.write_updated_conversation_state(ctx);
|
||||
}
|
||||
|
||||
/// Sets a live conversation's server token, and updates the mapping in the history
|
||||
/// model.
|
||||
pub fn set_server_conversation_token_for_conversation(
|
||||
@@ -1075,6 +1092,9 @@ impl BlocklistAIHistoryModel {
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: Some(source_conversation.autoexecute_override().into()),
|
||||
// The event cursor belongs to the source conversation's run; the
|
||||
// forked conversation will establish its own cursor.
|
||||
last_event_sequence: None,
|
||||
};
|
||||
let forked_conversation_id = AIConversationId::new();
|
||||
if let Err(e) = sqlite_sender.send(ModelEvent::UpdateMultiAgentConversation {
|
||||
@@ -1227,6 +1247,9 @@ impl BlocklistAIHistoryModel {
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: Some(conversation.autoexecute_override().into()),
|
||||
// The event cursor belongs to the source conversation's run; the
|
||||
// forked conversation will establish its own cursor.
|
||||
last_event_sequence: None,
|
||||
};
|
||||
|
||||
let forked_conversation_id = AIConversationId::new();
|
||||
|
||||
@@ -1132,6 +1132,7 @@ fn test_find_by_token_after_insert_forked_conversation_from_tasks() {
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
};
|
||||
let tasks = vec![warp_multi_agent_api::Task {
|
||||
id: "root-task".to_string(),
|
||||
|
||||
@@ -100,6 +100,9 @@ pub struct OrchestrationEventPoller {
|
||||
/// Monotonic counter for SSE connection generations. Ensures stale
|
||||
/// callbacks from replaced connections are discarded.
|
||||
next_sse_generation: u64,
|
||||
/// Consecutive failure count for the post-restore `get_ambient_agent_task`
|
||||
/// fetch (resets on success). Drives exponential backoff for retries.
|
||||
restore_fetch_failures: HashMap<AIConversationId, usize>,
|
||||
}
|
||||
|
||||
pub enum OrchestrationEventPollerEvent {
|
||||
@@ -126,6 +129,35 @@ impl OrchestrationEventPoller {
|
||||
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
|
||||
/// looking them up via `ServerApiProvider`. Lets unit tests inject a
|
||||
/// `MockAIClient` while still subscribing to `BlocklistAIHistoryModel`.
|
||||
#[cfg(test)]
|
||||
pub(super) fn new_with_clients_for_test(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
server_api: Arc<ServerApi>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> Self {
|
||||
let history_model = BlocklistAIHistoryModel::handle(ctx);
|
||||
ctx.subscribe_to_model(&history_model, |me, event, ctx| {
|
||||
me.handle_history_event(event, ctx);
|
||||
});
|
||||
Self {
|
||||
ai_client,
|
||||
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(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -187,6 +219,7 @@ impl OrchestrationEventPoller {
|
||||
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.
|
||||
self.sse_connections.remove(conversation_id);
|
||||
@@ -202,9 +235,235 @@ impl OrchestrationEventPoller {
|
||||
| BlocklistAIHistoryEvent::UpdatedTodoList { .. }
|
||||
| BlocklistAIHistoryEvent::UpdatedAutoexecuteOverride { .. }
|
||||
| BlocklistAIHistoryEvent::SplitConversation { .. }
|
||||
| BlocklistAIHistoryEvent::RestoredConversations { .. }
|
||||
| BlocklistAIHistoryEvent::UpdatedConversationMetadata { .. }
|
||||
| BlocklistAIHistoryEvent::UpdatedConversationArtifacts { .. } => {}
|
||||
BlocklistAIHistoryEvent::RestoredConversations {
|
||||
conversation_ids, ..
|
||||
} => {
|
||||
self.on_restored_conversations(conversation_ids.clone(), ctx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handles restoration of conversations on startup (or driver re-attach).
|
||||
///
|
||||
/// 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.
|
||||
fn on_restored_conversations(
|
||||
&mut self,
|
||||
conversation_ids: Vec<AIConversationId>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
// Orchestration v2 owns the events endpoints and the cursor model.
|
||||
// V1 conversations may carry a run_id but the v2-only event APIs
|
||||
// would return spurious 4xx responses, so skip restore entirely
|
||||
// when V2 is disabled.
|
||||
if !FeatureFlag::OrchestrationV2.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
for conv_id in conversation_ids {
|
||||
let (run_id, cursor, status, is_viewer) = {
|
||||
let history_model = BlocklistAIHistoryModel::as_ref(ctx);
|
||||
let Some(conversation) = history_model.conversation(&conv_id) else {
|
||||
continue;
|
||||
};
|
||||
let is_viewer = conversation.is_viewing_shared_session();
|
||||
let run_id = conversation.run_id();
|
||||
let cursor = conversation.last_event_sequence().unwrap_or(0);
|
||||
let status = conversation.status().clone();
|
||||
(run_id, cursor, status, is_viewer)
|
||||
};
|
||||
|
||||
// Shared-session viewers receive updates through session sharing;
|
||||
// polling here would re-inject events the session has already
|
||||
// processed.
|
||||
if is_viewer {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Initialize the in-memory cursor from the persisted SQLite value.
|
||||
// A later server `GET /agent/runs/{run_id}` response may advance
|
||||
// it to `max(SQLite, server)` before delivery starts.
|
||||
//
|
||||
// Note: a status transition arriving in the window before
|
||||
// finish_restore_fetch completes may trigger
|
||||
// start_event_delivery with only the SQLite cursor. This is
|
||||
// acceptable — worst case is one extra batch of duplicate
|
||||
// events.
|
||||
self.event_cursor.insert(conv_id, cursor);
|
||||
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
|
||||
// of run_ids to open against.
|
||||
if let Some(ref own) = run_id {
|
||||
self.watched_run_ids
|
||||
.entry(conv_id)
|
||||
.or_default()
|
||||
.insert(own.clone());
|
||||
}
|
||||
|
||||
// No run_id means we can't query the server for children or for
|
||||
// the canonical cursor. There's nothing more to do here; if a
|
||||
// run_id gets assigned later the standard self-registration path
|
||||
// will pick it up.
|
||||
let Some(run_id) = run_id else {
|
||||
self.maybe_start_delivery_after_restore(conv_id, &status, ctx);
|
||||
continue;
|
||||
};
|
||||
|
||||
let Ok(task_id) = run_id.parse::<crate::ai::ambient_agents::AmbientAgentTaskId>()
|
||||
else {
|
||||
log::warn!("could not parse run_id {run_id:?} for {conv_id:?}");
|
||||
self.maybe_start_delivery_after_restore(conv_id, &status, ctx);
|
||||
continue;
|
||||
};
|
||||
|
||||
self.spawn_restore_fetch(conv_id, task_id, cursor, ctx);
|
||||
}
|
||||
}
|
||||
|
||||
/// Issues `GET /agent/runs/{task_id}` and routes the result through
|
||||
/// `finish_restore_fetch`. Used both for the initial post-restore fetch
|
||||
/// and for backoff-driven retries.
|
||||
fn spawn_restore_fetch(
|
||||
&mut self,
|
||||
conv_id: AIConversationId,
|
||||
task_id: crate::ai::ambient_agents::AmbientAgentTaskId,
|
||||
sqlite_cursor: i64,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let ai_client = self.ai_client.clone();
|
||||
ctx.spawn(
|
||||
async move { ai_client.get_ambient_agent_task(&task_id).await },
|
||||
move |me, run_result, ctx| {
|
||||
me.finish_restore_fetch(conv_id, task_id, sqlite_cursor, run_result, ctx);
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
/// Completes the post-restore async fetch by merging the server cursor,
|
||||
/// installing the server-reported child run_ids, and — if the parent is
|
||||
/// `Success` — starting event delivery. On a server-fetch failure,
|
||||
/// schedules a retry with exponential backoff: V2 children always have a
|
||||
/// server-side `ai_tasks` row, so the server is the authoritative source
|
||||
/// for the watched run_id set, and any local fallback would be incomplete
|
||||
/// anyway. Without network connectivity event delivery wouldn't function,
|
||||
/// so retrying is the right behavior.
|
||||
fn finish_restore_fetch(
|
||||
&mut self,
|
||||
conv_id: AIConversationId,
|
||||
task_id: crate::ai::ambient_agents::AmbientAgentTaskId,
|
||||
sqlite_cursor: i64,
|
||||
run_result: anyhow::Result<crate::ai::ambient_agents::task::AmbientAgentTask>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
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
|
||||
// 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);
|
||||
return;
|
||||
}
|
||||
|
||||
// Reset the retry counter on success.
|
||||
self.restore_fetch_failures.remove(&conv_id);
|
||||
|
||||
// Merge the server cursor: use the max of SQLite and server
|
||||
// values so we don't re-deliver events the client already
|
||||
// acknowledged locally.
|
||||
let server_seq = task.last_event_sequence.unwrap_or(0);
|
||||
let merged = sqlite_cursor.max(server_seq);
|
||||
self.event_cursor.insert(conv_id, merged);
|
||||
|
||||
// The server response includes `children` inline on
|
||||
// `AmbientAgentTask`; this is the authoritative set of
|
||||
// direct child run_ids for the parent.
|
||||
//
|
||||
// Insert children and reconnect SSE once if any new run_ids
|
||||
// were added and a connection is already open (e.g. because a
|
||||
// status transition raced with this fetch and opened SSE with
|
||||
// only the parent's own run_id).
|
||||
let had_sse = self.sse_connections.contains_key(&conv_id);
|
||||
let watched = self.watched_run_ids.entry(conv_id).or_default();
|
||||
let mut any_new_children = false;
|
||||
for child in task.children {
|
||||
if watched.insert(child) {
|
||||
any_new_children = true;
|
||||
}
|
||||
}
|
||||
if any_new_children && had_sse {
|
||||
self.reconnect_sse(conv_id, ctx);
|
||||
}
|
||||
|
||||
let status = BlocklistAIHistoryModel::as_ref(ctx)
|
||||
.conversation(&conv_id)
|
||||
.map(|c| c.status().clone())
|
||||
.unwrap_or(ConversationStatus::Success);
|
||||
self.maybe_start_delivery_after_restore(conv_id, &status, ctx);
|
||||
}
|
||||
Err(err) => {
|
||||
log::warn!("Restore: get_agent_run failed for {conv_id:?}: {err:#}; will retry");
|
||||
self.start_restore_fetch_retry_timer(conv_id, task_id, sqlite_cursor, ctx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
fn start_restore_fetch_retry_timer(
|
||||
&mut self,
|
||||
conv_id: AIConversationId,
|
||||
task_id: crate::ai::ambient_agents::AmbientAgentTaskId,
|
||||
sqlite_cursor: i64,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let failures = self
|
||||
.restore_fetch_failures
|
||||
.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]);
|
||||
ctx.spawn(
|
||||
async move { Timer::after(backoff).await },
|
||||
move |me, _, ctx| {
|
||||
// The conversation may have been removed in the meantime;
|
||||
// if so, drop the retry. Otherwise re-issue the fetch.
|
||||
if !me.event_cursor.contains_key(&conv_id) {
|
||||
me.restore_fetch_failures.remove(&conv_id);
|
||||
return;
|
||||
}
|
||||
me.spawn_restore_fetch(conv_id, task_id, sqlite_cursor, ctx);
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
/// Starts event delivery for a restored conversation if the parent is
|
||||
/// currently `Success` and has at least one watched run_id. `InProgress`
|
||||
/// parents are deferred to `on_conversation_status_updated` once they
|
||||
/// next transition to `Success`.
|
||||
fn maybe_start_delivery_after_restore(
|
||||
&mut self,
|
||||
conv_id: AIConversationId,
|
||||
status: &ConversationStatus,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let has_watched = self
|
||||
.watched_run_ids
|
||||
.get(&conv_id)
|
||||
.is_some_and(|w| !w.is_empty());
|
||||
if !has_watched {
|
||||
return;
|
||||
}
|
||||
if matches!(status, ConversationStatus::Success) {
|
||||
self.start_event_delivery(conv_id, ctx);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -429,6 +688,40 @@ impl OrchestrationEventPoller {
|
||||
.unwrap_or(previous_cursor);
|
||||
self.event_cursor.insert(conversation_id, max_seq);
|
||||
|
||||
// Persist the cursor to SQLite so that after a restart we can resume
|
||||
// event delivery from this sequence number without re-delivering
|
||||
// events the parent has already acted on.
|
||||
BlocklistAIHistoryModel::handle(ctx).update(ctx, |model, ctx| {
|
||||
model.update_event_sequence(conversation_id, max_seq, ctx);
|
||||
});
|
||||
|
||||
// Also persist the cursor to the server so driver / cloud restarts
|
||||
// can resume without local SQLite state. Fire-and-forget: log on
|
||||
// failure, don't block event delivery. The server persists the
|
||||
// cursor on `ai_tasks.last_event_sequence`.
|
||||
let own_run_id = BlocklistAIHistoryModel::as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.and_then(|c| c.run_id());
|
||||
if let Some(run_id) = own_run_id {
|
||||
// TODO: consider debouncing this server write (see
|
||||
// specs/replay-agent-events-on-restore/TECH.md Risks).
|
||||
let ai_client = self.ai_client.clone();
|
||||
ctx.spawn(
|
||||
async move {
|
||||
ai_client
|
||||
.update_event_sequence_on_server(&run_id, max_seq)
|
||||
.await
|
||||
},
|
||||
move |_, result, _| {
|
||||
if let Err(err) = result {
|
||||
log::warn!(
|
||||
"Failed to persist event cursor to server for {conversation_id:?}: {err:#}"
|
||||
);
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// Track message IDs for server-side mark_delivered calls.
|
||||
let message_ids: Vec<String> = events
|
||||
.iter()
|
||||
|
||||
@@ -104,3 +104,499 @@ fn convert_lifecycle_events_maps_run_restarted() {
|
||||
Some(api::agent_event::lifecycle_event::Detail::InProgress(..))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ai_conversation_new_restored_preserves_last_event_sequence() {
|
||||
// Guards against regressions that drop the field when wiring the restore
|
||||
// path: a conversation restored with `last_event_sequence: Some(N)`
|
||||
// should expose it via `conversation.last_event_sequence()`.
|
||||
use crate::ai::agent::conversation::{AIConversation, AIConversationId};
|
||||
use crate::persistence::model::AgentConversationData;
|
||||
|
||||
let task = api::Task {
|
||||
id: "root".to_string(),
|
||||
messages: vec![api::Message {
|
||||
id: "m1".to_string(),
|
||||
task_id: "root".to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::AgentOutput(
|
||||
api::message::AgentOutput {
|
||||
text: "hi".to_string(),
|
||||
},
|
||||
)),
|
||||
request_id: String::new(),
|
||||
timestamp: None,
|
||||
}],
|
||||
dependencies: None,
|
||||
description: String::new(),
|
||||
summary: String::new(),
|
||||
server_data: String::new(),
|
||||
};
|
||||
let data = AgentConversationData {
|
||||
server_conversation_token: None,
|
||||
conversation_usage_metadata: None,
|
||||
reverted_action_ids: None,
|
||||
forked_from_server_conversation_token: None,
|
||||
artifacts_json: None,
|
||||
parent_agent_id: None,
|
||||
agent_name: None,
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: Some(42),
|
||||
};
|
||||
let conversation =
|
||||
AIConversation::new_restored(AIConversationId::new(), vec![task], Some(data))
|
||||
.expect("should restore");
|
||||
assert_eq!(conversation.last_event_sequence(), Some(42));
|
||||
}
|
||||
|
||||
// ---- Helpers for App-based poller tests ----
|
||||
|
||||
fn make_ambient_task_with_children(
|
||||
children: Vec<String>,
|
||||
) -> crate::ai::ambient_agents::AmbientAgentTask {
|
||||
let mut task = make_ambient_task_with_event_seq(None);
|
||||
task.children = children;
|
||||
task
|
||||
}
|
||||
|
||||
fn make_ambient_task_with_event_seq(
|
||||
last_event_sequence: Option<i64>,
|
||||
) -> crate::ai::ambient_agents::AmbientAgentTask {
|
||||
use chrono::Utc;
|
||||
crate::ai::ambient_agents::AmbientAgentTask {
|
||||
task_id: "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(),
|
||||
parent_run_id: None,
|
||||
title: "test".to_string(),
|
||||
state: crate::ai::ambient_agents::AmbientAgentTaskState::Succeeded,
|
||||
prompt: "prompt".to_string(),
|
||||
created_at: Utc::now(),
|
||||
started_at: Some(Utc::now()),
|
||||
updated_at: Utc::now(),
|
||||
status_message: None,
|
||||
source: None,
|
||||
session_id: None,
|
||||
session_link: None,
|
||||
creator: None,
|
||||
conversation_id: None,
|
||||
request_usage: None,
|
||||
agent_config_snapshot: None,
|
||||
artifacts: vec![],
|
||||
is_sandbox_running: false,
|
||||
last_event_sequence,
|
||||
children: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finish_restore_fetch_uses_server_cursor_when_sqlite_is_absent() {
|
||||
use crate::ai::agent::conversation::AIConversation;
|
||||
use crate::server::server_api::ai::MockAIClient;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
use std::sync::Arc;
|
||||
use warpui::App;
|
||||
|
||||
App::test((), |mut app| async move {
|
||||
let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true);
|
||||
|
||||
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[]));
|
||||
|
||||
// Restore a conversation with no SQLite cursor (`last_event_sequence:
|
||||
// None`). After the server fetch completes with `Some(42)` we expect
|
||||
// the in-memory cursor to be 42 (max(0, 42)).
|
||||
let conversation = AIConversation::new(false);
|
||||
let conversation_id = conversation.id();
|
||||
let terminal_view_id = warpui::EntityId::new();
|
||||
history_model.update(&mut app, |model, ctx| {
|
||||
model.restore_conversations(terminal_view_id, vec![conversation], ctx);
|
||||
});
|
||||
|
||||
let mock = MockAIClient::new();
|
||||
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)
|
||||
});
|
||||
|
||||
// Seed event_cursor as on_restored_conversations would before spawning
|
||||
// the async fetch. Without this the guard that detects mid-flight
|
||||
// conversation deletion would fire and return early.
|
||||
poller.update(&mut app, |me, _| {
|
||||
me.event_cursor.insert(conversation_id, 0);
|
||||
});
|
||||
|
||||
let task_id: crate::ai::ambient_agents::AmbientAgentTaskId =
|
||||
"550e8400-e29b-41d4-a716-446655440000".parse().unwrap();
|
||||
poller.update(&mut app, |me, ctx| {
|
||||
me.finish_restore_fetch(
|
||||
conversation_id,
|
||||
task_id,
|
||||
/* sqlite_cursor */ 0,
|
||||
Ok(make_ambient_task_with_event_seq(Some(42))),
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
|
||||
poller.read(&app, |me, _| {
|
||||
assert_eq!(me.event_cursor.get(&conversation_id).copied(), Some(42));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restored_inprogress_parent_defers_delivery_until_success() {
|
||||
use crate::ai::agent::conversation::{AIConversation, AIConversationId, ConversationStatus};
|
||||
use crate::server::server_api::ai::MockAIClient;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
use std::sync::Arc;
|
||||
use warpui::App;
|
||||
|
||||
App::test((), |mut app| async move {
|
||||
let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true);
|
||||
|
||||
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[]));
|
||||
|
||||
let mut conversation = AIConversation::new(false);
|
||||
// Use a parsable UUID-shaped run_id so the poller can construct
|
||||
// an `AmbientAgentTaskId` for the (mocked) server fetch.
|
||||
conversation.set_run_id("550e8400-e29b-41d4-a716-446655440100".to_string());
|
||||
let conversation_id: AIConversationId = conversation.id();
|
||||
let terminal_view_id = warpui::EntityId::new();
|
||||
history_model.update(&mut app, |model, ctx| {
|
||||
model.restore_conversations(terminal_view_id, vec![conversation], ctx);
|
||||
// The default status after restore is `InProgress` for live
|
||||
// conversations, but assert it explicitly to make the test
|
||||
// self-documenting.
|
||||
model.update_conversation_status(
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
ConversationStatus::InProgress,
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
|
||||
let mut mock = MockAIClient::new();
|
||||
// The async restore fetch may or may not complete during the test;
|
||||
// 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)
|
||||
});
|
||||
|
||||
// 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| {
|
||||
me.on_restored_conversations(vec![conversation_id], ctx);
|
||||
});
|
||||
poller.read(&app, |me, _| {
|
||||
assert_eq!(me.event_cursor.get(&conversation_id).copied(), Some(0));
|
||||
assert!(
|
||||
me.watched_run_ids
|
||||
.get(&conversation_id)
|
||||
.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).
|
||||
history_model.update(&mut app, |model, ctx| {
|
||||
model.update_conversation_status(
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
ConversationStatus::Success,
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
poller.read(&app, |me, _| {
|
||||
assert!(
|
||||
me.poll_in_flight.contains(&conversation_id),
|
||||
"Success transition with watched run_ids should start delivery"
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_poll_result_persists_max_seq_to_history_model() {
|
||||
use crate::ai::agent::conversation::{AIConversation, AIConversationId};
|
||||
use crate::persistence::ModelEvent;
|
||||
use crate::server::server_api::ai::MockAIClient;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
use crate::test_util::settings::initialize_settings_for_tests;
|
||||
use crate::{GlobalResourceHandles, GlobalResourceHandlesProvider};
|
||||
use std::sync::Arc;
|
||||
use warpui::App;
|
||||
|
||||
App::test((), |mut app| async move {
|
||||
let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true);
|
||||
|
||||
// `update_event_sequence` calls `write_updated_conversation_state`,
|
||||
// which reads `GeneralSettings`, `AppExecutionMode`, and the global
|
||||
// resource sender. Wire all of these up so the SQLite write can run.
|
||||
initialize_settings_for_tests(&mut app);
|
||||
let (sender, receiver) = std::sync::mpsc::sync_channel::<ModelEvent>(4);
|
||||
let mut global_resource_handles = GlobalResourceHandles::mock(&mut app);
|
||||
global_resource_handles.model_event_sender = Some(sender);
|
||||
app.add_singleton_model(|_| GlobalResourceHandlesProvider::new(global_resource_handles));
|
||||
|
||||
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[]));
|
||||
|
||||
let mut conversation = AIConversation::new(false);
|
||||
conversation.set_run_id("550e8400-e29b-41d4-a716-446655440200".to_string());
|
||||
let conversation_id: AIConversationId = conversation.id();
|
||||
let terminal_view_id = warpui::EntityId::new();
|
||||
history_model.update(&mut app, |model, ctx| {
|
||||
model.restore_conversations(terminal_view_id, vec![conversation], ctx);
|
||||
});
|
||||
|
||||
let mut mock = MockAIClient::new();
|
||||
// The fire-and-forget server PATCH should be issued; permissive Ok.
|
||||
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)
|
||||
});
|
||||
|
||||
// Build a poll batch with max sequence = 42. Use an unrecognized
|
||||
// event_type so `convert_lifecycle_events` returns empty and the
|
||||
// function early-exits before touching `OrchestrationEventService`
|
||||
// (which we did not register in this test App).
|
||||
let events = vec![
|
||||
AgentRunEvent {
|
||||
event_type: "unrecognized_event_type".to_string(),
|
||||
run_id: "some-other-run".to_string(),
|
||||
ref_id: None,
|
||||
execution_id: None,
|
||||
occurred_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
sequence: 17,
|
||||
},
|
||||
AgentRunEvent {
|
||||
event_type: "unrecognized_event_type".to_string(),
|
||||
run_id: "some-other-run".to_string(),
|
||||
ref_id: None,
|
||||
execution_id: None,
|
||||
occurred_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
sequence: 42,
|
||||
},
|
||||
];
|
||||
|
||||
poller.update(&mut app, |me, ctx| {
|
||||
me.handle_poll_result(
|
||||
conversation_id,
|
||||
/* self_run_id */ "some-other-run",
|
||||
/* previous_cursor */ 0,
|
||||
events,
|
||||
/* messages */ vec![],
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
|
||||
history_model.read(&app, |model, _| {
|
||||
let last_seq = model
|
||||
.conversation(&conversation_id)
|
||||
.and_then(|c| c.last_event_sequence());
|
||||
assert_eq!(
|
||||
last_seq,
|
||||
Some(42),
|
||||
"BlocklistAIHistoryModel.update_event_sequence must be called with max_seq"
|
||||
);
|
||||
});
|
||||
|
||||
// Drain at least one persistence event to confirm the SQLite write
|
||||
// path was triggered (sanity check for the side effect, not the
|
||||
// primary assertion).
|
||||
let _ = receiver.recv_timeout(std::time::Duration::from_secs(1));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finish_restore_fetch_no_ops_when_conversation_deleted_mid_flight() {
|
||||
// If the conversation is removed while the async fetch is in-flight, the
|
||||
// RemoveConversation handler clears event_cursor. finish_restore_fetch
|
||||
// uses the missing cursor as a sentinel and must not re-populate
|
||||
// watched_run_ids or event_cursor for the deleted conversation.
|
||||
use crate::ai::agent::conversation::AIConversation;
|
||||
use crate::server::server_api::ai::MockAIClient;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
use std::sync::Arc;
|
||||
use warpui::App;
|
||||
|
||||
App::test((), |mut app| async move {
|
||||
let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true);
|
||||
|
||||
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[]));
|
||||
|
||||
let mut conversation = AIConversation::new(false);
|
||||
conversation.set_run_id("550e8400-e29b-41d4-a716-446655440300".to_string());
|
||||
let conversation_id = conversation.id();
|
||||
let terminal_view_id = warpui::EntityId::new();
|
||||
history_model.update(&mut app, |model, ctx| {
|
||||
model.restore_conversations(terminal_view_id, vec![conversation], ctx);
|
||||
});
|
||||
|
||||
let mock = MockAIClient::new();
|
||||
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)
|
||||
});
|
||||
|
||||
// Seed cursor as on_restored_conversations would.
|
||||
poller.update(&mut app, |me, _| {
|
||||
me.event_cursor.insert(conversation_id, 0);
|
||||
});
|
||||
|
||||
// Simulate the RemoveConversation handler firing while the fetch is
|
||||
// in-flight: it clears event_cursor (and all other state).
|
||||
poller.update(&mut app, |me, _| {
|
||||
me.watched_run_ids.remove(&conversation_id);
|
||||
me.event_cursor.remove(&conversation_id);
|
||||
});
|
||||
|
||||
// The in-flight fetch now completes — with children.
|
||||
let task_id: crate::ai::ambient_agents::AmbientAgentTaskId =
|
||||
"550e8400-e29b-41d4-a716-446655440000".parse().unwrap();
|
||||
poller.update(&mut app, |me, ctx| {
|
||||
me.finish_restore_fetch(
|
||||
conversation_id,
|
||||
task_id,
|
||||
/* sqlite_cursor */ 0,
|
||||
Ok(make_ambient_task_with_children(vec![
|
||||
"child-run-1".to_string()
|
||||
])),
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
|
||||
poller.read(&app, |me, _| {
|
||||
assert!(
|
||||
!me.watched_run_ids.contains_key(&conversation_id),
|
||||
"watched_run_ids must not be repopulated for a deleted conversation"
|
||||
);
|
||||
assert!(
|
||||
!me.event_cursor.contains_key(&conversation_id),
|
||||
"event_cursor must not be repopulated for a deleted conversation"
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finish_restore_fetch_reconnects_sse_when_children_added_to_open_connection() {
|
||||
// When a status transition races with the restore fetch and opens SSE
|
||||
// before children are known, finish_restore_fetch must reconnect SSE
|
||||
// with the updated run_id set rather than leaving children unwatched.
|
||||
use crate::ai::agent::conversation::{AIConversation, ConversationStatus};
|
||||
use crate::server::server_api::ai::MockAIClient;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
use std::sync::Arc;
|
||||
use warpui::App;
|
||||
|
||||
App::test((), |mut app| async move {
|
||||
let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true);
|
||||
|
||||
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[]));
|
||||
|
||||
let own_run_id = "550e8400-e29b-41d4-a716-446655440400";
|
||||
let mut conversation = AIConversation::new(false);
|
||||
conversation.set_run_id(own_run_id.to_string());
|
||||
let conversation_id = conversation.id();
|
||||
let terminal_view_id = warpui::EntityId::new();
|
||||
history_model.update(&mut app, |model, ctx| {
|
||||
model.restore_conversations(terminal_view_id, vec![conversation], ctx);
|
||||
model.update_conversation_status(
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
ConversationStatus::InProgress,
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
|
||||
let mock = MockAIClient::new();
|
||||
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)
|
||||
});
|
||||
|
||||
// Seed the state on_restored_conversations would have set up, then
|
||||
// inject a fake open SSE connection (generation 0) simulating the
|
||||
// race: a status transition fired before the restore fetch completed.
|
||||
let (_, rx) = futures::channel::mpsc::unbounded::<SseStreamItem>();
|
||||
poller.update(&mut app, |me, _| {
|
||||
me.event_cursor.insert(conversation_id, 0);
|
||||
me.watched_run_ids
|
||||
.entry(conversation_id)
|
||||
.or_default()
|
||||
.insert(own_run_id.to_string());
|
||||
me.sse_connections.insert(
|
||||
conversation_id,
|
||||
SseConnectionState {
|
||||
event_receiver: rx,
|
||||
generation: 0,
|
||||
},
|
||||
);
|
||||
me.next_sse_generation = 1;
|
||||
});
|
||||
|
||||
// The restore fetch returns with a child run_id.
|
||||
let task_id: crate::ai::ambient_agents::AmbientAgentTaskId =
|
||||
"550e8400-e29b-41d4-a716-446655440000".parse().unwrap();
|
||||
poller.update(&mut app, |me, ctx| {
|
||||
me.finish_restore_fetch(
|
||||
conversation_id,
|
||||
task_id,
|
||||
/* sqlite_cursor */ 0,
|
||||
Ok(make_ambient_task_with_children(vec![
|
||||
"child-run-1".to_string()
|
||||
])),
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
|
||||
poller.read(&app, |me, _| {
|
||||
assert!(
|
||||
me.watched_run_ids
|
||||
.get(&conversation_id)
|
||||
.is_some_and(|w| w.contains("child-run-1")),
|
||||
"child run_id must be in watched set"
|
||||
);
|
||||
// The old generation-0 connection must have been replaced by a
|
||||
// new one with a higher generation, proving SSE was reconnected.
|
||||
let gen = me
|
||||
.sse_connections
|
||||
.get(&conversation_id)
|
||||
.map(|s| s.generation);
|
||||
assert!(
|
||||
gen.is_some_and(|g| g > 0),
|
||||
"SSE must be reconnected (new generation) after children are discovered; got gen={gen:?}"
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
@@ -357,6 +357,42 @@ impl OrchestrationEventService {
|
||||
} => {
|
||||
for conversation_id in conversation_ids {
|
||||
self.sync_conversation_status(*conversation_id, ctx);
|
||||
// Under V1 local lifecycle dispatch, child status
|
||||
// transitions are forwarded to the parent via
|
||||
// `lifecycle_subscription_routes`. That map is not
|
||||
// persisted, so re-register subscriptions for each
|
||||
// restored child whose parent is loaded locally so that
|
||||
// child status transitions continue to propagate after
|
||||
// a restart. V2 uses the server event log and does not
|
||||
// need this.
|
||||
if !FeatureFlag::OrchestrationV2.is_enabled() {
|
||||
let parent_agent_id = {
|
||||
let history_model = BlocklistAIHistoryModel::as_ref(ctx);
|
||||
let Some(child_conv) = history_model.conversation(conversation_id)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if !child_conv.is_child_agent_conversation() {
|
||||
continue;
|
||||
}
|
||||
child_conv
|
||||
.parent_conversation_id()
|
||||
.and_then(|pid| history_model.conversation(&pid))
|
||||
.and_then(|p| p.server_conversation_token())
|
||||
.map(|t| t.as_str().to_string())
|
||||
};
|
||||
if let Some(parent_agent_id) = parent_agent_id {
|
||||
// `None` event-type filter = subscribe to all
|
||||
// lifecycle types. The original filter (if any)
|
||||
// is not persisted; subscribing broader than the
|
||||
// original is acceptable per the tech spec.
|
||||
self.register_lifecycle_subscription(
|
||||
*conversation_id,
|
||||
parent_agent_id,
|
||||
None,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
BlocklistAIHistoryEvent::RemoveConversation {
|
||||
|
||||
@@ -357,3 +357,67 @@ fn test_lifecycle_event_type_from_proto_includes_cancelled_and_blocked() {
|
||||
api::LifecycleEventType::Blocked
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restored_v1_child_conversation_re_registers_lifecycle_subscription() {
|
||||
use crate::ai::agent::conversation::AIConversation;
|
||||
use warp_core::features::FeatureFlag;
|
||||
use warpui::{App, EntityId};
|
||||
|
||||
App::test((), |mut app| async move {
|
||||
// V1 path is gated on `!OrchestrationV2`.
|
||||
let _v1_guard = FeatureFlag::OrchestrationV2.override_enabled(false);
|
||||
|
||||
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[]));
|
||||
|
||||
// Build a parent conversation with a server token; under V1 the
|
||||
// parent's `server_conversation_token` is the agent identifier the
|
||||
// child subscribes to.
|
||||
let parent_token = "parent-token-v1";
|
||||
let mut parent_conversation = AIConversation::new(false);
|
||||
parent_conversation.set_server_conversation_token(parent_token.to_string());
|
||||
let parent_conversation_id = parent_conversation.id();
|
||||
|
||||
// Build a child conversation pointing at the parent.
|
||||
let mut child_conversation = AIConversation::new(false);
|
||||
child_conversation.set_parent_conversation_id(parent_conversation_id);
|
||||
let child_conversation_id = child_conversation.id();
|
||||
|
||||
let terminal_view_id = EntityId::new();
|
||||
history_model.update(&mut app, |model, ctx| {
|
||||
model.restore_conversations(
|
||||
terminal_view_id,
|
||||
vec![parent_conversation, child_conversation],
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
|
||||
// Drive the OrchestrationEventService through its standard
|
||||
// `handle_history_event` entry point; `restore_conversations` already
|
||||
// emitted `RestoredConversations`, so we replay it explicitly through
|
||||
// the service to keep this test independent of subscription wiring.
|
||||
let service = app.add_singleton_model(|_| OrchestrationEventService::default());
|
||||
service.update(&mut app, |svc, ctx| {
|
||||
svc.handle_history_event(
|
||||
&BlocklistAIHistoryEvent::RestoredConversations {
|
||||
terminal_view_id,
|
||||
conversation_ids: vec![parent_conversation_id, child_conversation_id],
|
||||
},
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
|
||||
service.read(&app, |svc, _| {
|
||||
let routes = svc
|
||||
.lifecycle_subscription_routes
|
||||
.get(&child_conversation_id)
|
||||
.expect("expected V1 lifecycle route to be registered for the child");
|
||||
assert_eq!(routes.len(), 1, "expected exactly one route");
|
||||
assert_eq!(routes[0].target_agent_id, parent_token);
|
||||
assert!(
|
||||
routes[0].subscribed_event_types.is_none(),
|
||||
"restore re-registers with `None` (subscribe to all event types)"
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
@@ -39,6 +39,8 @@ fn create_test_task(task_id: &str) -> AmbientAgentTask {
|
||||
agent_config_snapshot: None,
|
||||
artifacts: vec![],
|
||||
is_sandbox_running: false,
|
||||
last_event_sequence: None,
|
||||
children: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -130,6 +132,7 @@ fn test_from_task_includes_linked_directory_when_run_id_matches() {
|
||||
parent_conversation_id: None,
|
||||
run_id: Some(task_id.to_string()),
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -244,6 +247,7 @@ fn test_from_task_includes_linked_directory_when_server_token_matches() {
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
},
|
||||
);
|
||||
|
||||
|
||||
@@ -791,6 +791,39 @@ impl ServerApi {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Sends a PATCH request to a public API endpoint that returns no response body.
|
||||
async fn patch_public_api_unit<B>(&self, path: &str, body: &B) -> Result<()>
|
||||
where
|
||||
B: Serialize,
|
||||
{
|
||||
let auth_token = self
|
||||
.get_or_refresh_access_token()
|
||||
.await
|
||||
.context("Failed to get access token for API request")?;
|
||||
|
||||
let url = format!("{}/api/v1/{}", ChannelState::server_root_url(), path);
|
||||
|
||||
let mut request = self.client.patch(&url).json(body);
|
||||
if let Some(token) = auth_token.as_bearer_token() {
|
||||
request = request.bearer_auth(token);
|
||||
}
|
||||
|
||||
for (name, value) in self.ambient_agent_headers().await? {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
|
||||
let response = request
|
||||
.send()
|
||||
.await
|
||||
.with_context(|| format!("Failed to send API request to {url}"))?;
|
||||
|
||||
if response.status().is_success() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(Self::error_from_response(response).await)
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends an authenticated empty POST request to /client/login, which signals to the server
|
||||
/// that the user is logged in.
|
||||
pub async fn notify_login(&self) {
|
||||
|
||||
@@ -923,6 +923,16 @@ pub trait AIClient: 'static + Send + Sync {
|
||||
limit: i32,
|
||||
) -> anyhow::Result<Vec<AgentRunEvent>, anyhow::Error>;
|
||||
|
||||
/// Persists the latest observed event sequence number for a run on the
|
||||
/// server. Used to keep the server-side cursor in sync with the client so
|
||||
/// that driver/cloud restores can resume without replaying events the
|
||||
/// parent has already acted on.
|
||||
async fn update_event_sequence_on_server(
|
||||
&self,
|
||||
run_id: &str,
|
||||
sequence: i64,
|
||||
) -> anyhow::Result<(), anyhow::Error>;
|
||||
|
||||
async fn report_agent_event(
|
||||
&self,
|
||||
run_id: &str,
|
||||
@@ -1885,6 +1895,22 @@ impl AIClient for ServerApi {
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
async fn update_event_sequence_on_server(
|
||||
&self,
|
||||
run_id: &str,
|
||||
sequence: i64,
|
||||
) -> anyhow::Result<(), anyhow::Error> {
|
||||
#[derive(serde::Serialize)]
|
||||
struct UpdateBody {
|
||||
sequence: i64,
|
||||
}
|
||||
self.patch_public_api_unit(
|
||||
&format!("agent/runs/{run_id}/event-sequence"),
|
||||
&UpdateBody { sequence },
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn report_agent_event(
|
||||
&self,
|
||||
run_id: &str,
|
||||
|
||||
@@ -988,6 +988,7 @@ impl TerminalView {
|
||||
parent_conversation_id: None,
|
||||
run_id: None,
|
||||
autoexecute_override: None,
|
||||
last_event_sequence: None,
|
||||
};
|
||||
|
||||
match AIConversation::new_restored(conversation_id, tasks, Some(conversation_data)) {
|
||||
|
||||
@@ -44,6 +44,8 @@ fn task_with_run_time_and_credits() -> AmbientAgentTask {
|
||||
agent_config_snapshot: None,
|
||||
artifacts: vec![],
|
||||
is_sandbox_running: false,
|
||||
last_event_sequence: None,
|
||||
children: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user