REMOTE-2111 (3/3): Wire the periodic checkpoint coordinator into AgentDriver

app/src/ai/agent_sdk/driver.rs

@@ -313,6 +313,10 @@ pub struct AgentDriverOptions {
     pub snapshot_upload_timeout: Option<Duration>,
     /// Declarations script timeout override.
     pub snapshot_script_timeout: Option<Duration>,
+    /// 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<Duration>,
     /// 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<checkpoint_coordinator::CheckpointCoordinatorHandle>,
+
     /// 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,36 @@ 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(),
+                    // 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),
+                    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 +820,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 +864,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 +3833,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 +3858,26 @@ impl AgentDriver {
             return;
         }
 
+        // When the periodic checkpoint coordinator is active, it owns the entire
+        // 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(checkpoint_coordinator::finalize_budget(
+                    script_timeout,
+                    upload_timeout,
+                ))
+                .await;
+            return;
+        }
+
         let Ok((working_dir, client)) = spawner
             .spawn(|me, ctx| {
                 let client = ServerApiProvider::as_ref(ctx).get_harness_support_client();

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;

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<dyn HarnessSupportClient>,
@@ -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<dyn HarnessSupportClient>,
     generation: CheckpointGeneration,

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()),

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,15 +159,13 @@ 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,
     pub manifest_object: String,
     pub objects: Vec<String>,
 }
 
-#[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,

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::<FeatureFlag>()] =