From dd98a69a2255ec83c4a91a2ad9f399e5e922879e Mon Sep 17 00:00:00 2001 From: Joey Wang Date: Fri, 31 Jul 2026 10:10:59 -0400 Subject: [PATCH 1/2] Wire the periodic checkpoint coordinator into AgentDriver Instantiates and drives CheckpointCoordinatorHandle from AgentDriver's own lifecycle, completing REMOTE-2111 Phase 3: - AgentDriverOptions gains checkpoint_interval (override for the default 5-minute-plus-jitter cadence; primarily for tests). - AgentDriver spawns a coordinator when a cloud task_id is present, snapshot uploads aren't disabled, and both FeatureFlag::OzHandoff and the new FeatureFlag::PeriodicHandoffCheckpoints are enabled. - Driver finalization routes through the coordinator's finalize() (bounded by the existing snapshot upload timeout) when a coordinator is active, instead of the legacy one-shot end-of-run snapshot path. - warp_features: add PeriodicHandoffCheckpoints, off by default while the coordinator rolls out. Co-Authored-By: Oz --- app/src/ai/agent_sdk/driver.rs | 59 ++++++++++++++++++++++++++++++++- app/src/ai/agent_sdk/mod.rs | 1 + crates/warp_features/src/lib.rs | 6 ++++ 3 files changed, 65 insertions(+), 1 deletion(-) diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index 70524631e25..41fbd92efbb 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -313,6 +313,10 @@ pub struct AgentDriverOptions { pub snapshot_upload_timeout: Option, /// Declarations script timeout override. pub snapshot_script_timeout: Option, + /// Periodic checkpoint cadence override. Only used when + /// `FeatureFlag::PeriodicHandoffCheckpoints` is enabled; deliberately separate + /// from `snapshot_upload_timeout`/`snapshot_script_timeout` per-attempt budgets. + pub checkpoint_interval: Option, /// Skip the initial `StartFromAmbientRunPrompt` so the agent waits for a /// follow-up instead of hallucinating an empty turn. Sourced from the /// `--skip-initial-turn` CLI flag, which the worker emits when the @@ -378,6 +382,13 @@ pub struct AgentDriver { snapshot_upload_timeout: Duration, snapshot_script_timeout: Duration, + /// Periodic workspace-handoff checkpoint coordinator. `Some` only when + /// `FeatureFlag::OzHandoff` and `FeatureFlag::PeriodicHandoffCheckpoints` are both + /// enabled, the run has a cloud task id, and `--no-snapshot` was not set; + /// `None` otherwise, in which case `run_snapshot_upload` falls back to the legacy + /// one-shot upload path unchanged. + checkpoint_coordinator: Option, + /// Conversation ID this driver is running. Set at construction for /// resumed runs and on `ConversationServerTokenAssigned` for fresh /// runs; consumed by `unregister_streamer_consumer` at end of run. @@ -651,6 +662,7 @@ impl AgentDriver { snapshot_disabled, snapshot_upload_timeout, snapshot_script_timeout, + checkpoint_interval, skip_initial_turn, strict_mcp_startup, mcp_startup_timeout, @@ -759,6 +771,32 @@ impl AgentDriver { _ => None, }; + // Spawn the periodic checkpoint coordinator under the same gates as the + // declarations writer above, plus the dedicated rollout flag. `None` keeps + // `run_snapshot_upload` on the legacy one-shot upload path unchanged. + let checkpoint_coordinator = match task_id { + Some(id) + if FeatureFlag::OzHandoff.is_enabled() + && FeatureFlag::PeriodicHandoffCheckpoints.is_enabled() + && !snapshot_disabled_value => + { + let client = ServerApiProvider::as_ref(ctx).get_harness_support_client(); + Some(checkpoint_coordinator::CheckpointCoordinatorHandle::new( + client, + id, + working_dir.clone(), + ctx.spawner(), + checkpoint_interval + .unwrap_or(checkpoint_coordinator::DEFAULT_CHECKPOINT_INTERVAL), + snapshot_script_timeout + .unwrap_or(snapshot::DEFAULT_DECLARATIONS_SCRIPT_TIMEOUT), + snapshot_upload_timeout.unwrap_or(snapshot::DEFAULT_SNAPSHOT_UPLOAD_TIMEOUT), + ctx.background_executor(), + )) + } + _ => None, + }; + Ok(Self { terminal_driver, working_dir, @@ -778,6 +816,7 @@ impl AgentDriver { .unwrap_or(snapshot::DEFAULT_SNAPSHOT_UPLOAD_TIMEOUT), snapshot_script_timeout: snapshot_script_timeout .unwrap_or(snapshot::DEFAULT_DECLARATIONS_SCRIPT_TIMEOUT), + checkpoint_coordinator, run_conversation_id, parent_run_id: parent_run_id_for_self, third_party_harness_model_config, @@ -821,6 +860,7 @@ impl AgentDriver { snapshot_disabled: false, snapshot_upload_timeout: snapshot::DEFAULT_SNAPSHOT_UPLOAD_TIMEOUT, snapshot_script_timeout: snapshot::DEFAULT_DECLARATIONS_SCRIPT_TIMEOUT, + checkpoint_coordinator: None, run_conversation_id: None, parent_run_id: None, third_party_harness_model_config: None, @@ -3789,13 +3829,20 @@ impl AgentDriver { // Snapshot upload is only meaningful for cloud task runs, so short-circuit before // pulling the rest of the context onto this task. - let Ok((Some(task_id), snapshot_disabled, upload_timeout, script_timeout)) = spawner + let Ok(( + Some(task_id), + snapshot_disabled, + upload_timeout, + script_timeout, + checkpoint_coordinator, + )) = spawner .spawn(|me, _| { ( me.task_id, me.snapshot_disabled, me.snapshot_upload_timeout, me.snapshot_script_timeout, + me.checkpoint_coordinator.clone(), ) }) .await @@ -3807,6 +3854,16 @@ impl AgentDriver { return; } + // When the periodic checkpoint coordinator is active, it owns the entire + // end-of-run path: `finalize` regenerates declarations, runs one last + // best-effort attempt bounded by `upload_timeout`, and commits it as the + // selected checkpoint. This replaces the legacy one-shot upload below so + // there is exactly one end-of-run snapshot path, not two. + if let Some(coordinator) = checkpoint_coordinator { + coordinator.finalize(upload_timeout).await; + return; + } + let Ok((working_dir, client)) = spawner .spawn(|me, ctx| { let client = ServerApiProvider::as_ref(ctx).get_harness_support_client(); diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index 811c1cb522a..c7039b6b43f 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -1069,6 +1069,7 @@ impl AgentDriverRunner { .snapshot .snapshot_script_timeout .map(|duration| duration.into()), + checkpoint_interval: None, skip_initial_turn: args.skip_initial_turn, strict_mcp_startup: args.strict_mcp_startup, mcp_startup_timeout: args.mcp_startup_timeout.map(|duration| duration.into()), diff --git a/crates/warp_features/src/lib.rs b/crates/warp_features/src/lib.rs index 1612e769cea..1b5cd474a59 100644 --- a/crates/warp_features/src/lib.rs +++ b/crates/warp_features/src/lib.rs @@ -932,6 +932,12 @@ pub enum FeatureFlag { /// authenticated with the logged-in user's session token. No manual MCP /// setup or API key required. FactoryMcp, + + /// Enables periodic workspace-handoff checkpoints during a cloud agent run, + /// rather than only uploading a workspace snapshot once at end-of-run. + /// Requires `OzHandoff` to also be enabled; a no-op for local runs and when + /// `--no-snapshot` is set. Off by default while the coordinator rolls out. + PeriodicHandoffCheckpoints, } static FLAG_STATES: [AtomicBool; cardinality::()] = From d9bb9b832047e1363e373c63ba78cc1cd547b6b6 Mon Sep 17 00:00:00 2001 From: joeywangzr Date: Fri, 31 Jul 2026 15:43:39 +0000 Subject: [PATCH 2/2] Fix the finalize budget, hand the writer to the coordinator, drop scaffolding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups on the AgentDriver wiring. 1. Derive the finalize budget instead of passing `upload_timeout`. The coordinator's floor is `script_timeout + upload_timeout`, so passing `upload_timeout` alone made `remaining > floor` false for every possible configuration (defaults: 120s > 60s + 120s). Because this path `return`s past the legacy one-shot upload, enabling `PeriodicHandoffCheckpoints` meant no end-of-run snapshot at all: runs shorter than one checkpoint interval uploaded nothing, and longer runs lost everything since the last periodic tick. Use `checkpoint_coordinator::finalize_budget(script_timeout, upload_timeout)`, which owns the floor so the two cannot drift apart. 2. Give the coordinator the declarations writer. `run_snapshot_upload` flushes queued driver-side `file` appends before running the declarations script so none is in flight when the bash script starts appending to the same JSONL. The coordinator path skipped that, so checkpoints could miss the agent's most recent edits. The coordinator is gated on a superset of the writer's own conditions, so the handle is always present when the coordinator exists. 3. Remove the temporary `allow(dead_code)` scaffolding. Both earlier PRs in this stack noted their allows were "removable once that PR is merged on top of this one" — this is that PR. Removes the module-wide `#![allow(dead_code)]` in `checkpoint_coordinator` (which would otherwise permanently blind the lint over new async/networking code) plus the per-item allows in `snapshot` and `harness_support`, along with the now-stale comments explaining them. Verified with `cargo clippy -p warp --all-targets --tests -- -D warnings`. Co-Authored-By: Oz --- app/src/ai/agent_sdk/driver.rs | 20 ++++++++++++++++--- .../driver/checkpoint_coordinator.rs | 8 -------- app/src/ai/agent_sdk/driver/snapshot.rs | 9 --------- app/src/server/server_api/harness_support.rs | 7 ------- 4 files changed, 17 insertions(+), 27 deletions(-) diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index 41fbd92efbb..a43a1688d5d 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -785,6 +785,10 @@ impl AgentDriver { client, id, working_dir.clone(), + // Shared with the history subscription so every attempt can drain + // queued `file` appends before the declarations script runs, exactly + // as `run_snapshot_upload` does on the legacy path. + snapshot_file_writer.clone(), ctx.spawner(), checkpoint_interval .unwrap_or(checkpoint_coordinator::DEFAULT_CHECKPOINT_INTERVAL), @@ -3855,12 +3859,22 @@ impl AgentDriver { } // When the periodic checkpoint coordinator is active, it owns the entire - // end-of-run path: `finalize` regenerates declarations, runs one last - // best-effort attempt bounded by `upload_timeout`, and commits it as the + // end-of-run path: `finalize` drains the declarations writer, regenerates + // declarations, runs one last best-effort attempt, and commits it as the // selected checkpoint. This replaces the legacy one-shot upload below so // there is exactly one end-of-run snapshot path, not two. + // + // The budget must come from `finalize_budget`, not from `upload_timeout` alone: + // the coordinator's floor is `script_timeout + upload_timeout`, so a smaller + // budget silently skips the final attempt — and since this path `return`s past + // the legacy upload below, that would mean no end-of-run snapshot at all. if let Some(coordinator) = checkpoint_coordinator { - coordinator.finalize(upload_timeout).await; + coordinator + .finalize(checkpoint_coordinator::finalize_budget( + script_timeout, + upload_timeout, + )) + .await; return; } diff --git a/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs b/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs index 778f9e00b40..180540904c8 100644 --- a/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs +++ b/app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs @@ -16,14 +16,6 @@ //! 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; diff --git a/app/src/ai/agent_sdk/driver/snapshot.rs b/app/src/ai/agent_sdk/driver/snapshot.rs index 6c8bbd8c0c6..adefc3b2304 100644 --- a/app/src/ai/agent_sdk/driver/snapshot.rs +++ b/app/src/ai/agent_sdk/driver/snapshot.rs @@ -665,9 +665,6 @@ struct SnapshotOutcome { /// Outcome of one checkpoint attempt, where [`SnapshotOutcome`] only covers per-entry upload /// results within that attempt. -// The checkpoint pipeline below has no production caller until the periodic coordinator -// lands in a follow-up, stacked PR; the `allow(dead_code)`s go away with it. -#[allow(dead_code)] #[derive(Debug)] pub(super) enum CheckpointResult { /// `generation` is now the server's selected checkpoint. @@ -691,12 +688,10 @@ pub(super) enum CheckpointResult { /// [`SnapshotUploadMode`] for the server-side semantics. enum PipelineMode { Legacy, - #[allow(dead_code)] Checkpoint(CheckpointGeneration), } /// Disambiguates [`mint_generation`] calls landing in the same millisecond. -#[allow(dead_code)] static GENERATION_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); /// Mint a generation identifier, which satisfies [`CheckpointGeneration`]'s format by @@ -705,7 +700,6 @@ static GENERATION_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::Ato /// Call this exactly once per attempt, after that attempt's payload has been gathered. /// Re-uploading an already-gathered payload must reuse its generation; enforcing that is the /// caller's job (see the coordinator). -#[allow(dead_code)] pub(super) fn mint_generation() -> CheckpointGeneration { let millis = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -719,7 +713,6 @@ pub(super) fn mint_generation() -> CheckpointGeneration { /// /// Only the exact-set [`CommitSnapshotRequest`] needs this; everything earlier in the pipeline /// speaks logical names, and the server derives each presigned target's storage name itself. -#[allow(dead_code)] fn storage_name(generation: &CheckpointGeneration, logical: &str) -> String { format!("checkpoint_{}__{logical}", generation.as_str()) } @@ -977,7 +970,6 @@ async fn run_pipeline( /// Unlike [`upload_snapshot_from_declarations_file`], an unusable declarations file is /// [`CheckpointResult::Skipped`] rather than `None`, because the coordinator's state machine /// distinguishes "nothing to do" from "tried and failed". -#[allow(dead_code)] pub(super) async fn run_checkpoint_from_declarations_file( path: &Path, client: Arc, @@ -999,7 +991,6 @@ pub(super) async fn run_checkpoint_from_declarations_file( /// Upload and commit an already-gathered payload under `generation`. Split out so a caller /// re-running the exact same attempt can reuse both the payload and the generation. -#[allow(dead_code)] async fn run_checkpoint_pipeline( client: Arc, generation: CheckpointGeneration, diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index e1bbac47323..42fc341d6c6 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -132,9 +132,6 @@ impl CheckpointGeneration { /// Construct from a string the caller has already shaped to [`Self::is_valid`]. /// `snapshot::mint_generation` is the only production caller and satisfies it by /// construction, so the invariant is a debug assertion rather than a fallible return. - // Nothing calls the checkpoint API in this file yet -- the periodic coordinator that does - // lands in a follow-up, stacked PR, and takes the `allow(dead_code)`s with it. - #[allow(dead_code)] pub(crate) fn from_validated(value: String) -> Self { debug_assert!( Self::is_valid(&value), @@ -143,7 +140,6 @@ impl CheckpointGeneration { Self(value) } - #[allow(dead_code)] pub fn as_str(&self) -> &str { &self.0 } @@ -163,7 +159,6 @@ impl std::fmt::Display for CheckpointGeneration { /// /// Exact-set: the server persists `objects` verbatim as the commit marker and selection /// later returns exactly that set, not everything sharing the generation prefix. -#[allow(dead_code)] #[derive(Debug, Clone, serde::Serialize)] pub struct CommitSnapshotRequest { pub generation: String, @@ -171,7 +166,6 @@ pub struct CommitSnapshotRequest { pub objects: Vec, } -#[allow(dead_code)] #[derive(Debug, Clone, serde::Deserialize)] pub struct CommitSnapshotResponse { pub generation: String, @@ -347,7 +341,6 @@ pub trait HarnessSupportClient: 'static + Send + Sync { /// Only call this once every object in `request.objects` (including /// `request.manifest_object`) has uploaded successfully; the server verifies existence /// and per-attempt size limits and rejects the whole commit otherwise. - #[allow(dead_code)] async fn commit_snapshot( &self, request: &CommitSnapshotRequest,