From 5fae21c4eabe1e8e36c98bf9c111a0b4cdb3359a Mon Sep 17 00:00:00 2001 From: Isaac Yu Date: Mon, 26 Oct 2026 12:01:11 -0700 Subject: [PATCH 1/2] feat: agents idle while hosting an active preview url Co-Authored-By: Warp --- app/src/ai/agent_sdk/driver.rs | 118 ++++++-- .../ai/agent_sdk/driver/preview_retention.rs | 98 +++++++ .../driver/preview_retention_tests.rs | 272 ++++++++++++++++++ app/src/ai/agent_sdk/mod.rs | 2 + app/src/ai/agent_sdk/mod_tests.rs | 1 + app/src/server/server_api/ai.rs | 23 ++ app/src/server/server_api/ai_tests.rs | 68 ++++- 7 files changed, 551 insertions(+), 31 deletions(-) create mode 100644 app/src/ai/agent_sdk/driver/preview_retention.rs create mode 100644 app/src/ai/agent_sdk/driver/preview_retention_tests.rs diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index a8f41723a9b..483d58d8ba0 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -117,6 +117,7 @@ pub(crate) mod harness; mod harness_output_monitor; pub(crate) mod mcp_startup; pub(super) mod output; +mod preview_retention; mod prewarmed; mod snapshot; pub(crate) mod terminal; @@ -124,6 +125,7 @@ mod termination; use environment::PrepareEnvironmentError; pub(crate) use mcp_startup::MCP_SERVER_STARTUP_TIMEOUT; +use preview_retention::PreviewRetention; pub(crate) use snapshot::upload_snapshot_for_handoff; use terminal::TerminalDriverEvent; use termination::{Interrupt, InterruptWatch, RunEndCause}; @@ -278,6 +280,7 @@ pub(crate) struct IdleTimeoutSender { /// exact moment of commitment, rather than only after the completion reaches the model /// thread via the oneshot and any further async plumbing (QUALITY-1801). on_commit: Arc, + preview_retention: Option, } // Hand-written so cloning does not require `T: Clone`. Every field is a shared handle, so @@ -290,6 +293,7 @@ impl Clone for IdleTimeoutSender { pending: Arc::clone(&self.pending), wait: Arc::clone(&self.wait), on_commit: Arc::clone(&self.on_commit), + preview_retention: self.preview_retention.clone(), } } } @@ -303,6 +307,7 @@ impl IdleTimeoutSender { pending: Arc::new(Mutex::new(None)), wait: Arc::new(RealIdleWait), on_commit: Arc::new(|| {}), + preview_retention: None, } } @@ -318,6 +323,7 @@ impl IdleTimeoutSender { if let Ok(mut guard) = self.tx_cell.lock() && let Some(sender) = guard.take() { + self.generation.fetch_add(1, Ordering::SeqCst); (self.on_commit)(); let _ = sender.send(value); } @@ -325,34 +331,63 @@ impl IdleTimeoutSender { /// End the run after `timeout` by sending `value`, unless cancelled before then. pub(crate) fn end_run_after(&self, timeout: Duration, value: T) { + self.schedule_idle_timeout(timeout, value, false); + } + + fn schedule_idle_timeout(&self, timeout: Duration, value: T, retain_previews: bool) { // Increment the generation counter to invalidate any existing timers, // then capture the new generation for our timer to check against. - let current_gen = self.generation.fetch_add(1, Ordering::SeqCst) + 1; - let tx_cell = Arc::clone(&self.tx_cell); - let generation = Arc::clone(&self.generation); + let current_gen = { + let Ok(_guard) = self.tx_cell.lock() else { + return; + }; + self.generation.fetch_add(1, Ordering::SeqCst) + 1 + }; + let sender = self.clone(); let wait = Arc::clone(&self.wait); - let on_commit = Arc::clone(&self.on_commit); // Spawn a background thread that will complete the oneshot after the idle timeout, // unless a follow-up query resets the timer (by bumping the generation counter). thread::spawn(move || { wait.wait(timeout); - - // Check if our timer generation is still current. If not, a follow-up - // query or other activity has "cancelled" this timer by bumping the generation. - if generation.load(Ordering::SeqCst) != current_gen { + if retain_previews && let Some(retention) = sender.preview_retention.clone() { + let background = retention.background.clone(); + background + .spawn(async move { + sender + .finish_retained_idle(current_gen, value, retention) + .await; + }) + .detach(); + } else { + sender.commit_idle_timeout(current_gen, value); + } + }); + } + async fn finish_retained_idle(&self, generation: usize, value: T, retention: PreviewRetention) { + while self.generation.load(Ordering::SeqCst) == generation { + let delay = (retention.check)().await; + if self.generation.load(Ordering::SeqCst) != generation { return; } - if let Ok(mut guard) = tx_cell.lock() - && let Some(sender) = guard.take() - { - // Commit before sending: this is the only signal the model layer ever gets - // that a deferred window has elapsed, so it must land before the completion - // is observable at all (QUALITY-1801). - on_commit(); - let _ = sender.send(value); + match delay { + Some(delay) => (retention.wait)(delay).await, + None => { + self.commit_idle_timeout(generation, value); + return; + } } - }); + } + } + + fn commit_idle_timeout(&self, generation: usize, value: T) { + if let Ok(mut guard) = self.tx_cell.lock() + && self.generation.load(Ordering::SeqCst) == generation + && let Some(sender) = guard.take() + { + (self.on_commit)(); + let _ = sender.send(value); + } } /// Cancel any pending idle timers. @@ -363,7 +398,8 @@ impl IdleTimeoutSender { if let Ok(mut pending) = self.pending.lock() { *pending = None; } - if self.generation.load(Ordering::SeqCst) > 0 { + // Serialize invalidation with the final generation check and synchronous exit commitment. + if let Ok(_guard) = self.tx_cell.lock() { self.generation.fetch_add(1, Ordering::SeqCst); } } @@ -376,7 +412,7 @@ impl IdleTimeoutSender { self.cancel_idle_timeout(); self.end_run_now(value); } - Some(idle_timeout) => self.end_run_after(idle_timeout, value), + Some(idle_timeout) => self.schedule_idle_timeout(idle_timeout, value, true), None => self.end_run_now(value), } } @@ -388,10 +424,14 @@ impl IdleTimeoutSender { /// Re-arming replaces the recorded outcome, so a run that fails, resumes, and fails again /// exits reporting its most recent failure. fn arm_refreshable(&self, window: Duration, value: T) { + if window.is_zero() { + self.complete_with_optional_idle(Some(window), value); + return; + } if let Ok(mut pending) = self.pending.lock() { *pending = Some((value.clone(), window)); } - self.end_run_after(window, value); + self.schedule_idle_timeout(window, value, true); } /// Push an armed deadline out by its original window. Returns that window, or `None` if @@ -401,7 +441,7 @@ impl IdleTimeoutSender { let pending = self.pending.lock().ok()?; pending.clone()? }; - self.end_run_after(window, value); + self.schedule_idle_timeout(window, value, true); Some(window) } } @@ -654,6 +694,8 @@ pub struct AgentDriverOptions { pub environment_variables: HashMap, /// ID of the task being executed. pub task_id: Option, + /// Execution identity supplied by the cloud launch. + pub execution_id: Option, pub use_factory_repositories: bool, /// Server-owned experiments retained without client interpretation. pub experimental: Option>, @@ -719,6 +761,7 @@ struct RunConfiguration { preparation_file: Option, secrets: HashMap, task_id: Option, + execution_id: Option, use_factory_repositories: bool, experimental: Option>, parent_run_id: Option, @@ -755,6 +798,7 @@ fn split_run_options( secrets, environment_variables, task_id, + execution_id, use_factory_repositories, experimental, parent_run_id, @@ -874,6 +918,7 @@ fn split_run_options( preparation_file, secrets, task_id, + execution_id, use_factory_repositories, experimental, parent_run_id, @@ -943,6 +988,7 @@ pub struct AgentDriver { // The associated task ID for this agent run, if any. task_id: Option, + execution_id: Option, use_factory_repositories: bool, #[allow( dead_code, @@ -1329,6 +1375,25 @@ impl From for AgentDriverError { } impl AgentDriver { + fn idle_timeout_sender( + &self, + tx: oneshot::Sender, + ctx: &ModelContext, + ) -> IdleTimeoutSender { + let mut sender = IdleTimeoutSender::new(tx); + if preview_retention::is_namespace_sandbox(warp_isolation_platform::detect()) + && let Some(task_id) = self.task_id + { + sender.preview_retention = PreviewRetention::new( + task_id, + self.execution_id.as_deref(), + ServerApiProvider::as_ref(ctx).get_ai_client(), + ctx.background_executor(), + ); + } + sender + } + /// Creates the driver and terminal without running workspace setup or submitting a prompt. #[tracing::instrument(name = "AgentDriver::create", skip_all, err, fields( tags.cloud_agent = true, @@ -1433,6 +1498,7 @@ impl AgentDriver { _preparation_file: None, output_format: OutputFormat::default(), task_id: None, + execution_id: None, use_factory_repositories: false, experimental: None, team_scope: None, @@ -1479,6 +1545,7 @@ impl AgentDriver { preparation_file, secrets, task_id, + execution_id, use_factory_repositories, experimental, parent_run_id, @@ -1506,6 +1573,7 @@ impl AgentDriver { self._preparation_file = preparation_file; self.secrets = Arc::new(secrets); self.task_id = task_id; + self.execution_id = execution_id; self.use_factory_repositories = use_factory_repositories; self.factory_repository_contract = use_factory_repositories; self.experimental = experimental; @@ -2496,8 +2564,8 @@ impl AgentDriver { Self::report_failure_before_lingering(foreground, stage, error, window).await; let (tx, rx) = oneshot::channel::<()>(); - let controller = DebugWindowController::new(IdleTimeoutSender::new(tx)); let armed = foreground.spawn(move |me, ctx| { + let controller = DebugWindowController::new(me.idle_timeout_sender(tx, ctx)); me.arm_setup_failure_debug_window(controller, window, ctx); }); if let Err(error) = armed.await { @@ -2833,7 +2901,9 @@ impl AgentDriver { events: &SetupClientEventReporter, ) -> Result, AgentDriverError> { let (exit_tx, exit_rx) = oneshot::channel(); - let harness_exit = IdleTimeoutSender::new(exit_tx); + let harness_exit = foreground + .spawn(move |me, ctx| me.idle_timeout_sender(exit_tx, ctx)) + .await?; // Harnesses that report progress through the native conversation instead of a CLI agent // session have no hook plugin to install; the conversation's own status drives their @@ -3635,7 +3705,7 @@ impl AgentDriver { #[cfg(test)] let post_commit_gate = tests::test_post_commit_gate(); #[cfg_attr(not(test), allow(unused_mut))] - let mut run_exit = IdleTimeoutSender::new(internal_tx).with_on_commit({ + let mut run_exit = self.idle_timeout_sender(internal_tx, ctx).with_on_commit({ let committed_conversation_id = Arc::clone(&committed_conversation_id); move || { if let Ok(guard) = committed_conversation_id.lock() diff --git a/app/src/ai/agent_sdk/driver/preview_retention.rs b/app/src/ai/agent_sdk/driver/preview_retention.rs new file mode 100644 index 00000000000..dd84057c416 --- /dev/null +++ b/app/src/ai/agent_sdk/driver/preview_retention.rs @@ -0,0 +1,98 @@ +use std::future::Future; +use std::sync::Arc; +use std::time::Duration; + +use warp_isolation_platform::IsolationPlatformType; +use warpui::r#async::executor::Background; +use warpui::r#async::{BoxFuture, FutureExt as _, Timer}; + +use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::server::server_api::ai::{AIClient, ListRunPreviewsResponse}; + +const PREVIEW_CHECK_TIMEOUT: Duration = Duration::from_secs(5); +const PREVIEW_RECHECK_INTERVAL: Duration = Duration::from_secs(60); +pub(super) fn is_namespace_sandbox(platform: Option) -> bool { + match platform { + #[cfg(not(target_family = "wasm"))] + Some(IsolationPlatformType::Namespace) => true, + #[cfg(not(target_family = "wasm"))] + Some( + IsolationPlatformType::Docker + | IsolationPlatformType::DockerSandbox + | IsolationPlatformType::Kubernetes, + ) => false, + None => false, + } +} + +#[derive(Clone)] +pub(super) struct PreviewRetention { + pub background: Arc, + pub check: Arc BoxFuture<'static, Option> + Send + Sync>, + pub wait: Arc BoxFuture<'static, ()> + Send + Sync>, +} +async fn check_previews( + request: impl Future>, + expected_execution_id: Option, + timeout: Duration, +) -> Option { + match request.with_timeout(timeout).await { + Ok(Ok(response)) => retention_delay(&response, expected_execution_id), + Ok(Err(error)) => { + log::warn!("Preview idle retention check failed; allowing idle exit: {error:#}"); + None + } + Err(_) => { + log::warn!("Preview idle retention check timed out; allowing idle exit"); + None + } + } +} + +impl PreviewRetention { + pub fn new( + task_id: AmbientAgentTaskId, + execution_id: Option<&str>, + client: Arc, + background: Arc, + ) -> Option { + let expected_execution_id = match execution_id { + Some(raw) => Some(raw.parse::().ok().filter(|id| *id > 0)?), + None => None, + }; + Some(Self { + background, + check: Arc::new(move || { + let client = client.clone(); + Box::pin(async move { + check_previews( + client.list_run_previews(&task_id), + expected_execution_id, + PREVIEW_CHECK_TIMEOUT, + ) + .await + }) + }), + wait: Arc::new(|delay| { + Box::pin(async move { + Timer::after(delay).await; + }) + }), + }) + } +} + +fn retention_delay( + response: &ListRunPreviewsResponse, + expected_execution_id: Option, +) -> Option { + let execution_id = response.execution_id.filter(|id| *id > 0)?; + if expected_execution_id.is_some_and(|expected| expected != execution_id) { + return None; + } + (!response.previews.is_empty()).then_some(PREVIEW_RECHECK_INTERVAL) +} + +#[cfg(test)] +#[path = "preview_retention_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/driver/preview_retention_tests.rs b/app/src/ai/agent_sdk/driver/preview_retention_tests.rs new file mode 100644 index 00000000000..559ac18ef94 --- /dev/null +++ b/app/src/ai/agent_sdk/driver/preview_retention_tests.rs @@ -0,0 +1,272 @@ +use std::collections::VecDeque; +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use futures::channel::oneshot; +use futures::executor::block_on; +use futures::poll; +use serde_json::json; + +use super::super::IdleTimeoutSender; +use super::*; + +fn response(preview_count: usize) -> ListRunPreviewsResponse { + serde_json::from_value(json!({ + "execution_id": 42, + "previews": vec![json!({}); preview_count], + })) + .unwrap() +} + +#[test] +fn preview_retention_rejects_malformed_launch_execution_id() { + let task_id = "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(); + for raw in ["", "not-an-id", "0", "-1", "9223372036854775808"] { + assert!( + PreviewRetention::new( + task_id, + Some(raw), + Arc::new(crate::server::server_api::ai::MockAIClient::new()), + Arc::new(Background::new(1, |_| "preview-retention-test".into())), + ) + .is_none() + ); + } +} +#[test] +#[cfg(not(target_family = "wasm"))] +fn preview_retention_only_applies_to_namespace() { + assert!(!is_namespace_sandbox(None)); + assert!(!is_namespace_sandbox(Some(IsolationPlatformType::Docker))); + assert!(!is_namespace_sandbox(Some( + IsolationPlatformType::Kubernetes + ))); + assert!(is_namespace_sandbox(Some(IsolationPlatformType::Namespace))); + assert!(!is_namespace_sandbox(Some( + IsolationPlatformType::DockerSandbox + ))); +} + +fn retention( + check: impl Fn() -> BoxFuture<'static, Option> + Send + Sync + 'static, + wait: impl Fn(Duration) -> BoxFuture<'static, ()> + Send + Sync + 'static, +) -> PreviewRetention { + PreviewRetention { + background: Arc::new(Background::new(1, |_| "preview-retention-test".into())), + check: Arc::new(check), + wait: Arc::new(wait), + } +} + +#[test] +fn preview_retention_uses_server_filtered_presence_and_fixed_wait() { + for (preview_count, expected) in [ + (0, None), + (1, Some(Duration::from_secs(60))), + (3, Some(Duration::from_secs(60))), + ] { + assert_eq!( + retention_delay(&response(preview_count), Some(42)), + expected + ); + } + let mut no_execution = response(1); + no_execution.execution_id = None; + assert_eq!(retention_delay(&no_execution, None), None); + for execution_id in [Some(0), Some(-1), Some(43)] { + let mut other_execution = response(1); + other_execution.execution_id = execution_id; + assert_eq!(retention_delay(&other_execution, Some(42)), None); + } + assert_eq!( + retention_delay(&response(1), None), + Some(Duration::from_secs(60)) + ); +} + +#[test] +fn preview_retention_rechecks_after_wait_before_commit() { + block_on(async { + let (tx, mut rx) = oneshot::channel(); + let committed = Arc::new(AtomicUsize::new(0)); + let count = committed.clone(); + let sender = IdleTimeoutSender::new(tx).with_on_commit(move || { + count.fetch_add(1, Ordering::SeqCst); + }); + let checks = Arc::new(Mutex::new(VecDeque::from([ + Some(Duration::from_secs(60)), + None, + ]))); + let (release, gate) = oneshot::channel(); + let gate = Mutex::new(Some(gate)); + let retention = retention( + move || { + let result = checks.lock().unwrap().pop_front().unwrap(); + Box::pin(async move { result }) + }, + move |duration| { + assert_eq!(duration, Duration::from_secs(60)); + let gate = gate.lock().unwrap().take().unwrap(); + Box::pin(async move { + gate.await.unwrap(); + }) + }, + ); + let mut finish = Box::pin(sender.finish_retained_idle(0, 7, retention)); + assert!(poll!(finish.as_mut()).is_pending()); + assert_eq!(rx.try_recv().unwrap(), None); + assert_eq!(committed.load(Ordering::SeqCst), 0); + release.send(()).unwrap(); + finish.await; + assert_eq!(rx.try_recv().unwrap(), Some(7)); + assert_eq!(committed.load(Ordering::SeqCst), 1); + }); +} + +#[test] +fn preview_retention_followup_invalidates_in_flight_response() { + for response in [None, Some(Duration::from_secs(60))] { + block_on(async { + let (tx, mut rx) = oneshot::channel::(); + let committed = Arc::new(AtomicUsize::new(0)); + let count = committed.clone(); + let sender = IdleTimeoutSender::new(tx).with_on_commit(move || { + count.fetch_add(1, Ordering::SeqCst); + }); + let (release, gate) = oneshot::channel(); + let gate = Mutex::new(Some(gate)); + let retention = retention( + move || { + let gate = gate.lock().unwrap().take().unwrap(); + Box::pin(async move { gate.await.unwrap() }) + }, + |_| panic!("stale response must not retain a superseded turn"), + ); + let mut finish = Box::pin(sender.finish_retained_idle(0, 1, retention)); + assert!(poll!(finish.as_mut()).is_pending()); + sender.cancel_idle_timeout(); + release.send(response).unwrap(); + finish.await; + assert_eq!(rx.try_recv().unwrap(), None); + assert_eq!(committed.load(Ordering::SeqCst), 0); + sender.commit_idle_timeout(0, 2); + assert_eq!(rx.try_recv().unwrap(), None); + sender.end_run_now(3); + assert_eq!(rx.try_recv().unwrap(), Some(3)); + }); + } +} + +#[test] +fn preview_retention_explicit_exit_bypasses_in_flight_query() { + block_on(async { + let (tx, mut rx) = oneshot::channel(); + let sender = IdleTimeoutSender::new(tx); + let retention = retention(|| Box::pin(futures::future::pending()), |_| unreachable!()); + let mut finish = Box::pin(sender.finish_retained_idle(0, 1, retention)); + assert!(poll!(finish.as_mut()).is_pending()); + sender.end_run_now(2); + assert_eq!(rx.try_recv().unwrap(), Some(2)); + }); +} + +#[test] +fn preview_retention_none_zero_and_automatic_recovery_bypass_check() { + for timeout in [None, Some(Duration::ZERO)] { + let (tx, rx) = oneshot::channel(); + let mut sender = IdleTimeoutSender::new(tx); + sender.preview_retention = Some(retention( + || panic!("must bypass check"), + |_| unreachable!(), + )); + sender.complete_with_optional_idle(timeout, 5); + assert_eq!(block_on(rx).unwrap(), 5); + } + let (tx, rx) = oneshot::channel(); + let mut sender = IdleTimeoutSender::new(tx); + sender.preview_retention = Some(retention( + || panic!("must bypass check"), + |_| unreachable!(), + )); + sender.end_run_after(Duration::ZERO, 6); + assert_eq!(block_on(rx).unwrap(), 6); +} + +struct NoIdleWait; + +impl super::super::IdleWait for NoIdleWait { + fn wait(&self, _: Duration) {} +} + +struct SynchronizedIdleWait(std::sync::Barrier); + +impl super::super::IdleWait for SynchronizedIdleWait { + fn wait(&self, _: Duration) { + self.0.wait(); + } +} + +#[test] +fn preview_retention_refreshed_failure_window_checks_before_commit() { + let (tx, rx) = oneshot::channel(); + let checks = Arc::new(AtomicUsize::new(0)); + let checks_for_commit = checks.clone(); + let mut sender = IdleTimeoutSender::new(tx).with_on_commit(move || { + assert_eq!(checks_for_commit.load(Ordering::SeqCst), 1); + }); + let wait = Arc::new(SynchronizedIdleWait(std::sync::Barrier::new(3))); + sender.wait = wait.clone(); + sender.preview_retention = Some(retention( + move || { + checks.fetch_add(1, Ordering::SeqCst); + Box::pin(async { None }) + }, + |_| unreachable!(), + )); + sender.arm_refreshable(Duration::from_secs(10), 7); + assert_eq!(sender.refresh(), Some(Duration::from_secs(10))); + wait.0.wait(); + assert_eq!(block_on(rx).unwrap(), 7); +} + +#[test] +fn preview_retention_completion_and_failure_idle_check_before_commit() { + for failure_window in [false, true] { + let (tx, rx) = oneshot::channel(); + let checks = Arc::new(AtomicUsize::new(0)); + let checks_for_commit = checks.clone(); + let mut sender = IdleTimeoutSender::new(tx).with_on_commit(move || { + assert_eq!(checks_for_commit.load(Ordering::SeqCst), 1); + }); + sender.wait = Arc::new(NoIdleWait); + sender.preview_retention = Some(retention( + move || { + checks.fetch_add(1, Ordering::SeqCst); + Box::pin(async { None }) + }, + |_| unreachable!(), + )); + if failure_window { + sender.arm_refreshable(Duration::from_secs(10), 7); + } else { + sender.complete_with_optional_idle(Some(Duration::from_secs(10)), 7); + } + assert_eq!(block_on(rx).unwrap(), 7); + } +} +#[test] +fn preview_retention_query_failure_and_timeout_allow_exit() { + let failed = futures::future::ready(Err(anyhow::anyhow!("404 older server"))); + assert_eq!( + block_on(check_previews(failed, None, PREVIEW_CHECK_TIMEOUT)), + None + ); + assert_eq!( + block_on(check_previews( + futures::future::pending(), + None, + Duration::ZERO + )), + None + ); +} diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index 751b5c2330f..9b4fe45a810 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -725,6 +725,7 @@ fn build_execution_task_and_options( let options = AgentDriverOptions { working_dir: working_dir.clone(), task_id: Some(task_id), + execution_id: Some(execution_id.to_owned()), use_factory_repositories, experimental: None, parent_run_id: config.parent_run_id.map(|id| id.into_inner()), @@ -1533,6 +1534,7 @@ impl AgentDriverRunner { let driver_options = driver::AgentDriverOptions { working_dir: working_dir.clone(), task_id, + execution_id: args.execution_id.clone(), use_factory_repositories, experimental: None, parent_run_id: None, diff --git a/app/src/ai/agent_sdk/mod_tests.rs b/app/src/ai/agent_sdk/mod_tests.rs index 75f7dbc74b8..b756c92ec2e 100644 --- a/app/src/ai/agent_sdk/mod_tests.rs +++ b/app/src/ai/agent_sdk/mod_tests.rs @@ -650,6 +650,7 @@ pub(crate) fn agent_driver_options() -> AgentDriverOptions { AgentDriverOptions { working_dir: std::env::current_dir().unwrap(), task_id: None, + execution_id: None, use_factory_repositories: false, experimental: None, parent_run_id: None, diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index 06dc3ea2220..4eff91705d6 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -1003,6 +1003,15 @@ struct ListRunsResponse { runs: Vec, } +#[derive(Debug, serde::Deserialize)] +pub struct ListRunPreviewsResponse { + pub execution_id: Option, + pub previews: Vec, +} + +#[derive(Debug, serde::Deserialize)] +pub struct RunPreview {} + impl<'de> serde::Deserialize<'de> for ListRunsResponse { fn deserialize(deserializer: D) -> Result where @@ -1268,6 +1277,10 @@ pub struct UpdateMemoryResponse { #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] pub trait AIClient: 'static + Send + Sync { + async fn list_run_previews( + &self, + task_id: &AmbientAgentTaskId, + ) -> anyhow::Result; async fn generate_commands_from_natural_language( &self, prompt: String, @@ -2015,6 +2028,16 @@ fn convert_upload_field( #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl AIClient for ServerApi { + async fn list_run_previews( + &self, + task_id: &AmbientAgentTaskId, + ) -> anyhow::Result { + self.get_public_api_response_for_task(task_id, &format!("agent/runs/{task_id}/previews")) + .await? + .json::() + .await + .map_err(Into::into) + } async fn generate_commands_from_natural_language( &self, prompt: String, diff --git a/app/src/server/server_api/ai_tests.rs b/app/src/server/server_api/ai_tests.rs index b198b4ed237..f445fc76711 100644 --- a/app/src/server/server_api/ai_tests.rs +++ b/app/src/server/server_api/ai_tests.rs @@ -6,16 +6,19 @@ use serde_json::{Value, json}; use warp_graphql::ai::PlatformErrorCode; use warp_graphql::platform_error::{PlatformErrorInfo, PlatformErrorMessageFormat}; use warp_graphql::queries::execution_config::CodeForge; -use warp_server_client::base_client::{CLOUD_AGENT_ID_HEADER, TEAM_UID_HEADER}; +use warp_server_client::base_client::{ + AMBIENT_WORKLOAD_TOKEN_HEADER, CLOUD_AGENT_ID_HEADER, TEAM_UID_HEADER, +}; use super::super::ServerApi; use super::{ - AIClient, AgentMessageHeader, AgentRunEvent, AgentSource, AmbientAgentTaskState, Artifact, - ArtifactDownloadResponse, ArtifactType, CONNECTED_SELF_HOSTED_WORKERS_PATH, - ConnectedSelfHostedWorker, CreateAgentRequest, ExecutionLocation, ForkConversationResponse, - ListConnectedSelfHostedWorkersResponse, ListRunsResponse, PrepareAttachmentUploadsResponse, - ReadAgentMessageResponse, RunFollowupRequest, RunSortBy, RunSortOrder, SpawnAgentRequest, - TaskGitCredentialsError, TaskListFilter, TaskStatusUpdate, UploadFieldValue, UserQueryMode, + AIClient, AgentMessageHeader, AgentRunEvent, AgentSource, AmbientAgentTaskId, + AmbientAgentTaskState, Artifact, ArtifactDownloadResponse, ArtifactType, + CONNECTED_SELF_HOSTED_WORKERS_PATH, ConnectedSelfHostedWorker, CreateAgentRequest, + ExecutionLocation, ForkConversationResponse, ListConnectedSelfHostedWorkersResponse, + ListRunsResponse, PrepareAttachmentUploadsResponse, ReadAgentMessageResponse, + RunFollowupRequest, RunSortBy, RunSortOrder, SpawnAgentRequest, TaskGitCredentialsError, + TaskListFilter, TaskStatusUpdate, UploadFieldValue, UserQueryMode, agent_task_status_message_input, build_fork_conversation_url, build_list_agent_runs_url, build_run_followup_url, is_unknown_git_credential_schema_error, }; @@ -29,6 +32,57 @@ fn request_scope_for_team(team_uid: ServerId) -> RequestTeamScope { RequestTeamScope::from_scope(&TeamContextForOperation::new_for_test(team_uid)) } +#[test] +fn list_run_previews_uses_ambient_run_auth_and_ignores_preview_details() { + let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(); + let request = { + let mut server = warp_core::channel::ChannelState::mock_server(); + server.mock("GET", format!("/api/v1/agent/runs/{task_id}/previews").as_str()) + .match_header(CLOUD_AGENT_ID_HEADER, task_id.to_string().as_str()) + .match_header(AMBIENT_WORKLOAD_TOKEN_HEADER, "test-workload-token") + .with_status(200) + .with_body(r#"{"execution_id":42,"previews":[{"expires_at":"2026-10-22T18:00:00Z","port":3000}]}"#) + .expect(1) + .create() + }; + let api = execution_bootstrap_server_api(); + let response = block_on(api.list_run_previews(&task_id)).unwrap(); + assert_eq!(response.execution_id, Some(42)); + assert_eq!(response.previews.len(), 1); + request.assert(); +} + +#[test] +fn list_run_previews_no_execution_and_older_server() { + let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440001".parse().unwrap(); + for (status, body) in [ + (200, r#"{"previews":[]}"#), + (404, r#"{"error":"not found"}"#), + ] { + let request = { + let mut server = warp_core::channel::ChannelState::mock_server(); + server + .mock( + "GET", + format!("/api/v1/agent/runs/{task_id}/previews").as_str(), + ) + .with_status(status) + .with_body(body) + .expect(1) + .create() + }; + let result = block_on(execution_bootstrap_server_api().list_run_previews(&task_id)); + if status == 200 { + let response = result.unwrap(); + assert_eq!(response.execution_id, None); + assert!(response.previews.is_empty()); + } else { + assert!(result.is_err()); + } + request.assert(); + } +} + fn execution_bootstrap_response() -> Value { json!({ "data": { From 377305e072358f646632491096bbbd3abbbe564a Mon Sep 17 00:00:00 2001 From: Isaac Yu Date: Mon, 26 Oct 2026 14:14:49 -0700 Subject: [PATCH 2/2] fix: test cases --- app/src/ai/agent_sdk/driver.rs | 9 +++-- .../driver/preview_retention_tests.rs | 34 +++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index 483d58d8ba0..802fe1277e6 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -350,7 +350,10 @@ impl IdleTimeoutSender { // unless a follow-up query resets the timer (by bumping the generation counter). thread::spawn(move || { wait.wait(timeout); - if retain_previews && let Some(retention) = sender.preview_retention.clone() { + if retain_previews + && !timeout.is_zero() + && let Some(retention) = sender.preview_retention.clone() + { let background = retention.background.clone(); background .spawn(async move { @@ -424,10 +427,6 @@ impl IdleTimeoutSender { /// Re-arming replaces the recorded outcome, so a run that fails, resumes, and fails again /// exits reporting its most recent failure. fn arm_refreshable(&self, window: Duration, value: T) { - if window.is_zero() { - self.complete_with_optional_idle(Some(window), value); - return; - } if let Ok(mut pending) = self.pending.lock() { *pending = Some((value.clone(), window)); } diff --git a/app/src/ai/agent_sdk/driver/preview_retention_tests.rs b/app/src/ai/agent_sdk/driver/preview_retention_tests.rs index 559ac18ef94..e077a14549b 100644 --- a/app/src/ai/agent_sdk/driver/preview_retention_tests.rs +++ b/app/src/ai/agent_sdk/driver/preview_retention_tests.rs @@ -206,6 +206,40 @@ impl super::super::IdleWait for SynchronizedIdleWait { } } +#[test] +fn preview_retention_zero_failure_window_bypasses_check() { + let (tx, rx) = oneshot::channel(); + let mut sender = IdleTimeoutSender::new(tx); + sender.wait = Arc::new(NoIdleWait); + sender.preview_retention = Some(retention( + || panic!("zero window must bypass check"), + |_| unreachable!(), + )); + + sender.arm_refreshable(Duration::ZERO, 7); + + assert_eq!(block_on(rx).unwrap(), 7); +} + +#[test] +fn preview_retention_zero_failure_window_remains_refreshable_without_check() { + let (tx, mut rx) = oneshot::channel(); + let mut sender = IdleTimeoutSender::new(tx); + let wait = Arc::new(SynchronizedIdleWait(std::sync::Barrier::new(3))); + sender.wait = wait.clone(); + sender.preview_retention = Some(retention( + || panic!("refreshed zero window must bypass check"), + |_| unreachable!(), + )); + + sender.arm_refreshable(Duration::ZERO, 7); + + assert_eq!(rx.try_recv().unwrap(), None); + assert_eq!(sender.refresh(), Some(Duration::ZERO)); + wait.0.wait(); + assert_eq!(block_on(rx).unwrap(), 7); +} + #[test] fn preview_retention_refreshed_failure_window_checks_before_commit() { let (tx, rx) = oneshot::channel();