diff --git a/app/src/ai/agent/api/convert_conversation.rs b/app/src/ai/agent/api/convert_conversation.rs index b77e430cd..943e9e456 100644 --- a/app/src/ai/agent/api/convert_conversation.rs +++ b/app/src/ai/agent/api/convert_conversation.rs @@ -82,6 +82,7 @@ pub fn convert_conversation_data_to_ai_conversation( parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, @@ -97,6 +98,7 @@ pub fn convert_conversation_data_to_ai_conversation( parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, // TODO: Populate run_id from server metadata once it is exposed // in ServerAIConversationMetadata. For cloud conversations that // were spawned via the server API, the run_id is created at task diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 6fc03805e..d3fc408d7 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -355,6 +355,7 @@ impl AIConversation { parent_agent_id, agent_name, parent_conversation_id, + is_remote_child, run_id, autoexecute_override, last_event_sequence, @@ -380,6 +381,7 @@ impl AIConversation { let parent_conversation_id = data .parent_conversation_id .and_then(|id| AIConversationId::try_from(id).ok()); + let is_remote_child = data.is_remote_child; let run_id = data.run_id; let autoexecute_override = if FeatureFlag::RememberFastForwardState.is_enabled() { data.autoexecute_override @@ -399,6 +401,7 @@ impl AIConversation { parent_agent_id, agent_name, parent_conversation_id, + is_remote_child, run_id, autoexecute_override, last_event_sequence, @@ -413,6 +416,7 @@ impl AIConversation { None, None, None, + false, None, AIConversationAutoexecuteMode::default(), None, @@ -457,7 +461,7 @@ impl AIConversation { parent_agent_id, agent_name, parent_conversation_id, - is_remote_child: false, + is_remote_child, last_event_sequence, }) } @@ -809,7 +813,7 @@ impl AIConversation { /// 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() + self.parent_conversation_id.is_some() || self.parent_agent_id.is_some() } /// True iff this conversation knows about a parent agent — either via a @@ -2838,6 +2842,7 @@ impl AIConversation { parent_agent_id: self.parent_agent_id.clone(), agent_name: self.agent_name.clone(), parent_conversation_id: self.parent_conversation_id.map(|id| id.to_string()), + is_remote_child: self.is_remote_child, run_id: self.task_id.map(|id| id.to_string()), autoexecute_override: Some(self.autoexecute_override.into()), last_event_sequence: self.last_event_sequence, diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index d330f264c..3ba96b4f1 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/app/src/ai/agent/conversation_tests.rs @@ -121,6 +121,41 @@ fn restored_conversation_defaults_autoexecute_override_when_not_persisted() { ); } +#[test] +fn restored_conversation_uses_persisted_last_event_sequence() { + let conversation_data: AgentConversationData = + serde_json::from_str(r#"{"server_conversation_token":null,"last_event_sequence":42}"#) + .unwrap(); + + let conversation = restored_conversation(Some(conversation_data)); + + assert_eq!(conversation.last_event_sequence(), Some(42)); +} + +#[test] +fn restored_conversation_uses_persisted_remote_child_marker() { + let conversation_data: AgentConversationData = + serde_json::from_str(r#"{"server_conversation_token":null,"is_remote_child":true}"#) + .unwrap(); + + let conversation = restored_conversation(Some(conversation_data)); + + assert!(conversation.is_remote_child()); +} + +#[test] +fn child_conversation_detection_uses_parent_agent_id() { + let conversation_data: AgentConversationData = serde_json::from_str( + r#"{"server_conversation_token":null,"parent_agent_id":"parent-run-id"}"#, + ) + .unwrap(); + + let conversation = restored_conversation(Some(conversation_data)); + + assert!(conversation.is_child_agent_conversation()); + assert_eq!(conversation.parent_conversation_id(), None); +} + #[test] fn restored_conversation_defaults_unknown_persisted_autoexecute_override() { let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true); diff --git a/app/src/ai/agent_conversations_model_tests.rs b/app/src/ai/agent_conversations_model_tests.rs index 37ee5d715..e20fe1dd3 100644 --- a/app/src/ai/agent_conversations_model_tests.rs +++ b/app/src/ai/agent_conversations_model_tests.rs @@ -152,6 +152,7 @@ fn test_display_status_uses_matching_conversation_for_in_progress_task() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -204,6 +205,7 @@ fn test_display_status_updates_when_blocked_conversation_resumes() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -280,6 +282,7 @@ fn test_display_status_terminal_task_state_overrides_matching_conversation() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -332,6 +335,7 @@ fn test_status_filter_uses_display_status_for_task_backed_conversations() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -766,6 +770,7 @@ fn test_get_tasks_and_conversations_prefers_task_when_task_id_matches_conversati parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -823,6 +828,7 @@ fn test_get_tasks_and_conversations_prefers_task_when_server_token_matches() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, @@ -879,6 +885,7 @@ fn test_get_tasks_and_conversations_keeps_unrelated_tasks_and_conversations() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/app/src/ai/agent_events/message_hydrator.rs b/app/src/ai/agent_events/message_hydrator.rs index e73d78c35..7d5711181 100644 --- a/app/src/ai/agent_events/message_hydrator.rs +++ b/app/src/ai/agent_events/message_hydrator.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use std::time::Duration; +use crate::ai::ambient_agents::AmbientAgentTaskId; use anyhow::{anyhow, Context, Result}; #[cfg(not(target_family = "wasm"))] use futures::future::Either; @@ -9,6 +10,7 @@ use warpui::r#async::Timer; use crate::ai::agent::ReceivedMessageInput; use crate::server::server_api::ai::{AIClient, AgentRunEvent, ReadAgentMessageResponse}; +use crate::server::server_api::ServerApi; pub(crate) const DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(5); @@ -17,6 +19,8 @@ pub(crate) const DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT: Duration = Duration::from_ #[derive(Clone)] pub(crate) struct MessageHydrator { ai_client: Arc, + task_scoped_server_api: Option>, + task_id: Option, #[cfg_attr(target_family = "wasm", allow(dead_code))] fetch_timeout: Duration, } @@ -26,16 +30,40 @@ impl MessageHydrator { Self::with_fetch_timeout(ai_client, DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT) } + pub(crate) fn for_task(server_api: Arc, task_id: AmbientAgentTaskId) -> Self { + let ai_client: Arc = server_api.clone(); + Self { + ai_client, + task_scoped_server_api: Some(server_api), + task_id: Some(task_id), + fetch_timeout: DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT, + } + } + pub(crate) fn with_fetch_timeout( ai_client: Arc, fetch_timeout: Duration, ) -> Self { Self { ai_client, + task_scoped_server_api: None, + task_id: None, fetch_timeout, } } + async fn read_message(&self, message_id: &str) -> Result { + match (self.task_scoped_server_api.as_ref(), self.task_id) { + (Some(server_api), Some(task_id)) => { + server_api + .read_agent_message_for_task(&task_id, message_id) + .await + } + _ => self.ai_client.read_agent_message(message_id).await, + } + .with_context(|| format!("Failed to read agent message {message_id}")) + } + pub(crate) async fn hydrate_event_for_recipient( &self, event: &AgentRunEvent, @@ -55,6 +83,17 @@ impl MessageHydrator { return None; } }; + if message.body.is_empty() { + log::warn!( + "Hydrated empty-body agent message: message_id={} event_sequence={} recipient_run_id={} sender_run_id={} subject={:?} task_id={:?}", + message.message_id, + event.sequence, + recipient_run_id, + message.sender_run_id, + message.subject, + self.task_id.map(|task_id| task_id.to_string()) + ); + } Some(ReceivedMessageInput { message_id: message.message_id, @@ -70,15 +109,13 @@ impl MessageHydrator { &self, message_id: &str, ) -> Result { - let read_message = self.ai_client.read_agent_message(message_id); + let read_message = self.read_message(message_id); let timeout = Timer::after(self.fetch_timeout); futures::pin_mut!(read_message); futures::pin_mut!(timeout); match futures::future::select(read_message, timeout).await { - Either::Left((result, _)) => { - result.with_context(|| format!("Failed to read agent message {message_id}")) - } + Either::Left((result, _)) => result, Either::Right(_) => Err(anyhow!("Timed out reading agent message {message_id}")), } } @@ -88,10 +125,7 @@ impl MessageHydrator { &self, message_id: &str, ) -> Result { - self.ai_client - .read_agent_message(message_id) - .await - .with_context(|| format!("Failed to read agent message {message_id}")) + self.read_message(message_id).await } pub(crate) async fn read_message_from_event_with_timeout( @@ -105,10 +139,15 @@ impl MessageHydrator { } pub(crate) async fn mark_message_delivered(&self, message_id: &str) -> Result<()> { - self.ai_client - .mark_message_delivered(message_id) - .await - .with_context(|| format!("Failed to mark agent message {message_id} as delivered")) + match (self.task_scoped_server_api.as_ref(), self.task_id) { + (Some(server_api), Some(task_id)) => { + server_api + .mark_message_delivered_for_task(&task_id, message_id) + .await + } + _ => self.ai_client.mark_message_delivered(message_id).await, + } + .with_context(|| format!("Failed to mark agent message {message_id} as delivered")) } pub(crate) async fn mark_messages_delivered_best_effort<'a, I>( diff --git a/app/src/ai/agent_sdk/ambient.rs b/app/src/ai/agent_sdk/ambient.rs index 802e37292..9da79dfac 100644 --- a/app/src/ai/agent_sdk/ambient.rs +++ b/app/src/ai/agent_sdk/ambient.rs @@ -9,7 +9,7 @@ use crate::ai::ambient_agents::spawn::{ }; use crate::ai::ambient_agents::task::HarnessConfig; use crate::ai::ambient_agents::AmbientAgentTaskState; -use crate::ai::ambient_agents::{AgentConfigSnapshot, AmbientAgentTask}; +use crate::ai::ambient_agents::{AgentConfigSnapshot, AmbientAgentTask, AmbientAgentTaskId}; use crate::ai::artifacts::Artifact; use crate::auth::AuthStateProvider; use crate::server::server_api::ai::{ @@ -51,7 +51,7 @@ use crate::ai::agent_sdk::driver::attachments::{ use crate::cloud_object::model::persistence::CloudModel; use crate::server::ids::{ServerId, SyncId}; -use super::common::{EnvironmentChoice, ResolveConfigurationError}; +use super::common::{parse_ambient_task_id, EnvironmentChoice, ResolveConfigurationError}; const MAX_LINE_WIDTH: usize = 90; const STREAM_RETRY_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10]; @@ -643,17 +643,42 @@ impl AmbientAgentRunner { output_format: OutputFormat, ctx: &mut ModelContext, ) -> anyhow::Result<()> { - let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + let provider = ServerApiProvider::as_ref(ctx); + let ai_client = provider.get_ai_client(); + let server_api = provider.get(); + let scoped_task_id = task_id_for_message_send(&args.sender_run_id)?; let future = async move { - let response = ai_client - .send_agent_message(SendAgentMessageRequest { - to: args.to, - subject: args.subject, - body: args.body, - sender_run_id: args.sender_run_id, - }) - .await?; + let request = SendAgentMessageRequest { + to: args.to, + subject: args.subject, + body: args.body, + sender_run_id: args.sender_run_id, + }; + let log_context = SendAgentMessageLogContext::new(&request, scoped_task_id.as_ref()); + log_context.log_start(); + let send_message = async move { + match scoped_task_id { + Some(task_id) => { + server_api + .send_agent_message_for_task(&task_id, request) + .await + } + None => ai_client.send_agent_message(request).await, + } + }; + let response = match send_message.await { + Ok(response) => { + log_context.log_success(&response); + response + } + Err(err) => { + let err = err.context(log_context.error_context()); + log_context.log_error(&err); + eprintln!("{err:#}"); + return Err(err); + } + }; print_send_message_response(&response, output_format)?; Ok(()) }; @@ -661,25 +686,31 @@ impl AmbientAgentRunner { Ok(()) } + fn list_messages( &self, args: MessageListArgs, output_format: OutputFormat, ctx: &mut ModelContext, ) -> anyhow::Result<()> { - let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + let provider = ServerApiProvider::as_ref(ctx); + let ai_client = provider.get_ai_client(); + let server_api = provider.get(); let future = async move { - let messages = ai_client - .list_agent_messages( - &args.run_id, - ListAgentMessagesRequest { - unread_only: args.unread, - since: args.since, - limit: args.limit, - }, - ) - .await?; + let request = ListAgentMessagesRequest { + unread_only: args.unread, + since: args.since, + limit: args.limit, + }; + let messages = match task_id_from_run_id(&args.run_id) { + Some(task_id) => { + server_api + .list_agent_messages_for_task(&task_id, &args.run_id, request) + .await? + } + None => ai_client.list_agent_messages(&args.run_id, request).await?, + }; super::output::print_list(messages, output_format); Ok(()) }; @@ -711,10 +742,20 @@ impl AmbientAgentRunner { output_format: OutputFormat, ctx: &mut ModelContext, ) -> anyhow::Result<()> { - let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + let provider = ServerApiProvider::as_ref(ctx); + let ai_client = provider.get_ai_client(); + let server_api = provider.get(); + let scoped_task_id = task_id_from_oz_run_id_env()?; let future = async move { - let message = ai_client.read_agent_message(&args.message_id).await?; + let message = match scoped_task_id { + Some(task_id) => { + server_api + .read_agent_message_for_task(&task_id, &args.message_id) + .await? + } + None => ai_client.read_agent_message(&args.message_id).await?, + }; print_read_message_response(&message, output_format)?; Ok(()) }; @@ -729,10 +770,20 @@ impl AmbientAgentRunner { output_format: OutputFormat, ctx: &mut ModelContext, ) -> anyhow::Result<()> { - let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + let provider = ServerApiProvider::as_ref(ctx); + let ai_client = provider.get_ai_client(); + let server_api = provider.get(); + let scoped_task_id = task_id_from_oz_run_id_env()?; let future = async move { - ai_client.mark_message_delivered(&args.message_id).await?; + match scoped_task_id { + Some(task_id) => { + server_api + .mark_message_delivered_for_task(&task_id, &args.message_id) + .await? + } + None => ai_client.mark_message_delivered(&args.message_id).await?, + } print_mark_message_delivered_result(&args.message_id, output_format)?; Ok(()) }; @@ -936,6 +987,91 @@ fn write_stream_record(record: &T) -> anyhow::Result<()> { stdout.flush().context("unable to flush stdout")?; Ok(()) } + +fn task_id_from_run_id(run_id: &str) -> Option { + run_id.parse().ok() +} + +fn task_id_from_oz_run_id_env() -> anyhow::Result> { + match std::env::var(warp_cli::OZ_RUN_ID_ENV) { + Ok(run_id) => parse_ambient_task_id(&run_id, "Invalid OZ_RUN_ID").map(Some), + Err(std::env::VarError::NotPresent) => Ok(None), + Err(std::env::VarError::NotUnicode(_)) => Err(anyhow!( + "{} is set but is not valid Unicode", + warp_cli::OZ_RUN_ID_ENV + )), + } +} + +fn task_id_for_message_send(sender_run_id: &str) -> anyhow::Result> { + match task_id_from_run_id(sender_run_id) { + Some(task_id) => Ok(Some(task_id)), + None => task_id_from_oz_run_id_env(), + } +} + +#[derive(Debug, Clone)] +struct SendAgentMessageLogContext { + sender_run_id: String, + task_id: Option, + target_agent_ids: Vec, + subject: String, + body_len: usize, +} + +impl SendAgentMessageLogContext { + fn new(request: &SendAgentMessageRequest, task_id: Option<&AmbientAgentTaskId>) -> Self { + Self { + sender_run_id: request.sender_run_id.clone(), + task_id: task_id.map(|task_id| task_id.to_string()), + target_agent_ids: request.to.clone(), + subject: request.subject.clone(), + body_len: request.body.chars().count(), + } + } + + fn error_context(&self) -> String { + format!( + "Failed to send agent message (sender_run_id={:?}, task_id={:?}, target_agent_ids={:?})", + self.sender_run_id, self.task_id, self.target_agent_ids + ) + } + + fn log_start(&self) { + log::info!( + "Sending ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={}", + self.sender_run_id, + self.task_id, + self.target_agent_ids, + self.subject, + self.body_len + ); + } + + fn log_success(&self, response: &SendAgentMessageResponse) { + log::info!( + "Sent ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} message_ids={:?}", + self.sender_run_id, + self.task_id, + self.target_agent_ids, + self.subject, + self.body_len, + response.message_ids + ); + } + + fn log_error(&self, err: &anyhow::Error) { + log::warn!( + "Failed to send ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} error={err:#}", + self.sender_run_id, + self.task_id, + self.target_agent_ids, + self.subject, + self.body_len + ); + } +} + async fn watch_messages_forever( server_api: Arc, ai_client: Arc, @@ -943,15 +1079,25 @@ async fn watch_messages_forever( ) -> anyhow::Result<()> { let run_id = args.run_id; let watched_run_ids = vec![run_id.clone()]; + let scoped_task_id = task_id_from_run_id(&run_id); let mut last_seen_sequence = args.since_sequence; let mut initial_connect = true; let mut failures = 0usize; loop { - let mut stream = match server_api - .stream_agent_events(&watched_run_ids, last_seen_sequence) - .await - { + let stream_result = match scoped_task_id.as_ref() { + Some(task_id) => { + server_api + .stream_agent_events_for_task(task_id, &watched_run_ids, last_seen_sequence) + .await + } + None => { + server_api + .stream_agent_events(&watched_run_ids, last_seen_sequence) + .await + } + }; + let mut stream = match stream_result { Ok(stream) => { if !initial_connect { eprintln!( @@ -1007,8 +1153,15 @@ async fn watch_messages_forever( last_seen_sequence = event.sequence; continue; }; - - let message = match ai_client.read_agent_message(&message_id).await { + let message_result = match scoped_task_id.as_ref() { + Some(task_id) => { + server_api + .read_agent_message_for_task(task_id, &message_id) + .await + } + None => ai_client.read_agent_message(&message_id).await, + }; + let message = match message_result { Ok(message) => message, Err(err) => { failures += 1; diff --git a/app/src/ai/agent_sdk/ambient_tests.rs b/app/src/ai/agent_sdk/ambient_tests.rs index 4522cad6c..aa298fa89 100644 --- a/app/src/ai/agent_sdk/ambient_tests.rs +++ b/app/src/ai/agent_sdk/ambient_tests.rs @@ -1,6 +1,4 @@ -//! Unit tests for `filter_from_args`. Verifies the clap enums are faithfully translated into -//! `TaskListFilter` without dropping any fields. - +//! Unit tests for ambient agent CLI argument mapping and message helpers. use chrono::{TimeZone, Utc}; use warp_cli::json_filter::JsonOutput; @@ -12,6 +10,9 @@ use warp_cli::task::{ use super::*; use crate::server::server_api::ai::{ArtifactType, ExecutionLocation, RunSortBy, RunSortOrder}; +const TASK_ID: &str = "00000000-0000-0000-0000-000000000001"; +const OTHER_TASK_ID: &str = "00000000-0000-0000-0000-000000000002"; + /// A `ListTasksArgs` whose fields are all at their defaults. fn empty_args() -> ListTasksArgs { ListTasksArgs { @@ -167,3 +168,49 @@ fn every_field_maps_through() { assert_eq!(filter.sort_order, Some(RunSortOrder::Asc)); assert_eq!(filter.cursor.as_deref(), Some("abcd==")); } + +#[test] +fn task_id_from_run_id_accepts_task_uuid() { + let task_id = task_id_from_run_id(TASK_ID).expect("valid task id"); + + assert_eq!(task_id.to_string(), TASK_ID); +} + +#[test] +fn task_id_from_run_id_ignores_non_task_ids() { + assert!(task_id_from_run_id("local-child-run").is_none()); +} + +#[test] +#[serial_test::serial] +fn task_id_for_message_send_prefers_sender_run_id() { + std::env::set_var(warp_cli::OZ_RUN_ID_ENV, OTHER_TASK_ID); + let task_id = task_id_for_message_send(TASK_ID) + .expect("valid task id") + .expect("task id"); + std::env::remove_var(warp_cli::OZ_RUN_ID_ENV); + + assert_eq!(task_id.to_string(), TASK_ID); +} + +#[test] +#[serial_test::serial] +fn task_id_for_message_send_falls_back_to_oz_run_id() { + std::env::set_var(warp_cli::OZ_RUN_ID_ENV, TASK_ID); + let task_id = task_id_for_message_send("local-child-run") + .expect("valid env task id") + .expect("task id"); + std::env::remove_var(warp_cli::OZ_RUN_ID_ENV); + + assert_eq!(task_id.to_string(), TASK_ID); +} + +#[test] +#[serial_test::serial] +fn task_id_from_oz_run_id_env_rejects_invalid_value() { + std::env::set_var(warp_cli::OZ_RUN_ID_ENV, "not-a-task-id"); + let err = task_id_from_oz_run_id_env().expect_err("invalid task id"); + std::env::remove_var(warp_cli::OZ_RUN_ID_ENV); + + assert!(err.to_string().contains("Invalid OZ_RUN_ID")); +} diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index e09aaeea6..9926ac28a 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -22,7 +22,8 @@ use crate::ai::skills::{SkillManager, SkillWatcher}; use crate::ai::{ agent::conversation::AIConversationId, agent_sdk::driver::harness::{ - task_env_vars, HarnessKind, HarnessRunner, ResumePayload, SavePoint, ThirdPartyHarness, + task_env_vars, HarnessCleanupDisposition, HarnessKind, HarnessRunner, ResumePayload, + SavePoint, ThirdPartyHarness, }, }; use crate::terminal::cli_agent_sessions::plugin_manager::{ @@ -271,6 +272,10 @@ pub struct AgentDriver { // The conversation ID to continue (if provided). restored_conversation_id: Option, + /// If set, a third-party-harness conversation to resume. Consumed when + /// preparing the harness runner and cleared afterward. + resume_payload: Option, + /// Cloud providers set up within this driver session. cloud_providers: Vec>, @@ -282,11 +287,6 @@ pub struct AgentDriver { snapshot_upload_timeout: Duration, snapshot_script_timeout: Duration, - /// If set, a third-party-harness conversation to resume. Consumed by `prepare_harness` - /// when building the runner and taken back to `None` after use so subsequent runs start - /// fresh. - resume_payload: Option, - /// Conversation ID this driver is running. Set at construction for /// resumed runs and on `ConversationServerTokenAssigned` for fresh /// runs; consumed by `unregister_streamer_consumer` at end of run. @@ -653,6 +653,7 @@ impl AgentDriver { harness: None, idle_on_complete, restored_conversation_id, + resume_payload, cloud_providers, environment, snapshot_disabled: snapshot_disabled.unwrap_or(false), @@ -660,7 +661,6 @@ impl AgentDriver { .unwrap_or(snapshot::DEFAULT_SNAPSHOT_UPLOAD_TIMEOUT), snapshot_script_timeout: snapshot_script_timeout .unwrap_or(snapshot::DEFAULT_DECLARATIONS_SCRIPT_TIMEOUT), - resume_payload, run_conversation_id, parent_run_id: parent_run_id_for_self, }) @@ -1566,11 +1566,6 @@ impl AgentDriver { .await .map_err(|_| AgentDriverError::InvalidRuntimeState)?; harness.prepare_environment_config(&working_dir, system_prompt.as_deref(), &secrets)?; - - // Pull the resume payload off the driver so the harness runner can rehydrate any - // existing session/conversation state before launching its CLI. The payload variant - // is harness-specific; harnesses match on their own [`ResumePayload`] variant and - // ignore others. let resume = foreground .spawn(|me, _| me.resume_payload.take()) .await @@ -1634,14 +1629,31 @@ impl AgentDriver { // Final save after the command finishes. log::debug!("Triggering final save of harness conversation data"); - report_if_error!(runner + let final_save_succeeded = match runner .save_conversation(SavePoint::Final, foreground) .await - .context("Failed to save harness conversation (final)")); - report_if_error!(runner - .cleanup(foreground) + .context("Failed to save harness conversation (final)") + { + Ok(()) => true, + Err(err) => { + report_error!(err); + false + } + }; + let cleanup_disposition = if final_save_succeeded + && matches!(command_result.as_ref(), Ok(exit_code) if exit_code.was_successful()) + { + HarnessCleanupDisposition::PreserveResumptionStateIfSupported + } else { + HarnessCleanupDisposition::DropResumptionState + }; + if let Err(err) = runner + .cleanup(cleanup_disposition, foreground) .await - .context("Failed to clean up harness runtime state")); + .context("Failed to clean up harness runtime state") + { + report_error!(err); + } let exit_code = command_result?; log::debug!("Agent harness exited with status {exit_code}"); diff --git a/app/src/ai/agent_sdk/driver/error_classification_tests.rs b/app/src/ai/agent_sdk/driver/error_classification_tests.rs index 2f2f85e19..1d0769d48 100644 --- a/app/src/ai/agent_sdk/driver/error_classification_tests.rs +++ b/app/src/ai/agent_sdk/driver/error_classification_tests.rs @@ -99,6 +99,35 @@ fn environment_not_found_is_failed_with_resource_not_found() { ); } +#[test] +fn conversation_harness_mismatch_is_failed_with_env_setup() { + let (state, update) = classify_driver_error(&AgentDriverError::ConversationHarnessMismatch { + conversation_id: "conv-123".into(), + expected: "claude".into(), + got: "oz".into(), + }); + assert_eq!(state, AgentTaskState::Failed); + assert_eq!( + update.error_code, + Some(PlatformErrorCode::EnvironmentSetupFailed) + ); + assert!(update.message.contains("conv-123")); + assert!(update.message.contains("--harness claude")); +} + +#[test] +fn conversation_resume_state_missing_is_failed_with_resource_not_found() { + let (state, update) = + classify_driver_error(&AgentDriverError::ConversationResumeStateMissing { + harness: "claude".into(), + conversation_id: "conv-123".into(), + }); + assert_eq!(state, AgentTaskState::Failed); + assert_eq!(update.error_code, Some(PlatformErrorCode::ResourceNotFound)); + assert!(update.message.contains("conv-123")); + assert!(update.message.contains("claude")); +} + // --- ShareSessionFailed variants --- #[test] diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code.rs b/app/src/ai/agent_sdk/driver/harness/claude_code.rs index b6761eb61..d1e546c1d 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code.rs @@ -29,25 +29,31 @@ use super::claude_transcript::{ }; use super::json_utils::{read_json_file_or_default, write_json_file}; use super::{ - write_temp_file, HarnessRunner, ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness, + cli_agent_session_status, write_temp_file, HarnessCleanupDisposition, HarnessRunner, + ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness, }; mod parent_bridge; +mod wake_driver; #[cfg(test)] use super::super::OZ_MESSAGE_LISTENER_STATE_ROOT_ENV; -use parent_bridge::MessageBridge; #[cfg(test)] use parent_bridge::{ acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir, - parent_bridge_char_count, parent_bridge_hook_output_ack_file, parent_bridge_hook_output_file, - parent_bridge_root, parent_bridge_staged_message_path, parent_bridge_surfaced_message_path, - prepare_parent_bridge_hook_output, render_parent_bridge_message_block, - stage_parent_bridge_message, MessageBridgeHookOutput, MessageBridgeMessageRecord, - MESSAGE_BRIDGE_CONTEXT_PREAMBLE, + parent_bridge_char_count, parent_bridge_event_cursor_file, parent_bridge_hook_output_ack_file, + parent_bridge_hook_output_file, parent_bridge_root, parent_bridge_staged_message_path, + parent_bridge_surfaced_message_path, prepare_parent_bridge_hook_output, + read_parent_bridge_event_cursor, render_parent_bridge_message_block, + stage_parent_bridge_message, write_parent_bridge_event_cursor, MessageBridgeHookOutput, + MessageBridgeMessageRecord, MESSAGE_BRIDGE_CONTEXT_PREAMBLE, }; +use parent_bridge::{MessageBridge, MessageBridgeCleanupDisposition}; +#[cfg(test)] +use shell_words::quote as shell_quote; +#[cfg(test)] +use wake_driver::{ClaudeWakeRemoteContext, CLAUDE_WAKE_PROMPT_FILE_NAME}; pub(crate) struct ClaudeHarness; - #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl ThirdPartyHarness for ClaudeHarness { @@ -229,7 +235,6 @@ impl ClaudeHarnessRunner { // avoiding shell-quoting issues with complex content (e.g. skill instructions). let temp_file = write_temp_file("oz_prompt_", prompt)?; let prompt_path = temp_file.path().display().to_string(); - let (session_id, preexisting_conversation_id, resuming) = match resume { Some(ClaudeResumeInfo { conversation_id, @@ -362,9 +367,33 @@ impl ClaudeHarnessRunner { .await } - fn cleanup_parent_bridge(&self) -> Result<()> { + async fn should_preserve_parent_bridge( + &self, + cleanup_disposition: HarnessCleanupDisposition, + foreground: &ModelSpawner, + ) -> bool { + if !matches!( + cleanup_disposition, + HarnessCleanupDisposition::PreserveResumptionStateIfSupported + ) { + return false; + } + + !matches!( + cli_agent_session_status(&self.terminal_driver, foreground).await, + Some(crate::terminal::cli_agent_sessions::CLIAgentSessionStatus::Blocked { .. }) + | Some(crate::terminal::cli_agent_sessions::CLIAgentSessionStatus::InProgress) + ) + } + + fn cleanup_parent_bridge(&self, preserve_state: bool) -> Result<()> { if let Some(parent_bridge) = self.parent_bridge.as_ref() { - parent_bridge.cleanup()?; + let cleanup_disposition = if preserve_state { + MessageBridgeCleanupDisposition::PreserveState + } else { + MessageBridgeCleanupDisposition::RemoveState + }; + parent_bridge.cleanup(cleanup_disposition)?; } Ok(()) } @@ -414,7 +443,7 @@ impl HarnessRunner for ClaudeHarnessRunner { { Ok(command_handle) => command_handle, Err(err) => { - self.cleanup_parent_bridge() + self.cleanup_parent_bridge(false) .map_err(AgentDriverError::ConfigBuildFailed)?; return Err(err); } @@ -494,9 +523,16 @@ impl HarnessRunner for ClaudeHarnessRunner { Ok(()) } - async fn cleanup(&self, _foreground: &ModelSpawner) -> Result<()> { + async fn cleanup( + &self, + cleanup_disposition: HarnessCleanupDisposition, + foreground: &ModelSpawner, + ) -> Result<()> { self.flush_parent_bridge_acks().await?; - self.cleanup_parent_bridge() + let preserve_state = self + .should_preserve_parent_bridge(cleanup_disposition, foreground) + .await; + self.cleanup_parent_bridge(preserve_state) } } @@ -526,7 +562,6 @@ async fn upload_transcript( .with_context(|| format!("Failed to get transcript upload target for {conversation_id}"))?; upload_to_target(client.http_client(), &target, body).await } - fn prepare_claude_environment_config( working_dir: &Path, secrets: &HashMap, diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs index fe05e8771..cba5002ef 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs @@ -28,12 +28,15 @@ use crate::ai::agent_events::{ AgentEventDriverConfig, MessageHydrator, ServerApiAgentEventSource, }; use crate::ai::agent_sdk::driver::{AgentDriver, OZ_MESSAGE_LISTENER_STATE_ROOT_ENV}; +use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::server::server_api::ai::AIClient; use crate::server::server_api::ai::AgentRunEvent; use crate::server::server_api::ServerApi; const LEGACY_MESSAGE_LISTENER_STATE_ROOT_ENV: &str = "OZ_PARENT_STATE_ROOT"; const PARENT_BRIDGE_DEFAULT_STATE_ROOT: &str = ".claude-code/oz-parent-bridge"; const PARENT_BRIDGE_SURFACED_DIR_NAME: &str = "surfaced"; +const PARENT_BRIDGE_EVENT_CURSOR_FILE_NAME: &str = "event-cursor.json"; const PARENT_BRIDGE_HOOK_OUTPUT_FILE_NAME: &str = "pending-hook-output.json"; const PARENT_BRIDGE_HOOK_OUTPUT_ACK_FILE_NAME: &str = "pending-hook-output.ack"; const PARENT_BRIDGE_MAX_CONTEXT_CHARS_ENV: &str = "OZ_PARENT_MAX_CONTEXT_CHARS"; @@ -48,6 +51,16 @@ pub(super) struct MessageBridge { runtime: Mutex>, state_lock: AsyncMutex<()>, } + +pub(super) enum MessageBridgeCleanupDisposition { + RemoveState, + PreserveState, +} + +#[derive(Debug, PartialEq, Serialize, Deserialize)] +struct MessageBridgeEventCursor { + since_sequence: i64, +} struct MessageBridgeRuntime { task: SpawnedFutureHandle, } @@ -55,6 +68,7 @@ struct MessageBridgeRuntime { struct MessageBridgeEventConsumer { run_id: String, state_dir: PathBuf, + server_api: Arc, } #[cfg_attr(target_family = "wasm", async_trait(?Send))] @@ -91,6 +105,13 @@ impl AgentEventConsumer for MessageBridgeEventConsumer { Ok(AgentEventConsumerControlFlow::Continue) } + + async fn persist_cursor(&mut self, sequence: i64) -> anyhow::Result<()> { + write_parent_bridge_event_cursor(&self.state_dir, sequence)?; + self.server_api + .update_event_sequence_on_server(&self.run_id, sequence) + .await + } } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -130,6 +151,12 @@ struct SelectedMessageBridgeMessages { } impl MessageBridge { + fn hydrator(&self, server_api: Arc) -> MessageHydrator { + match self.run_id.parse::() { + Ok(task_id) => MessageHydrator::for_task(server_api, task_id), + Err(_) => MessageHydrator::new(server_api), + } + } pub(super) fn new(run_id: String, session_id: Uuid) -> Result { Ok(Self { run_id, @@ -177,8 +204,7 @@ impl MessageBridge { if !self.state_dir.exists() { return Ok(()); } - - let hydrator = MessageHydrator::new(server_api); + let hydrator = self.hydrator(server_api); let _guard = self.state_lock.lock().await; acknowledge_parent_bridge_hook_output(&hydrator, &self.state_dir).await?; prepare_parent_bridge_hook_output( @@ -193,16 +219,18 @@ impl MessageBridge { if !self.state_dir.exists() { return Ok(()); } - - let hydrator = MessageHydrator::new(server_api); + let hydrator = self.hydrator(server_api); let _guard = self.state_lock.lock().await; acknowledge_parent_bridge_hook_output(&hydrator, &self.state_dir).await } - pub(super) fn cleanup(&self) -> Result<()> { + pub(super) fn cleanup(&self, disposition: MessageBridgeCleanupDisposition) -> Result<()> { if let Some(runtime) = self.runtime.lock().take() { runtime.task.abort(); } + if matches!(disposition, MessageBridgeCleanupDisposition::PreserveState) { + return Ok(()); + } match fs::remove_dir_all(&self.state_dir) { Ok(()) => {} Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} @@ -249,6 +277,10 @@ pub(super) fn parent_bridge_hook_output_ack_file(state_dir: &Path) -> PathBuf { state_dir.join(PARENT_BRIDGE_HOOK_OUTPUT_ACK_FILE_NAME) } +pub(super) fn parent_bridge_event_cursor_file(state_dir: &Path) -> PathBuf { + state_dir.join(PARENT_BRIDGE_EVENT_CURSOR_FILE_NAME) +} + fn parent_bridge_message_path(dir: &Path, sequence: i64, message_id: &str) -> PathBuf { dir.join(format!("{sequence:020}-{message_id}.json")) } @@ -277,6 +309,28 @@ pub(super) fn ensure_parent_bridge_state_dir(state_dir: &Path) -> Result<()> { Ok(()) } +pub(super) fn read_parent_bridge_event_cursor(state_dir: &Path) -> Result { + let path = parent_bridge_event_cursor_file(state_dir); + if !path.exists() { + return Ok(0); + } + + let cursor = serde_json::from_slice::( + &fs::read(&path).with_context(|| format!("Failed to read {}", path.display()))?, + ) + .with_context(|| format!("Failed to parse {}", path.display()))?; + Ok(cursor.since_sequence) +} + +pub(super) fn write_parent_bridge_event_cursor(state_dir: &Path, sequence: i64) -> Result<()> { + write_parent_bridge_json_atomically( + &parent_bridge_event_cursor_file(state_dir), + &MessageBridgeEventCursor { + since_sequence: sequence, + }, + ) +} + pub(super) fn stage_parent_bridge_message( state_dir: &Path, record: &MessageBridgeMessageRecord, @@ -288,7 +342,7 @@ pub(super) fn stage_parent_bridge_message( Ok(()) } -fn parent_bridge_max_context_chars() -> usize { +pub(super) fn parent_bridge_max_context_chars() -> usize { std::env::var(PARENT_BRIDGE_MAX_CONTEXT_CHARS_ENV) .ok() .and_then(|value| value.trim().parse::().ok()) @@ -579,15 +633,49 @@ async fn run_parent_bridge_forever( state_dir: PathBuf, ) -> Result<()> { ensure_parent_bridge_state_dir(&state_dir)?; + let since_sequence = + read_parent_bridge_resume_cursor(server_api.as_ref(), &run_id, &state_dir).await?; // The shared driver keeps `since_sequence` in memory across its own retry - // loop, which is all this per-session bridge needs because the state dir is - // not reused across sessions. - let config = AgentEventDriverConfig::retry_forever(vec![run_id.clone()], 0); - let source = ServerApiAgentEventSource::new(server_api); - let mut consumer = MessageBridgeEventConsumer { run_id, state_dir }; + // loop and we also persist it inside the session state dir so dormant runs + // can resume without replaying already handled events. + let config = AgentEventDriverConfig::retry_forever(vec![run_id.clone()], since_sequence); + let source = ServerApiAgentEventSource::new(server_api.clone()); + let mut consumer = MessageBridgeEventConsumer { + run_id, + state_dir, + server_api, + }; run_agent_event_driver(source, config, &mut consumer).await } +async fn read_parent_bridge_resume_cursor( + server_api: &ServerApi, + run_id: &str, + state_dir: &Path, +) -> Result { + // The server cursor is the durable cross-client source of truth, but the + // bridge also keeps a local cursor for same-machine recovery. If Warp or + // Claude restarts after the bridge has staged events locally but before the + // server cursor update is visible, the local cursor prevents replaying + // messages already handed to this Claude session. + let local_sequence = read_parent_bridge_event_cursor(state_dir)?; + let Ok(task_id) = run_id.parse() else { + return Ok(local_sequence); + }; + + let server_sequence = match server_api.get_ambient_agent_task(&task_id).await { + Ok(task) => task.last_event_sequence.unwrap_or(0), + Err(err) => { + log::warn!( + "Failed to read server-backed event cursor for Claude message bridge run {run_id}: {err:#}" + ); + 0 + } + }; + + Ok(local_sequence.max(server_sequence)) +} + fn write_parent_bridge_json_atomically(path: &Path, value: &T) -> Result<()> { write_parent_bridge_bytes_atomically(path, &serde_json::to_vec(value)?) } diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs b/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs new file mode 100644 index 000000000..feb1dc3d2 --- /dev/null +++ b/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs @@ -0,0 +1,317 @@ +use std::collections::HashMap; +use std::ffi::{OsStr, OsString}; +use std::path::PathBuf; +use std::sync::Arc; + +use crate::ai::agent::conversation::{AIConversation, ConversationStatus}; +use anyhow::{Context, Result}; +use shell_words::quote as shell_quote; +use uuid::Uuid; +use warp_cli::agent::Harness; +use warp_graphql::ai::AgentTaskState; + +use crate::ai::agent_events::MessageHydrator; +use crate::ai::ambient_agents::{AmbientAgentTaskId, AmbientAgentTaskState}; +use crate::server::server_api::ai::AIClient; +use crate::server::server_api::harness_support::ResolvePromptRequest; +use crate::server::server_api::ServerApi; +use crate::terminal::CLIAgent; + +use super::super::claude_transcript::{ + claude_config_dir, write_envelope, write_session_index_entry, ClaudeTranscriptEnvelope, +}; +use super::super::task_env_vars; +use super::parent_bridge::{ + acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir, parent_bridge_root, +}; +use super::{claude_command, prepare_claude_environment_config, ClaudeHarness}; + +const CLAUDE_WAKE_PROMPT: &str = + "New lead-agent messages are available. Read the latest lead-agent updates and continue the task accordingly."; +pub(super) const CLAUDE_WAKE_PROMPT_FILE_NAME: &str = "wake-turn-prompt.txt"; +const CLAUDE_WAKE_EXTERNALLY_MANAGED_LISTENER_ENV_VARS: &[&str] = &[ + "OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY", + "OZ_PARENT_LISTENER_MANAGED_EXTERNALLY", +]; + +#[derive(Debug)] +pub(super) struct ClaudeWakeRemoteContext { + pub(super) session_id: Uuid, + pub(super) envelope: ClaudeTranscriptEnvelope, + pub(super) wake_prompt: String, +} + +struct ClaudeWakeCandidate { + task_id: AmbientAgentTaskId, + parent_run_id: Option, + working_dir: Option, +} + +impl ClaudeHarness { + pub(crate) async fn wake_dormant_session( + server_api: Arc, + conversation: AIConversation, + parent_conversation: Option, + working_dir: Option, + ) -> Result> { + let Some(candidate) = + Self::local_wake_candidate(&conversation, parent_conversation.as_ref(), working_dir) + else { + return Ok(None); + }; + let ClaudeWakeCandidate { + task_id, + parent_run_id, + working_dir, + } = candidate; + + let task = server_api.get_ambient_agent_task(&task_id).await?; + let harness = task + .agent_config_snapshot + .as_ref() + .and_then(|snapshot| snapshot.harness.as_ref()) + .map(|config| config.harness_type); + log::info!( + "Evaluating dormant Claude wake: task_id={task_id} server_task_state={:?} harness={harness:?}", + task.state + ); + if !is_local_wake_task_state_ready(task.state.clone()) || harness != Some(Harness::Claude) { + log::info!( + "Skipping dormant Claude wake: task_id={task_id} server_task_state={:?} harness={harness:?}", + task.state + ); + return Ok(None); + } + + let remote = Self::fetch_local_wake_remote_context(task_id, server_api.clone()).await?; + let command = Self::prepare_local_wake_command( + server_api.clone(), + task_id, + parent_run_id, + working_dir, + remote, + ) + .await?; + + log::info!("Reopening dormant Claude task before wake command: task_id={task_id}"); + server_api + .update_agent_task(task_id, Some(AgentTaskState::InProgress), None, None, None) + .await + .map_err(|err| { + anyhow::anyhow!( + "Failed to reopen dormant Claude task {task_id} before wake: {err:#}" + ) + })?; + log::info!("Reopened dormant Claude task before wake command: task_id={task_id}"); + + Ok(Some(command)) + } + + fn local_wake_candidate( + conversation: &AIConversation, + parent_conversation: Option<&AIConversation>, + working_dir: Option, + ) -> Option { + let conversation_id = conversation.id(); + if !matches!(conversation.status(), ConversationStatus::Success) { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=not_success status={:?}", + conversation.status() + ); + return None; + } + if !conversation.is_child_agent_conversation() || conversation.is_remote_child() { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=not_local_child is_child_agent_conversation={} is_remote_child={}", + conversation.is_child_agent_conversation(), + conversation.is_remote_child() + ); + return None; + } + let Some(task_id) = conversation.task_id() else { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=missing_task_id" + ); + return None; + }; + let parent_run_id = conversation + .parent_agent_id() + .map(str::to_owned) + .or_else(|| parent_conversation.and_then(AIConversation::run_id)); + + Some(ClaudeWakeCandidate { + task_id, + parent_run_id, + working_dir, + }) + } + + async fn fetch_local_wake_remote_context( + task_id: AmbientAgentTaskId, + server_api: Arc, + ) -> Result { + let resolved = server_api + .resolve_prompt_for_task( + &task_id, + ResolvePromptRequest { + skill: None, + attachments_dir: None, + }, + ) + .await + .with_context(|| format!("Failed to resolve Claude wake prompt for task {task_id}"))?; + let bytes = server_api + .fetch_transcript_for_task(&task_id) + .await + .with_context(|| format!("Failed to fetch Claude transcript for task {task_id}"))?; + let envelope: ClaudeTranscriptEnvelope = + serde_json::from_slice(&bytes).with_context(|| { + format!("Failed to deserialize Claude transcript for wake task {task_id}") + })?; + let wake_prompt = match resolved.resumption_prompt { + Some(resumption_prompt) if !resumption_prompt.is_empty() => { + format!( + "{resumption_prompt} + +{CLAUDE_WAKE_PROMPT}" + ) + } + _ => CLAUDE_WAKE_PROMPT.to_string(), + }; + Ok(ClaudeWakeRemoteContext { + session_id: envelope.uuid, + envelope, + wake_prompt, + }) + } + + pub(super) async fn prepare_local_wake_command( + server_api: Arc, + task_id: AmbientAgentTaskId, + parent_run_id: Option, + working_dir: Option, + mut remote: ClaudeWakeRemoteContext, + ) -> Result { + let working_dir = working_dir.unwrap_or_else(|| remote.envelope.cwd.clone()); + prepare_claude_environment_config(&working_dir, &HashMap::new()) + .context("Failed to prepare Claude environment for wake")?; + + remote.envelope.cwd = working_dir.clone(); + let config_root = claude_config_dir().context("Failed to resolve Claude config dir")?; + write_envelope(&remote.envelope, &config_root) + .context("Failed to rehydrate Claude transcript for wake")?; + if let Err(error) = write_session_index_entry(remote.session_id, &working_dir, &config_root) + { + log::warn!("Failed to update Claude sessions-index.json for wake: {error:#}"); + } + + let state_dir = parent_bridge_root()?.join(remote.session_id.to_string()); + ensure_parent_bridge_state_dir(&state_dir)?; + let hydrator = MessageHydrator::for_task(server_api, task_id); + acknowledge_parent_bridge_hook_output(&hydrator, &state_dir).await?; + let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME); + std::fs::write(&prompt_path, remote.wake_prompt.as_bytes()) + .with_context(|| format!("Failed to write {}", prompt_path.display()))?; + + let command = claude_command( + CLIAgent::Claude.command_prefix(), + &remote.session_id, + &prompt_path.display().to_string(), + None, + true, + ); + let env_vars = local_wake_task_env_vars(Some(&task_id), parent_run_id.as_deref()); + + Ok(prefix_command_with_env_vars(command, env_vars)) + } +} + +fn local_wake_task_env_vars( + task_id: Option<&AmbientAgentTaskId>, + parent_run_id: Option<&str>, +) -> HashMap { + let mut env_vars = task_env_vars(task_id, parent_run_id, Harness::Claude); + // The local wake command is executed directly in the existing child + // terminal, not through `AgentDriver::run_harness`, so Warp does not start + // `MessageBridge` for this resumed Claude process. Leave the listener in + // the Claude plugin's self-managed mode; otherwise the hook waits for + // state files that no managed bridge is producing and the wake message is + // never surfaced to Claude. + for env_name in CLAUDE_WAKE_EXTERNALLY_MANAGED_LISTENER_ENV_VARS { + env_vars.remove(OsStr::new(env_name)); + } + env_vars +} + +fn is_local_wake_task_state_ready(state: AmbientAgentTaskState) -> bool { + match state { + AmbientAgentTaskState::Succeeded => true, + // The local conversation status is already gated on `Success` before + // this function is called. The server task update is fire-and-forget, + // so it can still report `InProgress` for a short window after the + // local Claude process has actually stopped. Treat that stale server + // state as wakeable for local children. + AmbientAgentTaskState::InProgress => true, + AmbientAgentTaskState::Queued + | AmbientAgentTaskState::Pending + | AmbientAgentTaskState::Claimed + | AmbientAgentTaskState::Failed + | AmbientAgentTaskState::Error + | AmbientAgentTaskState::Blocked + | AmbientAgentTaskState::Cancelled + | AmbientAgentTaskState::Unknown => false, + } +} + +fn prefix_command_with_env_vars(command: String, env_vars: HashMap) -> String { + if env_vars.is_empty() { + return command; + } + + let mut env_pairs = env_vars + .into_iter() + .map(|(key, value)| { + ( + key.to_string_lossy().into_owned(), + value.to_string_lossy().into_owned(), + ) + }) + .collect::>(); + env_pairs.sort_unstable_by(|(left, _), (right, _)| left.cmp(right)); + + let assignments = env_pairs + .into_iter() + .map(|(key, value)| format!("{key}={}", shell_quote(&value))) + .collect::>() + .join(" "); + + format!("env {assignments} {command}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn local_wake_task_state_ready_allows_success_and_stale_in_progress() { + assert!(is_local_wake_task_state_ready( + AmbientAgentTaskState::Succeeded + )); + assert!(is_local_wake_task_state_ready( + AmbientAgentTaskState::InProgress + )); + + for state in [ + AmbientAgentTaskState::Queued, + AmbientAgentTaskState::Pending, + AmbientAgentTaskState::Claimed, + AmbientAgentTaskState::Failed, + AmbientAgentTaskState::Error, + AmbientAgentTaskState::Blocked, + AmbientAgentTaskState::Cancelled, + AmbientAgentTaskState::Unknown, + ] { + assert!(!is_local_wake_task_state_ready(state)); + } + } +} diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs b/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs index 71342cf28..9c9a94447 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs @@ -6,10 +6,14 @@ use std::sync::Arc; use tempfile::TempDir; use uuid::Uuid; +use warp_cli::{OZ_HARNESS_ENV, OZ_PARENT_RUN_ID_ENV, OZ_RUN_ID_ENV}; use super::*; use crate::ai::agent_events::MessageHydrator; +use crate::ai::agent_sdk::driver::harness::claude_transcript::encode_cwd; +use crate::ai::agent_sdk::driver::OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV; use crate::server::server_api::ai::{MockAIClient, ReadAgentMessageResponse}; +use crate::server::server_api::ServerApiProvider; fn sample_parent_bridge_message( sequence: i64, @@ -91,6 +95,32 @@ fn claude_command_pipes_prompt_path() { ); } +#[test] +fn write_session_index_entry_creates_expected_entry() { + let tmp = TempDir::new().unwrap(); + let cwd = Path::new("/my/project"); + let session_id = Uuid::new_v4(); + + write_session_index_entry(session_id, cwd, tmp.path()).unwrap(); + + let index_path = tmp.path().join("sessions-index.json"); + let index: Value = serde_json::from_slice(&fs::read(index_path).unwrap()).unwrap(); + let session_key = session_id.to_string(); + let entry = &index[&session_key]; + let encoded = encode_cwd(cwd); + + assert_eq!(entry["sessionId"], Value::String(session_key.clone())); + assert_eq!( + entry["cwd"], + Value::String(cwd.to_string_lossy().into_owned()) + ); + assert_eq!(entry["projectPath"], Value::String(encoded.clone())); + assert_eq!( + entry["transcriptPath"], + Value::String(format!("projects/{encoded}/{session_id}.jsonl")) + ); +} + #[test] #[serial_test::serial] fn parent_bridge_root_prefers_environment_override() { @@ -119,6 +149,77 @@ fn stage_parent_bridge_message_writes_message_record() { assert!(staged_record.sender_run_id.is_empty()); } +#[tokio::test] +async fn parent_bridge_event_cursor_defaults_to_zero_when_missing() { + let tmp = TempDir::new().unwrap(); + let state_dir = tmp.path().join("session-123"); + ensure_parent_bridge_state_dir(&state_dir).unwrap(); + + assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 0); + assert!(!parent_bridge_event_cursor_file(&state_dir).exists()); +} + +#[tokio::test] +async fn parent_bridge_event_cursor_round_trips() { + let tmp = TempDir::new().unwrap(); + let state_dir = tmp.path().join("session-123"); + ensure_parent_bridge_state_dir(&state_dir).unwrap(); + + write_parent_bridge_event_cursor(&state_dir, 42).unwrap(); + + assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 42); + assert!(parent_bridge_event_cursor_file(&state_dir).exists()); +} + +#[test] +#[serial_test::serial] +fn message_bridge_cleanup_preserves_state_for_wakeable_runs() { + let tmp = TempDir::new().unwrap(); + std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, tmp.path()); + + let session_id = Uuid::new_v4(); + let bridge = MessageBridge::new("run-123".to_string(), session_id).unwrap(); + let state_dir = tmp.path().join(session_id.to_string()); + ensure_parent_bridge_state_dir(&state_dir).unwrap(); + let record = sample_staged_parent_bridge_message(42, "msg-123"); + stage_parent_bridge_message(&state_dir, &record).unwrap(); + write_parent_bridge_event_cursor(&state_dir, 42).unwrap(); + + bridge + .cleanup(MessageBridgeCleanupDisposition::PreserveState) + .unwrap(); + std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV); + + assert!(state_dir.exists()); + assert!(parent_bridge_staged_message_path(&state_dir, 42, "msg-123").exists()); + assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 42); +} + +#[test] +#[serial_test::serial] +fn message_bridge_cleanup_removes_state_for_non_wakeable_runs() { + let tmp = TempDir::new().unwrap(); + std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, tmp.path()); + + let session_id = Uuid::new_v4(); + let bridge = MessageBridge::new("run-123".to_string(), session_id).unwrap(); + let state_dir = tmp.path().join(session_id.to_string()); + ensure_parent_bridge_state_dir(&state_dir).unwrap(); + stage_parent_bridge_message( + &state_dir, + &sample_staged_parent_bridge_message(42, "msg-123"), + ) + .unwrap(); + write_parent_bridge_event_cursor(&state_dir, 42).unwrap(); + + bridge + .cleanup(MessageBridgeCleanupDisposition::RemoveState) + .unwrap(); + std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV); + + assert!(!state_dir.exists()); +} + #[tokio::test] async fn prepare_parent_bridge_hook_output_moves_selected_messages_to_surfaced_dir() { let tmp = TempDir::new().unwrap(); @@ -543,6 +644,86 @@ fn resolve_suffix_returns_none_for_short_key() { assert_eq!(resolve_anthropic_api_key_suffix(&secrets), None); } +#[test] +#[serial_test::serial] +fn prepare_local_wake_command_rehydrates_transcript_with_self_managed_listener() { + let home_dir = TempDir::new().unwrap(); + let claude_config_dir = TempDir::new().unwrap(); + let bridge_state_root = TempDir::new().unwrap(); + let working_dir = home_dir.path().join("workspace/project"); + fs::create_dir_all(&working_dir).unwrap(); + + std::env::set_var("HOME", home_dir.path()); + std::env::set_var("CLAUDE_CONFIG_DIR", claude_config_dir.path()); + std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, bridge_state_root.path()); + + let session_id = Uuid::new_v4(); + let remote = ClaudeWakeRemoteContext { + session_id, + envelope: ClaudeTranscriptEnvelope { + cwd: Path::new("/stale/cwd").to_path_buf(), + uuid: session_id, + claude_version: None, + entries: vec![serde_json::json!({"type": "assistant", "text": "done"})], + subagents: HashMap::new(), + todos: HashMap::new(), + }, + wake_prompt: "resume prompt\n\nwake prompt".to_string(), + }; + let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440010".parse().unwrap(); + let parent_run_id = "parent-run-456".to_string(); + + let command = futures::executor::block_on(ClaudeHarness::prepare_local_wake_command( + ServerApiProvider::new_for_test().get(), + task_id, + Some(parent_run_id.clone()), + Some(working_dir.clone()), + remote, + )) + .unwrap(); + + let state_dir = bridge_state_root.path().join(session_id.to_string()); + let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME); + + assert!(command.contains("--resume")); + assert!(command.starts_with("env ")); + assert!(command.contains(&session_id.to_string())); + assert!(command.contains(CLAUDE_WAKE_PROMPT_FILE_NAME)); + assert!(command.contains(&format!( + "{OZ_RUN_ID_ENV}={}", + shell_quote(&task_id.to_string()) + ))); + assert!(command.contains(&format!( + "{OZ_PARENT_RUN_ID_ENV}={}", + shell_quote(&parent_run_id) + ))); + assert!(command.contains(&format!("{OZ_HARNESS_ENV}={}", shell_quote("claude")))); + assert!(!command.contains(OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV)); + assert!(!command.contains("OZ_PARENT_LISTENER_MANAGED_EXTERNALLY")); + assert_eq!( + fs::read_to_string(&prompt_path).unwrap(), + "resume prompt\n\nwake prompt" + ); + assert!(!parent_bridge_hook_output_file(&state_dir).exists()); + + let restored_envelope = + read_envelope(session_id, &working_dir, claude_config_dir.path()).unwrap(); + assert_eq!(restored_envelope.cwd, working_dir); + assert_eq!( + restored_envelope.entries, + vec![serde_json::json!({"type": "assistant", "text": "done"})] + ); + assert!(home_dir.path().join(".claude.json").exists()); + assert!(claude_config_dir + .path() + .join(CLAUDE_SETTINGS_FILE_NAME) + .exists()); + + std::env::remove_var("HOME"); + std::env::remove_var("CLAUDE_CONFIG_DIR"); + std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV); +} + #[test] fn resolve_suffix_returns_none_for_short_anthropic_api_key() { let secrets = HashMap::from([( diff --git a/app/src/ai/agent_sdk/driver/harness/codex.rs b/app/src/ai/agent_sdk/driver/harness/codex.rs index 37ccc1da6..6f7bfb2d8 100644 --- a/app/src/ai/agent_sdk/driver/harness/codex.rs +++ b/app/src/ai/agent_sdk/driver/harness/codex.rs @@ -311,6 +311,7 @@ fn write_codex_auth_json(path: &Path, auth: &CodexAuthDotJson) -> Result<()> { { use std::io::Write as _; use std::os::unix::fs::OpenOptionsExt; + use std::os::unix::fs::PermissionsExt; let mut file = fs::OpenOptions::new() .write(true) .create(true) @@ -318,6 +319,8 @@ fn write_codex_auth_json(path: &Path, auth: &CodexAuthDotJson) -> Result<()> { .mode(0o600) .open(path) .with_context(|| format!("Failed to open {} for writing", path.display()))?; + file.set_permissions(fs::Permissions::from_mode(0o600)) + .with_context(|| format!("Failed to set permissions on {}", path.display()))?; file.write_all(&bytes) .with_context(|| format!("Failed to write {}", path.display()))?; } diff --git a/app/src/ai/agent_sdk/driver/harness/gemini.rs b/app/src/ai/agent_sdk/driver/harness/gemini.rs index c3ecda57e..02aa12348 100644 --- a/app/src/ai/agent_sdk/driver/harness/gemini.rs +++ b/app/src/ai/agent_sdk/driver/harness/gemini.rs @@ -22,7 +22,10 @@ use crate::terminal::CLIAgent; use super::super::terminal::{CommandHandle, TerminalDriver}; use super::super::{AgentDriver, AgentDriverError}; use super::json_utils::{read_json_file_or_default, write_json_file}; -use super::{write_temp_file, HarnessRunner, ResumePayload, SavePoint, ThirdPartyHarness}; +use super::{ + write_temp_file, HarnessCleanupDisposition, HarnessRunner, ResumePayload, SavePoint, + ThirdPartyHarness, +}; pub(crate) struct GeminiHarness; @@ -215,6 +218,14 @@ impl HarnessRunner for GeminiHarnessRunner { ) .await } + + async fn cleanup( + &self, + _cleanup_disposition: HarnessCleanupDisposition, + _foreground: &ModelSpawner, + ) -> Result<()> { + Ok(()) + } } fn prepare_gemini_environment_config( diff --git a/app/src/ai/agent_sdk/driver/harness/mod.rs b/app/src/ai/agent_sdk/driver/harness/mod.rs index cbb083520..fdbcd2e3a 100644 --- a/app/src/ai/agent_sdk/driver/harness/mod.rs +++ b/app/src/ai/agent_sdk/driver/harness/mod.rs @@ -39,7 +39,6 @@ pub(crate) mod claude_transcript; mod codex; mod gemini; mod json_utils; - pub(crate) use claude_code::ClaudeHarness; use claude_transcript::ClaudeResumeInfo; use codex::CodexHarness; @@ -167,9 +166,9 @@ pub(crate) fn harness_kind(harness: Harness) -> Result Ok(HarnessKind::Oz), Harness::Claude => Ok(HarnessKind::ThirdParty(Box::new(ClaudeHarness))), + Harness::Codex => Ok(HarnessKind::ThirdParty(Box::new(CodexHarness))), Harness::OpenCode => Ok(HarnessKind::Unsupported(Harness::OpenCode)), Harness::Gemini => Ok(HarnessKind::ThirdParty(Box::new(GeminiHarness))), - Harness::Codex => Ok(HarnessKind::ThirdParty(Box::new(CodexHarness))), Harness::Unknown => Err(AgentDriverError::InvalidRuntimeState), } } @@ -325,6 +324,18 @@ pub(crate) enum SavePoint { PostTurn, } +/// Controls how much harness-owned state should survive cleanup after the CLI +/// exits. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum HarnessCleanupDisposition { + /// Tear down all harness-owned resume and wake state. + DropResumptionState, + /// The harness exited cleanly and its final save completed, so wake/resume + /// state may be preserved if the harness-specific runtime also considers + /// the run complete. + PreserveResumptionStateIfSupported, +} + /// Stateful per-run representation of an external harness produced /// by [`ThirdPartyHarness::build_runner`]. /// @@ -362,7 +373,11 @@ pub(crate) trait HarnessRunner: Send + Sync { } /// Clean up any harness-owned background state after the harness exits. - async fn cleanup(&self, _foreground: &ModelSpawner) -> Result<()> { + async fn cleanup( + &self, + _cleanup_disposition: HarnessCleanupDisposition, + _foreground: &ModelSpawner, + ) -> Result<()> { Ok(()) } } @@ -373,20 +388,29 @@ pub(crate) async fn has_running_cli_agent( terminal_driver: &ModelHandle, foreground: &ModelSpawner, ) -> bool { + matches!( + cli_agent_session_status(terminal_driver, foreground).await, + Some(CLIAgentSessionStatus::InProgress) + ) +} + +/// Returns the tracked CLI agent session status for the terminal, if any. +pub(crate) async fn cli_agent_session_status( + terminal_driver: &ModelHandle, + foreground: &ModelSpawner, +) -> Option { let driver = terminal_driver.clone(); - let Ok(running) = foreground + foreground .spawn(move |_, ctx| { let terminal_view_id = driver.as_ref(ctx).terminal_view().id(); CLIAgentSessionsModel::handle(ctx) .as_ref(ctx) .session(terminal_view_id) - .is_some_and(|s| s.status == CLIAgentSessionStatus::InProgress) + .map(|session| session.status.clone()) }) .await - else { - return false; - }; - running + .ok() + .flatten() } /// Create a [`NamedTempFile`] with the given prefix and write `content` into it. diff --git a/app/src/ai/agent_sdk/driver/snapshot_tests.rs b/app/src/ai/agent_sdk/driver/snapshot_tests.rs index 322fb6b44..7c4b42542 100644 --- a/app/src/ai/agent_sdk/driver/snapshot_tests.rs +++ b/app/src/ai/agent_sdk/driver/snapshot_tests.rs @@ -86,6 +86,10 @@ impl HarnessSupportClient for TestClient { unimplemented!("not used by upload_snapshot_from_declarations_file") } + async fn fetch_transcript(&self) -> Result { + unimplemented!("not used by upload_snapshot_from_declarations_file") + } + async fn get_block_snapshot_upload_target( &self, _conversation_id: &AIConversationId, @@ -138,10 +142,6 @@ impl HarnessSupportClient for TestClient { Ok(targets) } - async fn fetch_transcript(&self) -> Result { - unimplemented!("not used by upload_snapshot_from_declarations_file") - } - fn http_client(&self) -> &http_client::Client { &self.http } diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index 54e693cbe..397cd304c 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -466,6 +466,24 @@ fn build_server_side_task( Ok((config, task)) } +fn reconcile_task_harness( + task_id: &str, + selected_harness: &mut Harness, + task_harness: Harness, +) -> Result { + if *selected_harness == Harness::Oz { + *selected_harness = task_harness; + } else if task_harness != *selected_harness { + return Err(AgentDriverError::TaskHarnessMismatch { + task_id: task_id.to_string(), + expected: task_harness.to_string(), + got: selected_harness.to_string(), + }); + } + + harness_kind(*selected_harness) +} + /// Resolve a `Prompt` to a plain string. fn resolve_prompt(prompt: &Prompt, ctx: &AppContext) -> Result { match prompt { @@ -570,21 +588,23 @@ impl AgentDriverRunner { // Pull relevant variables out of args before moving it into the closure. let share_requests = args.share.share.clone(); let bedrock_inference_role = args.bedrock_inference_role.clone(); + let has_task_id = args.task_id.is_some(); let args_harness = args.harness; - // `--conversation` path (user-invoked local resume): validate before any task side // effects so mismatches fail fast. The `--task-id` path derives its conversation id // from the server-side task metadata inside `build_driver_options_and_task`. Both // can currently be passed together (the worker server-side appends `--conversation` // alongside `--task-id` for Slack/Linear followups); when both are set, the explicit // `--conversation` value wins via the merge below. - if let Some(conversation_id) = args.conversation.as_deref() { - common::fetch_and_validate_conversation_harness( - server_api.clone(), - conversation_id, - args_harness, - ) - .await?; + if !has_task_id { + if let Some(conversation_id) = args.conversation.as_deref() { + common::fetch_and_validate_conversation_harness( + server_api.clone(), + conversation_id, + args_harness, + ) + .await?; + } } let resume_conversation_id = args.conversation.clone(); @@ -847,7 +867,6 @@ impl AgentDriverRunner { .await?; None }; - // Resolve environment and cloud providers. Self::resolve_environment(foreground, environment_id, &mut driver_options).await?; @@ -1055,13 +1074,11 @@ impl AgentDriverRunner { // task is linked to an existing conversation, since task harness and conversation harness // always match (the task spawned the conversation). if let Some(task_harness) = task_harness { - if task_harness != driver_options.selected_harness { - return Err(AgentDriverError::TaskHarnessMismatch { - task_id: task_id_str, - expected: task_harness.to_string(), - got: driver_options.selected_harness.to_string(), - }); - } + task.harness = reconcile_task_harness( + &task_id_str, + &mut driver_options.selected_harness, + task_harness, + )?; } // Set the task ID on the ServerApi so it's sent with all subsequent requests. @@ -1133,7 +1150,6 @@ impl AgentDriverRunner { "Failed to convert conversation data to AIConversation".into(), ) })?; - Ok(Some(driver::ResumeOptions::Oz(Box::new( ConversationRestorationInNewPaneType::Historical { conversation, diff --git a/app/src/ai/agent_sdk/mod_tests.rs b/app/src/ai/agent_sdk/mod_tests.rs index 4091d9343..4870d9968 100644 --- a/app/src/ai/agent_sdk/mod_tests.rs +++ b/app/src/ai/agent_sdk/mod_tests.rs @@ -1,12 +1,14 @@ +use super::{command_requires_auth, command_to_telemetry_event, reconcile_task_harness}; use serde_json::json; use warp_cli::{ + agent::Harness, artifact::{ArtifactCommand, DownloadArtifactArgs, GetArtifactArgs, UploadArtifactArgs}, task::{MessageCommand, MessageSendArgs, MessageWatchArgs, TaskCommand}, CliCommand, }; use warp_core::telemetry::TelemetryEvent; -use super::{command_requires_auth, command_to_telemetry_event}; +const TASK_ID: &str = "00000000-0000-0000-0000-000000000001"; #[test] fn logout_does_not_require_auth() { @@ -128,6 +130,38 @@ fn run_message_send_telemetry_defaults_to_unknown_harness() { assert_eq!(event.payload(), Some(json!({ "harness": "unknown" }))); } +#[test] +fn reconcile_task_harness_adopts_task_harness_when_cli_uses_default() { + let mut selected_harness = Harness::Oz; + let harness = reconcile_task_harness(TASK_ID, &mut selected_harness, Harness::Claude) + .expect("default harness should adopt task harness"); + + assert_eq!(selected_harness, Harness::Claude); + assert_eq!(harness.harness(), Harness::Claude); +} + +#[test] +fn reconcile_task_harness_allows_matching_explicit_harness() { + let mut selected_harness = Harness::Claude; + let harness = reconcile_task_harness(TASK_ID, &mut selected_harness, Harness::Claude) + .expect("matching harness should succeed"); + + assert_eq!(selected_harness, Harness::Claude); + assert_eq!(harness.harness(), Harness::Claude); +} + +#[test] +fn reconcile_task_harness_rejects_explicit_mismatch() { + let mut selected_harness = Harness::Gemini; + let err = reconcile_task_harness(TASK_ID, &mut selected_harness, Harness::Claude) + .expect_err("mismatched harness should fail"); + + assert_eq!(selected_harness, Harness::Gemini); + assert!(err.to_string().contains("Task")); + assert!(err.to_string().contains("--harness gemini")); + assert!(err.to_string().contains("claude")); +} + #[test] #[serial_test::serial] fn run_message_watch_telemetry_defaults_to_unknown_harness() { diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 5681dd3ce..0ef183d53 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -418,6 +418,9 @@ impl BlocklistAIActionExecutor { id: Option, ctx: &mut ModelContext, ) { + self.send_message_executor.update(ctx, |executor, _| { + executor.set_ambient_agent_task_id(id); + }); self.request_computer_use_executor .update(ctx, |executor, _| { executor.set_ambient_agent_task_id(id); diff --git a/app/src/ai/blocklist/action_model/execute/send_message.rs b/app/src/ai/blocklist/action_model/execute/send_message.rs index 9818371ce..6e42c09a0 100644 --- a/app/src/ai/blocklist/action_model/execute/send_message.rs +++ b/app/src/ai/blocklist/action_model/execute/send_message.rs @@ -1,9 +1,19 @@ +#[cfg(not(target_family = "wasm"))] +use anyhow::anyhow; +#[cfg(not(target_family = "wasm"))] +use futures::future::Either; use futures::{future::BoxFuture, FutureExt}; -use warpui::{Entity, ModelContext, SingletonEntity}; +#[cfg(not(target_family = "wasm"))] +use std::time::Duration; +#[cfg(not(target_family = "wasm"))] +use warpui::r#async::Timer; +use warpui::{AppContext, Entity, ModelContext, SingletonEntity}; use crate::ai::agent::{ - AIAgentAction, AIAgentActionResultType, AIAgentActionType, SendMessageToAgentResult, + conversation::AIConversationId, AIAgentAction, AIAgentActionResultType, AIAgentActionType, + SendMessageToAgentResult, }; +use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::history_model::BlocklistAIHistoryModel; use crate::ai::blocklist::orchestration_events::{OrchestrationEventService, SendMessageResult}; use crate::ai::blocklist::telemetry::{ @@ -11,18 +21,112 @@ use crate::ai::blocklist::telemetry::{ TeamAgentCommunicationFailureReason, TeamAgentCommunicationKind, TeamAgentCommunicationTransport, TeamAgentOrchestrationVersion, }; -use crate::server::server_api::ai::SendAgentMessageRequest; +use crate::server::server_api::ai::{SendAgentMessageRequest, SendAgentMessageResponse}; use crate::server::server_api::ServerApiProvider; use warp_core::features::FeatureFlag; use warp_core::send_telemetry_from_ctx; use super::{ActionExecution, AnyActionExecution, ExecuteActionInput, PreprocessActionInput}; -pub struct SendMessageToAgentExecutor; +#[cfg(not(target_family = "wasm"))] +const SEND_AGENT_MESSAGE_TIMEOUT: Duration = Duration::from_secs(15); + +pub struct SendMessageToAgentExecutor { + ambient_agent_task_id: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SendMessageTaskResolution { + ConversationTask, + AmbientTaskFallback, + NoTaskContext, +} + +fn sender_run_id_and_task_id_for_send( + conversation_id: AIConversationId, + ambient_agent_task_id: Option, + ctx: &AppContext, +) -> ( + String, + Option, + SendMessageTaskResolution, +) { + let conversation = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id); + let conversation_task_id = conversation.and_then(|conversation| conversation.task_id()); + let (task_id, task_resolution) = match (conversation_task_id, ambient_agent_task_id) { + (Some(task_id), _) => (Some(task_id), SendMessageTaskResolution::ConversationTask), + (None, Some(task_id)) => ( + Some(task_id), + SendMessageTaskResolution::AmbientTaskFallback, + ), + (None, None) => (None, SendMessageTaskResolution::NoTaskContext), + }; + let sender_run_id = conversation + .and_then(|conversation| conversation.run_id()) + .or_else(|| task_id.map(|task_id| task_id.to_string())) + .unwrap_or_default(); + (sender_run_id, task_id, task_resolution) +} + +#[cfg(not(target_family = "wasm"))] +async fn send_agent_message_with_timeout( + server_api: std::sync::Arc, + ai_client: std::sync::Arc, + task_id: Option, + request: SendAgentMessageRequest, +) -> anyhow::Result { + let task_id_for_timeout = task_id.map(|task_id| task_id.to_string()); + let send_message = async move { + match task_id { + Some(task_id) => { + server_api + .send_agent_message_for_task(&task_id, request) + .await + } + None => ai_client.send_agent_message(request).await, + } + }; + let timeout = Timer::after(SEND_AGENT_MESSAGE_TIMEOUT); + futures::pin_mut!(send_message); + futures::pin_mut!(timeout); + + match futures::future::select(send_message, timeout).await { + Either::Left((result, _)) => result, + Either::Right(_) => Err(anyhow!( + "Timed out sending orchestration message{}", + task_id_for_timeout + .map(|task_id| format!(" for task {task_id}")) + .unwrap_or_default() + )), + } +} + +#[cfg(target_family = "wasm")] +async fn send_agent_message_with_timeout( + server_api: std::sync::Arc, + ai_client: std::sync::Arc, + task_id: Option, + request: SendAgentMessageRequest, +) -> anyhow::Result { + match task_id { + Some(task_id) => { + server_api + .send_agent_message_for_task(&task_id, request) + .await + } + None => ai_client.send_agent_message(request).await, + } +} impl SendMessageToAgentExecutor { pub fn new() -> Self { - Self + Self { + ambient_agent_task_id: None, + } + } + + pub fn set_ambient_agent_task_id(&mut self, id: Option) { + self.ambient_agent_task_id = id; } pub(super) fn should_autoexecute( @@ -57,15 +161,21 @@ impl SendMessageToAgentExecutor { let message_body = message.clone(); if FeatureFlag::OrchestrationV2.is_enabled() { - let sender_run_id = BlocklistAIHistoryModel::as_ref(ctx) - .conversation(&conversation_id) - .and_then(|c| c.run_id()) - .map(|s| s.to_string()) - .unwrap_or_default(); + let (sender_run_id, task_id, task_resolution) = sender_run_id_and_task_id_for_send( + conversation_id, + self.ambient_agent_task_id, + ctx, + ); let log_addresses = addresses.clone(); let log_subject = subject.clone(); let log_sender_run_id = sender_run_id.clone(); + let log_task_id = task_id.map(|task_id| task_id.to_string()); + let log_body_len = message_body.chars().count(); + let server_api = ServerApiProvider::as_ref(ctx).get(); let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + log::info!( + "Sending orchestration message: conversation_id={conversation_id:?} resolution={task_resolution:?} sender_run_id={log_sender_run_id:?} task_id={log_task_id:?} target_agent_ids={log_addresses:?} subject={log_subject:?} body_len={log_body_len}" + ); let request = SendAgentMessageRequest { to: addresses, subject, @@ -73,11 +183,16 @@ impl SendMessageToAgentExecutor { sender_run_id, }; return ActionExecution::new_async( - async move { ai_client.send_agent_message(request).await }, + async move { + send_agent_message_with_timeout(server_api, ai_client, task_id, request).await + }, move |result, ctx| match result { Ok(response) => { let message_id = response.message_ids.into_iter().next().unwrap_or_default(); + log::info!( + "Sent orchestration message: conversation_id={conversation_id:?} resolution={task_resolution:?} sender_run_id={log_sender_run_id:?} task_id={log_task_id:?} target_agent_ids={log_addresses:?} subject={log_subject:?} body_len={log_body_len} message_id={message_id:?}" + ); AIAgentActionResultType::SendMessageToAgent( SendMessageToAgentResult::Success { message_id }, ) @@ -103,7 +218,7 @@ impl SendMessageToAgentExecutor { ctx ); log::warn!( - "Failed to send child-agent message via server API: conversation_id={conversation_id:?} sender_run_id={log_sender_run_id:?} target_agent_ids={log_addresses:?} subject={log_subject:?} error={err:#}" + "Failed to send child-agent message via server API: conversation_id={conversation_id:?} resolution={task_resolution:?} sender_run_id={log_sender_run_id:?} task_id={log_task_id:?} target_agent_ids={log_addresses:?} subject={log_subject:?} body_len={log_body_len} error={err:#}" ); AIAgentActionResultType::SendMessageToAgent( SendMessageToAgentResult::Error(error_message), @@ -145,3 +260,7 @@ impl Default for SendMessageToAgentExecutor { impl Entity for SendMessageToAgentExecutor { type Event = (); } + +#[cfg(test)] +#[path = "send_message_tests.rs"] +mod tests; diff --git a/app/src/ai/blocklist/action_model/execute/send_message_tests.rs b/app/src/ai/blocklist/action_model/execute/send_message_tests.rs new file mode 100644 index 000000000..a0e2ac50c --- /dev/null +++ b/app/src/ai/blocklist/action_model/execute/send_message_tests.rs @@ -0,0 +1,60 @@ +use super::*; +use crate::ai::blocklist::BlocklistAIHistoryModel; +use warpui::{App, EntityId}; + +#[test] +fn sender_run_id_and_task_id_for_send_falls_back_to_ambient_task_id() { + App::test((), |mut app| async move { + let terminal_view_id = EntityId::new(); + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let conversation_id = history_model.update(&mut app, |history_model, ctx| { + history_model.start_new_conversation(terminal_view_id, false, false, ctx) + }); + let ambient_task_id = "11111111-1111-1111-1111-111111111111" + .parse() + .expect("valid ambient task id"); + + let (sender_run_id, task_id, task_resolution) = app.read(|ctx| { + sender_run_id_and_task_id_for_send(conversation_id, Some(ambient_task_id), ctx) + }); + + assert_eq!(sender_run_id, ambient_task_id.to_string()); + assert_eq!(task_id, Some(ambient_task_id)); + assert_eq!( + task_resolution, + SendMessageTaskResolution::AmbientTaskFallback + ); + }); +} + +#[test] +fn sender_run_id_and_task_id_for_send_prefers_conversation_task_id() { + App::test((), |mut app| async move { + let terminal_view_id = EntityId::new(); + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let conversation_id = history_model.update(&mut app, |history_model, ctx| { + history_model.start_new_conversation(terminal_view_id, false, false, ctx) + }); + let conversation_task_id = "22222222-2222-2222-2222-222222222222" + .parse() + .expect("valid conversation task id"); + let ambient_task_id = "33333333-3333-3333-3333-333333333333" + .parse() + .expect("valid ambient task id"); + + history_model.update(&mut app, |history_model, _| { + history_model + .conversation_mut(&conversation_id) + .expect("conversation exists") + .set_task_id(conversation_task_id); + }); + + let (sender_run_id, task_id, task_resolution) = app.read(|ctx| { + sender_run_id_and_task_id_for_send(conversation_id, Some(ambient_task_id), ctx) + }); + + assert_eq!(sender_run_id, conversation_task_id.to_string()); + assert_eq!(task_id, Some(conversation_task_id)); + assert_eq!(task_resolution, SendMessageTaskResolution::ConversationTask); + }); +} diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 09fdbaef7..4147e98b4 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -30,6 +30,8 @@ use crate::ai::agent::{ PassiveSuggestionTriggerType, RunningCommand, }; use crate::ai::agent::{DocumentContentAttachmentSource, FileContext}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::ClaudeHarness; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, @@ -51,7 +53,7 @@ use crate::global_resource_handles::GlobalResourceHandlesProvider; use crate::network::NetworkStatus; use crate::notebooks::editor::model::FileLinkResolutionContext; use crate::persistence::ModelEvent; -use crate::server::server_api::AIApiError; +use crate::server::server_api::{AIApiError, ServerApiProvider}; use crate::terminal::model::block::{ formatted_terminal_contents_for_input, BlockId, CURSOR_MARKER, }; @@ -71,12 +73,17 @@ use parking_lot::FairMutex; use pending_response_streams::PendingResponseStreams; use session_sharing_protocol::common::ParticipantId; use std::collections::{HashMap, HashSet}; +#[cfg(not(target_family = "wasm"))] +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use warp_core::assertions::safe_assert; use warp_multi_agent_api::{message, Task, ToolType}; use warpui::r#async::{SpawnedFutureHandle, Timer}; +use super::orchestration_event_streamer::{ + OrchestrationEventStreamer, OrchestrationEventStreamerEvent, +}; use super::orchestration_events::{OrchestrationEventService, OrchestrationEventServiceEvent}; use warpui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; @@ -162,6 +169,10 @@ pub enum BlocklistAIControllerEvent { filename: Option, }, + ExecuteLocalHarnessCommand { + command: String, + }, + FreeTierLimitCheckTriggered, } @@ -312,6 +323,9 @@ pub struct BlocklistAIController { /// Pending auto-resume tasks that are waiting for network connectivity. /// These should be cancelled when a new request is sent for the same conversation. pending_auto_resume_handles: HashMap, + /// Pending dormant Claude wake preparations for success-idle child conversations. + #[cfg_attr(target_family = "wasm", allow(dead_code))] + pending_local_claude_wakes: HashMap, /// Passive conversations explicitly requested to follow up after actions complete. pending_passive_follow_ups: HashSet, /// Passive suggestion results that should be included with the next request @@ -352,6 +366,21 @@ enum FollowUpTrigger { UserRequested, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LocalClaudeWakeTrigger { + PendingEvents, + WakeOnlyStream, +} + +impl LocalClaudeWakeTrigger { + fn requires_pending_events(self) -> bool { + match self { + Self::PendingEvents => true, + Self::WakeOnlyStream => false, + } + } +} + struct InputQuery { which_task: WhichTask, input_query: InputQueryType, @@ -537,6 +566,14 @@ impl BlocklistAIController { me.handle_pending_events_ready(*conversation_id, ctx); }); } + if FeatureFlag::OrchestrationV2.is_enabled() { + let streamer = OrchestrationEventStreamer::handle(ctx); + ctx.subscribe_to_model(&streamer, move |me, event, ctx| { + let OrchestrationEventStreamerEvent::DormantClaudeWakeReady { conversation_id } = + event; + me.handle_dormant_claude_wake_ready(*conversation_id, ctx); + }); + } Self { input_model, context_model, @@ -550,6 +587,7 @@ impl BlocklistAIController { ambient_agent_task_id: None, attachments_download_dir: None, pending_auto_resume_handles: HashMap::new(), + pending_local_claude_wakes: HashMap::new(), pending_passive_follow_ups: HashSet::new(), pending_passive_suggestion_results: HashMap::new(), } @@ -598,8 +636,15 @@ impl BlocklistAIController { .remove(&conversation_id) .unwrap_or_default(); - let cancellation_reason = - self.cancel_active_conversation_for_follow_up(conversation_id, ctx); + let ai_history_model = BlocklistAIHistoryModel::as_ref(ctx); + let active_conversation_id = ai_history_model.active_conversation_id(self.terminal_view_id); + let cancellation_reason = CancellationReason::FollowUpSubmitted { + is_for_same_conversation: active_conversation_id + .is_some_and(|id| id == conversation_id), + }; + if let Some(active_conversation_id) = active_conversation_id { + self.cancel_conversation_progress(active_conversation_id, cancellation_reason, ctx); + } if let Some(slash_command_request) = SlashCommandRequest::from_query(query.as_str()) { slash_command_request.send_request(self, is_queued_prompt, ctx); @@ -1193,40 +1238,9 @@ impl BlocklistAIController { slash_command: SlashCommandRequest, ctx: &mut ModelContext, ) { - // Slash commands are a fresh user turn; mirror `send_query`'s - // cancel-and-resend so we don't trip `send_request_input`'s in-flight - // invariant. - if let Some(conversation_id) = slash_command.conversation_id(self, ctx) { - self.cancel_active_conversation_for_follow_up(conversation_id, ctx); - } slash_command.send_request(self, /*is_queued_prompt*/ false, ctx); } - /// Cancel any in-flight progress on the active conversation in preparation - /// for sending a follow-up turn that will land on `target_conversation_id`. - /// Without this pre-cancel, [`Self::send_request_input`] would trip its - /// in-flight invariant when the new turn re-uses an existing conversation. - /// - /// Returns the [`CancellationReason::FollowUpSubmitted`] reason used so - /// callers can reuse it for downstream side effects (e.g. cancelling - /// pending actions on the target conversation). - fn cancel_active_conversation_for_follow_up( - &mut self, - target_conversation_id: AIConversationId, - ctx: &mut ModelContext, - ) -> CancellationReason { - let active_conversation_id = - BlocklistAIHistoryModel::as_ref(ctx).active_conversation_id(self.terminal_view_id); - let reason = CancellationReason::FollowUpSubmitted { - is_for_same_conversation: active_conversation_id - .is_some_and(|id| id == target_conversation_id), - }; - if let Some(active_conversation_id) = active_conversation_id { - self.cancel_conversation_progress(active_conversation_id, reason, ctx); - } - reason - } - /// Same as [`Self::send_slash_command_request`] but marks the emitted `SentRequest` /// event as a queued prompt submission so UI subscribers (e.g. the input editor) /// don't clear the input buffer on the auto-send. @@ -1484,32 +1498,197 @@ impl BlocklistAIController { self.pending_passive_follow_ups.remove(&conversation_id); } - /// Handles the EventsReady signal. Checks readiness, drains - /// pending events from the service, and injects them into the conversation. - fn handle_pending_events_ready( - &mut self, + fn conversation_ready_for_pending_events( + &self, conversation_id: AIConversationId, - ctx: &mut ModelContext, - ) { + ctx: &ModelContext, + ) -> bool { let owns = BlocklistAIHistoryModel::as_ref(ctx) .all_live_conversations_for_terminal_view(self.terminal_view_id) - .any(|c| c.id() == conversation_id); - if !owns { - return; + .any(|conversation| conversation.id() == conversation_id); + let has_active_stream = self + .in_flight_response_streams + .has_active_stream_for_conversation(conversation_id, ctx); + let Some(conversation) = + BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) + else { + log::info!( + "Pending events are not ready: conversation_id={conversation_id:?} reason=conversation_missing owns_conversation={owns} has_active_stream={has_active_stream}" + ); + return false; + }; + let is_success = matches!(conversation.status(), ConversationStatus::Success); + if !owns || has_active_stream || !is_success { + log::info!( + "Pending events are not ready: conversation_id={conversation_id:?} owns_conversation={owns} has_active_stream={has_active_stream} status={:?}", + conversation.status() + ); + return false; } + true + } + + #[cfg(target_family = "wasm")] + fn maybe_prepare_local_claude_wake( + &mut self, + _conversation_id: AIConversationId, + _trigger: LocalClaudeWakeTrigger, + _ctx: &mut ModelContext, + ) -> bool { + false + } + + #[cfg(not(target_family = "wasm"))] + fn maybe_prepare_local_claude_wake( + &mut self, + conversation_id: AIConversationId, + trigger: LocalClaudeWakeTrigger, + ctx: &mut ModelContext, + ) -> bool { if self - .in_flight_response_streams - .has_active_stream_for_conversation(conversation_id, ctx) + .pending_local_claude_wakes + .contains_key(&conversation_id) { - return; + log::info!("Dormant Claude wake already pending: conversation_id={conversation_id:?}"); + return true; + } + if trigger.requires_pending_events() { + let has_pending_events = OrchestrationEventService::handle(ctx) + .update(ctx, |svc, _| svc.has_pending_events(conversation_id)); + if !has_pending_events { + return false; + } } - // Only drain when the conversation is actually idle. - let is_success = BlocklistAIHistoryModel::as_ref(ctx) - .conversation(&conversation_id) - .is_some_and(|c| matches!(c.status(), ConversationStatus::Success)); - if !is_success { + if !self.conversation_ready_for_pending_events(conversation_id, ctx) { + return false; + } + + let history_model = BlocklistAIHistoryModel::as_ref(ctx); + let Some(conversation) = history_model.conversation(&conversation_id).cloned() else { + log::info!( + "Skipping dormant Claude wake preparation: conversation_id={conversation_id:?} reason=conversation_missing" + ); + return false; + }; + let parent_conversation = conversation + .parent_conversation_id() + .and_then(|parent_conversation_id| history_model.conversation(&parent_conversation_id)) + .cloned(); + let working_dir = self + .active_session + .as_ref(ctx) + .current_working_directory() + .cloned() + .map(PathBuf::from); + let task_id = conversation.task_id(); + + let server_api = ServerApiProvider::as_ref(ctx).get(); + let handle = ctx.spawn( + async move { + log::info!( + "Preparing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id:?}" + ); + ClaudeHarness::wake_dormant_session( + server_api.clone(), + conversation, + parent_conversation, + working_dir, + ) + .await + }, + move |me, result, ctx| { + me.pending_local_claude_wakes.remove(&conversation_id); + match result { + Ok(Some(command)) => { + log::info!( + "Executing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id:?}" + ); + BlocklistAIHistoryModel::handle(ctx).update( + ctx, + |history_model, ctx| { + history_model.update_conversation_status( + me.terminal_view_id, + conversation_id, + ConversationStatus::InProgress, + ctx, + ); + }, + ); + ctx.emit(BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { + command, + }); + } + Ok(None) => { + match trigger { + LocalClaudeWakeTrigger::PendingEvents => { + log::info!( + "Falling back to generic pending-event injection after dormant Claude wake eligibility check: conversation_id={conversation_id:?} task_id={task_id:?}" + ); + me.inject_pending_events_for_request(conversation_id, ctx); + } + LocalClaudeWakeTrigger::WakeOnlyStream => { + log::info!( + "Retrying wake-only dormant Claude eligibility check: conversation_id={conversation_id:?} task_id={task_id:?}" + ); + me.schedule_dormant_claude_wake_ready_retry(conversation_id, ctx); + } + } + } + Err(err) => { + log::warn!( + "Failed to prepare dormant Claude wake command for {conversation_id:?} task_id={task_id:?}: {err:#}" + ); + match trigger { + LocalClaudeWakeTrigger::PendingEvents => { + me.schedule_pending_events_ready_retry(conversation_id, ctx); + } + LocalClaudeWakeTrigger::WakeOnlyStream => { + me.schedule_dormant_claude_wake_ready_retry(conversation_id, ctx); + } + } + } + } + }, + ); + self.pending_local_claude_wakes + .insert(conversation_id, handle); + true + } + + fn schedule_pending_events_ready_retry( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + ctx.spawn( + async move { Timer::after(Duration::from_secs(2)).await }, + move |me, _, ctx| { + me.handle_pending_events_ready(conversation_id, ctx); + }, + ); + } + + fn schedule_dormant_claude_wake_ready_retry( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + ctx.spawn( + async move { Timer::after(Duration::from_secs(2)).await }, + move |me, _, ctx| { + me.handle_dormant_claude_wake_ready(conversation_id, ctx); + }, + ); + } + + fn inject_pending_events_for_request( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if !self.conversation_ready_for_pending_events(conversation_id, ctx) { return; } @@ -1546,6 +1725,44 @@ impl BlocklistAIController { } } + /// Handles the EventsReady signal. Checks readiness, drains + /// pending events from the service, and injects them into the conversation. + fn handle_pending_events_ready( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if !self.conversation_ready_for_pending_events(conversation_id, ctx) { + return; + } + + if self.maybe_prepare_local_claude_wake( + conversation_id, + LocalClaudeWakeTrigger::PendingEvents, + ctx, + ) { + return; + } + + self.inject_pending_events_for_request(conversation_id, ctx); + } + + fn handle_dormant_claude_wake_ready( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if !self.maybe_prepare_local_claude_wake( + conversation_id, + LocalClaudeWakeTrigger::WakeOnlyStream, + ctx, + ) { + log::info!( + "Ignoring dormant Claude wake-ready signal: conversation_id={conversation_id:?}" + ); + } + } + pub fn resume_conversation( &mut self, conversation_id: AIConversationId, diff --git a/app/src/ai/blocklist/controller/slash_command.rs b/app/src/ai/blocklist/controller/slash_command.rs index 15f6ddc8e..60aa0348f 100644 --- a/app/src/ai/blocklist/controller/slash_command.rs +++ b/app/src/ai/blocklist/controller/slash_command.rs @@ -6,8 +6,8 @@ use warpui::{AppContext, ModelContext, SingletonEntity}; use crate::{ ai::{ agent::{ - conversation::AIConversationId, AIAgentContext, AIAgentInput, CloneRepositoryURL, - EntrypointType, RequestMetadata, + conversation::AIConversationId, AIAgentContext, AIAgentInput, CancellationReason, + CloneRepositoryURL, EntrypointType, RequestMetadata, }, blocklist::agent_view::AgentViewEntryOrigin, }, @@ -90,6 +90,8 @@ impl SlashCommandRequest { if inputs.is_empty() { return; } + let active_conversation_id = BlocklistAIHistoryModel::as_ref(ctx) + .active_conversation_id(controller.terminal_view_id); // If no existing conversation, create a new one. // When AgentView is enabled, enter agent view which creates the conversation @@ -114,6 +116,18 @@ impl SlashCommandRequest { return; }; + let cancellation_reason = CancellationReason::FollowUpSubmitted { + is_for_same_conversation: active_conversation_id + .is_some_and(|id| id == conversation_id), + }; + if let Some(active_conversation_id) = active_conversation_id { + controller.cancel_conversation_progress( + active_conversation_id, + cancellation_reason, + ctx, + ); + } + let Some(conversation) = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) else { diff --git a/app/src/ai/blocklist/history_model.rs b/app/src/ai/blocklist/history_model.rs index eb0910aad..d407a3145 100644 --- a/app/src/ai/blocklist/history_model.rs +++ b/app/src/ai/blocklist/history_model.rs @@ -1090,10 +1090,9 @@ impl BlocklistAIHistoryModel { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, 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(); @@ -1245,10 +1244,9 @@ impl BlocklistAIHistoryModel { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, 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, }; diff --git a/app/src/ai/blocklist/history_model_test.rs b/app/src/ai/blocklist/history_model_test.rs index 6dc2a4f48..9c7e9df65 100644 --- a/app/src/ai/blocklist/history_model_test.rs +++ b/app/src/ai/blocklist/history_model_test.rs @@ -879,6 +879,50 @@ fn test_toggle_autoexecute_override_persists_updated_conversation_state() { }); } +#[test] +fn test_update_event_sequence_persists_updated_conversation_state() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + + let (sender, receiver) = std::sync::mpsc::sync_channel(1); + 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 terminal_view_id = EntityId::new(); + + let conversation_id = history_model.update(&mut app, |history_model, ctx| { + history_model.start_new_conversation(terminal_view_id, false, false, ctx) + }); + + history_model.update(&mut app, |history_model, ctx| { + history_model.update_event_sequence(conversation_id, 42, ctx); + }); + + let event = receiver.recv_timeout(Duration::from_secs(1)).unwrap(); + + let ModelEvent::UpdateMultiAgentConversation { + conversation_id: persisted_conversation_id, + conversation_data, + .. + } = event + else { + panic!("expected UpdateMultiAgentConversation event"); + }; + + assert_eq!(persisted_conversation_id, conversation_id.to_string()); + assert_eq!(conversation_data.last_event_sequence, Some(42)); + + history_model.read(&app, |history_model, _| { + let conversation = history_model + .conversation(&conversation_id) + .expect("conversation should exist"); + assert_eq!(conversation.last_event_sequence(), Some(42)); + }); + }); +} + #[test] fn test_find_by_token_after_merge_cloud_metadata() { App::test((), |mut app| async move { @@ -1130,6 +1174,7 @@ fn test_find_by_token_after_insert_forked_conversation_from_tasks() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/app/src/ai/blocklist/orchestration_event_streamer.rs b/app/src/ai/blocklist/orchestration_event_streamer.rs index 0676172d8..97fc9c386 100644 --- a/app/src/ai/blocklist/orchestration_event_streamer.rs +++ b/app/src/ai/blocklist/orchestration_event_streamer.rs @@ -4,13 +4,14 @@ use super::orchestration_events::{ OrchestrationEventService, PendingEvent, PendingEventDetail, }; use crate::ai::agent::{ - conversation::AIConversationId, AIAgentExchangeId, AIAgentOutputMessageType, - ReceivedMessageInput, + conversation::{AIAgentHarness, AIConversationId, ConversationStatus}, + AIAgentExchangeId, AIAgentOutputMessageType, ReceivedMessageInput, }; use crate::ai::agent_events::{ run_agent_event_driver, AgentEventConsumer, AgentEventConsumerControlFlow, AgentEventDriverConfig, MessageHydrator, ServerApiAgentEventSource, }; +use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::server::server_api::ai::{AIClient, AgentRunEvent}; use crate::server::server_api::{ServerApi, ServerApiProvider}; use anyhow::anyhow; @@ -20,9 +21,10 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::Duration; use uuid::Uuid; +use warp_cli::agent::Harness; use warp_core::features::FeatureFlag; use warp_multi_agent_api as api; -use warpui::r#async::Timer; +use warpui::r#async::{SpawnedFutureHandle, Timer}; use warpui::{ Entity, EntityId, GetSingletonModelHandle, ModelContext, SingletonEntity, UpdateModel, }; @@ -51,6 +53,46 @@ struct SseForwardingConsumer { tx: mpsc::UnboundedSender, self_run_id: String, hydrator: MessageHydrator, + hydrate_new_messages: bool, +} + +/// State for a wake-only listener. Unlike `SseConnectionState`, this listener +/// never forwards or persists events; it stops on the first event and asks the +/// controller to cold-start the dormant Claude run so the parent bridge can +/// take over delivery. +struct WakeConnectionState { + generation: u64, + task: SpawnedFutureHandle, +} + +struct DormantClaudeWakeConsumer { + run_id: String, + wake_sequence: Option, +} + +impl DormantClaudeWakeConsumer { + fn new(run_id: String) -> Self { + Self { + run_id, + wake_sequence: None, + } + } +} + +#[cfg_attr(target_family = "wasm", async_trait(?Send))] +#[cfg_attr(not(target_family = "wasm"), async_trait)] +impl AgentEventConsumer for DormantClaudeWakeConsumer { + async fn on_event( + &mut self, + event: AgentRunEvent, + ) -> anyhow::Result { + if event.run_id != self.run_id { + return Ok(AgentEventConsumerControlFlow::Continue); + } + + self.wake_sequence = Some(event.sequence); + Ok(AgentEventConsumerControlFlow::Stop) + } } #[cfg_attr(target_family = "wasm", async_trait(?Send))] @@ -60,10 +102,13 @@ impl AgentEventConsumer for SseForwardingConsumer { &mut self, event: AgentRunEvent, ) -> anyhow::Result { - let fetched_message = self - .hydrator - .hydrate_event_for_recipient(&event, &self.self_run_id) - .await; + let fetched_message = if self.hydrate_new_messages { + self.hydrator + .hydrate_event_for_recipient(&event, &self.self_run_id) + .await + } else { + None + }; self.tx .unbounded_send(SseStreamItem { @@ -98,8 +143,18 @@ struct ConversationStreamState { /// model id for `agent_sdk`) that need events delivered to this /// conversation. consumers: HashSet, + /// Execution harness from the task row, when available. Local harness + /// child conversations are created before they have server conversation + /// metadata, so this lets us recognize dormant local Claude children + /// without relying on `ServerAIConversationMetadata`. + harness: Option, /// Active SSE connection, if one is open. sse_connection: Option, + /// Active wake-only listener for dormant local Claude children, if one is + /// open. This is separate from generic SSE because generic delivery would + /// hydrate messages and advance the shared cursor before Claude's parent + /// bridge can consume them. + wake_connection: Option, /// Consecutive `get_ambient_agent_task` failure count for the /// post-restore retry loop; resets on success. restore_fetch_failures: usize, @@ -123,13 +178,22 @@ pub struct OrchestrationEventStreamer { /// Monotonic counter for SSE connection generations. Ensures stale /// callbacks from replaced connections are discarded. next_sse_generation: u64, + /// Monotonic counter for wake-only listener generations. Ensures stale + /// callbacks from replaced listeners are discarded. + next_wake_generation: u64, } pub enum OrchestrationEventStreamerEvent { - // Reserved for future use (e.g., status signals to the controller). + DormantClaudeWakeReady { conversation_id: AIConversationId }, } impl OrchestrationEventStreamer { + fn message_hydrator_for_run_id(&self, run_id: &str) -> MessageHydrator { + match run_id.parse::() { + Ok(task_id) => MessageHydrator::for_task(self.server_api.clone(), task_id), + Err(_) => MessageHydrator::new(self.ai_client.clone()), + } + } pub fn new(ctx: &mut ModelContext) -> Self { let provider = ServerApiProvider::as_ref(ctx); let ai_client = provider.get_ai_client(); @@ -143,6 +207,7 @@ impl OrchestrationEventStreamer { server_api, streams: HashMap::new(), next_sse_generation: 0, + next_wake_generation: 0, } } @@ -164,6 +229,7 @@ impl OrchestrationEventStreamer { server_api, streams: HashMap::new(), next_sse_generation: 0, + next_wake_generation: 0, } } @@ -192,6 +258,7 @@ impl OrchestrationEventStreamer { // If the server-token event fired before this registration, pick // up the now-available child role here. self.ensure_self_run_id_watched(conversation_id, ctx); + self.spawn_task_harness_fetch_if_needed(conversation_id, ctx); self.reevaluate_eligibility(conversation_id, ctx); } @@ -281,6 +348,12 @@ impl OrchestrationEventStreamer { } => { self.on_restored_conversations(conversation_ids.clone(), ctx); } + BlocklistAIHistoryEvent::UpdatedConversationStatus { + conversation_id, .. + } + | BlocklistAIHistoryEvent::UpdatedConversationMetadata { + conversation_id, .. + } => self.reevaluate_eligibility(*conversation_id, ctx), BlocklistAIHistoryEvent::StartedNewConversation { .. } | BlocklistAIHistoryEvent::CreatedSubtask { .. } | BlocklistAIHistoryEvent::UpgradedTask { .. } @@ -292,8 +365,6 @@ impl OrchestrationEventStreamer { | BlocklistAIHistoryEvent::UpdatedTodoList { .. } | BlocklistAIHistoryEvent::UpdatedAutoexecuteOverride { .. } | BlocklistAIHistoryEvent::SplitConversation { .. } - | BlocklistAIHistoryEvent::UpdatedConversationStatus { .. } - | BlocklistAIHistoryEvent::UpdatedConversationMetadata { .. } | BlocklistAIHistoryEvent::UpdatedConversationArtifacts { .. } => {} } } @@ -303,11 +374,58 @@ impl OrchestrationEventStreamer { conversation_id: AIConversationId, ctx: &mut ModelContext, ) { + self.spawn_task_harness_fetch_if_needed(conversation_id, ctx); if self.ensure_self_run_id_watched(conversation_id, ctx) { self.reevaluate_eligibility(conversation_id, ctx); } } + fn spawn_task_harness_fetch_if_needed( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if self + .streams + .get(&conversation_id) + .is_some_and(|stream| stream.harness.is_some()) + { + return; + } + let Some(run_id) = self.self_run_id(conversation_id, ctx) else { + return; + }; + let Ok(task_id) = run_id.parse::() else { + return; + }; + let local_cursor = self + .streams + .get(&conversation_id) + .map(|stream| stream.event_cursor) + .unwrap_or(0); + let ai_client = self.ai_client.clone(); + ctx.spawn( + async move { ai_client.get_ambient_agent_task(&task_id).await }, + move |me, result, ctx| { + let task = match result { + Ok(task) => task, + Err(err) => { + log::warn!( + "Failed to fetch task harness for {conversation_id:?} task_id={task_id}: {err:#}" + ); + return; + } + }; + if let Some(stream) = me.streams.get_mut(&conversation_id) { + stream.harness = agent_task_harness(&task).or(stream.harness); + stream.event_cursor = + local_cursor.max(task.last_event_sequence.unwrap_or(0)); + } + me.reevaluate_eligibility(conversation_id, ctx); + }, + ); + } + /// Inserts `self_run_id` into the conversation's watched set if the /// conversation has any orchestration role (child or parent) and is /// not a passive remote-run view. Returns whether anything was @@ -404,7 +522,8 @@ impl OrchestrationEventStreamer { .retain(|id| !confirmed_ids.contains(id)); } - let hydrator = MessageHydrator::new(self.ai_client.clone()); + let hydrator = + self.message_hydrator_for_run_id(conversation.run_id().as_deref().unwrap_or_default()); ctx.spawn( async move { hydrator @@ -440,7 +559,11 @@ impl OrchestrationEventStreamer { // Dropping the SSE receiver causes the driver task's next send // to fail and exit; the drain timer's `is_current` check then // no-ops on its next tick. - self.streams.remove(&conversation_id); + if let Some(mut stream) = self.streams.remove(&conversation_id) { + if let Some(connection) = stream.wake_connection.take() { + connection.task.abort(); + } + } // Prune the removed conversation's run_id from every other // tracked conversation's watched set, then re-evaluate eligibility @@ -582,6 +705,7 @@ impl OrchestrationEventStreamer { // Reset the retry counter on success. stream.restore_fetch_failures = 0; + stream.harness = agent_task_harness(&task).or(stream.harness); // Merge the server cursor: use the max of SQLite and // server values so we don't re-deliver events the @@ -706,6 +830,29 @@ impl OrchestrationEventStreamer { .is_some_and(|c| c.is_viewing_shared_session() || c.is_remote_child()) } + fn should_skip_sse_for_dormant_local_claude_child( + &self, + conversation_id: AIConversationId, + ctx: &warpui::AppContext, + ) -> bool { + let Some(conversation) = + BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) + else { + return false; + }; + conversation.is_child_agent_conversation() + && !conversation.is_remote_child() + && matches!(conversation.status(), ConversationStatus::Success) + && (conversation + .server_metadata() + .is_some_and(|metadata| metadata.harness == AIAgentHarness::ClaudeCode) + || self + .streams + .get(&conversation_id) + .and_then(|stream| stream.harness) + .is_some_and(|harness| harness == Harness::Claude)) + } + /// True iff this conversation should currently hold an SSE connection. /// A subscription is needed only when there is an active consumer in /// this process (an open agent view or an agent_sdk driver) AND the @@ -718,12 +865,33 @@ impl OrchestrationEventStreamer { if self.is_remote_run_view(conversation_id, ctx) { return false; } + if self.should_skip_sse_for_dormant_local_claude_child(conversation_id, ctx) { + log::info!( + "Skipping generic SSE delivery for dormant local Claude child {conversation_id:?}; parent bridge will deliver wake events" + ); + return false; + } let has_parent = BlocklistAIHistoryModel::as_ref(ctx) .conversation(&conversation_id) .is_some_and(|c| c.has_parent_agent()); has_parent || self.is_parent_agent_conversation(conversation_id, ctx) } + /// True iff this conversation should hold the wake-only listener used for + /// dormant local Claude children. Generic SSE intentionally stays closed + /// for these conversations so it cannot hydrate messages or advance the + /// server cursor before Claude's parent bridge starts. + fn is_dormant_claude_wake_listener_eligible( + &self, + conversation_id: AIConversationId, + ctx: &warpui::AppContext, + ) -> bool { + self.has_active_consumer(conversation_id) + && !self.is_remote_run_view(conversation_id, ctx) + && self.should_skip_sse_for_dormant_local_claude_child(conversation_id, ctx) + && self.self_run_id(conversation_id, ctx).is_some() + } + /// Returns the list of run_ids to subscribe to for `conversation_id`. /// Includes both the conversation's own `self_run_id` (when it is a /// child) and any registered child run_ids (when the conversation @@ -760,6 +928,131 @@ impl OrchestrationEventStreamer { (false, true) => self.teardown_sse(conversation_id, ctx), (false, false) => {} } + + let wake_eligible = self.is_dormant_claude_wake_listener_eligible(conversation_id, ctx); + let wake_connected = self + .streams + .get(&conversation_id) + .is_some_and(|s| s.wake_connection.is_some()); + + match (wake_eligible, wake_connected) { + (true, false) => self.start_dormant_claude_wake_listener(conversation_id, ctx), + (true, true) => {} + (false, true) => self.teardown_dormant_claude_wake_listener(conversation_id), + (false, false) => {} + } + } + + /// Opens a wake-only listener for a dormant local Claude child. The + /// listener observes the child's run_id, stops on the first event, and + /// emits a controller signal without enqueueing any event data or + /// persisting any cursor. The Claude parent bridge will consume the event + /// after the CLI has been relaunched. + fn start_dormant_claude_wake_listener( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + let Some(run_id) = self.self_run_id(conversation_id, ctx) else { + return; + }; + + let local_cursor = self + .streams + .get(&conversation_id) + .map(|s| s.event_cursor) + .unwrap_or(0); + let generation = self.next_wake_generation; + self.next_wake_generation += 1; + + log::info!( + "Opening dormant Claude wake listener for {conversation_id:?} \ + (gen={generation}, run_id={run_id:?}, since={local_cursor})" + ); + + let ai_client = self.ai_client.clone(); + let source = ServerApiAgentEventSource::new(self.server_api.clone()); + let task_run_id = run_id.clone(); + let handle = ctx.spawn( + async move { + let since_sequence = resolve_dormant_claude_wake_cursor( + ai_client, + task_run_id.clone(), + local_cursor, + ) + .await; + let config = AgentEventDriverConfig::retry_forever( + vec![task_run_id.clone()], + since_sequence, + ); + let mut consumer = DormantClaudeWakeConsumer::new(task_run_id); + run_agent_event_driver(source, config, &mut consumer).await?; + Ok::<_, anyhow::Error>(consumer.wake_sequence) + }, + move |me, result, ctx| { + let is_current = me + .streams + .get(&conversation_id) + .and_then(|s| s.wake_connection.as_ref()) + .is_some_and(|c| c.generation == generation); + if !is_current { + return; + } + + if let Some(stream) = me.streams.get_mut(&conversation_id) { + stream.wake_connection = None; + } + + match result { + Ok(Some(sequence)) => { + log::info!( + "Dormant Claude wake listener observed event for \ + {conversation_id:?} at sequence {sequence}" + ); + ctx.emit(OrchestrationEventStreamerEvent::DormantClaudeWakeReady { + conversation_id, + }); + } + Ok(None) => { + log::warn!( + "Dormant Claude wake listener stopped for {conversation_id:?} \ + without observing an event" + ); + if me.is_dormant_claude_wake_listener_eligible(conversation_id, ctx) { + me.start_dormant_claude_wake_listener(conversation_id, ctx); + } + } + Err(err) => { + log::warn!( + "Dormant Claude wake listener failed for {conversation_id:?} \ + (gen={generation}): {err:#}" + ); + if me.is_dormant_claude_wake_listener_eligible(conversation_id, ctx) { + me.start_dormant_claude_wake_listener(conversation_id, ctx); + } + } + } + }, + ); + + let stream = self.streams.entry(conversation_id).or_default(); + stream.wake_connection = Some(WakeConnectionState { + generation, + task: handle, + }); + } + + fn teardown_dormant_claude_wake_listener(&mut self, conversation_id: AIConversationId) { + if let Some(stream) = self.streams.get_mut(&conversation_id) { + if let Some(connection) = stream.wake_connection.take() { + log::info!( + "Tearing down dormant Claude wake listener for {conversation_id:?} \ + (gen={})", + connection.generation + ); + connection.task.abort(); + } + } } /// Opens a long-lived SSE connection for `conversation_id`. Events @@ -781,7 +1074,6 @@ impl OrchestrationEventStreamer { .unwrap_or(0); let server_api = self.server_api.clone(); - let ai_client = self.ai_client.clone(); let self_run_id = self.self_run_id(conversation_id, ctx).unwrap_or_default(); @@ -789,10 +1081,8 @@ impl OrchestrationEventStreamer { let generation = self.next_sse_generation; self.next_sse_generation += 1; - self.streams - .entry(conversation_id) - .or_default() - .sse_connection = Some(SseConnectionState { + let stream = self.streams.entry(conversation_id).or_default(); + stream.sse_connection = Some(SseConnectionState { event_receiver: rx, generation, }); @@ -804,7 +1094,7 @@ impl OrchestrationEventStreamer { let config = AgentEventDriverConfig::retry_forever(run_ids.clone(), cursor); let source = ServerApiAgentEventSource::new(server_api); - let hydrator = MessageHydrator::new(ai_client); + let hydrator = self.message_hydrator_for_run_id(&self_run_id); ctx.spawn( async move { @@ -812,6 +1102,7 @@ impl OrchestrationEventStreamer { tx, self_run_id, hydrator, + hydrate_new_messages: true, }; run_agent_event_driver(source, config, &mut consumer).await }, @@ -963,10 +1254,9 @@ impl OrchestrationEventStreamer { } // Track message IDs for server-side mark_delivered calls. - let message_ids: Vec = events + let message_ids: Vec = messages .iter() - .filter(|e| e.event_type == "new_message" && e.run_id == self_run_id) - .filter_map(|e| e.ref_id.clone()) + .map(|message| message.message_id.clone()) .collect(); if !message_ids.is_empty() { self.streams @@ -981,7 +1271,7 @@ impl OrchestrationEventStreamer { return; } - let pending = build_pending_events(messages, lifecycle_events); + let pending = build_pending_events(&events, messages, lifecycle_events); OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { svc.enqueue_event_batch(conversation_id, pending, ctx); }); @@ -1022,6 +1312,35 @@ impl Entity for OrchestrationEventStreamer { impl SingletonEntity for OrchestrationEventStreamer {} +async fn resolve_dormant_claude_wake_cursor( + ai_client: Arc, + run_id: String, + local_cursor: i64, +) -> i64 { + let Ok(task_id) = run_id.parse::() else { + return local_cursor; + }; + + match ai_client.get_ambient_agent_task(&task_id).await { + Ok(task) => local_cursor.max(task.last_event_sequence.unwrap_or(0)), + Err(err) => { + log::warn!( + "Failed to read server cursor for dormant Claude wake listener \ + run {run_id}: {err:#}; using local cursor {local_cursor}" + ); + local_cursor + } + } +} + +fn agent_task_harness(task: &crate::ai::ambient_agents::task::AmbientAgentTask) -> Option { + task.agent_config_snapshot + .as_ref() + .and_then(|snapshot| snapshot.harness.as_ref()) + .map(|config| config.harness_type) + .filter(|harness| *harness != Harness::Unknown) +} + fn parse_occurred_at(s: &str) -> prost_types::Timestamp { chrono::DateTime::parse_from_rfc3339(s) .map(|dt| prost_types::Timestamp { @@ -1083,20 +1402,32 @@ fn convert_lifecycle_events(events: &[AgentRunEvent], self_run_id: &str) -> Vec< } fn build_pending_events( + events: &[AgentRunEvent], messages: Vec, lifecycle_events: Vec, ) -> Vec { let mut pending = Vec::with_capacity(messages.len() + lifecycle_events.len()); for msg in &messages { + let metadata = events + .iter() + .find(|event| { + event.event_type == "new_message" + && event.ref_id.as_deref() == Some(msg.message_id.as_str()) + }) + .map(|event| (event.sequence, event.occurred_at.clone())); + let (sequence, occurred_at) = + metadata.unwrap_or_else(|| (0, chrono::Utc::now().to_rfc3339())); pending.push(PendingEvent { event_id: msg.message_id.clone(), source_agent_id: msg.sender_agent_id.clone(), attempt_count: 0, detail: PendingEventDetail::Message { + sequence, message_id: msg.message_id.clone(), addresses: msg.addresses.clone(), subject: msg.subject.clone(), message_body: msg.message_body.clone(), + occurred_at, }, }); } diff --git a/app/src/ai/blocklist/orchestration_event_streamer_tests.rs b/app/src/ai/blocklist/orchestration_event_streamer_tests.rs index fdf6eb145..58c326f0d 100644 --- a/app/src/ai/blocklist/orchestration_event_streamer_tests.rs +++ b/app/src/ai/blocklist/orchestration_event_streamer_tests.rs @@ -1,6 +1,6 @@ use super::*; use crate::ai::agent_events::{ - agent_event_backoff, agent_event_failures_exceeded_threshold, + agent_event_backoff, agent_event_failures_exceeded_threshold, AgentEventConsumerControlFlow, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS, }; @@ -142,6 +142,7 @@ fn ai_conversation_new_restored_preserves_last_event_sequence() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: Some(42), @@ -190,6 +191,320 @@ fn make_ambient_task_with_event_seq( } } +fn make_server_metadata_with_harness( + harness: AIAgentHarness, +) -> crate::ai::agent::conversation::ServerAIConversationMetadata { + use crate::ai::agent::api::ServerConversationToken; + use crate::cloud_object::{Revision, ServerMetadata, ServerPermissions}; + use crate::persistence::model::ConversationUsageMetadata; + use crate::server::ids::ServerId; + use chrono::Utc; + + crate::ai::agent::conversation::ServerAIConversationMetadata { + title: "test".to_string(), + working_directory: None, + harness, + usage: ConversationUsageMetadata { + was_summarized: false, + context_window_usage: 0.0, + credits_spent: 0.0, + credits_spent_for_last_block: None, + token_usage: vec![], + tool_usage_metadata: Default::default(), + }, + metadata: ServerMetadata { + uid: ServerId::default(), + revision: Revision::now(), + metadata_last_updated_ts: Utc::now().into(), + trashed_ts: None, + folder_id: None, + is_welcome_object: false, + creator_uid: None, + last_editor_uid: None, + current_editor_uid: None, + }, + permissions: ServerPermissions::mock_personal(), + ambient_agent_task_id: None, + server_conversation_token: ServerConversationToken::new("server-token".to_string()), + artifacts: vec![], + } +} + +#[test] +fn dormant_local_claude_child_skips_generic_sse_but_allows_wake_listener() { + 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 parent_id = AIConversation::new(false).id(); + let mut conversation = AIConversation::new(false); + let run_id = "550e8400-e29b-41d4-a716-446655440610".to_string(); + conversation.set_run_id(run_id.clone()); + conversation.set_parent_conversation_id(parent_id); + conversation.set_server_metadata(make_server_metadata_with_harness( + AIAgentHarness::ClaudeCode, + )); + 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::Success, + ctx, + ); + }); + + let mock = MockAIClient::new(); + let ai_client: Arc = 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, |me, _| { + let stream = me.streams.entry(conversation_id).or_default(); + stream.consumers.insert(warpui::EntityId::new()); + stream.watched_run_ids.insert(run_id); + }); + + streamer.read(&app, |me, ctx| { + assert!( + !me.is_eligible(conversation_id, ctx), + "generic SSE must stay closed for dormant local Claude children" + ); + assert!( + me.is_dormant_claude_wake_listener_eligible(conversation_id, ctx), + "wake-only listener should open for dormant local Claude children" + ); + }); + }); +} + +#[test] +fn dormant_local_claude_child_uses_task_harness_when_server_metadata_missing() { + 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 warp_cli::agent::Harness; + 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 parent_id = AIConversation::new(false).id(); + let mut conversation = AIConversation::new(false); + let run_id = "550e8400-e29b-41d4-a716-446655440611".to_string(); + conversation.set_run_id(run_id.clone()); + conversation.set_parent_conversation_id(parent_id); + 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::Success, + ctx, + ); + }); + + let mock = MockAIClient::new(); + let ai_client: Arc = 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, |me, _| { + let stream = me.streams.entry(conversation_id).or_default(); + stream.consumers.insert(warpui::EntityId::new()); + stream.watched_run_ids.insert(run_id); + }); + + streamer.read(&app, |me, ctx| { + assert!( + me.is_eligible(conversation_id, ctx), + "generic SSE should remain eligible before the task harness is known" + ); + assert!( + !me.is_dormant_claude_wake_listener_eligible(conversation_id, ctx), + "wake-only listener should wait until the task harness identifies Claude" + ); + }); + + streamer.update(&mut app, |me, _| { + me.streams + .get_mut(&conversation_id) + .expect("stream exists") + .harness = Some(Harness::Claude); + }); + + streamer.read(&app, |me, ctx| { + assert!( + !me.is_eligible(conversation_id, ctx), + "generic SSE must close after task metadata identifies a dormant local Claude child" + ); + assert!( + me.is_dormant_claude_wake_listener_eligible(conversation_id, ctx), + "wake-only listener should open based on cached task harness even without server metadata" + ); + }); + }); +} +#[tokio::test] +async fn dormant_claude_wake_consumer_stops_on_first_target_event() { + let mut consumer = DormantClaudeWakeConsumer::new("target-run".to_string()); + + let ignored_event = AgentRunEvent { + event_type: "new_message".to_string(), + run_id: "other-run".to_string(), + ref_id: Some("message-1".to_string()), + execution_id: None, + occurred_at: "2026-01-01T00:00:00Z".to_string(), + sequence: 7, + }; + assert_eq!( + consumer.on_event(ignored_event).await.unwrap(), + AgentEventConsumerControlFlow::Continue + ); + assert_eq!(consumer.wake_sequence, None); + + // The wake consumer uses the default no-op cursor persistence hook; it + // should not persist SQLite or server cursors while waiting to wake Claude. + consumer.persist_cursor(7).await.unwrap(); + + let target_event = AgentRunEvent { + event_type: "new_message".to_string(), + run_id: "target-run".to_string(), + ref_id: Some("message-2".to_string()), + execution_id: None, + occurred_at: "2026-01-01T00:00:01Z".to_string(), + sequence: 8, + }; + assert_eq!( + consumer.on_event(target_event).await.unwrap(), + AgentEventConsumerControlFlow::Stop + ); + assert_eq!(consumer.wake_sequence, Some(8)); +} + +#[test] +fn restored_conversations_skip_v2_streaming_when_orchestration_v2_disabled() { + 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(false); + + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); + + let mut conversation = AIConversation::new(false); + conversation.set_run_id("550e8400-e29b-41d4-a716-446655440500".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 = 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, |me, ctx| { + me.on_restored_conversations(vec![conversation_id], ctx); + }); + + streamer.read(&app, |me, _| { + assert!( + me.streams.is_empty(), + "V2-disabled restore must not initialize stream state" + ); + }); + }); +} + +#[test] +fn build_pending_events_preserves_message_sequence_and_timestamp() { + let occurred_at = "2026-01-02T03:04:05Z"; + let pending = build_pending_events( + &[AgentRunEvent { + event_type: "new_message".to_string(), + run_id: "sender-run".to_string(), + ref_id: Some("message-123".to_string()), + execution_id: None, + occurred_at: occurred_at.to_string(), + sequence: 77, + }], + vec![ReceivedMessageInput { + message_id: "message-123".to_string(), + sender_agent_id: "sender-agent".to_string(), + addresses: vec!["recipient-agent".to_string()], + subject: "subject".to_string(), + message_body: "body".to_string(), + }], + vec![], + ); + + assert_eq!(pending.len(), 1); + let detail = &pending[0].detail; + let PendingEventDetail::Message { + sequence, + message_id, + occurred_at: event_occurred_at, + .. + } = detail + else { + panic!("expected pending message event"); + }; + assert_eq!(*sequence, 77); + assert_eq!(message_id, "message-123"); + assert_eq!(event_occurred_at, occurred_at); +} + +#[tokio::test] +async fn sse_forwarding_consumer_skips_message_hydration_when_disabled() { + use futures::StreamExt; + + let (tx, mut rx) = futures::channel::mpsc::unbounded(); + let mut ai_client = crate::server::server_api::ai::MockAIClient::new(); + ai_client.expect_read_agent_message().times(0); + let ai_client: Arc = Arc::new(ai_client); + let hydrator = MessageHydrator::new(ai_client); + let mut consumer = SseForwardingConsumer { + tx, + self_run_id: "child-run".to_string(), + hydrator, + hydrate_new_messages: false, + }; + let event = make_run_event("new_message", "child-run", Some("message-123")); + + consumer.on_event(event).await.unwrap(); + + let item = rx.next().await.expect("expected forwarded event"); + assert_eq!(item.event.ref_id.as_deref(), Some("message-123")); + assert!(item.fetched_message.is_none()); +} #[test] fn finish_restore_fetch_uses_server_cursor_when_sqlite_is_absent() { use crate::ai::agent::conversation::AIConversation; diff --git a/app/src/ai/blocklist/orchestration_events.rs b/app/src/ai/blocklist/orchestration_events.rs index ebec82d8d..6bb8c0803 100644 --- a/app/src/ai/blocklist/orchestration_events.rs +++ b/app/src/ai/blocklist/orchestration_events.rs @@ -56,10 +56,14 @@ impl LifecycleEventDetailStage { #[derive(Debug, Clone)] pub enum PendingEventDetail { Message { + #[cfg_attr(not(test), allow(dead_code))] + sequence: i64, message_id: String, addresses: Vec, subject: String, message_body: String, + #[cfg_attr(not(test), allow(dead_code))] + occurred_at: String, }, Lifecycle { event: api::AgentEvent, @@ -357,42 +361,7 @@ 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, - ); - } - } + self.restore_v1_lifecycle_subscription(*conversation_id, ctx); } } BlocklistAIHistoryEvent::RemoveConversation { @@ -451,6 +420,36 @@ impl OrchestrationEventService { .insert(conversation_id, conversation.status().clone()); } + fn restore_v1_lifecycle_subscription( + &mut self, + source_conversation_id: AIConversationId, + ctx: &ModelContext, + ) { + if FeatureFlag::OrchestrationV2.is_enabled() { + return; + } + + let target_agent_id = { + let history_model = BlocklistAIHistoryModel::as_ref(ctx); + let Some(conversation) = history_model.conversation(&source_conversation_id) else { + return; + }; + if !conversation.is_child_agent_conversation() { + return; + } + + conversation + .parent_conversation_id() + .and_then(|parent_id| history_model.conversation(&parent_id)) + .and_then(|parent| parent.orchestration_agent_id()) + .or_else(|| conversation.parent_agent_id().map(str::to_string)) + }; + + if let Some(target_agent_id) = target_agent_id { + self.register_lifecycle_subscription(source_conversation_id, target_agent_id, None); + } + } + fn on_conversation_status_updated( &mut self, conversation_id: AIConversationId, @@ -747,6 +746,7 @@ impl OrchestrationEventService { // We keep `message_id` stable across targets so dedupe/threading can reason // about a single message delivered to multiple recipients. let message_id = Uuid::new_v4().to_string(); + let occurred_at = chrono::Utc::now().to_rfc3339(); for (_, target_conversation_id) in resolved_targets { let event_id = Uuid::new_v4().to_string(); @@ -756,10 +756,12 @@ impl OrchestrationEventService { source_agent_id: sender_agent_id.to_string(), attempt_count: 0, detail: PendingEventDetail::Message { + sequence: 0, message_id: message_id.clone(), addresses: target_agent_ids.to_vec(), subject: subject.clone(), message_body: message_body.clone(), + occurred_at: occurred_at.clone(), }, }; self.pending_events @@ -890,6 +892,12 @@ impl OrchestrationEventService { ctx.emit(OrchestrationEventServiceEvent::EventsReady { conversation_id }); } + pub fn has_pending_events(&self, conversation_id: AIConversationId) -> bool { + self.pending_events + .get(&conversation_id) + .is_some_and(|events| !events.is_empty()) + } + /// Drain and return all pending events for a conversation. fn drain_pending_events(&mut self, conversation_id: &AIConversationId) -> Vec { self.pending_events @@ -933,10 +941,12 @@ impl OrchestrationEventService { for event in &deliverable { match &event.detail { PendingEventDetail::Message { + sequence: _, message_id, addresses, subject, message_body, + occurred_at: _, } => messages.push(ReceivedMessageInput { message_id: message_id.clone(), sender_agent_id: event.source_agent_id.clone(), diff --git a/app/src/ai/blocklist/orchestration_events_tests.rs b/app/src/ai/blocklist/orchestration_events_tests.rs index c3ba0a52e..30055c0db 100644 --- a/app/src/ai/blocklist/orchestration_events_tests.rs +++ b/app/src/ai/blocklist/orchestration_events_tests.rs @@ -1,7 +1,10 @@ #![allow(deprecated)] use super::*; +use crate::ai::blocklist::BlocklistAIHistoryModel; use std::collections::HashSet; +use warp_core::features::FeatureFlag; use warp_multi_agent_api as api; +use warpui::{App, EntityId}; // Helper for constructing lifecycle pending events with minimal boilerplate. // Tests use this to focus on queue/coalescing behavior rather than payload setup. @@ -71,10 +74,12 @@ fn message_pending_event(event_id: &str) -> PendingEvent { source_agent_id: "sender".to_string(), attempt_count: 0, detail: PendingEventDetail::Message { + sequence: 0, message_id: "message-1".to_string(), addresses: vec!["target".to_string()], subject: "subject".to_string(), message_body: "body".to_string(), + occurred_at: "2026-01-01T00:00:00Z".to_string(), }, } } @@ -269,10 +274,12 @@ fn test_did_event_round_trip_through_server_matches_message_event_by_message_id( source_agent_id: "sender".to_string(), attempt_count: 0, detail: PendingEventDetail::Message { + sequence: 0, message_id: "message-1".to_string(), addresses: vec!["target".to_string()], subject: "subject".to_string(), message_body: "body".to_string(), + occurred_at: "2026-01-01T00:00:00Z".to_string(), }, }; @@ -359,65 +366,79 @@ fn test_lifecycle_event_type_from_proto_includes_cancelled_and_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}; +fn test_has_pending_events_tracks_any_event_kind() { + let conversation_id = crate::ai::agent::conversation::AIConversationId::new(); + let mut service = OrchestrationEventService::new_without_subscriptions(); + assert!(!service.has_pending_events(conversation_id)); + service.pending_events.insert( + conversation_id, + vec![ + lifecycle_pending_event( + "lifecycle-1", + "child-a", + api::LifecycleEventType::InProgress, + 0, + ), + message_pending_event("message-event-1"), + lifecycle_pending_event( + "lifecycle-2", + "child-b", + api::LifecycleEventType::Succeeded, + 0, + ), + message_pending_event("message-event-2"), + ], + ); + assert!(service.has_pending_events(conversation_id)); + service.pending_events.remove(&conversation_id); + assert!(!service.has_pending_events(conversation_id)); +} +#[test] +fn test_restored_v1_child_reregisters_lifecycle_subscription() { 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 _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(false); 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( + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let service = app.add_model(|_| OrchestrationEventService::new_without_subscriptions()); + + let (_parent_conversation_id, child_conversation_id) = + history_model.update(&mut app, |history_model, ctx| { + let parent_conversation_id = + history_model.start_new_conversation(terminal_view_id, false, false, ctx); + history_model.set_server_conversation_token_for_conversation( + parent_conversation_id, + "parent-token".to_string(), + ); + let child_conversation_id = history_model.start_new_child_conversation( + terminal_view_id, + "child".to_string(), + parent_conversation_id, + ctx, + ); + history_model.set_server_conversation_token_for_conversation( + child_conversation_id, + "child-token".to_string(), + ); + (parent_conversation_id, child_conversation_id) + }); + + service.update(&mut app, |service, ctx| { + service.handle_history_event( &BlocklistAIHistoryEvent::RestoredConversations { terminal_view_id, - conversation_ids: vec![parent_conversation_id, child_conversation_id], + conversation_ids: vec![child_conversation_id], }, ctx, ); - }); - service.read(&app, |svc, _| { - let routes = svc + let routes = service .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)" - ); + .expect("restored V1 child should have a lifecycle subscription"); + assert_eq!(routes.len(), 1); + assert_eq!(routes[0].target_agent_id, "parent-token"); + assert_eq!(routes[0].subscribed_event_types, None); }); }); } diff --git a/app/src/ai/conversation_details_panel_tests.rs b/app/src/ai/conversation_details_panel_tests.rs index 001898a40..aecfd76a8 100644 --- a/app/src/ai/conversation_details_panel_tests.rs +++ b/app/src/ai/conversation_details_panel_tests.rs @@ -130,6 +130,7 @@ fn test_from_task_includes_linked_directory_when_run_id_matches() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.to_string()), autoexecute_override: None, last_event_sequence: None, @@ -245,6 +246,7 @@ fn test_from_task_includes_linked_directory_when_server_token_matches() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/app/src/pane_group/child_agent.rs b/app/src/pane_group/child_agent.rs index ed633684d..6d678115c 100644 --- a/app/src/pane_group/child_agent.rs +++ b/app/src/pane_group/child_agent.rs @@ -1,5 +1,7 @@ -use std::{collections::HashMap, ffi::OsString}; +use std::{collections::HashMap, ffi::OsString, path::PathBuf}; +use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::ai::attachment_utils::attachments_download_dir; use warpui::{EntityId, SingletonEntity, ViewContext, ViewHandle}; use crate::ai::agent::conversation::{AIConversationId, ConversationStatus}; @@ -15,6 +17,31 @@ pub(crate) struct HiddenChildAgentConversation { pub terminal_view_id: EntityId, pub conversation_id: AIConversationId, } +#[derive(Clone, Debug)] +pub(crate) struct HiddenChildAgentTaskContext { + pub task_id: AmbientAgentTaskId, + pub working_dir: Option, +} + +pub(crate) fn apply_hidden_child_agent_task_context( + terminal_view: &ViewHandle, + task_context: &HiddenChildAgentTaskContext, + ctx: &mut ViewContext, +) { + let task_id = task_context.task_id; + let working_dir = task_context.working_dir.clone(); + + terminal_view.update(ctx, move |terminal_view, ctx| { + terminal_view + .ai_controller() + .update(ctx, |controller, ctx| { + controller.set_ambient_agent_task_id(Some(task_id), ctx); + if let Some(working_dir) = working_dir.as_deref() { + controller.set_attachments_download_dir(attachments_download_dir(working_dir)); + } + }); + }); +} fn propagate_parent_agent_settings( group: &PaneGroup, @@ -72,6 +99,7 @@ pub(crate) fn create_hidden_child_agent_conversation( name: String, parent_conversation_id: AIConversationId, env_vars: HashMap, + task_context: Option, ctx: &mut ViewContext, ) -> Option { let new_pane_id = @@ -84,6 +112,9 @@ pub(crate) fn create_hidden_child_agent_conversation( let terminal_view_id = new_terminal_view.id(); propagate_parent_agent_settings(group, parent_pane_id, terminal_view_id, ctx); + if let Some(task_context) = task_context.as_ref() { + apply_hidden_child_agent_task_context(&new_terminal_view, task_context, ctx); + } let conversation_id = start_new_child_conversation(terminal_view_id, name, parent_conversation_id, ctx); @@ -117,6 +148,7 @@ fn create_error_child_agent_conversation_context( name.clone(), parent_conversation_id, HashMap::new(), + None, ctx, ) { return Some((Some(terminal_view), terminal_view_id, conversation_id)); diff --git a/app/src/pane_group/mod.rs b/app/src/pane_group/mod.rs index 3cc4772ae..81d865a84 100644 --- a/app/src/pane_group/mod.rs +++ b/app/src/pane_group/mod.rs @@ -170,6 +170,7 @@ pub mod focus_state; pub mod pane; pub mod tree; pub mod working_directories; +use child_agent::{apply_hidden_child_agent_task_context, HiddenChildAgentTaskContext}; use focus_state::PaneGroupFocusState; @@ -212,6 +213,25 @@ const MINIMUM_PANE_SIZE: f32 = 50.; const MINIMUM_PANE_SIZE_UDI: f32 = 190.; const KEYBOARD_RESIZE_DELTA: f32 = 10.; +type AmbientAgentViewModelHandle = + ModelHandle; + +trait AmbientAgentViewModelHandleExt<'a> { + fn into_optional_handle(self) -> Option<&'a AmbientAgentViewModelHandle>; +} + +impl<'a> AmbientAgentViewModelHandleExt<'a> for &'a AmbientAgentViewModelHandle { + fn into_optional_handle(self) -> Option<&'a AmbientAgentViewModelHandle> { + Some(self) + } +} + +impl<'a> AmbientAgentViewModelHandleExt<'a> for Option<&'a AmbientAgentViewModelHandle> { + fn into_optional_handle(self) -> Option<&'a AmbientAgentViewModelHandle> { + self + } +} + fn get_minimum_pane_size(app: &AppContext) -> f32 { use crate::settings::InputSettings; if InputSettings::as_ref(app).is_universal_developer_input_enabled(app) { @@ -3083,10 +3103,75 @@ impl PaneGroup { ctx: &mut ViewContext, ) { let child_id = child_conversation.id(); + if child_conversation.is_remote_child() { + let Some(task_id) = child_conversation.task_id() else { + log::warn!( + "Cannot restore remote child conversation {child_id:?} without a task ID" + ); + return; + }; + + let new_pane_id = + self.insert_ambient_agent_pane_hidden_for_child_agent(parent_pane_id, ctx); + + if let Some(new_terminal_view) = self.terminal_view_from_pane_id(new_pane_id, ctx) { + let mut restored = false; + new_terminal_view.update(ctx, |terminal_view, ctx| { + terminal_view.restore_conversation_after_view_creation( + RestoredAIConversation::new(child_conversation), + true, + ctx, + ); + terminal_view.enter_agent_view( + None, + Some(child_id), + AgentViewEntryOrigin::CloudAgent, + ctx, + ); + let Some(ambient_agent_view_model) = terminal_view + .ambient_agent_view_model() + .into_optional_handle() + .cloned() + else { + return; + }; + ambient_agent_view_model.update(ctx, |model, ctx| { + model.set_conversation_id(Some(child_id)); + model.enter_viewing_existing_session(task_id, ctx); + }); + restored = true; + }); + if restored { + self.child_agent_panes.insert(child_id, new_pane_id.into()); + } else { + log::error!( + "Failed to restore remote child agent pane {child_id:?}: missing ambient agent view model" + ); + self.discard_pane(new_pane_id.into(), ctx); + } + } else { + log::error!("Failed to get terminal view for remote child agent pane {child_id:?}"); + self.discard_pane(new_pane_id.into(), ctx); + } + return; + } + let child_task_context = + child_conversation + .task_id() + .map(|task_id| HiddenChildAgentTaskContext { + task_id, + working_dir: child_conversation + .current_working_directory() + .or_else(|| child_conversation.initial_working_directory()) + .map(PathBuf::from), + }); let new_pane_id = self.insert_terminal_pane_hidden_for_child_agent(parent_pane_id, HashMap::new(), ctx); if let Some(new_terminal_view) = self.terminal_view_from_pane_id(new_pane_id, ctx) { + if let Some(task_context) = child_task_context.as_ref() { + apply_hidden_child_agent_task_context(&new_terminal_view, task_context, ctx); + } new_terminal_view.update(ctx, |terminal_view, ctx| { terminal_view.restore_conversation_after_view_creation( RestoredAIConversation::new(child_conversation), diff --git a/app/src/pane_group/mod_tests.rs b/app/src/pane_group/mod_tests.rs index aa63d7d59..756dc0f6d 100644 --- a/app/src/pane_group/mod_tests.rs +++ b/app/src/pane_group/mod_tests.rs @@ -2,9 +2,18 @@ use crate::terminal::cli_agent_sessions::CLIAgentSessionsModel; use crate::{ ai::{ active_agent_views_model::ActiveAgentViewsModel, + agent::{ + conversation::{AIConversation, AIConversationId}, + PassiveSuggestionTrigger, + }, agent_conversations_model::AgentConversationsModel, ambient_agents::github_auth_notifier::GitHubAuthNotifier, - blocklist::BlocklistAIHistoryModel, + ambient_agents::AmbientAgentTaskId, + blocklist::{ + orchestration_event_streamer::OrchestrationEventStreamer, + orchestration_events::OrchestrationEventService, + task_status_sync_model::TaskStatusSyncModel, BlocklistAIHistoryModel, + }, document::ai_document_model::AIDocumentModel, execution_profiles::profiles::AIExecutionProfilesModel, llms::LLMPreferences, @@ -18,6 +27,7 @@ use crate::{ AIRequestUsageModel, }, auth::auth_manager::AuthManager, + changelog_model::ChangelogModel, cloud_object::model::persistence::CloudModel, context_chips::prompt::Prompt, experiments, @@ -64,8 +74,11 @@ use crate::{ use repo_metadata::RepoMetadataModel; use repo_metadata::{repositories::DetectedRepositories, watcher::DirectoryWatcher}; use std::collections::HashMap; +use uuid::Uuid; +use warp_core::features::FeatureFlag; use watcher::HomeDirectoryWatcher; +use super::child_agent::{create_hidden_child_agent_conversation, HiddenChildAgentTaskContext}; use super::*; use crate::terminal::resizable_data::ResizableData; use ai::{ @@ -84,6 +97,7 @@ fn initialize_app(app: &mut App) { initialize_settings_for_tests(app); app.add_singleton_model(|_ctx| ServerApiProvider::new_for_test()); + app.add_singleton_model(|ctx| ChangelogModel::new(ServerApiProvider::as_ref(ctx).get())); app.add_singleton_model(|_| AuthStateProvider::new_for_test()); app.add_singleton_model(AppTelemetryContextProvider::new_context_provider); app.add_singleton_model(AuthManager::new_for_test); @@ -124,7 +138,13 @@ fn initialize_app(app: &mut App) { app.add_singleton_model(TerminalKeybindings::new); app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); app.add_singleton_model(|_| CLIAgentSessionsModel::new()); + app.add_singleton_model(OrchestrationEventService::new); + app.add_singleton_model(TaskStatusSyncModel::new); + if FeatureFlag::OrchestrationV2.is_enabled() { + app.add_singleton_model(OrchestrationEventStreamer::new); + } app.add_singleton_model(|_| ActiveAgentViewsModel::new()); + app.add_singleton_model(crate::ai::blocklist::BlocklistAIPermissions::new); app.add_singleton_model(AgentNotificationsModel::new); app.add_singleton_model(|ctx| { AIExecutionProfilesModel::new(&crate::LaunchMode::new_for_unit_test(), ctx) @@ -222,6 +242,72 @@ fn new_notebook(ctx: &mut ViewContext) -> ViewHandle { ctx.add_typed_action_view(NotebookView::new) } +fn new_ambient_agent_task_id() -> AmbientAgentTaskId { + Uuid::new_v4().to_string().parse().unwrap() +} + +fn start_parent_conversation( + panes: &PaneGroup, + parent_pane_id: PaneId, + ctx: &mut ViewContext, +) -> AIConversationId { + let parent_terminal_view_id = panes + .terminal_view_from_pane_id(parent_pane_id, ctx) + .expect("parent pane should have a terminal view") + .id(); + + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_conversation(parent_terminal_view_id, false, false, ctx) + }) +} + +fn request_ambient_agent_task_id_for_hidden_child( + panes: &PaneGroup, + child_conversation_id: AIConversationId, + child_pane_id: PaneId, + ctx: &mut ViewContext, +) -> Option { + let terminal_view = panes + .terminal_view_from_pane_id(child_pane_id, ctx) + .expect("child pane should have a terminal view"); + let ai_controller = terminal_view.as_ref(ctx).ai_controller().clone(); + + ai_controller.update(ctx, |controller, ctx| { + controller + .build_passive_suggestions_request_params( + Some(child_conversation_id), + PassiveSuggestionTrigger::FilesChanged, + vec![], + ctx, + ) + .expect("child pane should build passive suggestion request params") + .1 + .ambient_agent_task_id + }) +} + +fn ambient_child_session_state( + panes: &PaneGroup, + child_pane_id: PaneId, + ctx: &mut ViewContext, +) -> (Option, bool, Option) { + let terminal_view = panes + .terminal_view_from_pane_id(child_pane_id, ctx) + .expect("child pane should have a terminal view"); + let terminal_view_ref = terminal_view.as_ref(ctx); + let active_conversation_id = terminal_view_ref.active_conversation_id(ctx); + let ambient_model = terminal_view_ref + .ambient_agent_view_model() + .expect("child pane should have an ambient agent model") + .as_ref(ctx); + + ( + ambient_model.task_id(), + ambient_model.is_agent_running(), + active_conversation_id, + ) +} + struct PreAttachReturnsFalsePane { pane_id: PaneId, pane_configuration: ModelHandle, @@ -376,6 +462,131 @@ fn test_insert_hidden_child_agent_pane_keeps_focus_and_active_session() { }); } +#[test] +fn test_hidden_child_creation_applies_ambient_task_id_to_controller() { + let _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(true); + + App::test((), |mut app| async move { + initialize_app(&mut app); + let pane_group = mock_pane_group(&mut app, Default::default()); + + pane_group.update(&mut app, |panes, ctx| { + let parent_pane_id = get_newly_created_pane_id(panes, &[]); + let parent_conversation_id = start_parent_conversation(panes, parent_pane_id, ctx); + let task_id = new_ambient_agent_task_id(); + + let child = create_hidden_child_agent_conversation( + panes, + parent_pane_id, + "Agent 1".to_string(), + parent_conversation_id, + HashMap::new(), + Some(HiddenChildAgentTaskContext { + task_id, + working_dir: None, + }), + ctx, + ) + .expect("fresh hidden child conversation should be created"); + + let child_pane_id = panes + .child_agent_panes + .get(&child.conversation_id) + .copied() + .expect("fresh hidden child pane should be tracked"); + + assert_eq!( + request_ambient_agent_task_id_for_hidden_child( + panes, + child.conversation_id, + child_pane_id, + ctx, + ), + Some(task_id) + ); + }); + }); +} + +#[test] +fn test_restored_hidden_child_pane_reapplies_ambient_task_id_to_controller() { + let _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(true); + + App::test((), |mut app| async move { + initialize_app(&mut app); + let pane_group = mock_pane_group(&mut app, Default::default()); + + pane_group.update(&mut app, |panes, ctx| { + let parent_pane_id = get_newly_created_pane_id(panes, &[]); + let parent_conversation_id = start_parent_conversation(panes, parent_pane_id, ctx); + let task_id = new_ambient_agent_task_id(); + + let mut child_conversation = AIConversation::new(false); + child_conversation.set_parent_conversation_id(parent_conversation_id); + child_conversation.set_task_id(task_id); + let child_conversation_id = child_conversation.id(); + + panes.create_hidden_child_agent_pane(child_conversation, parent_pane_id, ctx); + + let child_pane_id = panes + .child_agent_panes + .get(&child_conversation_id) + .copied() + .expect("restored hidden child pane should be tracked"); + + assert_eq!( + request_ambient_agent_task_id_for_hidden_child( + panes, + child_conversation_id, + child_pane_id, + ctx, + ), + Some(task_id) + ); + }); + }); +} + +#[test] +fn test_restored_remote_hidden_child_pane_enters_existing_ambient_session() { + let _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(true); + + App::test((), |mut app| async move { + initialize_app(&mut app); + let pane_group = mock_pane_group(&mut app, Default::default()); + + pane_group.update(&mut app, |panes, ctx| { + let parent_pane_id = get_newly_created_pane_id(panes, &[]); + let parent_conversation_id = start_parent_conversation(panes, parent_pane_id, ctx); + let task_id = new_ambient_agent_task_id(); + + let mut child_conversation = AIConversation::new(false); + child_conversation.set_parent_conversation_id(parent_conversation_id); + child_conversation.set_task_id(task_id); + child_conversation.mark_as_remote_child(); + let child_conversation_id = child_conversation.id(); + + panes.create_hidden_child_agent_pane(child_conversation, parent_pane_id, ctx); + + let child_pane_id = panes + .child_agent_panes + .get(&child_conversation_id) + .copied() + .expect("restored remote hidden child pane should be tracked"); + + let (ambient_task_id, is_agent_running, active_conversation_id) = + ambient_child_session_state(panes, child_pane_id, ctx); + + assert_eq!(ambient_task_id, Some(task_id)); + assert!( + is_agent_running, + "remote child restore should view the existing ambient session" + ); + assert_eq!(active_conversation_id, Some(child_conversation_id)); + }); + }); +} + #[test] fn test_active_session_id_reset_on_last_pane_close() { App::test((), |mut app| async move { diff --git a/app/src/pane_group/pane/local_harness_launch.rs b/app/src/pane_group/pane/local_harness_launch.rs index d1535e8c6..abd6f4f9e 100644 --- a/app/src/pane_group/pane/local_harness_launch.rs +++ b/app/src/pane_group/pane/local_harness_launch.rs @@ -89,7 +89,6 @@ pub(super) async fn prepare_local_harness_child_launch( let command = match harness { Harness::Oz => unreachable!("normalize_local_child_harness filters out Oz"), Harness::Unknown => unreachable!("normalize_local_child_harness filters out Unknown"), - Harness::Codex => unreachable!("normalize_local_child_harness filters out Codex"), Harness::Claude => { let working_dir = startup_directory .or_else(|| std::env::current_dir().ok()) @@ -126,6 +125,7 @@ pub(super) async fn prepare_local_harness_child_launch( .map_err(|error: AgentDriverError| error.to_string())?; build_local_opencode_child_command(&prompt) } + Harness::Codex => unreachable!("normalize_local_child_harness filters out Codex"), Harness::Gemini => unreachable!("normalize_local_child_harness filters out Gemini"), }; diff --git a/app/src/pane_group/pane/terminal_pane.rs b/app/src/pane_group/pane/terminal_pane.rs index 1b697f36d..b857dd832 100644 --- a/app/src/pane_group/pane/terminal_pane.rs +++ b/app/src/pane_group/pane/terminal_pane.rs @@ -56,6 +56,8 @@ use crate::{ #[cfg(feature = "local_fs")] use crate::ai::blocklist::BlocklistAIHistoryEvent; #[cfg(not(target_family = "wasm"))] +use crate::pane_group::child_agent::HiddenChildAgentTaskContext; +#[cfg(not(target_family = "wasm"))] use crate::server::server_api::ServerApiProvider; use warp_core::execution_mode::AppExecutionMode; @@ -1129,6 +1131,7 @@ fn handle_terminal_view_event( request.name, request.parent_conversation_id, HashMap::new(), + None, ctx, ) { register_legacy_local_lifecycle_subscription( @@ -1164,6 +1167,7 @@ fn handle_terminal_view_event( } => { let startup_directory = group.startup_path_for_new_session(Some(terminal_pane_id), ctx); + let launch_startup_directory = startup_directory.clone(); let ai_client = ServerApiProvider::handle(ctx).as_ref(ctx).get_ai_client(); let parent_pane_id = pane_id; let request_name = request.name.clone(); @@ -1183,7 +1187,7 @@ fn handle_terminal_view_event( harness_type, parent_run_id, shell_type, - startup_directory, + launch_startup_directory, ai_client, ) .await @@ -1207,6 +1211,10 @@ fn handle_terminal_view_event( request_name.clone(), parent_conversation_id, env_vars, + Some(HiddenChildAgentTaskContext { + task_id, + working_dir: startup_directory.clone(), + }), ctx, ) { BlocklistAIHistoryModel::handle(ctx).update( @@ -1223,16 +1231,16 @@ fn handle_terminal_view_event( ); new_terminal_view.update(ctx, |terminal_view, ctx| { - terminal_view.execute_command_or_set_pending( - &command, - ctx, - ); terminal_view.enter_agent_view( None, Some(conversation_id), AgentViewEntryOrigin::ChildAgent, ctx, ); + terminal_view.execute_command_or_set_pending( + &command, + ctx, + ); }); } else { create_error_child_agent_conversation( diff --git a/app/src/server/server_api.rs b/app/src/server/server_api.rs index 0c31d7d41..ecc7de31c 100644 --- a/app/src/server/server_api.rs +++ b/app/src/server/server_api.rs @@ -490,6 +490,16 @@ impl ServerApi { .collect()) } + async fn ambient_agent_headers_for_task( + &self, + task_id: &AmbientAgentTaskId, + ) -> Result> { + let mut headers = self.ambient_agent_headers().await?; + headers.retain(|(name, _)| *name != CLOUD_AGENT_ID_HEADER); + headers.push((CLOUD_AGENT_ID_HEADER, task_id.to_string())); + Ok(headers) + } + fn create_oauth_client() -> self::auth::OAuth2Client { let server_root = Url::parse(&ChannelState::server_root_url()).expect("Server root URL must be valid"); @@ -711,6 +721,40 @@ impl ServerApi { Ok(request.eventsource()) } + pub async fn stream_agent_events_for_task( + &self, + task_id: &AmbientAgentTaskId, + run_ids: &[String], + since_sequence: i64, + ) -> Result { + debug_assert!(!run_ids.is_empty(), "run_ids must not be empty"); + let auth_token = self + .get_or_refresh_access_token() + .await + .context("Failed to get access token for SSE stream")?; + + let run_ids_param: String = run_ids + .iter() + .map(|id| format!("run_ids[]={}", urlencoding::encode(id))) + .collect::>() + .join("&"); + let url = format!( + "{}/api/v1/agent/events/stream?{run_ids_param}&since={since_sequence}", + ChannelState::rtc_http_url() + ); + + let mut request = self.client.get(&url); + if let Some(token) = auth_token.as_bearer_token() { + request = request.bearer_auth(token); + } + + for (name, value) in self.ambient_agent_headers_for_task(task_id).await? { + request = request.header(name, value); + } + + Ok(request.eventsource()) + } + /// Sends a POST request to a public API endpoint and returns the raw response on success. async fn post_public_api_response( &self, diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index ca0d64f2b..0b3fe2c28 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -1008,6 +1008,73 @@ fn into_file_artifact_record( } } +impl ServerApi { + pub(crate) async fn send_agent_message_for_task( + &self, + task_id: &AmbientAgentTaskId, + request: SendAgentMessageRequest, + ) -> anyhow::Result { + let response = self + .post_public_api_response_for_task(task_id, "agent/messages", &request) + .await?; + let response = response.json::().await?; + Ok(response) + } + + #[cfg_attr(target_family = "wasm", allow(dead_code))] + pub(crate) async fn list_agent_messages_for_task( + &self, + task_id: &AmbientAgentTaskId, + run_id: &str, + request: ListAgentMessagesRequest, + ) -> anyhow::Result, anyhow::Error> { + let mut params = vec![format!("limit={}", request.limit)]; + if request.unread_only { + params.push("unread=true".to_string()); + } + if let Some(since) = request.since { + params.push(format!("since={}", urlencoding::encode(&since))); + } + + let path = format!("agent/messages/{run_id}?{}", params.join("&")); + let response = self + .get_public_api_response_for_task(task_id, &path) + .await?; + let response = response.json::>().await?; + Ok(response) + } + + pub(crate) async fn mark_message_delivered_for_task( + &self, + task_id: &AmbientAgentTaskId, + message_id: &str, + ) -> anyhow::Result<(), anyhow::Error> { + self.post_public_api_response_for_task( + task_id, + &format!("agent/messages/{message_id}/delivered"), + &(), + ) + .await?; + Ok(()) + } + + pub(crate) async fn read_agent_message_for_task( + &self, + task_id: &AmbientAgentTaskId, + message_id: &str, + ) -> anyhow::Result { + let response = self + .post_public_api_response_for_task( + task_id, + &format!("agent/messages/{message_id}/read"), + &(), + ) + .await?; + let response = response.json::().await?; + Ok(response) + } +} + #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl AIClient for ServerApi { @@ -1926,6 +1993,7 @@ impl AIClient for ServerApi { struct UpdateBody { sequence: i64, } + self.patch_public_api_unit( &format!("agent/runs/{run_id}/event-sequence"), &UpdateBody { sequence }, diff --git a/app/src/server/server_api/ai_test.rs b/app/src/server/server_api/ai_test.rs index 1a859dbf7..d38aaea21 100644 --- a/app/src/server/server_api/ai_test.rs +++ b/app/src/server/server_api/ai_test.rs @@ -1,6 +1,9 @@ use chrono::TimeZone; use chrono::Utc; +use futures::executor::block_on; +use super::super::auth::CLOUD_AGENT_ID_HEADER; +use super::super::ServerApi; use super::{ build_list_agent_runs_url, build_run_followup_url, AgentMessageHeader, AgentRunEvent, AgentSource, AmbientAgentTaskState, Artifact, ArtifactDownloadResponse, ArtifactType, @@ -9,6 +12,27 @@ use super::{ }; use crate::notebooks::NotebookId; +#[test] +fn ambient_agent_headers_for_task_overrides_existing_cloud_agent_header() { + let server_api = ServerApi::new_for_test(); + let ambient_task_id = "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(); + let task_scoped_id = "123e4567-e89b-12d3-a456-426614174000".parse().unwrap(); + + server_api.set_ambient_agent_task_id(Some(ambient_task_id)); + + let cloud_agent_headers: Vec<_> = + block_on(server_api.ambient_agent_headers_for_task(&task_scoped_id)) + .unwrap() + .into_iter() + .filter(|(name, _)| *name == CLOUD_AGENT_ID_HEADER) + .collect(); + + assert_eq!( + cloud_agent_headers, + vec![(CLOUD_AGENT_ID_HEADER, task_scoped_id.to_string())] + ); +} + #[test] fn test_deserialize_file_artifact_download_response() { let json = r#"{ diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index 7f0a982c5..301f6fd4a 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -12,7 +12,9 @@ use super::ServerApi; use crate::ai::agent::conversation::AIConversationId; #[cfg(not(target_family = "wasm"))] use crate::ai::agent_sdk::retry::with_bounded_retry; +use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::artifacts::Artifact; +use crate::server::server_api::auth::AuthClient; /// A presigned upload target returned by the server. #[derive(Debug, Clone, serde::Deserialize)] @@ -168,6 +170,119 @@ pub trait HarnessSupportClient: 'static + Send + Sync { fn http_client(&self) -> &http_client::Client; } +impl ServerApi { + pub(crate) async fn get_public_api_response_for_task( + &self, + task_id: &AmbientAgentTaskId, + path: &str, + ) -> Result { + let auth_token = self + .get_or_refresh_access_token() + .await + .context("Failed to get access token for API request")?; + + let url = format!("{}/api/v1/{}", crate::ChannelState::server_root_url(), path); + + let mut request = self.client.get(&url); + if let Some(token) = auth_token.as_bearer_token() { + request = request.bearer_auth(token); + } + + for (name, value) in self.ambient_agent_headers_for_task(task_id).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(response) + } else { + Err(Self::error_from_response(response).await) + } + } + + pub(crate) async fn post_public_api_response_for_task( + &self, + task_id: &AmbientAgentTaskId, + path: &str, + body: &B, + ) -> Result + where + B: serde::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/{}", crate::ChannelState::server_root_url(), path); + + let mut request = self.client.post(&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_for_task(task_id).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(response) + } else { + Err(Self::error_from_response(response).await) + } + } + + pub(crate) async fn resolve_prompt_for_task( + &self, + task_id: &AmbientAgentTaskId, + request: ResolvePromptRequest, + ) -> Result { + let response = self + .post_public_api_response_for_task(task_id, "harness-support/resolve-prompt", &request) + .await?; + let url = response.url().clone(); + response + .json::() + .await + .with_context(|| format!("Failed to deserialize response from {url}")) + } + + pub(crate) async fn fetch_transcript_for_task( + &self, + task_id: &AmbientAgentTaskId, + ) -> Result { + #[cfg(not(target_family = "wasm"))] + { + with_bounded_retry("fetch task-scoped harness-support transcript", || async { + let response = self + .get_public_api_response_for_task(task_id, "harness-support/transcript") + .await?; + response + .bytes() + .await + .context("Failed to read task-scoped harness-support transcript body") + }) + .await + } + #[cfg(target_family = "wasm")] + { + let _ = task_id; + unreachable!( + "fetch_transcript_for_task is not supported on wasm; agent_sdk is not built on this target" + ); + } + } +} + #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl HarnessSupportClient for ServerApi { diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index 26b715863..3fe41af41 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -4693,6 +4693,9 @@ impl TerminalView { if let BlocklistAIControllerEvent::SentRequest { model_id, .. } = event { self.maybe_insert_aws_bedrock_login_banner(model_id, ctx); } + if let BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { command } = event { + self.execute_command_or_set_pending(command, ctx); + } if let BlocklistAIControllerEvent::FinishedReceivingOutput { conversation_id, .. } = event @@ -11729,6 +11732,32 @@ impl TerminalView { } } + fn child_conversation_id_for_cli_status_updates( + &self, + ctx: &AppContext, + ) -> Option { + if let Some(conversation_id) = BlocklistAIHistoryModel::as_ref(ctx) + .active_conversation(self.view_id) + .and_then(|conversation| { + conversation + .is_child_agent_conversation() + .then_some(conversation.id()) + }) + { + return Some(conversation_id); + } + + let mut child_conversation_ids = BlocklistAIHistoryModel::as_ref(ctx) + .all_live_conversations_for_terminal_view(self.view_id) + .filter(|conversation| conversation.is_child_agent_conversation()) + .map(|conversation| conversation.id()); + let child_conversation_id = child_conversation_ids.next()?; + child_conversation_ids + .next() + .is_none() + .then_some(child_conversation_id) + } + /// If the startup auto-open setting is enabled, auto-opens rich input for a /// CLI agent session. Called after creating a command-detected session or /// registering a listener so rich input is shown immediately. @@ -11812,14 +11841,7 @@ impl TerminalView { return; } - let active_child_conversation_id = BlocklistAIHistoryModel::as_ref(ctx) - .active_conversation(self.view_id) - .and_then(|conversation| { - conversation - .is_child_agent_conversation() - .then_some(conversation.id()) - }); - if let Some(conversation_id) = active_child_conversation_id { + if let Some(conversation_id) = self.child_conversation_id_for_cli_status_updates(ctx) { BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { history_model.update_conversation_status( self.view_id, diff --git a/app/src/terminal/view/load_ai_conversation.rs b/app/src/terminal/view/load_ai_conversation.rs index bb4c506b2..f99e45019 100644 --- a/app/src/terminal/view/load_ai_conversation.rs +++ b/app/src/terminal/view/load_ai_conversation.rs @@ -998,6 +998,7 @@ impl TerminalView { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/app/src/terminal/view_test.rs b/app/src/terminal/view_test.rs index 73047e3d6..3c1f9794b 100644 --- a/app/src/terminal/view_test.rs +++ b/app/src/terminal/view_test.rs @@ -4305,6 +4305,88 @@ fn cli_session_status_updates_active_child_conversation() { }) } +#[test] +fn cli_session_status_updates_single_child_conversation_without_agent_view() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let _agent_view = FeatureFlag::AgentView.override_enabled(true); + + let terminal = add_window_with_terminal(&mut app, None); + + let child_conversation_id = terminal.update(&mut app, |view, ctx| { + let parent_conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_conversation(view.view_id, false, false, ctx) + }); + let child_conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_child_conversation( + view.view_id, + "Agent 2".to_string(), + parent_conversation_id, + ctx, + ) + }); + + CLIAgentSessionsModel::handle(ctx).update(ctx, |sessions, ctx| { + sessions.set_session( + view.view_id, + CLIAgentSession { + agent: CLIAgent::Claude, + status: CLIAgentSessionStatus::InProgress, + session_context: CLIAgentSessionContext::default(), + input_state: CLIAgentInputState::Closed, + should_auto_toggle_input: false, + listener: None, + remote_host: None, + plugin_version: None, + draft_text: None, + custom_command_prefix: None, + }, + ctx, + ); + }); + + child_conversation_id + }); + + terminal.read(&app, |_view, ctx| { + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&child_conversation_id) + .expect("child conversation should exist"); + assert_eq!(conversation.status(), &ConversationStatus::InProgress); + }); + + terminal.update(&mut app, |view, ctx| { + CLIAgentSessionsModel::handle(ctx).update(ctx, |sessions, ctx| { + sessions.update_from_event( + view.view_id, + &CLIAgentEvent { + v: 1, + agent: CLIAgent::Claude, + event: CLIAgentEventType::Stop, + session_id: None, + cwd: None, + project: None, + payload: CLIAgentEventPayload { + response: Some("Done".to_owned()), + ..Default::default() + }, + }, + ctx, + ); + }); + }); + + terminal.read(&app, |_view, ctx| { + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&child_conversation_id) + .expect("child conversation should exist"); + assert_eq!(conversation.status(), &ConversationStatus::Success); + }); + }) +} + #[test] fn manual_dismiss_disables_auto_toggle_for_session() { App::test((), |mut app| async move { diff --git a/crates/http_client/src/lib.rs b/crates/http_client/src/lib.rs index f10875d09..936925b8b 100644 --- a/crates/http_client/src/lib.rs +++ b/crates/http_client/src/lib.rs @@ -203,16 +203,16 @@ impl Client { ) } - pub fn patch(&self, url: U) -> RequestBuilder<'_> { + pub fn put(&self, url: U) -> RequestBuilder<'_> { self.builder( - self.wrapped.patch(url.clone()), + self.wrapped.put(url.clone()), Self::include_warp_http_headers(url), ) } - pub fn put(&self, url: U) -> RequestBuilder<'_> { + pub fn patch(&self, url: U) -> RequestBuilder<'_> { self.builder( - self.wrapped.put(url.clone()), + self.wrapped.patch(url.clone()), Self::include_warp_http_headers(url), ) } diff --git a/crates/persistence/src/model.rs b/crates/persistence/src/model.rs index f99220c60..e5069c619 100644 --- a/crates/persistence/src/model.rs +++ b/crates/persistence/src/model.rs @@ -1005,6 +1005,10 @@ impl<'de> Deserialize<'de> for PersistedAutoexecuteMode { }) } } +fn is_false(value: &bool) -> bool { + !*value +} + // Serializes to `conversation_data` column in `agent_conversations`. #[derive(Debug, Serialize, Deserialize, Clone)] pub struct AgentConversationData { @@ -1029,6 +1033,10 @@ pub struct AgentConversationData { /// The local conversation ID of the parent conversation. #[serde(default, skip_serializing_if = "Option::is_none")] pub parent_conversation_id: Option, + /// True when this conversation is a parent-side placeholder for a child + /// agent executing on a remote worker. + #[serde(default, skip_serializing_if = "is_false")] + pub is_remote_child: bool, /// The server-assigned run identifier (`ai_tasks.id`) for v2 orchestration. /// For local agents this arrives via StreamInit; for cloud agents it will /// come from SpawnAgentResponse once the local→cloud spawn path is wired. @@ -1347,6 +1355,7 @@ mod tests { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: Some(42), @@ -1356,6 +1365,27 @@ mod tests { assert_eq!(roundtripped.last_event_sequence, Some(42)); } + #[test] + fn agent_conversation_data_roundtrips_remote_child_marker() { + 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, + is_remote_child: true, + run_id: None, + autoexecute_override: None, + last_event_sequence: None, + }; + let json = serde_json::to_string(&data).expect("serialize"); + let roundtripped: AgentConversationData = serde_json::from_str(&json).expect("deserialize"); + assert!(roundtripped.is_remote_child); + } + #[test] fn agent_conversation_data_deserializes_legacy_payload_without_last_event_sequence() { // Legacy rows persisted before this feature landed omit the field @@ -1364,6 +1394,7 @@ mod tests { let data: AgentConversationData = serde_json::from_str(legacy_json).expect("legacy rows must deserialize"); assert_eq!(data.last_event_sequence, None); + assert!(!data.is_remote_child); } #[test] @@ -1377,6 +1408,7 @@ mod tests { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/crates/warp_cli/src/agent.rs b/crates/warp_cli/src/agent.rs index 3d4de67c1..155dccfa6 100644 --- a/crates/warp_cli/src/agent.rs +++ b/crates/warp_cli/src/agent.rs @@ -284,6 +284,10 @@ pub struct RunAgentArgs { #[command(flatten)] pub snapshot: SnapshotArgs, /// Identifier for the task that spawned this agent, used to report progress. + /// + /// When `--conversation` is omitted, the conversation id is read off the server-side + /// task metadata. Some worker follow-up call sites still pass both flags, so keep + /// accepting the compatibility shape until all producers have been updated. #[arg(long = "task-id", hide = true, conflicts_with_all = ["prompt", "saved_prompt", "file"])] pub task_id: Option, diff --git a/crates/warp_cli/src/lib_tests.rs b/crates/warp_cli/src/lib_tests.rs index 7ee73def0..f64ec3a18 100644 --- a/crates/warp_cli/src/lib_tests.rs +++ b/crates/warp_cli/src/lib_tests.rs @@ -1401,6 +1401,30 @@ fn agent_run_cloud_accepts_snapshot_flags() { ); } +#[test] +fn agent_run_accepts_task_id_with_conversation_for_worker_followups() { + let args = Args::try_parse_from([ + "warp", + "agent", + "run", + "--task-id", + "task-123", + "--conversation", + "conv-123", + ]) + .unwrap(); + + let Some(Command::CommandLine(boxed_cmd)) = args.command else { + panic!("Expected `warp agent run` command"); + }; + let CliCommand::Agent(AgentCommand::Run(run_args)) = boxed_cmd.as_ref() else { + panic!("Expected `warp agent run` command"); + }; + + assert_eq!(run_args.task_id.as_deref(), Some("task-123")); + assert_eq!(run_args.conversation.as_deref(), Some("conv-123")); +} + #[test] fn agent_run_cloud_accepts_computer_use_flag() { let args = Args::try_parse_from([