From 0b6d1263ea2b3dffd69e29de73871931d8d5b83c Mon Sep 17 00:00:00 2001 From: Joey Wang Date: Thu, 30 Jul 2026 23:49:44 -0400 Subject: [PATCH 1/2] Add periodic checkpoint coordinator state machine Adds the Idle -> Due -> InFlight -> Idle scheduling loop (with Finalizing -> Stopped reachable from any state) that will drive periodic workspace-handoff checkpoints, per REMOTE-2111 Phase 3. - checkpoint_coordinator.rs: CheckpointCoordinatorHandle spawns the loop and exposes finalize() for graceful shutdown. The timer only ever moves Idle to Due; Due polls a safe-boundary predicate (conversation not mid-turn) before entering InFlight, which runs exactly one attempt via snapshot::run_checkpoint_from_declarations_file on a background task so a racing finalize() can bound how long it waits without ever stranding the attempt. finalize() starts at most one more attempt (skipped below the gather/upload timeout floor) or awaits an already-in-flight one -- never both -- bounded end to end by its budget. - terminal/view.rs: add TerminalView::ai_action_model(), which is_safe_boundary uses to check for unfinished actions on the conversation. This module has no production caller yet -- AgentDriver wires CheckpointCoordinatorHandle::new into its spawn/finalize lifecycle in a follow-up, stacked PR -- so it is marked allow(dead_code) at the module level with a comment explaining why; removable once that PR merges on top of this one. Co-Authored-By: Oz --- app/src/ai/agent_sdk/driver.rs | 1 + .../driver/checkpoint_coordinator.rs | 456 +++++++++++++ .../driver/checkpoint_coordinator_tests.rs | 600 ++++++++++++++++++ app/src/terminal/view.rs | 7 + 4 files changed, 1064 insertions(+) create mode 100644 app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs create mode 100644 app/src/ai/agent_sdk/driver/checkpoint_coordinator_tests.rs diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index e3741bd2bd0..70524631e25 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -99,6 +99,7 @@ use crate::terminal::view::ConversationRestorationInNewPaneType; pub(crate) mod attachments; #[cfg(feature = "local_fs")] pub(crate) mod cache_setup; +mod checkpoint_coordinator; pub(crate) mod cloud_provider; pub(crate) mod environment; mod error_classification; diff --git a/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs b/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs new file mode 100644 index 00000000000..050ffb26196 --- /dev/null +++ b/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs @@ -0,0 +1,456 @@ +//! Periodic workspace-handoff checkpoint coordinator (REMOTE-2111 Phase 3). +//! +//! Drives the five-state machine described in `docs/remote-2111-checkpoint-spec.md` +//! (warp-server): `Idle -> Due -> InFlight -> Idle` on the periodic path, with +//! `Finalizing -> Stopped` reachable from `Idle`, `Due`, or `InFlight` via +//! [`CheckpointCoordinatorHandle::finalize`]. The timer only ever moves `Idle` to +//! `Due`; all gather/upload/commit work happens through +//! `super::snapshot::run_checkpoint_from_declarations_file`, reusing the same +//! declarations file and gather/upload pipeline as the legacy end-of-run snapshot. +//! +//! Safe-boundary gating ("only touch the filesystem/network when the conversation +//! isn't mid-turn") is implemented as a bounded poll of `AgentDriver`'s own state via +//! its `ModelSpawner`, rather than a push subscription: `AgentDriver` already reads +//! exactly the state needed (`run_conversation_id`, the terminal view's action model) +//! through this same read-only, spawner-based pattern used by `run_snapshot_upload`. +//! This trades a small amount of latency (up to [`SAFE_BOUNDARY_POLL_INTERVAL`]) for +//! avoiding new push-subscription wiring through the UI model graph. + +// This module has no production caller yet -- `AgentDriver` wires +// `CheckpointCoordinatorHandle::new` into its spawn/finalize lifecycle in a +// follow-up, stacked PR. Until then, `CheckpointCoordinatorHandle::new` (and +// everything it reaches) is unreachable outside `#[cfg(test)]`, which reaches the +// same code through `new_for_test`. This module-level allow is temporary and should +// be removable once that PR merges on top of this one. +#![allow(dead_code)] + +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use futures::FutureExt as _; +use futures::future::BoxFuture; +use instant::Instant; +use rand::Rng as _; +use tokio::sync::{mpsc, oneshot}; +use warpui::r#async::executor::Background; +use warpui::r#async::{FutureExt as _, Timer}; +use warpui::{ModelSpawner, SingletonEntity}; + +use super::AgentDriver; +use super::snapshot::{self, CheckpointResult}; +use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::ai::blocklist::BlocklistAIHistoryModel; +use crate::server::server_api::harness_support::{CheckpointGeneration, HarnessSupportClient}; + +/// A safe-boundary predicate, decoupled from `ModelSpawner` so +/// [`coordinator_loop`] can be exercised in isolation by tests. Production code builds +/// this from [`is_safe_boundary`]; tests supply a directly-controllable closure. +type BoundaryCheck = Arc BoxFuture<'static, bool> + Send + Sync>; + +/// Default cadence between the start of one checkpoint attempt and the timer firing +/// again, absent an override on `AgentDriverOptions`. Deliberately not +/// `HARNESS_SAVE_INTERVAL` (30s) -- see the spec's "5 minutes plus jitter" cadence. +pub(super) const DEFAULT_CHECKPOINT_INTERVAL: Duration = Duration::from_secs(5 * 60); +/// Upper bound on additive jitter, so agents scheduled at the same time don't all +/// checkpoint in lockstep. +const CHECKPOINT_JITTER: Duration = Duration::from_secs(30); +/// How often the `Due` state re-checks whether the conversation is at a safe boundary. +const SAFE_BOUNDARY_POLL_INTERVAL: Duration = Duration::from_secs(2); + +/// A request to finalize, carrying the deadline by which the coordinator must ack so +/// shutdown can proceed. +struct FinalizeRequest { + deadline: Instant, + ack: oneshot::Sender<()>, +} + +/// Handle used by `AgentDriver` to request finalization of the periodic checkpoint +/// coordinator. Cloneable and fire-and-forget: dropping every handle without calling +/// [`finalize`](Self::finalize) simply leaves the coordinator running periodic +/// attempts until the process exits. +#[derive(Clone)] +pub(super) struct CheckpointCoordinatorHandle { + finalize_tx: mpsc::UnboundedSender, +} + +impl CheckpointCoordinatorHandle { + /// Spawn the coordinator task on `background` and return a handle to it. + #[allow(clippy::too_many_arguments)] + pub(super) fn new( + client: Arc, + task_id: AmbientAgentTaskId, + working_dir: PathBuf, + spawner: ModelSpawner, + interval: Duration, + script_timeout: Duration, + upload_timeout: Duration, + background: Arc, + ) -> Self { + let boundary_check: BoundaryCheck = Arc::new(move || { + let spawner = spawner.clone(); + Box::pin(async move { is_safe_boundary(&spawner).await }) + }); + Self::spawn_with_boundary_check( + client, + task_id, + working_dir, + boundary_check, + interval, + CHECKPOINT_JITTER, + script_timeout, + upload_timeout, + background, + ) + } + + /// Test-facing constructor that bypasses `ModelSpawner` (and so the full + /// UI framework) by taking the safe-boundary predicate directly, and disables jitter + /// (production jitter is bounded by [`CHECKPOINT_JITTER`], up to 30s, which would + /// otherwise make tests using a short `interval` flaky/slow). + #[cfg(test)] + #[allow(clippy::too_many_arguments)] + fn new_for_test( + client: Arc, + task_id: AmbientAgentTaskId, + working_dir: PathBuf, + boundary_check: BoundaryCheck, + interval: Duration, + script_timeout: Duration, + upload_timeout: Duration, + background: Arc, + ) -> Self { + Self::spawn_with_boundary_check( + client, + task_id, + working_dir, + boundary_check, + interval, + Duration::ZERO, + script_timeout, + upload_timeout, + background, + ) + } + + #[allow(clippy::too_many_arguments)] + fn spawn_with_boundary_check( + client: Arc, + task_id: AmbientAgentTaskId, + working_dir: PathBuf, + boundary_check: BoundaryCheck, + interval: Duration, + jitter: Duration, + script_timeout: Duration, + upload_timeout: Duration, + background: Arc, + ) -> Self { + let (finalize_tx, finalize_rx) = mpsc::unbounded_channel(); + let loop_background = background.clone(); + background + .spawn(coordinator_loop( + client, + task_id, + working_dir, + boundary_check, + interval, + jitter, + script_timeout, + upload_timeout, + loop_background, + finalize_rx, + )) + .detach(); + Self { finalize_tx } + } + + /// Request finalization: run at most one more checkpoint attempt if none is + /// already in flight (skipped if `budget` doesn't exceed the gather/upload + /// floor), or await an already-in-flight attempt instead -- never both -- then + /// stop the coordinator. Bounded by `budget` end to end. Safe to call at most + /// once; safe to never call. + pub(super) async fn finalize(&self, budget: Duration) { + let (ack_tx, ack_rx) = oneshot::channel(); + let request = FinalizeRequest { + deadline: Instant::now() + budget, + ack: ack_tx, + }; + if self.finalize_tx.send(request).is_err() { + // Coordinator task already exited; nothing to wait for. + return; + } + // The coordinator always acks well within `budget` (either immediately, when + // below the floor, or after its own internally bounded attempt). This extra + // bound is defense-in-depth so a coordinator bug cannot wedge shutdown. + let _ = tokio::time::timeout(budget, ack_rx).await; + } +} + +/// Add up to `jitter` of additive random delay to `interval` so many agents scheduled at once +/// don't checkpoint in lockstep. Production always passes [`CHECKPOINT_JITTER`]; tests pass +/// `Duration::ZERO` for determinism. +fn jittered_interval(interval: Duration, jitter: Duration) -> Duration { + let jitter_ms = u64::try_from(jitter.as_millis()).unwrap_or(u64::MAX); + let extra = if jitter_ms == 0 { + 0 + } else { + rand::thread_rng().gen_range(0..=jitter_ms) + }; + interval + Duration::from_millis(extra) +} + +/// Run one checkpoint attempt to completion: regenerate declarations, then gather, +/// upload, and commit. Bounded by `upload_timeout`; `script_timeout` separately +/// bounds only the declarations-script sub-step (matching the legacy pipeline). +async fn run_one_attempt( + client: Arc, + task_id: AmbientAgentTaskId, + working_dir: PathBuf, + script_timeout: Duration, + upload_timeout: Duration, +) -> CheckpointResult { + snapshot::run_declarations_script(&working_dir, &task_id, script_timeout).await; + let path = snapshot::resolve_declarations_path(Some(&task_id)); + match snapshot::run_checkpoint_from_declarations_file(&path, client) + .with_timeout(upload_timeout) + .await + { + Ok(result) => result, + Err(_) => CheckpointResult::Failed { + generation: None, + reason: format!("checkpoint attempt exceeded {upload_timeout:?} upload timeout"), + }, + } +} + +/// Spawn one attempt on `background` and return a receiver that resolves with its +/// result once the attempt completes. The spawned task runs to completion +/// independently of whether anything ever reads from the receiver, so a caller that +/// stops waiting (e.g. because a shutdown budget elapsed) cannot strand the attempt +/// or cause it to be silently abandoned mid-upload -- it simply keeps running in the +/// background and, if it succeeds, still commits. +fn start_attempt( + client: Arc, + task_id: AmbientAgentTaskId, + working_dir: PathBuf, + script_timeout: Duration, + upload_timeout: Duration, + background: &Background, +) -> oneshot::Receiver { + let (tx, rx) = oneshot::channel(); + background + .spawn(async move { + let result = + run_one_attempt(client, task_id, working_dir, script_timeout, upload_timeout).await; + let _ = tx.send(result); + }) + .detach(); + rx +} + +/// Query `AgentDriver` (via its spawner) for whether the conversation is currently at +/// a safe boundary: either yielded via `wait_for_events`, or simply not mid-turn with +/// no pending/running actions. Returns `true` (safe) if the driver has no +/// conversation yet, if its conversation can no longer be found, or if the driver +/// itself has been dropped -- in each case there is nothing left to interrupt. +async fn is_safe_boundary(spawner: &ModelSpawner) -> bool { + spawner + .spawn(|driver, ctx| { + let Some(conversation_id) = driver.run_conversation_id else { + return true; + }; + let Some(status) = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .map(|conversation| conversation.status().clone()) + else { + return true; + }; + if status.is_waiting_for_events() { + return true; + } + if status.is_in_progress() || status.is_transient_error() { + return false; + } + let terminal_view = driver + .terminal_driver + .as_ref(ctx) + .terminal_view() + .as_ref(ctx); + !terminal_view + .ai_action_model() + .as_ref(ctx) + .has_unfinished_actions_for_conversation(conversation_id) + }) + .await + .unwrap_or(true) +} + +/// Handle a finalize request received while no attempt is currently in flight: start +/// exactly one best-effort attempt only if `budget` exceeds the gather/upload floor, +/// bound it by the remaining budget, then ack. +async fn finalize_with_new_attempt( + request: FinalizeRequest, + client: Arc, + task_id: AmbientAgentTaskId, + working_dir: PathBuf, + script_timeout: Duration, + upload_timeout: Duration, +) { + let floor = script_timeout + upload_timeout; + let remaining = request.deadline.saturating_duration_since(Instant::now()); + if remaining > floor { + log::info!( + "Starting final checkpoint attempt at shutdown (remaining budget {remaining:?})" + ); + let attempt = run_one_attempt(client, task_id, working_dir, script_timeout, upload_timeout); + if tokio::time::timeout(remaining, attempt).await.is_err() { + log::warn!("Final checkpoint attempt did not complete within {remaining:?}"); + } + } else { + log::info!( + "Skipping final checkpoint attempt: remaining shutdown budget {remaining:?} \ + is below the {floor:?} floor" + ); + } + let _ = request.ack.send(()); +} + +/// Handle a finalize request received while an attempt started by the periodic +/// timer is already in flight: never start a second attempt -- just await the +/// existing one, bounded by the remaining budget, then ack. +async fn finalize_with_in_flight_attempt( + request: FinalizeRequest, + result_rx: oneshot::Receiver, +) { + let remaining = request.deadline.saturating_duration_since(Instant::now()); + match tokio::time::timeout(remaining, result_rx).await { + Ok(Ok(result)) => { + log::info!("In-flight checkpoint attempt resolved during finalization: {result:?}"); + } + Ok(Err(_)) => { + log::warn!("In-flight checkpoint attempt's result channel dropped without a result"); + } + Err(_) => { + // The spawned attempt keeps running in the background regardless; we + // just stop waiting for it so shutdown can proceed within budget. + log::warn!( + "In-flight checkpoint attempt did not resolve within the remaining \ + {remaining:?} shutdown budget; continuing shutdown without it" + ); + } + } + let _ = request.ack.send(()); +} + +/// The coordinator's main loop. `Idle` and `Due` are collapsed into the top of the +/// loop body: the timer is the only thing that ever moves `Idle` to `Due`, and `Due` +/// then polls the safe-boundary predicate. `InFlight` runs the attempt on a +/// background task (via [`start_attempt`]) specifically so a finalize request racing +/// in can bound how long it waits without ever stranding the attempt itself. +#[allow(clippy::too_many_arguments)] +async fn coordinator_loop( + client: Arc, + task_id: AmbientAgentTaskId, + working_dir: PathBuf, + boundary_check: BoundaryCheck, + interval: Duration, + jitter: Duration, + script_timeout: Duration, + upload_timeout: Duration, + background: Arc, + mut finalize_rx: mpsc::UnboundedReceiver, +) { + loop { + // --- Idle: wait for the next (jittered) tick or a finalize request. --- + futures::select! { + _ = Timer::after(jittered_interval(interval, jitter)).fuse() => {} + request = finalize_rx.recv().fuse() => { + let Some(request) = request else { return }; + finalize_with_new_attempt( + request, + client.clone(), + task_id, + working_dir.clone(), + script_timeout, + upload_timeout, + ) + .await; + return; + } + } + + // --- Due: poll for a safe boundary, staying responsive to finalize. Checked + // immediately on entry (not only after the first poll interval elapses) so an + // already-safe conversation doesn't pay needless latency. + let mut at_boundary = boundary_check().await; + while !at_boundary { + futures::select! { + _ = Timer::after(SAFE_BOUNDARY_POLL_INTERVAL).fuse() => { + at_boundary = boundary_check().await; + } + request = finalize_rx.recv().fuse() => { + let Some(request) = request else { return }; + finalize_with_new_attempt( + request, + client.clone(), + task_id, + working_dir.clone(), + script_timeout, + upload_timeout, + ) + .await; + return; + } + } + } + + // --- InFlight: run exactly one attempt, never overlapping another. --- + let mut result_rx = start_attempt( + client.clone(), + task_id, + working_dir.clone(), + script_timeout, + upload_timeout, + &background, + ); + futures::select! { + result = (&mut result_rx).fuse() => { + match result { + Ok(CheckpointResult::Committed { generation }) => { + log::info!( + "Periodic checkpoint committed: generation={}", + generation.as_str() + ); + } + Ok(CheckpointResult::Skipped) => { + log::info!("Periodic checkpoint skipped: no usable declarations"); + } + Ok(CheckpointResult::Failed { generation, reason }) => { + log::warn!( + "Periodic checkpoint attempt failed (generation={:?}): {reason}", + generation.as_ref().map(CheckpointGeneration::as_str) + ); + } + Err(_) => { + log::warn!( + "Periodic checkpoint attempt's result channel dropped without a result" + ); + } + } + // Success, skip, or failure: return to Idle and wait a full interval + // before the next attempt either way. The periodic timer itself + // (rather than a distinct short backoff) is the retry mechanism for + // failures too, matching the spec's "no RPO/RTO SLO" framing. + } + request = finalize_rx.recv().fuse() => { + let Some(request) = request else { return }; + finalize_with_in_flight_attempt(request, result_rx).await; + return; + } + } + } +} + +#[cfg(test)] +#[path = "checkpoint_coordinator_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/driver/checkpoint_coordinator_tests.rs b/app/src/ai/agent_sdk/driver/checkpoint_coordinator_tests.rs new file mode 100644 index 00000000000..a9cdd41be82 --- /dev/null +++ b/app/src/ai/agent_sdk/driver/checkpoint_coordinator_tests.rs @@ -0,0 +1,600 @@ +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::str::FromStr; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use anyhow::Result; +use async_trait::async_trait; +use instant::Instant; +use mockito::{Matcher, Server}; +use tempfile::TempDir; +use tokio::sync::oneshot; +use uuid::Uuid; + +use super::*; +use crate::ai::agent::conversation::AIConversationId; +use crate::ai::agent_sdk::test_support::build_test_http_client; +use crate::ai::artifacts::Artifact; +use crate::server::server_api::harness_support::{ + CommitSnapshotRequest, CommitSnapshotResponse, ReportArtifactResponse, ResolvePromptRequest, + ResolvedHarnessPrompt, SnapshotUploadRequest, UploadTarget, +}; + +fn fresh_task_id() -> AmbientAgentTaskId { + AmbientAgentTaskId::from_str(&Uuid::new_v4().to_string()).unwrap() +} + +fn test_background() -> Arc { + Arc::new(Background::new(1, |_| { + "checkpoint-coordinator-test".to_string() + })) +} + +/// Poll `condition` until it's true or `timeout` elapses, then assert it held. +async fn wait_for(mut condition: impl FnMut() -> bool, timeout: Duration) { + let deadline = Instant::now() + timeout; + while !condition() && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!(condition(), "condition not met within {timeout:?}"); +} + +/// Writes a single-file declarations JSONL at the default per-task-id path that +/// `run_one_attempt` resolves internally (it has no path-override hook), and removes the +/// per-task directory on drop so parallel tests never collide (each uses a fresh task id). +struct DeclarationsFixture { + dir: PathBuf, +} + +impl DeclarationsFixture { + fn new(task_id: &AmbientAgentTaskId, declared_file: &Path) -> Self { + let path = snapshot::resolve_declarations_path(Some(task_id)); + let dir = path.parent().unwrap().to_path_buf(); + std::fs::create_dir_all(&dir).unwrap(); + let mut contents = serde_json::json!({ + "version": 1, + "kind": "file", + "path": declared_file.to_string_lossy(), + }) + .to_string(); + contents.push('\n'); + std::fs::write(&path, contents).unwrap(); + Self { dir } + } +} + +impl Drop for DeclarationsFixture { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.dir); + } +} + +/// A `HarnessSupportClient` whose `get_snapshot_upload_targets` always fails immediately +/// (no HTTP involved), so `run_one_attempt` resolves to `Failed` cheaply. Used for tests that +/// only care about coordinator timing/gating, not upload success. `commit_snapshot` panics: +/// it must never be reached when uploads always fail. +struct FailingUploadTargetsClient { + http: http_client::Client, + call_count: AtomicUsize, +} + +impl FailingUploadTargetsClient { + fn new() -> Arc { + Arc::new(Self { + http: build_test_http_client(), + call_count: AtomicUsize::new(0), + }) + } +} + +#[async_trait] +impl HarnessSupportClient for FailingUploadTargetsClient { + async fn create_external_conversation(&self, _format: &str) -> Result { + unimplemented!("not used by the coordinator") + } + async fn get_transcript_upload_target( + &self, + _conversation_id: &AIConversationId, + ) -> Result { + unimplemented!("not used by the coordinator") + } + async fn get_block_snapshot_upload_target( + &self, + _conversation_id: &AIConversationId, + ) -> Result { + unimplemented!("not used by the coordinator") + } + async fn resolve_prompt( + &self, + _request: ResolvePromptRequest, + ) -> Result { + unimplemented!("not used by the coordinator") + } + async fn report_artifact(&self, _artifact: &Artifact) -> Result { + unimplemented!("not used by the coordinator") + } + async fn notify_user(&self, _message: &str) -> Result<()> { + unimplemented!("not used by the coordinator") + } + async fn finish_task(&self, _success: bool, _summary: &str) -> Result<()> { + unimplemented!("not used by the coordinator") + } + async fn report_clean_shutdown(&self) -> Result<()> { + unimplemented!("not used by the coordinator") + } + async fn report_error_shutdown( + &self, + _error_category: String, + _error_message: String, + ) -> Result<()> { + unimplemented!("not used by the coordinator") + } + async fn get_snapshot_upload_targets( + &self, + _request: &SnapshotUploadRequest, + ) -> Result> { + self.call_count.fetch_add(1, Ordering::SeqCst); + anyhow::bail!("simulated get_snapshot_upload_targets failure") + } + async fn commit_snapshot( + &self, + _request: &CommitSnapshotRequest, + ) -> Result { + panic!("commit_snapshot must never be called when uploads always fail"); + } + async fn fetch_transcript(&self) -> Result { + unimplemented!("not used by the coordinator") + } + fn http_client(&self) -> &http_client::Client { + &self.http + } +} + +/// A `HarnessSupportClient` that serves real presigned-URL uploads against a `mockito` server +/// and records call counts, with an optional artificial delay before `commit_snapshot` +/// resolves (used to keep an attempt "in flight" for finalization tests). +struct RecordingClient { + server_base_url: String, + http: http_client::Client, + commit_delay: Duration, + fail_commit: bool, + upload_calls: AtomicUsize, + commit_calls: AtomicUsize, +} + +impl RecordingClient { + fn new(server_base_url: String) -> Arc { + Self::with_options(server_base_url, Duration::ZERO, false) + } + + fn with_commit_delay(server_base_url: String, delay: Duration) -> Arc { + Self::with_options(server_base_url, delay, false) + } + + fn with_options( + server_base_url: String, + commit_delay: Duration, + fail_commit: bool, + ) -> Arc { + Arc::new(Self { + server_base_url, + http: build_test_http_client(), + commit_delay, + fail_commit, + upload_calls: AtomicUsize::new(0), + commit_calls: AtomicUsize::new(0), + }) + } +} + +#[async_trait] +impl HarnessSupportClient for RecordingClient { + async fn create_external_conversation(&self, _format: &str) -> Result { + unimplemented!("not used by the coordinator") + } + async fn get_transcript_upload_target( + &self, + _conversation_id: &AIConversationId, + ) -> Result { + unimplemented!("not used by the coordinator") + } + async fn get_block_snapshot_upload_target( + &self, + _conversation_id: &AIConversationId, + ) -> Result { + unimplemented!("not used by the coordinator") + } + async fn resolve_prompt( + &self, + _request: ResolvePromptRequest, + ) -> Result { + unimplemented!("not used by the coordinator") + } + async fn report_artifact(&self, _artifact: &Artifact) -> Result { + unimplemented!("not used by the coordinator") + } + async fn notify_user(&self, _message: &str) -> Result<()> { + unimplemented!("not used by the coordinator") + } + async fn finish_task(&self, _success: bool, _summary: &str) -> Result<()> { + unimplemented!("not used by the coordinator") + } + async fn report_clean_shutdown(&self) -> Result<()> { + unimplemented!("not used by the coordinator") + } + async fn report_error_shutdown( + &self, + _error_category: String, + _error_message: String, + ) -> Result<()> { + unimplemented!("not used by the coordinator") + } + async fn get_snapshot_upload_targets( + &self, + request: &SnapshotUploadRequest, + ) -> Result> { + self.upload_calls.fetch_add(1, Ordering::SeqCst); + Ok(request + .files + .iter() + .map(|f| UploadTarget { + url: format!("{}/upload/{}", self.server_base_url, f.filename), + method: "PUT".to_string(), + headers: HashMap::new(), + fields: Vec::new(), + }) + .collect()) + } + async fn commit_snapshot( + &self, + request: &CommitSnapshotRequest, + ) -> Result { + self.commit_calls.fetch_add(1, Ordering::SeqCst); + if !self.commit_delay.is_zero() { + tokio::time::sleep(self.commit_delay).await; + } + if self.fail_commit { + anyhow::bail!("simulated commit_snapshot failure"); + } + Ok(CommitSnapshotResponse { + generation: request.generation.clone(), + }) + } + async fn fetch_transcript(&self) -> Result { + unimplemented!("not used by the coordinator") + } + fn http_client(&self) -> &http_client::Client { + &self.http + } +} + +// ------------------------------------------------------------------------------------------------ +// Pure helpers. +// ------------------------------------------------------------------------------------------------ + +#[test] +fn jittered_interval_adds_bounded_nonnegative_jitter() { + let base = Duration::from_secs(60); + for _ in 0..50 { + let jittered = jittered_interval(base, CHECKPOINT_JITTER); + assert!(jittered >= base, "jitter must never shrink the interval"); + assert!( + jittered <= base + CHECKPOINT_JITTER, + "jitter must stay within the configured bound" + ); + } +} + +#[test] +fn jittered_interval_zero_jitter_is_a_no_op() { + let base = Duration::from_millis(5); + assert_eq!(jittered_interval(base, Duration::ZERO), base); +} + +// ------------------------------------------------------------------------------------------------ +// finalize_with_new_attempt / finalize_with_in_flight_attempt. +// ------------------------------------------------------------------------------------------------ + +#[tokio::test] +async fn finalize_with_new_attempt_below_floor_skips_the_attempt_entirely() { + let client = FailingUploadTargetsClient::new(); + let task_id = fresh_task_id(); + let (ack_tx, ack_rx) = oneshot::channel(); + let request = FinalizeRequest { + deadline: Instant::now() + Duration::from_millis(50), + ack: ack_tx, + }; + // floor = script_timeout + upload_timeout = 2s, comfortably above the 50ms budget. + finalize_with_new_attempt( + request, + client.clone(), + task_id, + PathBuf::from("/tmp"), + Duration::from_secs(1), + Duration::from_secs(1), + ) + .await; + assert_eq!( + client.call_count.load(Ordering::SeqCst), + 0, + "a below-floor budget must skip the attempt without calling the client at all" + ); + assert!(ack_rx.await.is_ok(), "ack must still be sent"); +} + +#[tokio::test] +async fn finalize_with_new_attempt_above_floor_starts_and_commits() { + let mut server = Server::new_async().await; + let _uploads = server + .mock("PUT", Matcher::Regex("^/upload/.+$".to_string())) + .with_status(200) + .create_async() + .await; + + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let client = RecordingClient::new(server.url()); + let (ack_tx, ack_rx) = oneshot::channel(); + let request = FinalizeRequest { + deadline: Instant::now() + Duration::from_secs(10), + ack: ack_tx, + }; + finalize_with_new_attempt( + request, + client.clone(), + task_id, + tempdir.path().to_path_buf(), + Duration::from_millis(10), + Duration::from_secs(5), + ) + .await; + assert_eq!( + client.commit_calls.load(Ordering::SeqCst), + 1, + "an above-floor budget must run exactly one attempt to completion" + ); + assert!(ack_rx.await.is_ok()); +} + +#[tokio::test] +async fn finalize_with_in_flight_attempt_stops_waiting_once_the_budget_elapses() { + let (result_tx, result_rx) = oneshot::channel::(); + let (ack_tx, ack_rx) = oneshot::channel(); + let request = FinalizeRequest { + deadline: Instant::now() + Duration::from_millis(50), + ack: ack_tx, + }; + let start = Instant::now(); + finalize_with_in_flight_attempt(request, result_rx).await; + assert!( + start.elapsed() < Duration::from_secs(2), + "must stop waiting once the budget elapses, not block indefinitely" + ); + assert!( + ack_rx.await.is_ok(), + "ack must still be sent after timing out" + ); + drop(result_tx); +} + +// ------------------------------------------------------------------------------------------------ +// start_attempt: no stranded work when the caller stops waiting. +// ------------------------------------------------------------------------------------------------ + +#[tokio::test] +async fn start_attempt_runs_to_completion_even_if_the_receiver_is_dropped() { + let client = FailingUploadTargetsClient::new(); + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let background = test_background(); + let rx = start_attempt( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + Duration::from_secs(5), + Duration::from_secs(5), + &background, + ); + drop(rx); // The caller stops waiting immediately; the spawned task must still run. + + wait_for( + || client.call_count.load(Ordering::SeqCst) >= 1, + Duration::from_secs(2), + ) + .await; +} + +// ------------------------------------------------------------------------------------------------ +// Full coordinator_loop state-machine behavior, via the test-only boundary-check injection. +// ------------------------------------------------------------------------------------------------ + +#[tokio::test] +async fn coordinator_defers_until_safe_boundary_then_commits() { + let mut server = Server::new_async().await; + let _uploads = server + .mock("PUT", Matcher::Regex("^/upload/.+$".to_string())) + .with_status(200) + .create_async() + .await; + + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let client = RecordingClient::new(server.url()); + let boundary_calls = Arc::new(AtomicUsize::new(0)); + let boundary_calls_for_closure = boundary_calls.clone(); + // False for the first two checks (simulating a mutating tool in flight / not yet + // WaitingForEvents), true from the third check onward. + let boundary_check: BoundaryCheck = Arc::new(move || { + let calls = boundary_calls_for_closure.clone(); + Box::pin(async move { calls.fetch_add(1, Ordering::SeqCst) >= 2 }) + }); + + let handle = CheckpointCoordinatorHandle::new_for_test( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + boundary_check, + Duration::from_millis(5), + Duration::from_secs(5), + Duration::from_secs(5), + test_background(), + ); + + wait_for( + || client.commit_calls.load(Ordering::SeqCst) >= 1, + Duration::from_secs(10), + ) + .await; + assert!( + boundary_calls.load(Ordering::SeqCst) >= 3, + "expected the coordinator to keep polling until the boundary became safe" + ); + drop(handle); +} + +#[tokio::test] +async fn coordinator_returns_to_idle_and_waits_a_full_interval_before_the_next_attempt() { + let mut server = Server::new_async().await; + let _uploads = server + .mock("PUT", Matcher::Regex("^/upload/.+$".to_string())) + .with_status(200) + .create_async() + .await; + + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let client = RecordingClient::new(server.url()); + let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { true })); + let interval = Duration::from_millis(400); + + let handle = CheckpointCoordinatorHandle::new_for_test( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + always_safe, + interval, + Duration::from_secs(5), + Duration::from_secs(5), + test_background(), + ); + + wait_for( + || client.commit_calls.load(Ordering::SeqCst) >= 1, + Duration::from_secs(5), + ) + .await; + + // Well within `interval` after the first commit, no second attempt should have started. + tokio::time::sleep(interval / 4).await; + assert_eq!( + client.commit_calls.load(Ordering::SeqCst), + 1, + "success must return to Idle and wait a full interval before the next attempt" + ); + + // After the interval (plus the immediate safe-boundary check) elapses, a second attempt runs. + wait_for( + || client.commit_calls.load(Ordering::SeqCst) >= 2, + Duration::from_secs(5), + ) + .await; + drop(handle); +} + +#[tokio::test] +async fn coordinator_finalize_awaits_the_in_flight_attempt_without_a_second_commit() { + let mut server = Server::new_async().await; + let _uploads = server + .mock("PUT", Matcher::Regex("^/upload/.+$".to_string())) + .with_status(200) + .create_async() + .await; + + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + // A slow commit keeps the periodic attempt "in flight" long enough for finalize() to race + // in while it's still running. + let client = RecordingClient::with_commit_delay(server.url(), Duration::from_millis(500)); + let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { true })); + + let handle = CheckpointCoordinatorHandle::new_for_test( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + always_safe, + Duration::from_millis(5), + Duration::from_secs(5), + Duration::from_secs(5), + test_background(), + ); + + // Give the coordinator time to enter InFlight and call commit_snapshot (which is now + // sleeping for commit_delay), but not enough time for it to resolve. + wait_for( + || client.commit_calls.load(Ordering::SeqCst) >= 1, + Duration::from_secs(2), + ) + .await; + + // Finalize with ample budget: it must await the already-in-flight attempt rather than + // starting a new one. + handle.finalize(Duration::from_secs(5)).await; + + assert_eq!( + client.commit_calls.load(Ordering::SeqCst), + 1, + "finalization must never start a second attempt while one is already in flight" + ); +} + +#[tokio::test] +async fn coordinator_finalize_from_idle_skips_attempt_below_floor() { + let client = FailingUploadTargetsClient::new(); + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + // No declarations fixture: the client would panic on any upload-targets/commit call, so + // this also proves no attempt is made. + let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { true })); + + let handle = CheckpointCoordinatorHandle::new_for_test( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + always_safe, + Duration::from_secs(600), // never fires during this test + Duration::from_secs(5), + Duration::from_secs(5), + test_background(), + ); + + // floor = 10s; a 100ms budget is well below it. + handle.finalize(Duration::from_millis(100)).await; + + assert_eq!( + client.call_count.load(Ordering::SeqCst), + 0, + "a below-floor finalize budget from Idle must skip the attempt entirely" + ); +} diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index c4325282041..02c7bae9683 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -7906,6 +7906,13 @@ impl TerminalView { &self.ai_input_model } + /// Used by [`crate::ai::agent_sdk::driver::checkpoint_coordinator`] to check whether the + /// terminal's conversation has any pending or running actions before starting a periodic + /// checkpoint attempt. + pub fn ai_action_model(&self) -> &ModelHandle { + &self.ai_action_model + } + pub fn agent_view_controller(&self) -> &ModelHandle { &self.agent_view_controller } From 6735cd69542c5129a1ff67d4a5db81bc9a9187ec Mon Sep 17 00:00:00 2001 From: joeywangzr Date: Fri, 31 Jul 2026 15:32:23 +0000 Subject: [PATCH 2/2] Make the coordinator's finalize reachable and its loop terminable Review follow-ups on the periodic checkpoint coordinator. 1. Expose `finalize_budget` so the finalize floor is reachable. `finalize_with_new_attempt` gates on `remaining > script_timeout + upload_timeout`. The wiring PR passes `upload_timeout` alone as the budget, so that comparison is false for *every* configuration (defaults: 120s > 60s + 120s), and because the coordinator owns the whole end-of-run path once enabled, the result is no end-of-run snapshot at all. Add `finalize_budget(script_timeout, upload_timeout)` and document that callers must derive the budget from it, so `AgentDriver` cannot drift out of sync with the floor. Downgrade the skip log to WARN, since reaching it means the end-of-run checkpoint was dropped. Also give `finalize`'s outer ack timeout slack over the inner deadline. Bounding both on exactly `budget` meant the "defense-in-depth" bound would routinely preempt an attempt the coordinator was legitimately still finishing, rather than only catching a bug. 2. Stop the loop when the driver is gone. `is_safe_boundary` mapped `Err(ModelDropped)` to "safe", so a dropped `AgentDriver` left the coordinator gathering and uploading the whole workspace every interval forever. Nothing else stops the loop, and `run_snapshot_upload` returns without finalizing under exactly the same condition (spawner round trip fails, or `OzHandoff` is off). Replace the boolean predicate with a `Boundary` tri-state so `DriverGone` returns from `coordinator_loop` instead of being conflated with `Safe`. 3. Checkpoint blocked conversations, and cap boundary deferral. `Blocked` is backed by a pending action, so it fell through to `has_unfinished_actions_for_conversation` and reported busy forever: a run parked on user approval (often hours) polled every 2s and never checkpointed, though nothing was mutating the workspace. Treat `Blocked` and `is_done` as quiescent alongside `WaitingForEvents`. Add `MAX_BOUNDARY_DEFERRAL` (10m) so a long turn can no longer starve the feature entirely. A slightly-inconsistent checkpoint beats none, since the previous committed generation is only replaced on success. 4. Keep boundary polling responsive to finalize. The boundary check was awaited outside `select!`, so a stalled model task queue could wedge shutdown behind an unbounded await. `wait_for_safe_boundary` now races it against the finalize channel. 5. Flush the declarations writer before every attempt. `AgentDriver::run_snapshot_upload` flushes pending driver-side `file` appends before running the declarations script, precisely so no write is in flight when the bash script starts appending. The coordinator did not, so both periodic and final checkpoints could miss the agent's most recent edits and race the script on the shared append-only JSONL. Thread an optional `DeclarationsWriterHandle` through and flush it in `run_one_attempt`. 6. Correct `DEFAULT_CHECKPOINT_INTERVAL`'s doc: the cadence is measured from attempt completion, not attempt start. Tests: the derived budget clears the floor (and the old pairing provably cannot); finalize during boundary polling still runs an attempt; a gone driver stops the loop without attempting; the loop stops after finalize; and a failed attempt returns to Idle and retries on the next tick. Co-Authored-By: Oz --- .../driver/checkpoint_coordinator.rs | 248 ++++++++++++---- .../driver/checkpoint_coordinator_tests.rs | 273 +++++++++++++++++- 2 files changed, 467 insertions(+), 54 deletions(-) diff --git a/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs b/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs index 050ffb26196..778f9e00b40 100644 --- a/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs +++ b/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs @@ -38,25 +38,67 @@ use warpui::r#async::{FutureExt as _, Timer}; use warpui::{ModelSpawner, SingletonEntity}; use super::AgentDriver; -use super::snapshot::{self, CheckpointResult}; +use super::snapshot::{self, CheckpointResult, DeclarationsWriterHandle}; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::BlocklistAIHistoryModel; use crate::server::server_api::harness_support::{CheckpointGeneration, HarnessSupportClient}; +/// Whether the conversation can tolerate a checkpoint attempt right now. +/// +/// `DriverGone` is deliberately distinct from `Busy`: a dropped `AgentDriver` used to be +/// reported as "safe", which left the coordinator gathering and uploading the whole +/// workspace every interval forever, since nothing else ever stops the loop. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Boundary { + /// Not mid-turn: safe to touch the filesystem and network. + Safe, + /// Mid-turn or actions in flight; retry after [`SAFE_BOUNDARY_POLL_INTERVAL`]. + Busy, + /// The `AgentDriver` this coordinator serves no longer exists, so there is nothing + /// left to checkpoint for and the loop must stop. + DriverGone, +} + /// A safe-boundary predicate, decoupled from `ModelSpawner` so /// [`coordinator_loop`] can be exercised in isolation by tests. Production code builds /// this from [`is_safe_boundary`]; tests supply a directly-controllable closure. -type BoundaryCheck = Arc BoxFuture<'static, bool> + Send + Sync>; +type BoundaryCheck = Arc BoxFuture<'static, Boundary> + Send + Sync>; -/// Default cadence between the start of one checkpoint attempt and the timer firing -/// again, absent an override on `AgentDriverOptions`. Deliberately not -/// `HARNESS_SAVE_INTERVAL` (30s) -- see the spec's "5 minutes plus jitter" cadence. +/// Default cadence between the end of one checkpoint attempt and the timer firing again, +/// absent an override on `AgentDriverOptions`. Deliberately not `HARNESS_SAVE_INTERVAL` +/// (30s) -- see the spec's "5 minutes plus jitter" cadence. +/// +/// Note this is measured from attempt *completion*, not attempt start: the timer is only +/// restarted once the `InFlight` state resolves, so back-to-back attempts can never +/// overlap. pub(super) const DEFAULT_CHECKPOINT_INTERVAL: Duration = Duration::from_secs(5 * 60); /// Upper bound on additive jitter, so agents scheduled at the same time don't all /// checkpoint in lockstep. const CHECKPOINT_JITTER: Duration = Duration::from_secs(30); /// How often the `Due` state re-checks whether the conversation is at a safe boundary. const SAFE_BOUNDARY_POLL_INTERVAL: Duration = Duration::from_secs(2); +/// How long `Due` will wait for a safe boundary before checkpointing anyway. +/// +/// Without a cap, a single long turn (or a conversation parked in a state the predicate +/// never calls safe) starves the feature entirely -- exactly the long-running case periodic +/// checkpoints exist for. A slightly-inconsistent checkpoint is strictly better than none, +/// since the previous committed generation is only replaced on success. +const MAX_BOUNDARY_DEFERRAL: Duration = Duration::from_secs(10 * 60); +/// Slack added on top of the per-attempt floor by [`finalize_budget`], and to the +/// belt-and-braces bound in [`CheckpointCoordinatorHandle::finalize`], so neither the +/// coordinator's ack round trip nor the outer timeout eats into the attempt's own budget. +const FINALIZE_ACK_SLACK: Duration = Duration::from_secs(10); + +/// The shutdown budget a caller must grant [`CheckpointCoordinatorHandle::finalize`] for a +/// final attempt to be possible. +/// +/// Exposed so `AgentDriver` cannot drift out of sync with the floor enforced in +/// [`finalize_with_new_attempt`]. Passing anything at or below `script_timeout + +/// upload_timeout` silently skips the final attempt -- and because the coordinator owns the +/// whole end-of-run path, that means no end-of-run snapshot at all. +pub(super) fn finalize_budget(script_timeout: Duration, upload_timeout: Duration) -> Duration { + script_timeout + upload_timeout + FINALIZE_ACK_SLACK +} /// A request to finalize, carrying the deadline by which the coordinator must ack so /// shutdown can proceed. @@ -81,6 +123,7 @@ impl CheckpointCoordinatorHandle { client: Arc, task_id: AmbientAgentTaskId, working_dir: PathBuf, + declarations_writer: Option, spawner: ModelSpawner, interval: Duration, script_timeout: Duration, @@ -95,6 +138,7 @@ impl CheckpointCoordinatorHandle { client, task_id, working_dir, + declarations_writer, boundary_check, interval, CHECKPOINT_JITTER, @@ -124,6 +168,7 @@ impl CheckpointCoordinatorHandle { client, task_id, working_dir, + None, boundary_check, interval, Duration::ZERO, @@ -138,6 +183,7 @@ impl CheckpointCoordinatorHandle { client: Arc, task_id: AmbientAgentTaskId, working_dir: PathBuf, + declarations_writer: Option, boundary_check: BoundaryCheck, interval: Duration, jitter: Duration, @@ -152,6 +198,7 @@ impl CheckpointCoordinatorHandle { client, task_id, working_dir, + declarations_writer, boundary_check, interval, jitter, @@ -169,6 +216,11 @@ impl CheckpointCoordinatorHandle { /// floor), or await an already-in-flight attempt instead -- never both -- then /// stop the coordinator. Bounded by `budget` end to end. Safe to call at most /// once; safe to never call. + /// + /// Callers should derive `budget` from [`finalize_budget`] rather than passing a + /// per-attempt timeout directly: the floor in [`finalize_with_new_attempt`] is + /// `script_timeout + upload_timeout`, so a smaller budget silently skips the final + /// attempt. pub(super) async fn finalize(&self, budget: Duration) { let (ack_tx, ack_rx) = oneshot::channel(); let request = FinalizeRequest { @@ -181,8 +233,11 @@ impl CheckpointCoordinatorHandle { } // The coordinator always acks well within `budget` (either immediately, when // below the floor, or after its own internally bounded attempt). This extra - // bound is defense-in-depth so a coordinator bug cannot wedge shutdown. - let _ = tokio::time::timeout(budget, ack_rx).await; + // bound is defense-in-depth so a coordinator bug cannot wedge shutdown -- hence + // the added slack: bounding on exactly `budget` would routinely preempt the + // attempt the coordinator is legitimately still finishing, instead of only + // catching a bug. + let _ = tokio::time::timeout(budget + FINALIZE_ACK_SLACK, ack_rx).await; } } @@ -199,16 +254,25 @@ fn jittered_interval(interval: Duration, jitter: Duration) -> Duration { interval + Duration::from_millis(extra) } -/// Run one checkpoint attempt to completion: regenerate declarations, then gather, -/// upload, and commit. Bounded by `upload_timeout`; `script_timeout` separately -/// bounds only the declarations-script sub-step (matching the legacy pipeline). +/// Run one checkpoint attempt to completion: drain pending declarations writes, +/// regenerate declarations, then gather, upload, and commit. Bounded by `upload_timeout`; +/// `script_timeout` separately bounds only the declarations-script sub-step (matching the +/// legacy pipeline). async fn run_one_attempt( client: Arc, task_id: AmbientAgentTaskId, working_dir: PathBuf, + declarations_writer: Option, script_timeout: Duration, upload_timeout: Duration, ) -> CheckpointResult { + // Drain queued driver-side `file` appends before the bash script starts appending its + // own `repo` entries to the same append-only JSONL. `AgentDriver::run_snapshot_upload` + // does this on the legacy path for exactly this reason; without it a checkpoint can + // both miss the agent's most recent edits and race the script on the shared file. + if let Some(writer) = &declarations_writer { + writer.flush().await; + } snapshot::run_declarations_script(&working_dir, &task_id, script_timeout).await; let path = snapshot::resolve_declarations_path(Some(&task_id)); match snapshot::run_checkpoint_from_declarations_file(&path, client) @@ -229,10 +293,12 @@ async fn run_one_attempt( /// stops waiting (e.g. because a shutdown budget elapsed) cannot strand the attempt /// or cause it to be silently abandoned mid-upload -- it simply keeps running in the /// background and, if it succeeds, still commits. +#[allow(clippy::too_many_arguments)] fn start_attempt( client: Arc, task_id: AmbientAgentTaskId, working_dir: PathBuf, + declarations_writer: Option, script_timeout: Duration, upload_timeout: Duration, background: &Background, @@ -240,8 +306,15 @@ fn start_attempt( let (tx, rx) = oneshot::channel(); background .spawn(async move { - let result = - run_one_attempt(client, task_id, working_dir, script_timeout, upload_timeout).await; + let result = run_one_attempt( + client, + task_id, + working_dir, + declarations_writer, + script_timeout, + upload_timeout, + ) + .await; let _ = tx.send(result); }) .detach(); @@ -249,50 +322,119 @@ fn start_attempt( } /// Query `AgentDriver` (via its spawner) for whether the conversation is currently at -/// a safe boundary: either yielded via `wait_for_events`, or simply not mid-turn with -/// no pending/running actions. Returns `true` (safe) if the driver has no -/// conversation yet, if its conversation can no longer be found, or if the driver -/// itself has been dropped -- in each case there is nothing left to interrupt. -async fn is_safe_boundary(spawner: &ModelSpawner) -> bool { - spawner +/// a safe boundary. +/// +/// [`Boundary::Safe`] when the driver has no conversation yet, when its conversation can no +/// longer be found, or when the conversation is quiescent -- there is nothing to interrupt. +/// [`Boundary::DriverGone`] when the driver itself has been dropped, which stops the loop +/// rather than being conflated with "safe" and checkpointing forever. +async fn is_safe_boundary(spawner: &ModelSpawner) -> Boundary { + let result = spawner .spawn(|driver, ctx| { let Some(conversation_id) = driver.run_conversation_id else { - return true; + return Boundary::Safe; }; let Some(status) = BlocklistAIHistoryModel::as_ref(ctx) .conversation(&conversation_id) .map(|conversation| conversation.status().clone()) else { - return true; + return Boundary::Safe; }; - if status.is_waiting_for_events() { - return true; + // Quiescent states, checked before the pending-action sweep below. + // + // `Blocked` in particular *is* backed by a pending action, so letting it fall + // through would report `Busy` forever: a run parked on user approval (often for + // hours) would poll every couple of seconds and never checkpoint, even though + // nothing is mutating the workspace. That is precisely when a checkpoint is + // most valuable. + if status.is_waiting_for_events() || status.is_blocked() || status.is_done() { + return Boundary::Safe; } if status.is_in_progress() || status.is_transient_error() { - return false; + return Boundary::Busy; } let terminal_view = driver .terminal_driver .as_ref(ctx) .terminal_view() .as_ref(ctx); - !terminal_view + if terminal_view .ai_action_model() .as_ref(ctx) .has_unfinished_actions_for_conversation(conversation_id) + { + Boundary::Busy + } else { + Boundary::Safe + } }) - .await - .unwrap_or(true) + .await; + result.unwrap_or(Boundary::DriverGone) +} + +/// How the `Due` state resolved. +enum DueOutcome { + /// Proceed to `InFlight`. + Safe, + /// A finalize request arrived while waiting; the caller owns it. + Finalize(FinalizeRequest), + /// The coordinator should stop: the driver is gone, or every handle was dropped. + Stop, +} + +/// Poll for a safe boundary, staying responsive to finalize throughout. +/// +/// The boundary check is itself awaited inside `select!` rather than before it: it is a +/// round trip through `AgentDriver`'s model task queue, and a stalled queue must not be able +/// to wedge shutdown behind an unbounded await. +async fn wait_for_safe_boundary( + boundary_check: &BoundaryCheck, + finalize_rx: &mut mpsc::UnboundedReceiver, +) -> DueOutcome { + let due_since = Instant::now(); + loop { + // Checked immediately on entry (not only after the first poll interval elapses) so + // an already-safe conversation doesn't pay needless latency. + let boundary = futures::select! { + boundary = boundary_check().fuse() => boundary, + request = finalize_rx.recv().fuse() => { + return request.map_or(DueOutcome::Stop, DueOutcome::Finalize); + } + }; + match boundary { + Boundary::Safe => return DueOutcome::Safe, + Boundary::DriverGone => { + log::info!("AgentDriver is gone; stopping the periodic checkpoint coordinator"); + return DueOutcome::Stop; + } + Boundary::Busy => {} + } + if due_since.elapsed() >= MAX_BOUNDARY_DEFERRAL { + log::warn!( + "Conversation has not reached a safe boundary in {MAX_BOUNDARY_DEFERRAL:?}; \ + checkpointing anyway rather than skipping this cycle entirely" + ); + return DueOutcome::Safe; + } + futures::select! { + _ = Timer::after(SAFE_BOUNDARY_POLL_INTERVAL).fuse() => {} + request = finalize_rx.recv().fuse() => { + return request.map_or(DueOutcome::Stop, DueOutcome::Finalize); + } + } + } } /// Handle a finalize request received while no attempt is currently in flight: start /// exactly one best-effort attempt only if `budget` exceeds the gather/upload floor, /// bound it by the remaining budget, then ack. +#[allow(clippy::too_many_arguments)] async fn finalize_with_new_attempt( request: FinalizeRequest, client: Arc, task_id: AmbientAgentTaskId, working_dir: PathBuf, + declarations_writer: Option, script_timeout: Duration, upload_timeout: Duration, ) { @@ -302,12 +444,21 @@ async fn finalize_with_new_attempt( log::info!( "Starting final checkpoint attempt at shutdown (remaining budget {remaining:?})" ); - let attempt = run_one_attempt(client, task_id, working_dir, script_timeout, upload_timeout); + let attempt = run_one_attempt( + client, + task_id, + working_dir, + declarations_writer, + script_timeout, + upload_timeout, + ); if tokio::time::timeout(remaining, attempt).await.is_err() { log::warn!("Final checkpoint attempt did not complete within {remaining:?}"); } } else { - log::info!( + // Callers must size the budget with `finalize_budget`; anything at or below the + // floor lands here and produces no end-of-run checkpoint at all. + log::warn!( "Skipping final checkpoint attempt: remaining shutdown budget {remaining:?} \ is below the {floor:?} floor" ); @@ -352,6 +503,7 @@ async fn coordinator_loop( client: Arc, task_id: AmbientAgentTaskId, working_dir: PathBuf, + declarations_writer: Option, boundary_check: BoundaryCheck, interval: Duration, jitter: Duration, @@ -371,6 +523,7 @@ async fn coordinator_loop( client.clone(), task_id, working_dir.clone(), + declarations_writer.clone(), script_timeout, upload_timeout, ) @@ -379,29 +532,23 @@ async fn coordinator_loop( } } - // --- Due: poll for a safe boundary, staying responsive to finalize. Checked - // immediately on entry (not only after the first poll interval elapses) so an - // already-safe conversation doesn't pay needless latency. - let mut at_boundary = boundary_check().await; - while !at_boundary { - futures::select! { - _ = Timer::after(SAFE_BOUNDARY_POLL_INTERVAL).fuse() => { - at_boundary = boundary_check().await; - } - request = finalize_rx.recv().fuse() => { - let Some(request) = request else { return }; - finalize_with_new_attempt( - request, - client.clone(), - task_id, - working_dir.clone(), - script_timeout, - upload_timeout, - ) - .await; - return; - } + // --- Due: poll for a safe boundary, staying responsive to finalize. --- + match wait_for_safe_boundary(&boundary_check, &mut finalize_rx).await { + DueOutcome::Safe => {} + DueOutcome::Finalize(request) => { + finalize_with_new_attempt( + request, + client.clone(), + task_id, + working_dir.clone(), + declarations_writer.clone(), + script_timeout, + upload_timeout, + ) + .await; + return; } + DueOutcome::Stop => return, } // --- InFlight: run exactly one attempt, never overlapping another. --- @@ -409,6 +556,7 @@ async fn coordinator_loop( client.clone(), task_id, working_dir.clone(), + declarations_writer.clone(), script_timeout, upload_timeout, &background, diff --git a/app/src/ai/agent_sdk/driver/checkpoint_coordinator_tests.rs b/app/src/ai/agent_sdk/driver/checkpoint_coordinator_tests.rs index a9cdd41be82..554191ca9f1 100644 --- a/app/src/ai/agent_sdk/driver/checkpoint_coordinator_tests.rs +++ b/app/src/ai/agent_sdk/driver/checkpoint_coordinator_tests.rs @@ -297,6 +297,26 @@ fn jittered_interval_zero_jitter_is_a_no_op() { // finalize_with_new_attempt / finalize_with_in_flight_attempt. // ------------------------------------------------------------------------------------------------ +#[test] +fn finalize_budget_exceeds_the_attempt_floor_for_the_driver_defaults() { + // Regression: `AgentDriver::run_snapshot_upload` used to pass `upload_timeout` alone as + // the finalize budget, but the floor is `script_timeout + upload_timeout`. The + // `remaining > floor` gate was therefore false for *every* possible configuration, so + // the final attempt was always skipped -- and since the coordinator owns the whole + // end-of-run path, that meant no end-of-run snapshot at all. + let script = snapshot::DEFAULT_DECLARATIONS_SCRIPT_TIMEOUT; + let upload = snapshot::DEFAULT_SNAPSHOT_UPLOAD_TIMEOUT; + assert!( + finalize_budget(script, upload) > script + upload, + "the budget callers derive must exceed the floor `finalize_with_new_attempt` enforces" + ); + // The old, broken pairing, asserted explicitly so nobody reintroduces it. + assert!( + upload <= script + upload, + "passing the upload timeout alone can never clear the floor" + ); +} + #[tokio::test] async fn finalize_with_new_attempt_below_floor_skips_the_attempt_entirely() { let client = FailingUploadTargetsClient::new(); @@ -312,6 +332,7 @@ async fn finalize_with_new_attempt_below_floor_skips_the_attempt_entirely() { client.clone(), task_id, PathBuf::from("/tmp"), + None, Duration::from_secs(1), Duration::from_secs(1), ) @@ -324,6 +345,49 @@ async fn finalize_with_new_attempt_below_floor_skips_the_attempt_entirely() { assert!(ack_rx.await.is_ok(), "ack must still be sent"); } +#[tokio::test] +async fn finalize_with_the_derived_budget_runs_the_attempt() { + // The companion to the unit test above: a budget built by `finalize_budget` from the + // same timeouts the coordinator holds must clear the floor and actually commit. + let mut server = Server::new_async().await; + let _uploads = server + .mock("PUT", Matcher::Regex("^/upload/.+$".to_string())) + .with_status(200) + .create_async() + .await; + + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let script_timeout = Duration::from_millis(10); + let upload_timeout = Duration::from_secs(5); + let client = RecordingClient::new(server.url()); + let (ack_tx, ack_rx) = oneshot::channel(); + let request = FinalizeRequest { + deadline: Instant::now() + finalize_budget(script_timeout, upload_timeout), + ack: ack_tx, + }; + finalize_with_new_attempt( + request, + client.clone(), + task_id, + tempdir.path().to_path_buf(), + None, + script_timeout, + upload_timeout, + ) + .await; + assert_eq!( + client.commit_calls.load(Ordering::SeqCst), + 1, + "a budget derived from `finalize_budget` must run exactly one final attempt" + ); + assert!(ack_rx.await.is_ok()); +} + #[tokio::test] async fn finalize_with_new_attempt_above_floor_starts_and_commits() { let mut server = Server::new_async().await; @@ -350,6 +414,7 @@ async fn finalize_with_new_attempt_above_floor_starts_and_commits() { client.clone(), task_id, tempdir.path().to_path_buf(), + None, Duration::from_millis(10), Duration::from_secs(5), ) @@ -401,6 +466,7 @@ async fn start_attempt_runs_to_completion_even_if_the_receiver_is_dropped() { client.clone(), task_id, tempdir.path().to_path_buf(), + None, Duration::from_secs(5), Duration::from_secs(5), &background, @@ -440,7 +506,13 @@ async fn coordinator_defers_until_safe_boundary_then_commits() { // WaitingForEvents), true from the third check onward. let boundary_check: BoundaryCheck = Arc::new(move || { let calls = boundary_calls_for_closure.clone(); - Box::pin(async move { calls.fetch_add(1, Ordering::SeqCst) >= 2 }) + Box::pin(async move { + if calls.fetch_add(1, Ordering::SeqCst) >= 2 { + Boundary::Safe + } else { + Boundary::Busy + } + }) }); let handle = CheckpointCoordinatorHandle::new_for_test( @@ -482,7 +554,7 @@ async fn coordinator_returns_to_idle_and_waits_a_full_interval_before_the_next_a let _decl = DeclarationsFixture::new(&task_id, &file_path); let client = RecordingClient::new(server.url()); - let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { true })); + let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { Boundary::Safe })); let interval = Duration::from_millis(400); let handle = CheckpointCoordinatorHandle::new_for_test( @@ -537,7 +609,7 @@ async fn coordinator_finalize_awaits_the_in_flight_attempt_without_a_second_comm // A slow commit keeps the periodic attempt "in flight" long enough for finalize() to race // in while it's still running. let client = RecordingClient::with_commit_delay(server.url(), Duration::from_millis(500)); - let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { true })); + let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { Boundary::Safe })); let handle = CheckpointCoordinatorHandle::new_for_test( client.clone(), @@ -569,6 +641,199 @@ async fn coordinator_finalize_awaits_the_in_flight_attempt_without_a_second_comm ); } +#[tokio::test] +async fn coordinator_stops_when_the_driver_is_gone() { + // Regression: `is_safe_boundary` used to map a dropped `AgentDriver` to "safe", so the + // loop kept gathering and uploading the whole workspace every interval forever. Nothing + // else ever stops it -- `run_snapshot_upload` returns without finalizing when the + // spawner round trip fails, which is the same condition. + let client = FailingUploadTargetsClient::new(); + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let boundary_calls = Arc::new(AtomicUsize::new(0)); + let boundary_calls_for_closure = boundary_calls.clone(); + let driver_gone: BoundaryCheck = Arc::new(move || { + let calls = boundary_calls_for_closure.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Boundary::DriverGone + }) + }); + + let handle = CheckpointCoordinatorHandle::new_for_test( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + driver_gone, + Duration::from_millis(5), + Duration::from_secs(5), + Duration::from_secs(5), + test_background(), + ); + + wait_for( + || boundary_calls.load(Ordering::SeqCst) >= 1, + Duration::from_secs(5), + ) + .await; + // Well past several intervals: the loop must have returned rather than kept ticking. + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!( + boundary_calls.load(Ordering::SeqCst), + 1, + "a gone driver must stop the loop, not be treated as a safe boundary" + ); + assert_eq!( + client.call_count.load(Ordering::SeqCst), + 0, + "no attempt may run once the driver is gone" + ); + drop(handle); +} + +#[tokio::test] +async fn coordinator_finalize_while_waiting_for_a_safe_boundary_still_runs_an_attempt() { + // The `Due` state polls indefinitely while the conversation is busy. A finalize request + // arriving in that window must be serviced (with a new attempt, since none is in flight) + // rather than waiting for a boundary that may never come. + let mut server = Server::new_async().await; + let _uploads = server + .mock("PUT", Matcher::Regex("^/upload/.+$".to_string())) + .with_status(200) + .create_async() + .await; + + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let boundary_calls = Arc::new(AtomicUsize::new(0)); + let boundary_calls_for_closure = boundary_calls.clone(); + let never_safe: BoundaryCheck = Arc::new(move || { + let calls = boundary_calls_for_closure.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Boundary::Busy + }) + }); + + let script_timeout = Duration::from_millis(10); + let upload_timeout = Duration::from_secs(5); + let client = RecordingClient::new(server.url()); + let handle = CheckpointCoordinatorHandle::new_for_test( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + never_safe, + Duration::from_millis(5), + script_timeout, + upload_timeout, + test_background(), + ); + + // Wait until the coordinator is definitely parked in `Due`. + wait_for( + || boundary_calls.load(Ordering::SeqCst) >= 1, + Duration::from_secs(5), + ) + .await; + + handle + .finalize(finalize_budget(script_timeout, upload_timeout)) + .await; + + assert_eq!( + client.commit_calls.load(Ordering::SeqCst), + 1, + "finalize arriving during boundary polling must still run one final attempt" + ); +} + +#[tokio::test] +async fn coordinator_stops_running_attempts_after_finalize() { + // `finalize` is documented to stop the coordinator. Prove the loop actually returns + // instead of continuing to tick in the background after shutdown has moved on. + let mut server = Server::new_async().await; + let _uploads = server + .mock("PUT", Matcher::Regex("^/upload/.+$".to_string())) + .with_status(200) + .create_async() + .await; + + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let script_timeout = Duration::from_millis(10); + let upload_timeout = Duration::from_secs(5); + let client = RecordingClient::new(server.url()); + let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { Boundary::Safe })); + let interval = Duration::from_millis(50); + + let handle = CheckpointCoordinatorHandle::new_for_test( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + always_safe, + interval, + script_timeout, + upload_timeout, + test_background(), + ); + + handle + .finalize(finalize_budget(script_timeout, upload_timeout)) + .await; + let after_finalize = client.commit_calls.load(Ordering::SeqCst); + + // Many intervals' worth of time; a still-running loop would commit again. + tokio::time::sleep(interval * 10).await; + assert_eq!( + client.commit_calls.load(Ordering::SeqCst), + after_finalize, + "the coordinator must stop after finalize rather than keep checkpointing" + ); +} + +#[tokio::test] +async fn coordinator_failed_attempt_returns_to_idle_and_retries_next_interval() { + // Failure is retried by the ordinary periodic timer rather than a distinct backoff, so a + // failing attempt must neither wedge the loop nor spin. + let client = FailingUploadTargetsClient::new(); + let task_id = fresh_task_id(); + let tempdir = TempDir::new().unwrap(); + let file_path = tempdir.path().join("note.txt"); + std::fs::write(&file_path, b"hi").unwrap(); + let _decl = DeclarationsFixture::new(&task_id, &file_path); + + let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { Boundary::Safe })); + let handle = CheckpointCoordinatorHandle::new_for_test( + client.clone(), + task_id, + tempdir.path().to_path_buf(), + always_safe, + Duration::from_millis(20), + Duration::from_millis(10), + Duration::from_secs(5), + test_background(), + ); + + wait_for( + || client.call_count.load(Ordering::SeqCst) >= 2, + Duration::from_secs(10), + ) + .await; + drop(handle); +} + #[tokio::test] async fn coordinator_finalize_from_idle_skips_attempt_below_floor() { let client = FailingUploadTargetsClient::new(); @@ -576,7 +841,7 @@ async fn coordinator_finalize_from_idle_skips_attempt_below_floor() { let tempdir = TempDir::new().unwrap(); // No declarations fixture: the client would panic on any upload-targets/commit call, so // this also proves no attempt is made. - let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { true })); + let always_safe: BoundaryCheck = Arc::new(|| Box::pin(async { Boundary::Safe })); let handle = CheckpointCoordinatorHandle::new_for_test( client.clone(),