Skip to content

Commit d9bb9b8

Browse files
joeywangzroz-agent
andcommitted
Fix the finalize budget, hand the writer to the coordinator, drop scaffolding
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 <oz-agent@warp.dev>
1 parent dd98a69 commit d9bb9b8

4 files changed

Lines changed: 17 additions & 27 deletions

File tree

app/src/ai/agent_sdk/driver.rs

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -785,6 +785,10 @@ impl AgentDriver {
785785
client,
786786
id,
787787
working_dir.clone(),
788+
// Shared with the history subscription so every attempt can drain
789+
// queued `file` appends before the declarations script runs, exactly
790+
// as `run_snapshot_upload` does on the legacy path.
791+
snapshot_file_writer.clone(),
788792
ctx.spawner(),
789793
checkpoint_interval
790794
.unwrap_or(checkpoint_coordinator::DEFAULT_CHECKPOINT_INTERVAL),
@@ -3855,12 +3859,22 @@ impl AgentDriver {
38553859
}
38563860

38573861
// When the periodic checkpoint coordinator is active, it owns the entire
3858-
// end-of-run path: `finalize` regenerates declarations, runs one last
3859-
// best-effort attempt bounded by `upload_timeout`, and commits it as the
3862+
// end-of-run path: `finalize` drains the declarations writer, regenerates
3863+
// declarations, runs one last best-effort attempt, and commits it as the
38603864
// selected checkpoint. This replaces the legacy one-shot upload below so
38613865
// there is exactly one end-of-run snapshot path, not two.
3866+
//
3867+
// The budget must come from `finalize_budget`, not from `upload_timeout` alone:
3868+
// the coordinator's floor is `script_timeout + upload_timeout`, so a smaller
3869+
// budget silently skips the final attempt — and since this path `return`s past
3870+
// the legacy upload below, that would mean no end-of-run snapshot at all.
38623871
if let Some(coordinator) = checkpoint_coordinator {
3863-
coordinator.finalize(upload_timeout).await;
3872+
coordinator
3873+
.finalize(checkpoint_coordinator::finalize_budget(
3874+
script_timeout,
3875+
upload_timeout,
3876+
))
3877+
.await;
38643878
return;
38653879
}
38663880

app/src/ai/agent_sdk/driver/checkpoint_coordinator.rs

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,6 @@
1616
//! This trades a small amount of latency (up to [`SAFE_BOUNDARY_POLL_INTERVAL`]) for
1717
//! avoiding new push-subscription wiring through the UI model graph.
1818
19-
// This module has no production caller yet -- `AgentDriver` wires
20-
// `CheckpointCoordinatorHandle::new` into its spawn/finalize lifecycle in a
21-
// follow-up, stacked PR. Until then, `CheckpointCoordinatorHandle::new` (and
22-
// everything it reaches) is unreachable outside `#[cfg(test)]`, which reaches the
23-
// same code through `new_for_test`. This module-level allow is temporary and should
24-
// be removable once that PR merges on top of this one.
25-
#![allow(dead_code)]
26-
2719
use std::path::PathBuf;
2820
use std::sync::Arc;
2921
use std::time::Duration;

app/src/ai/agent_sdk/driver/snapshot.rs

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -665,9 +665,6 @@ struct SnapshotOutcome {
665665

666666
/// Outcome of one checkpoint attempt, where [`SnapshotOutcome`] only covers per-entry upload
667667
/// results within that attempt.
668-
// The checkpoint pipeline below has no production caller until the periodic coordinator
669-
// lands in a follow-up, stacked PR; the `allow(dead_code)`s go away with it.
670-
#[allow(dead_code)]
671668
#[derive(Debug)]
672669
pub(super) enum CheckpointResult {
673670
/// `generation` is now the server's selected checkpoint.
@@ -691,12 +688,10 @@ pub(super) enum CheckpointResult {
691688
/// [`SnapshotUploadMode`] for the server-side semantics.
692689
enum PipelineMode {
693690
Legacy,
694-
#[allow(dead_code)]
695691
Checkpoint(CheckpointGeneration),
696692
}
697693

698694
/// Disambiguates [`mint_generation`] calls landing in the same millisecond.
699-
#[allow(dead_code)]
700695
static GENERATION_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
701696

702697
/// 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
705700
/// Call this exactly once per attempt, after that attempt's payload has been gathered.
706701
/// Re-uploading an already-gathered payload must reuse its generation; enforcing that is the
707702
/// caller's job (see the coordinator).
708-
#[allow(dead_code)]
709703
pub(super) fn mint_generation() -> CheckpointGeneration {
710704
let millis = std::time::SystemTime::now()
711705
.duration_since(std::time::UNIX_EPOCH)
@@ -719,7 +713,6 @@ pub(super) fn mint_generation() -> CheckpointGeneration {
719713
///
720714
/// Only the exact-set [`CommitSnapshotRequest`] needs this; everything earlier in the pipeline
721715
/// speaks logical names, and the server derives each presigned target's storage name itself.
722-
#[allow(dead_code)]
723716
fn storage_name(generation: &CheckpointGeneration, logical: &str) -> String {
724717
format!("checkpoint_{}__{logical}", generation.as_str())
725718
}
@@ -977,7 +970,6 @@ async fn run_pipeline(
977970
/// Unlike [`upload_snapshot_from_declarations_file`], an unusable declarations file is
978971
/// [`CheckpointResult::Skipped`] rather than `None`, because the coordinator's state machine
979972
/// distinguishes "nothing to do" from "tried and failed".
980-
#[allow(dead_code)]
981973
pub(super) async fn run_checkpoint_from_declarations_file(
982974
path: &Path,
983975
client: Arc<dyn HarnessSupportClient>,
@@ -999,7 +991,6 @@ pub(super) async fn run_checkpoint_from_declarations_file(
999991

1000992
/// Upload and commit an already-gathered payload under `generation`. Split out so a caller
1001993
/// re-running the exact same attempt can reuse both the payload and the generation.
1002-
#[allow(dead_code)]
1003994
async fn run_checkpoint_pipeline(
1004995
client: Arc<dyn HarnessSupportClient>,
1005996
generation: CheckpointGeneration,

app/src/server/server_api/harness_support.rs

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -132,9 +132,6 @@ impl CheckpointGeneration {
132132
/// Construct from a string the caller has already shaped to [`Self::is_valid`].
133133
/// `snapshot::mint_generation` is the only production caller and satisfies it by
134134
/// construction, so the invariant is a debug assertion rather than a fallible return.
135-
// Nothing calls the checkpoint API in this file yet -- the periodic coordinator that does
136-
// lands in a follow-up, stacked PR, and takes the `allow(dead_code)`s with it.
137-
#[allow(dead_code)]
138135
pub(crate) fn from_validated(value: String) -> Self {
139136
debug_assert!(
140137
Self::is_valid(&value),
@@ -143,7 +140,6 @@ impl CheckpointGeneration {
143140
Self(value)
144141
}
145142

146-
#[allow(dead_code)]
147143
pub fn as_str(&self) -> &str {
148144
&self.0
149145
}
@@ -163,15 +159,13 @@ impl std::fmt::Display for CheckpointGeneration {
163159
///
164160
/// Exact-set: the server persists `objects` verbatim as the commit marker and selection
165161
/// later returns exactly that set, not everything sharing the generation prefix.
166-
#[allow(dead_code)]
167162
#[derive(Debug, Clone, serde::Serialize)]
168163
pub struct CommitSnapshotRequest {
169164
pub generation: String,
170165
pub manifest_object: String,
171166
pub objects: Vec<String>,
172167
}
173168

174-
#[allow(dead_code)]
175169
#[derive(Debug, Clone, serde::Deserialize)]
176170
pub struct CommitSnapshotResponse {
177171
pub generation: String,
@@ -347,7 +341,6 @@ pub trait HarnessSupportClient: 'static + Send + Sync {
347341
/// Only call this once every object in `request.objects` (including
348342
/// `request.manifest_object`) has uploaded successfully; the server verifies existence
349343
/// and per-attempt size limits and rejects the whole commit otherwise.
350-
#[allow(dead_code)]
351344
async fn commit_snapshot(
352345
&self,
353346
request: &CommitSnapshotRequest,

0 commit comments

Comments
 (0)