Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 28 additions & 16 deletions app/src/ai/blocklist/history_model.rs
Original file line number Diff line number Diff line change
Expand Up @@ -327,6 +327,8 @@ pub struct BlocklistAIHistoryModel {

/// In-flight optimistic conversation rename state keyed by conversation.
in_flight_conversation_renames: HashMap<AIConversationId, InFlightConversationRename>,
/// Conversations with a server metadata request in progress.
in_flight_server_metadata_fetches: HashSet<AIConversationId>,

#[cfg(feature = "local_fs")]
db_connection: Option<Arc<Mutex<SqliteConnection>>>,
Expand Down Expand Up @@ -2123,6 +2125,12 @@ impl BlocklistAIHistoryModel {
.unwrap()
.as_str()
.to_string();
if !self
.in_flight_server_metadata_fetches
.insert(conversation_id)
{
return;
}

let server_api = ServerApiProvider::as_ref(ctx).get_ai_client();
ctx.spawn(
Expand All @@ -2131,25 +2139,29 @@ impl BlocklistAIHistoryModel {
.list_ai_conversation_metadata(Some(vec![server_token]))
.await
},
move |model, result, ctx| match result {
Ok(mut metadata_list) if !metadata_list.is_empty() => {
if let Some(metadata) = metadata_list.pop() {
model.set_server_metadata_for_conversation(
conversation_id,
metadata,
ctx,
move |model, result, ctx| {
model
.in_flight_server_metadata_fetches
.remove(&conversation_id);
match result {
Ok(mut metadata_list) if !metadata_list.is_empty() => {
if let Some(metadata) = metadata_list.pop() {
model.set_server_metadata_for_conversation(
conversation_id,
metadata,
ctx,
);
}
}
Ok(_) => {
log::warn!("No metadata returned for conversation {conversation_id}");
}
Err(e) => {
log::warn!(
"Failed to fetch metadata for conversation {conversation_id}: {e:#}"
);
}
}
Ok(_) => {
log::warn!("No metadata returned for conversation {}", conversation_id);
}
Err(e) => {
log::warn!(
"Failed to fetch metadata for conversation {}: {e:#}",
conversation_id
);
}
},
);
}
Expand Down
43 changes: 43 additions & 0 deletions app/src/ai/blocklist/history_model_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,13 +33,15 @@ use crate::ai::blocklist::controller::RequestInput;
use crate::ai::llms::LLMId;
use crate::auth::AuthStateProvider;
use crate::cloud_object::{Owner, Revision, ServerMetadata, ServerPermissions};
use crate::features::FeatureFlag;
use crate::input_suggestions::HistoryInputSuggestion;
use crate::persistence::ModelEvent;
use crate::persistence::model::{
AgentConversation, AgentConversationData, AgentConversationRecord, AgentConversationSummary,
PersistedAutoexecuteMode,
};
use crate::server::ids::ServerId;
use crate::server::server_api::ServerApiProvider;
use crate::server::telemetry::context_provider::AppTelemetryContextProvider;
use crate::terminal::model::session::SessionId;
use crate::test_util::ai_agent_tasks::create_api_task;
Expand Down Expand Up @@ -4375,6 +4377,47 @@ fn hydrate_remote_child_placeholder_with_cloud_transcript_preserves_placeholder_
});
}

#[test]
fn repeated_stream_completions_share_one_in_flight_metadata_fetch() {
let _cloud_conversations = FeatureFlag::CloudConversations.override_enabled(true);

App::test((), |mut app| async move {
initialize_history_persistence_for_tests(&mut app);
app.add_singleton_model(|_| ServerApiProvider::new_for_test());
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test());
let terminal_surface_id = EntityId::new();
let conversation_id = history_model.update(&mut app, |history, ctx| {
let conversation_id =
history.start_new_conversation(terminal_surface_id, false, false, false, ctx);
history.set_server_conversation_token_for_conversation(
conversation_id,
"metadata-fetch-token".to_string(),
);
conversation_id
});
let stream_id = ResponseStreamId::new_for_test();

history_model.update(&mut app, |history, ctx| {
history.mark_response_stream_completed_successfully(
&stream_id,
conversation_id,
terminal_surface_id,
ctx,
);
history.mark_response_stream_completed_successfully(
&stream_id,
conversation_id,
terminal_surface_id,
ctx,
);
assert_eq!(
history.in_flight_server_metadata_fetches,
HashSet::from([conversation_id])
);
});
});
}

// --- conversation_output_status_from_conversation ---

/// Builds a conversation with one in-flight exchange, completes it with the
Expand Down
27 changes: 22 additions & 5 deletions app/src/ai/blocklist/orchestration_event_streamer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,8 @@ struct ConversationStreamState {
/// metadata, so this lets us recognize dormant local Claude children
/// without relying on `ServerAIConversationMetadata`.
harness: Option<Harness>,
/// Whether a task request to resolve the execution harness is in progress.
harness_fetch_in_flight: bool,
/// Active SSE connection, if one is open.
sse_connection: Option<SseConnectionState>,
/// Active wake-only listener for dormant local Claude children, if one is
Expand Down Expand Up @@ -262,6 +264,8 @@ struct OrchestratorStreamState {
/// cursor, so a replay does not generate spurious `ChildSpawned` events
/// for already-known children.
seeded: bool,
/// `true` while the cold-start REST seed request is in progress.
seed_fetch_in_flight: bool,
/// Observer-mode child tracker for this orchestrator family. `None` until
/// the family drain creates one on the first batch it handles.
tracker: Option<OrchestrationChildTracker>,
Expand Down Expand Up @@ -1236,7 +1240,11 @@ impl OrchestrationEventStreamer {
entry
.consumers
.insert(consumer_id, orchestrator_placeholder_conv_id);
needs_seed = !entry.seeded && entry.sse_connection.is_none();
needs_seed =
!entry.seeded && !entry.seed_fetch_in_flight && entry.sse_connection.is_none();
if needs_seed {
entry.seed_fetch_in_flight = true;
}
}
// Hydrate the orchestrator placeholder's persisted cursor into the
// per-orchestrator entry so a restart-from-disk picks up where the
Expand Down Expand Up @@ -1378,13 +1386,14 @@ impl OrchestrationEventStreamer {
result: anyhow::Result<Vec<crate::ai::ambient_agents::task::AmbientAgentTask>>,
ctx: &mut ModelContext<Self>,
) {
if !self.viewer_mode_orchestrators.contains_key(&parent_task_id) {
let Some(entry) = self.viewer_mode_orchestrators.get_mut(&parent_task_id) else {
log::warn!(
"[orch-viewer-streamer] ancestor seed fetch completed but viewer-mode entry \
for parent_task_id={parent_task_id} is gone; dropping"
);
return;
};
entry.seed_fetch_in_flight = false;
match result {
Ok(tasks) => {
let tasks_received = tasks.len();
Expand All @@ -1403,8 +1412,9 @@ impl OrchestrationEventStreamer {
continue;
}
let run_id = task.task_id.to_string();
entry.known_children.insert(run_id.clone());
seeded_run_ids.push(run_id);
if entry.known_children.insert(run_id.clone()) {
seeded_run_ids.push(run_id);
}
if let Some(seq) = task.last_event_sequence {
seed = seed.max(seq);
}
Expand Down Expand Up @@ -1724,7 +1734,7 @@ impl OrchestrationEventStreamer {
if self
.streams
.get(&conversation_id)
.is_some_and(|stream| stream.harness.is_some())
.is_some_and(|stream| stream.harness.is_some() || stream.harness_fetch_in_flight)
{
return;
}
Expand All @@ -1739,10 +1749,17 @@ impl OrchestrationEventStreamer {
.get(&conversation_id)
.map(|stream| stream.event_cursor)
.unwrap_or(0);
self.streams
.entry(conversation_id)
.or_default()
.harness_fetch_in_flight = true;
let ai_client = self.ai_client.clone();
ctx.spawn(
async move { ai_client.get_ambient_agent_task(&task_id).await },
move |me, result, ctx| {
if let Some(stream) = me.streams.get_mut(&conversation_id) {
stream.harness_fetch_in_flight = false;
}
let task = match result {
Ok(task) => task,
Err(err) => {
Expand Down
109 changes: 109 additions & 0 deletions app/src/ai/blocklist/orchestration_event_streamer_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -623,6 +623,39 @@ fn make_server_metadata_with_harness(
}
}

#[test]
fn repeated_harness_fetch_attempts_share_one_in_flight_request() {
App::test((), |mut app| async move {
let history_model =
app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], vec![], &[]));
let run_id = "550e8400-e29b-41d4-a716-446655440620";
let mut conversation = AIConversation::new(false, false);
conversation.set_run_id(run_id.to_string());
let conversation_id = conversation.id();
history_model.update(&mut app, |history, ctx| {
history.restore_conversations(warpui::EntityId::new(), vec![conversation], ctx);
});

let mut mock = MockAIClient::new();
mock.expect_get_ambient_agent_task()
.times(1)
.returning(|_| Err(anyhow::anyhow!("fetch observed")));
let ai_client: Arc<dyn AIClient> = Arc::new(mock);
let server_api = ServerApiProvider::new_for_test().get();
let streamer = app.add_singleton_model(|ctx| {
OrchestrationEventStreamer::new_with_clients_for_test(ai_client, server_api, ctx)
});

streamer.update(&mut app, |streamer, ctx| {
streamer.spawn_task_harness_fetch_if_needed(conversation_id, ctx);
streamer.spawn_task_harness_fetch_if_needed(conversation_id, ctx);
});
for _ in 0..3 {
futures_lite::future::yield_now().await;
}
});
}

#[test]
fn dormant_local_claude_child_skips_generic_sse_but_allows_wake_listener() {
use std::sync::Arc;
Expand Down Expand Up @@ -1698,6 +1731,43 @@ fn make_parent_task_id_for_test(byte: u8) -> AmbientAgentTaskId {
let s = uuid.to_string();
s.parse().expect("valid task id")
}
#[test]
fn repeated_viewer_registration_starts_one_ancestor_seed_fetch() {
App::test((), |mut app| async move {
app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], vec![], &[]));

let mut mock = MockAIClient::new();
mock.expect_list_ambient_agent_tasks()
.times(1)
.returning(|_, _| Err(anyhow::anyhow!("fetch observed")));
let ai_client: Arc<dyn AIClient> = Arc::new(mock);
let server_api = ServerApiProvider::new_for_test().get();
let streamer = app.add_singleton_model(|ctx| {
OrchestrationEventStreamer::new_with_clients_for_test(ai_client, server_api, ctx)
});
let parent_task_id = make_parent_task_id_for_test(0xa0);
let placeholder_id = AIConversation::new(true, false).id();
let consumer_id = warpui::EntityId::new();

streamer.update(&mut app, |streamer, ctx| {
streamer.register_viewer_mode_consumer(
parent_task_id,
placeholder_id,
consumer_id,
ctx,
);
streamer.register_viewer_mode_consumer(
parent_task_id,
placeholder_id,
consumer_id,
ctx,
);
});
for _ in 0..3 {
futures_lite::future::yield_now().await;
}
});
}

#[test]
fn is_known_child_dedupes_per_parent_after_first_observation() {
Expand Down Expand Up @@ -2249,6 +2319,45 @@ fn finish_ancestor_seed_fetch_emits_child_spawned_for_each_seeded_child() {
});
}

#[test]
fn repeated_ancestor_seed_results_do_not_rebroadcast_known_children() {
App::test((), |mut app| async move {
app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], vec![], &[]));

let ai_client: Arc<dyn AIClient> = Arc::new(MockAIClient::new());
let server_api = ServerApiProvider::new_for_test().get();
let streamer = app.add_singleton_model(|ctx| {
OrchestrationEventStreamer::new_with_clients_for_test(ai_client, server_api, ctx)
});
let parent_task_id = make_parent_task_id_for_test(0xd4);
let child_task_id = make_parent_task_id_for_test(0xd5);
streamer.update(&mut app, |streamer, _| {
streamer
.viewer_mode_orchestrators
.entry(parent_task_id)
.or_default();
});
let captured_spawns = capture_child_spawns(&mut app, &streamer);

streamer.update(&mut app, |streamer, ctx| {
streamer.finish_ancestor_seed_fetch(
parent_task_id,
Ok(vec![make_ambient_task_with_task_id(child_task_id, None)]),
ctx,
);
streamer.finish_ancestor_seed_fetch(
parent_task_id,
Ok(vec![make_ambient_task_with_task_id(child_task_id, None)]),
ctx,
);
});

assert_eq!(
captured_spawns.lock().as_slice(),
&[(parent_task_id, child_task_id.to_string())]
);
});
}
#[test]
fn register_viewer_mode_consumer_replays_known_children_for_later_panes() {
// Regression for the late-arriving-consumer arm of the same bug: the
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -308,7 +308,9 @@ impl OrchestrationViewerModel {
}
return;
}
self.metadata_fetches.insert(task_id);
if !self.metadata_fetches.insert(task_id) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ [IMPORTANT] This bug fix is reachable at the model level and metadata_fetch_dispatch_count is already available below, but the PR adds no regression test for duplicate triggers before the first request completes. Please add coverage that calls this path twice before completion and asserts only one metadata request is dispatched.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added tests

return;
}
#[cfg(test)]
{
self.metadata_fetch_dispatch_count += 1;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -448,6 +448,26 @@ fn status_change_does_not_refetch_metadata_after_materialization() {
});
});
}
#[test]
fn repeated_child_discovery_starts_one_legacy_metadata_fetch() {
let _unified_stack = FeatureFlag::OrchestrationUnifiedStack.override_enabled(false);
App::test((), |mut app| async move {
let fixture = setup(&mut app);

fixture.model.update(&mut app, |model, ctx| {
model.handle_child_spawned(CHILD_A_TASK_ID.to_string(), ctx);
model.handle_child_spawned(CHILD_A_TASK_ID.to_string(), ctx);
});

fixture.model.read(&app, |model, _| {
assert_eq!(model.metadata_fetch_dispatch_count, 1);
assert_eq!(
model.metadata_fetches,
HashSet::from([task_id(CHILD_A_TASK_ID)])
);
});
});
}

// ---- Pending-metadata poll --------------------------------------------------

Expand Down
2 changes: 1 addition & 1 deletion script/wasm/bundle
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ mkdir -p "$EXTRAS_DIR"
# N.B. The bundled outputs will always be warp.js and warp_bg.wasm.
if [[ $RELEASE_CHANNEL = "local" ]]; then
WARP_BIN="warp"
FEATURES="$FEATURES,remote_tty"
FEATURES="$FEATURES"
elif [[ $RELEASE_CHANNEL = "dev" ]]; then
WARP_BIN="dev"
elif [[ $RELEASE_CHANNEL = "preview" ]]; then
Expand Down
Loading