diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index 70524631e25..a43a1688d5d 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,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(); 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/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/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, 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::()] =