From 1e98b54ed02ce6fb444f1b8ec591f4bb6e62026b Mon Sep 17 00:00:00 2001 From: Joey Wang Date: Fri, 31 Jul 2026 09:55:33 -0400 Subject: [PATCH 1/4] Add checkpoint upload/commit pipeline mechanics Adds the storage mechanics for one periodic-handoff checkpoint attempt (REMOTE-2111), as a self-contained module the (forthcoming, stacked) periodic checkpoint coordinator will drive: - harness_support.rs: SnapshotUploadMode (legacy|checkpoint) and generation on SnapshotUploadRequest; CheckpointGeneration (a validated, client-minted generation identifier); CommitSnapshotRequest/ Response; HarnessSupportClient::commit_snapshot. - snapshot.rs: CheckpointResult (Committed/Skipped/Failed outcome of one attempt); PipelineMode selects legacy vs. checkpoint-mode uploads; mint_generation (per-attempt, collision-free generation IDs); storage_name (generation-prefixed object naming); and the run_checkpoint_from_declarations_file / run_checkpoint_pipeline entry points: gather, upload every blob plus the manifest in checkpoint mode, then commit the exact object set that landed. Commit is withheld entirely (never partial) if the manifest or any required blob fails to upload, or if upload-target allocation fails. The whole pipeline has no production caller yet -- the coordinator that drives it periodically lands in a follow-up, stacked PR -- so the new items are marked #[allow(dead_code)] with a comment explaining why; each annotation is removable once that PR merges on top of this one. Co-Authored-By: Oz --- app/src/ai/agent_sdk/driver/snapshot.rs | 223 ++++++++- app/src/ai/agent_sdk/driver/snapshot_tests.rs | 456 +++++++++++++++++- app/src/server/server_api/harness_support.rs | 138 ++++++ 3 files changed, 806 insertions(+), 11 deletions(-) diff --git a/app/src/ai/agent_sdk/driver/snapshot.rs b/app/src/ai/agent_sdk/driver/snapshot.rs index 74d8e23f893..121c7fc69e1 100644 --- a/app/src/ai/agent_sdk/driver/snapshot.rs +++ b/app/src/ai/agent_sdk/driver/snapshot.rs @@ -48,7 +48,8 @@ use crate::server::server_api::ai::{ UploadLocalHandoffSnapshotRequest, }; use crate::server::server_api::harness_support::{ - HarnessSupportClient, SnapshotFileInfo, SnapshotUploadRequest, UploadTarget, upload_to_target, + CheckpointGeneration, CommitSnapshotRequest, HarnessSupportClient, SnapshotFileInfo, + SnapshotUploadRequest, UploadTarget, upload_to_target, }; /// Default path of the declarations file when neither the env var override nor a task ID @@ -210,8 +211,9 @@ pub(super) async fn run_declarations_script( /// /// Reads `$OZ_SNAPSHOT_DECLARATIONS_FILE` for the operator/test override, then delegates to /// [`resolve_declarations_path_with_override`] so tests can exercise the pure logic without -/// racing on the shared env var. -fn resolve_declarations_path(task_id: Option<&AmbientAgentTaskId>) -> PathBuf { +/// racing on the shared env var. `pub(super)` so `checkpoint_coordinator` can resolve the same +/// path used by the declarations writer and by [`run_declarations_script`]. +pub(super) fn resolve_declarations_path(task_id: Option<&AmbientAgentTaskId>) -> PathBuf { resolve_declarations_path_with_override(task_id, std::env::var_os(DECLARATIONS_PATH_ENV_VAR)) } @@ -651,6 +653,87 @@ struct SnapshotOutcome { manifest_uploaded: bool, } +/// Outcome of one checkpoint attempt, as opposed to [`SnapshotOutcome`] which only tracks +/// per-entry upload results within a single attempt. +// The whole checkpoint pipeline below (through `run_checkpoint_pipeline`) has no +// production caller yet -- the periodic checkpoint coordinator that drives it lands +// in a follow-up, stacked PR. `#[allow(dead_code)]` is temporary and should be +// removable once that PR is merged on top of this one. +#[allow(dead_code)] +#[derive(Debug)] +pub(super) enum CheckpointResult { + /// Every required object (blobs plus manifest) for `generation` uploaded successfully + /// and the exact-set commit call succeeded; `generation` is now the server's selected + /// checkpoint. + Committed { generation: CheckpointGeneration }, + /// There were no usable declarations to checkpoint (declarations file missing, empty, + /// or containing no valid entries). No generation was minted and no network calls + /// beyond reading local state were made. + Skipped, + /// A required upload (a non-cap-skipped blob, or the manifest), the upload-target + /// allocation, or the commit call itself failed. `generation` is `None` only when the + /// attempt was cut off before a generation was even minted (e.g. an external timeout + /// wrapping the whole attempt). Any minted generation's objects (if uploaded) are left + /// as uncommitted debris in storage; the server's existing marker (if any) is untouched. + Failed { + generation: Option, + reason: String, + }, +} + +/// Selects which upload-accounting path the shared gather/upload pipeline uses for a given +/// attempt. See `SnapshotUploadMode` (`crate::server::server_api::harness_support`) for the +/// server-side semantics. +enum PipelineMode { + /// One-shot end-of-run upload: unprefixed object names, counted against the + /// execution's cumulative lifetime attachment quota. + Legacy, + /// Periodic or finalization checkpoint attempt: the server stores each requested file + /// as `checkpoint___` and does not charge the cumulative quota. + // Not constructed until the coordinator PR (see the allow(dead_code) note above). + #[allow(dead_code)] + Checkpoint(CheckpointGeneration), +} + +/// Monotonic disambiguator for [`mint_generation`] so two attempts minted within the same +/// millisecond (e.g. in tests, or on a very fast retry) never collide. +#[allow(dead_code)] +static GENERATION_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +/// Mint a new checkpoint generation identifier. +/// +/// Must be called exactly once per checkpoint *attempt*, and only after that attempt's +/// payload has been gathered ("frozen") — retrying the same already-gathered payload (e.g. +/// after a transient upload failure) must reuse the previously minted generation rather than +/// calling this again; any newly gathered payload always mints a fresh one. Enforcing that +/// distinction is the caller's responsibility (see the coordinator in +/// `checkpoint_coordinator.rs`). +/// +/// Format: `-`. This satisfies the server's +/// `[A-Za-z0-9._-]{1,128}` charset and never contains the reserved `__` separator. +#[allow(dead_code)] +pub(super) fn mint_generation() -> CheckpointGeneration { + let millis = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + let counter = GENERATION_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + CheckpointGeneration::from_validated(format!("{millis}-{counter}")) +} + +/// Compute the generation-prefixed storage object name for a logical filename (a blob or the +/// manifest), matching the server's `checkpoint___` convention. +/// +/// Used only when assembling the exact-set [`CommitSnapshotRequest`] after upload — the +/// *logical* name (produced by [`unique_filename`]) is what flows through gather, manifest +/// building, and the upload-targets request; the server itself derives the storage name for +/// each presigned upload target from that logical filename plus the request's `generation` +/// field, so no client-side renaming is needed before that point. +#[allow(dead_code)] +fn storage_name(generation: &CheckpointGeneration, logical: &str) -> String { + format!("checkpoint_{}__{logical}", generation.as_str()) +} + // --- Manifest schema --- #[derive(serde::Serialize)] @@ -887,6 +970,7 @@ async fn run_pipeline( upload_gathered_snapshot( client, + &PipelineMode::Legacy, manifest_filename, upload_files, repos, @@ -896,6 +980,125 @@ async fn run_pipeline( .await } +/// Run one checkpoint attempt from the declarations file at `path`: read declarations, gather +/// the payload, mint a generation for it, upload every blob plus the manifest in checkpoint +/// mode, and commit the exact set that landed. Never panics; all failure modes are reported +/// via the returned [`CheckpointResult`] (and, for unexpected failures, `report_error!`). +/// +/// Unlike [`upload_snapshot_from_declarations_file`], a missing/empty/unusable declarations +/// file is reported as [`CheckpointResult::Skipped`] rather than `None`, since the coordinator +/// needs to distinguish "nothing to do" from "tried and failed" to drive its state machine. +#[allow(dead_code)] +pub(super) async fn run_checkpoint_from_declarations_file( + path: &Path, + client: Arc, +) -> CheckpointResult { + log::info!("Checkpoint attempt starting from {}", path.display()); + let Some(declarations) = read_and_parse_declarations(path) else { + return CheckpointResult::Skipped; + }; + let declarations = drop_files_covered_by_repos(declarations); + if declarations.is_empty() { + log::info!("Checkpoint declarations empty after de-duplication; skipping attempt"); + return CheckpointResult::Skipped; + } + let gathered = gather_snapshot_entries(declarations).await; + // The generation is minted here, once the gathered payload (blob contents, manifest + // stubs) is frozen for this attempt — see `mint_generation`'s contract. + let generation = mint_generation(); + run_checkpoint_pipeline(client, generation, gathered).await +} + +/// Upload and commit an already-gathered payload under `generation`. Split out from +/// [`run_checkpoint_from_declarations_file`] so a caller retrying the exact same attempt (as +/// opposed to gathering fresh) can reuse both the payload and the generation. +#[allow(dead_code)] +async fn run_checkpoint_pipeline( + client: Arc, + generation: CheckpointGeneration, + gathered: GatheredSnapshot, +) -> CheckpointResult { + let GatheredSnapshot { + manifest_filename, + upload_files, + repos, + files, + pre_upload_entries, + } = gathered; + + let outcome = upload_gathered_snapshot( + client.clone(), + &PipelineMode::Checkpoint(generation.clone()), + manifest_filename.clone(), + upload_files, + repos, + files, + pre_upload_entries, + ) + .await; + log::info!( + "Checkpoint attempt generation={generation} pending commit", + generation = generation.as_str() + ); + let Some(outcome) = outcome else { + return CheckpointResult::Failed { + generation: Some(generation), + reason: "failed to allocate upload targets or serialize manifest".to_string(), + }; + }; + log_snapshot_outcome(&outcome); + + if !outcome.manifest_uploaded { + return CheckpointResult::Failed { + generation: Some(generation), + reason: "manifest failed to upload".to_string(), + }; + } + if outcome + .entries + .iter() + .any(|e| e.status == EntryStatus::Failed) + { + return CheckpointResult::Failed { + generation: Some(generation), + reason: "one or more required blobs failed to upload".to_string(), + }; + } + + // Exact-set commit: the manifest object plus every blob whose own upload actually + // succeeded (cap-skipped, gather-failed, and read-failed entries are never included, + // matching the server's exact-set contract). + let manifest_object = storage_name(&generation, &manifest_filename); + let mut objects: Vec = outcome + .entries + .iter() + .filter(|e| e.status == EntryStatus::Uploaded && e.label != manifest_filename) + .map(|e| storage_name(&generation, &e.label)) + .collect(); + objects.push(manifest_object.clone()); + + let commit_request = CommitSnapshotRequest { + generation: generation.as_str().to_string(), + manifest_object, + objects, + }; + match client.commit_snapshot(&commit_request).await { + Ok(response) => { + log::info!("Checkpoint committed: generation={}", response.generation); + CheckpointResult::Committed { generation } + } + Err(e) => { + let e = e.context("Failed to commit checkpoint snapshot"); + let reason = format!("{e:#}"); + report_error!(e); + CheckpointResult::Failed { + generation: Some(generation), + reason, + } + } + } +} + struct GatheredSnapshot { manifest_filename: String, upload_files: Vec, @@ -954,6 +1157,7 @@ async fn gather_snapshot_entries(declarations: Vec) -> Gathere async fn upload_gathered_snapshot( client: Arc, + mode: &PipelineMode, manifest_filename: String, mut upload_files: Vec, mut repos: Vec, @@ -990,12 +1194,13 @@ async fn upload_gathered_snapshot( let mut target_map: HashMap = HashMap::new(); for chunk in file_infos.chunks(UPLOAD_BATCH_SIZE) { - let targets = match client - .get_snapshot_upload_targets(&SnapshotUploadRequest { - files: chunk.to_vec(), - }) - .await - { + let request = match mode { + PipelineMode::Legacy => SnapshotUploadRequest::legacy(chunk.to_vec()), + PipelineMode::Checkpoint(generation) => { + SnapshotUploadRequest::checkpoint(generation.clone(), chunk.to_vec()) + } + }; + let targets = match client.get_snapshot_upload_targets(&request).await { Ok(t) => t, Err(e) => { // Pipeline-abort: route through report_error! so Sentry captures the structured diff --git a/app/src/ai/agent_sdk/driver/snapshot_tests.rs b/app/src/ai/agent_sdk/driver/snapshot_tests.rs index f305880f0df..a6791000016 100644 --- a/app/src/ai/agent_sdk/driver/snapshot_tests.rs +++ b/app/src/ai/agent_sdk/driver/snapshot_tests.rs @@ -1,7 +1,7 @@ use std::fs; #[cfg(all(unix, not(target_os = "macos")))] use std::os::unix::ffi::OsStringExt as _; -use std::sync::Arc; +use std::sync::{Arc, Mutex as StdMutex}; use async_trait::async_trait; use command::blocking::Command as BlockingCommand; @@ -14,7 +14,8 @@ 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::{ - ReportArtifactResponse, ResolvePromptRequest, ResolvedHarnessPrompt, + CommitSnapshotResponse, ReportArtifactResponse, ResolvePromptRequest, ResolvedHarnessPrompt, + SnapshotUploadMode, }; // ------------------------------------------------------------------------------------------------ @@ -42,6 +43,14 @@ struct TestClient { /// alignment contract, the trailing files in the request end up with no target and are /// marked `skipped` downstream. drop_trailing_targets: usize, + /// Whether `commit_snapshot` should return an error, simulating a server-side commit + /// rejection (e.g. a missing required object). + fail_commit: bool, + /// Every `get_snapshot_upload_targets` request received, in order, for checkpoint-mode + /// assertions (mode/generation used, plain logical filenames). + upload_requests: Arc>>, + /// Every `commit_snapshot` request received, in order, for exact-set assertions. + commit_requests: Arc>>, } impl TestClient { @@ -51,6 +60,9 @@ impl TestClient { http: build_test_http_client(), fail_get_targets: false, drop_trailing_targets: 0, + fail_commit: false, + upload_requests: Arc::new(StdMutex::new(Vec::new())), + commit_requests: Arc::new(StdMutex::new(Vec::new())), }) } @@ -60,6 +72,9 @@ impl TestClient { http: build_test_http_client(), fail_get_targets: true, drop_trailing_targets: 0, + fail_commit: false, + upload_requests: Arc::new(StdMutex::new(Vec::new())), + commit_requests: Arc::new(StdMutex::new(Vec::new())), }) } @@ -69,8 +84,31 @@ impl TestClient { http: build_test_http_client(), fail_get_targets: false, drop_trailing_targets: drop_trailing, + fail_commit: false, + upload_requests: Arc::new(StdMutex::new(Vec::new())), + commit_requests: Arc::new(StdMutex::new(Vec::new())), }) } + + fn new_failing_commit(server_base_url: String) -> Arc { + Arc::new(Self { + server_base_url, + http: build_test_http_client(), + fail_get_targets: false, + drop_trailing_targets: 0, + fail_commit: true, + upload_requests: Arc::new(StdMutex::new(Vec::new())), + commit_requests: Arc::new(StdMutex::new(Vec::new())), + }) + } + + fn upload_requests(&self) -> Vec { + self.upload_requests.lock().unwrap().clone() + } + + fn commit_requests(&self) -> Vec { + self.commit_requests.lock().unwrap().clone() + } } #[async_trait] @@ -132,6 +170,7 @@ impl HarnessSupportClient for TestClient { &self, request: &SnapshotUploadRequest, ) -> Result> { + self.upload_requests.lock().unwrap().push(request.clone()); if self.fail_get_targets { anyhow::bail!("simulated get_snapshot_upload_targets failure"); } @@ -155,6 +194,19 @@ impl HarnessSupportClient for TestClient { Ok(targets) } + async fn commit_snapshot( + &self, + request: &CommitSnapshotRequest, + ) -> Result { + self.commit_requests.lock().unwrap().push(request.clone()); + if self.fail_commit { + anyhow::bail!("simulated commit_snapshot failure"); + } + Ok(CommitSnapshotResponse { + generation: request.generation.clone(), + }) + } + fn http_client(&self) -> &http_client::Client { &self.http } @@ -1463,3 +1515,403 @@ fn e2e_repo_plus_inside_and_outside_files_filters_overlap() { file_mock.assert(); manifest_mock.assert(); } + +// ------------------------------------------------------------------------------------------------ +// REMOTE-2111: checkpoint (periodic handoff) pipeline. +// ------------------------------------------------------------------------------------------------ + +#[test] +fn mint_generation_produces_unique_charset_valid_ids() { + let a = mint_generation(); + let b = mint_generation(); + assert_ne!(a.as_str(), b.as_str(), "successive generations must differ"); + for generation in [&a, &b] { + let s = generation.as_str(); + assert!(!s.is_empty() && s.len() <= 128, "length out of bounds: {s}"); + assert!( + s.chars() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | '_')), + "generation contains disallowed characters: {s}" + ); + assert!(!s.contains("__"), "generation must not contain '__': {s}"); + } +} + +#[test] +fn storage_name_prefixes_logical_name_with_generation() { + let generation = CheckpointGeneration::new_for_test("1700000000000-0"); + assert_eq!( + storage_name(&generation, "snapshot_state.json"), + "checkpoint_1700000000000-0__snapshot_state.json" + ); +} + +#[test] +fn checkpoint_commits_generation_prefixed_storage_names_while_upload_targets_use_plain_names() { + let tempdir = snaptest_tempdir(); + let file_path = tempdir.path().join("note.txt"); + fs::write(&file_path, b"hello").unwrap(); + let decl_dir = snaptest_tempdir(); + let declarations_path = write_declarations(decl_dir.path(), &[], &[&file_path]); + + let mut server = Server::new(); + let file_mock = server + .mock("PUT", upload_path("note\\.txt")) + .with_status(200) + .expect(1) + .create(); + let manifest_mock = server + .mock("PUT", upload_path("snapshot_state\\.json")) + .with_status(200) + .expect(1) + .create(); + + let client = TestClient::new(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let CheckpointResult::Committed { generation } = result else { + panic!("expected Committed, got {result:?}"); + }; + + // Upload-target requests use checkpoint mode with the plain logical filename; the server + // (not the client) derives the storage name from `generation` + `filename`. + let upload_requests = client.upload_requests(); + assert!(!upload_requests.is_empty()); + for request in &upload_requests { + assert_eq!(request.mode, SnapshotUploadMode::Checkpoint); + assert_eq!(request.generation.as_deref(), Some(generation.as_str())); + } + assert!( + upload_requests + .iter() + .any(|r| r.files.iter().any(|f| f.filename == "note.txt")), + "expected an upload-targets request naming the plain logical filename" + ); + + // The commit request is the only place storage names appear, and they must be + // generation-prefixed. + let commit_requests = client.commit_requests(); + assert_eq!(commit_requests.len(), 1, "exactly one commit call expected"); + let commit = &commit_requests[0]; + assert_eq!(commit.generation, generation.as_str()); + let expected_manifest = format!("checkpoint_{}__snapshot_state.json", generation.as_str()); + assert_eq!(commit.manifest_object, expected_manifest); + assert!(commit.objects.contains(&commit.manifest_object)); + assert!( + commit + .objects + .contains(&format!("checkpoint_{}__note.txt", generation.as_str())), + "expected note.txt's storage name in commit objects: {:?}", + commit.objects + ); + file_mock.assert(); + manifest_mock.assert(); +} + +#[test] +fn checkpoint_withholds_commit_when_a_required_blob_fails() { + // A blob upload fails permanently (404 is a non-retryable status). The exact-set contract + // requires withholding the commit entirely rather than committing a partial set. The + // manifest mock must return success here so this test isolates the blob-failure branch: + // without it, an unmocked manifest PUT would also fail against mockito's default + // (non-2xx) response for unmatched routes, and this test would pass for the wrong reason + // (falling into the "manifest failed to upload" branch instead of the intended one). + let tempdir = snaptest_tempdir(); + let file_path = tempdir.path().join("bad.txt"); + fs::write(&file_path, b"will-fail").unwrap(); + let decl_dir = snaptest_tempdir(); + let declarations_path = write_declarations(decl_dir.path(), &[], &[&file_path]); + + let mut server = Server::new(); + let file_mock = server + .mock("PUT", upload_path("bad\\.txt")) + .with_status(404) + .expect(1) + .create(); + let manifest_mock = server + .mock("PUT", upload_path("snapshot_state\\.json")) + .with_status(200) + .create(); + + let client = TestClient::new(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let CheckpointResult::Failed { reason, .. } = result else { + panic!("expected Failed, got {result:?}"); + }; + assert!( + reason.contains("blob"), + "expected the blob-failure reason specifically, got: {reason}" + ); + assert!( + client.commit_requests().is_empty(), + "commit must be withheld when a required blob fails" + ); + file_mock.assert(); + drop(manifest_mock); +} + +#[test] +fn checkpoint_manifest_upload_failure_withholds_commit() { + // The manifest itself fails to upload while the only blob succeeds. Losing the manifest + // means there is no rehydration catalogue even if blobs landed, so commit must still be + // withheld -- this exercises a distinct branch from the blob-failure test above. + let tempdir = snaptest_tempdir(); + let file_path = tempdir.path().join("ok.txt"); + fs::write(&file_path, b"fine").unwrap(); + let decl_dir = snaptest_tempdir(); + let declarations_path = write_declarations(decl_dir.path(), &[], &[&file_path]); + + let mut server = Server::new(); + let file_mock = server + .mock("PUT", upload_path("ok\\.txt")) + .with_status(200) + .create(); + // The upload path retries transient-looking failures with bounded retry before giving + // up, so a persistently failing manifest upload is attempted more than once. + let manifest_mock = server + .mock("PUT", upload_path("snapshot_state\\.json")) + .with_status(500) + .expect_at_least(1) + .create(); + + let client = TestClient::new(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let CheckpointResult::Failed { reason, .. } = result else { + panic!("expected Failed, got {result:?}"); + }; + assert!( + reason.contains("manifest"), + "expected the manifest-failure reason specifically, got: {reason}" + ); + assert!( + client.commit_requests().is_empty(), + "commit must be withheld when the manifest fails to upload" + ); + manifest_mock.assert(); + drop(file_mock); +} + +#[test] +fn checkpoint_target_allocation_failure_skips_commit() { + // The server refuses to allocate upload targets at all. No blobs or manifest are ever + // uploaded, and commit must never be attempted. + let tempdir = snaptest_tempdir(); + let file_path = tempdir.path().join("note.txt"); + fs::write(&file_path, b"hello").unwrap(); + let decl_dir = snaptest_tempdir(); + let declarations_path = write_declarations(decl_dir.path(), &[], &[&file_path]); + + let server = Server::new(); + let client = TestClient::new_failing_get_targets(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let CheckpointResult::Failed { reason, .. } = result else { + panic!("expected Failed, got {result:?}"); + }; + assert!( + reason.contains("allocate"), + "expected the target-allocation-failure reason specifically, got: {reason}" + ); + assert!( + client.commit_requests().is_empty(), + "commit must never be attempted when upload targets can't be allocated" + ); +} + +#[test] +fn checkpoint_commits_despite_cap_skipped_entries() { + // Declaring more files than the per-run cap should still commit the kept subset; cap-skipped + // entries must never appear in the exact-set commit request. + let tempdir = snaptest_tempdir(); + let decl_dir = snaptest_tempdir(); + let declared_count = MAX_SNAPSHOT_FILES_PER_RUN + 1; + let file_paths: Vec = (0..declared_count) + .map(|i| { + let path = tempdir.path().join(format!("file_{i:03}.txt")); + fs::write(&path, format!("content-{i}").as_bytes()).unwrap(); + path + }) + .collect(); + let file_refs: Vec<&Path> = file_paths.iter().map(|p| p.as_path()).collect(); + let declarations_path = write_declarations(decl_dir.path(), &[], &file_refs); + + let mut server = Server::new(); + let upload_mock = server + .mock("PUT", upload_path(r".+")) + .with_status(200) + .create(); + + let client = TestClient::new(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let CheckpointResult::Committed { generation } = result else { + panic!("expected Committed, got {result:?}"); + }; + + let commit_requests = client.commit_requests(); + assert_eq!(commit_requests.len(), 1); + let commit = &commit_requests[0]; + // Kept blobs (cap - 1, since the manifest reserves a slot) + the manifest itself. + let expected_objects = (MAX_SNAPSHOT_FILES_PER_RUN - 1) + 1; + assert_eq!( + commit.objects.len(), + expected_objects, + "cap-skipped entries must be excluded from the exact-set commit: {:?}", + commit.objects + ); + assert!(commit.objects.contains(&commit.manifest_object)); + let prefix = format!("checkpoint_{}__", generation.as_str()); + assert!(commit.objects.iter().all(|o| o.starts_with(&prefix))); + drop(upload_mock); +} + +#[test] +fn checkpoint_skips_when_declarations_file_missing() { + let tempdir = snaptest_tempdir(); + let missing = tempdir.path().join("does-not-exist.txt"); + let server = Server::new(); + let client = TestClient::new(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &missing, + client.clone(), + )); + assert!(matches!(result, CheckpointResult::Skipped)); + assert!(client.commit_requests().is_empty()); + assert!(client.upload_requests().is_empty()); +} + +#[test] +fn checkpoint_clean_repo_commits_manifest_only() { + let tempdir = snaptest_tempdir(); + init_git_repo(tempdir.path(), false); + let decl_dir = snaptest_tempdir(); + let declarations_path = write_declarations(decl_dir.path(), &[tempdir.path()], &[]); + + let mut server = Server::new(); + // No blob mock — a clean repo produces no patch, so only the manifest should upload. + let manifest_mock = server + .mock("PUT", upload_path("snapshot_state\\.json")) + .with_status(200) + .expect(1) + .create(); + + let client = TestClient::new(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let CheckpointResult::Committed { .. } = result else { + panic!("expected Committed, got {result:?}"); + }; + let commit_requests = client.commit_requests(); + assert_eq!(commit_requests.len(), 1); + assert_eq!( + commit_requests[0].objects, + vec![commit_requests[0].manifest_object.clone()], + "a clean repo should commit only the manifest object" + ); + manifest_mock.assert(); +} + +#[test] +fn checkpoint_commit_failure_reports_failed_result() { + let tempdir = snaptest_tempdir(); + init_git_repo(tempdir.path(), false); + let decl_dir = snaptest_tempdir(); + let declarations_path = write_declarations(decl_dir.path(), &[tempdir.path()], &[]); + + let mut server = Server::new(); + let manifest_mock = server + .mock("PUT", upload_path("snapshot_state\\.json")) + .with_status(200) + .expect(1) + .create(); + + let client = TestClient::new_failing_commit(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + assert!( + matches!(result, CheckpointResult::Failed { .. }), + "expected Failed, got {result:?}" + ); + assert_eq!( + client.commit_requests().len(), + 1, + "commit should still be attempted exactly once" + ); + manifest_mock.assert(); +} + +#[test] +fn checkpoint_new_gather_mints_a_fresh_generation_each_time() { + // Two independent checkpoint attempts (each a fresh gather) must never reuse a generation. + let tempdir = snaptest_tempdir(); + init_git_repo(tempdir.path(), false); + let decl_dir = snaptest_tempdir(); + let declarations_path = write_declarations(decl_dir.path(), &[tempdir.path()], &[]); + + let mut server = Server::new(); + let manifest_mock = server + .mock("PUT", upload_path("snapshot_state\\.json")) + .with_status(200) + .expect(2) + .create(); + + let client = TestClient::new(server.url()); + let rt = Runtime::new().unwrap(); + let first = rt.block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let second = rt.block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let ( + CheckpointResult::Committed { + generation: gen_one, + }, + CheckpointResult::Committed { + generation: gen_two, + }, + ) = (first, second) + else { + panic!("expected both attempts to commit"); + }; + assert_ne!( + gen_one.as_str(), + gen_two.as_str(), + "each fresh gather must mint a new generation" + ); + manifest_mock.assert(); +} diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index 9aabb260e22..70156f7b213 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -53,12 +53,129 @@ pub enum UploadFieldValue { ContentData, } +/// Selects how the server accounts for a `SnapshotUploadRequest`'s uploads. +/// +/// `Legacy` (the default) uses unprefixed object names and counts uploads against +/// the execution's cumulative lifetime attachment quota, matching today's one-shot +/// end-of-run snapshot. `Checkpoint` signs generation-prefixed object names and does +/// not consume that cumulative quota; the server enforces per-attempt limits instead +/// when the generation is committed via [`HarnessSupportClient::commit_snapshot`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize)] +#[serde(rename_all = "snake_case")] +pub enum SnapshotUploadMode { + #[default] + Legacy, + Checkpoint, +} + /// Request body for upload-snapshot upload targets. #[derive(Debug, Clone, serde::Serialize)] pub struct SnapshotUploadRequest { + /// Upload accounting mode. Omitted (default) is equivalent to `legacy` on the + /// server; see [`SnapshotUploadMode`]. + #[serde(skip_serializing_if = "is_default_mode")] + pub mode: SnapshotUploadMode, + /// Required when `mode` is [`SnapshotUploadMode::Checkpoint`]. Identifies the + /// checkpoint attempt; every requested file is uploaded by the server as + /// `checkpoint___`. Ignored for `legacy` mode. + #[serde(skip_serializing_if = "Option::is_none")] + pub generation: Option, pub files: Vec, } +fn is_default_mode(mode: &SnapshotUploadMode) -> bool { + *mode == SnapshotUploadMode::default() +} + +impl SnapshotUploadRequest { + /// Build a legacy-mode request, matching today's one-shot end-of-run upload. + pub fn legacy(files: Vec) -> Self { + Self { + mode: SnapshotUploadMode::Legacy, + generation: None, + files, + } + } + + /// Build a checkpoint-mode request for the given generation. + pub fn checkpoint(generation: CheckpointGeneration, files: Vec) -> Self { + Self { + mode: SnapshotUploadMode::Checkpoint, + generation: Some(generation.into_inner()), + files, + } + } +} + +/// A checkpoint generation identifier minted by the client for one checkpoint attempt. +/// +/// Must match the server's `[A-Za-z0-9._-]{1,128}` format and must not contain the +/// reserved `__` separator (validated by +/// [`crate::ai::agent_sdk::driver::snapshot::mint_generation`], the only production +/// constructor). Storage object basenames are `checkpoint___`; +/// the generation is a GCS keying detail only and must never leak into agent-visible +/// paths or restore commands. +#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize)] +#[serde(transparent)] +pub struct CheckpointGeneration(String); + +impl CheckpointGeneration { + /// Wrap an already-validated generation string. Exposed for tests; production + /// code should go through `snapshot::mint_generation` instead. + /// + /// The only caller (`driver::snapshot`'s test module) is itself excluded on + /// Windows (snapshot upload is cloud-agent-only and Linux-only), so this must + /// be gated the same way or it is dead code under Windows clippy/test builds. + #[cfg(all(test, not(windows)))] + pub(crate) fn new_for_test(value: impl Into) -> Self { + Self(value.into()) + } + + /// Construct from a pre-validated string. Crate-visible so `driver::snapshot` can + /// mint generations without duplicating this type. + // Only called by `snapshot::mint_generation`, which is itself unused until the + // periodic checkpoint coordinator (a follow-up, stacked PR) lands. + #[allow(dead_code)] + pub(crate) fn from_validated(value: String) -> Self { + Self(value) + } + + #[allow(dead_code)] + pub fn as_str(&self) -> &str { + &self.0 + } + + fn into_inner(self) -> String { + self.0 + } +} + +impl std::fmt::Display for CheckpointGeneration { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0) + } +} + +/// Request body for committing a fully uploaded checkpoint generation. Exact-set: the +/// server persists `objects` verbatim as the commit marker and later selection returns +/// exactly that set, never every object sharing the generation prefix. See +/// `docs/remote-2111-checkpoint-spec.md` (warp-server) for the full protocol. +// Not constructed until the periodic checkpoint coordinator (a follow-up, stacked PR) +// starts calling `HarnessSupportClient::commit_snapshot`. +#[allow(dead_code)] +#[derive(Debug, Clone, serde::Serialize)] +pub struct CommitSnapshotRequest { + pub generation: String, + pub manifest_object: String, + pub objects: Vec, +} + +#[allow(dead_code)] +#[derive(Debug, Clone, serde::Deserialize)] +pub struct CommitSnapshotResponse { + pub generation: String, +} + /// Describes a single file in a snapshot upload request. #[derive(Debug, Clone, serde::Serialize)] pub struct SnapshotFileInfo { @@ -224,6 +341,19 @@ pub trait HarnessSupportClient: 'static + Send + Sync { request: &SnapshotUploadRequest, ) -> Result>; + /// Commit a fully uploaded checkpoint generation for the active execution's exact + /// object set. Must only be called after every object named in `request.objects` + /// (including `request.manifest_object`) has itself uploaded successfully; the + /// server verifies existence and per-attempt size limits before this becomes the + /// selected checkpoint. + // Not called until the periodic checkpoint coordinator (a follow-up, stacked PR) + // lands. + #[allow(dead_code)] + async fn commit_snapshot( + &self, + request: &CommitSnapshotRequest, + ) -> Result; + /// Download the raw third-party harness transcript bytes for the current task's /// conversation. /// @@ -457,6 +587,14 @@ impl HarnessSupportClient for ServerApi { Ok(response.uploads) } + async fn commit_snapshot( + &self, + request: &CommitSnapshotRequest, + ) -> Result { + self.post_public_api("harness-support/commit-snapshot", request) + .await + } + async fn fetch_transcript(&self) -> Result { #[cfg(not(target_family = "wasm"))] { From a682589e42d54a1db7a0abb9f9fa6bbb4b6ce022 Mon Sep 17 00:00:00 2001 From: joeywangzr Date: Fri, 31 Jul 2026 15:16:15 +0000 Subject: [PATCH 2/4] Harden checkpoint mechanics: missing upload targets and unsafe filenames MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups on the checkpoint upload/commit pipeline. 1. Withhold the commit when the server omits a blob's upload target. `upload_gathered_snapshot` already anticipates a short `upload-snapshot` response and only warns; the target-less blob then reached `upload_entry` and was recorded as `EntryStatus::Skipped`. The checkpoint gate rejected only `Failed`, so the attempt committed a silently smaller object set and made it the selected checkpoint — discarding a previously complete one. `Skipped` conflated two very different causes. Split out `EntryStatus::NoTarget` for "server returned no presigned target" and treat it as fatal for a checkpoint attempt, leaving `Skipped` to mean only the deliberate `MAX_SNAPSHOT_FILES_PER_RUN` cap. Both still surface in the manifest as `skipped` so rehydration consumers keep a stable status vocabulary; the distinguishing detail stays in `error`. The legacy end-of-run path is unchanged in behavior. 2. Sanitize agent-controlled filenames before they reach storage names. `gather_repo` sanitized its filename component but `gather_file` used the raw basename, which flows into `checkpoint___`. `__` is documented as the reserved separator and the charset as `[A-Za-z0-9._-]`, but nothing enforced that for the filename half — and basenames come from agent-created files. `a__b.txt`, or `checkpoint_1700000000000-0__evil.txt`, produced an ambiguous storage name that either fails the server's existence check at commit time (losing the whole checkpoint) or lands under a different generation. `sanitize_name_component` now collapses to the server charset and squashes `_` runs so `__` can never appear; `gather_file` and `sanitize_filename_component` both route through it, and `storage_name` debug-asserts the invariant. 3. Smaller fixes. - `CheckpointGeneration::from_validated` now actually validates (debug assertion) instead of only claiming to in its name and docs. - Move the "pending commit" log after the outcome check so it no longer fires for attempts that failed to allocate targets. - Correct the `CheckpointResult::Failed { generation: None }` doc: an external timeout also reports `None` after a generation was minted, so `None` must not be read as "nothing landed in storage". Tests: commit is withheld when the server omits a blob target (chunked so the truncation hits a blob rather than the always-last manifest); sanitization never yields `__` or out-of-charset bytes; and an end-to-end guard that a hostile basename cannot produce an ambiguous storage name. Co-Authored-By: Oz --- app/src/ai/agent_sdk/driver/snapshot.rs | 129 ++++++++++---- app/src/ai/agent_sdk/driver/snapshot_tests.rs | 163 +++++++++++++++++- app/src/server/server_api/harness_support.rs | 21 +++ 3 files changed, 277 insertions(+), 36 deletions(-) diff --git a/app/src/ai/agent_sdk/driver/snapshot.rs b/app/src/ai/agent_sdk/driver/snapshot.rs index 121c7fc69e1..50ac72bf2ca 100644 --- a/app/src/ai/agent_sdk/driver/snapshot.rs +++ b/app/src/ai/agent_sdk/driver/snapshot.rs @@ -584,7 +584,18 @@ struct SnapshotUploadFile { enum EntryStatus { Uploaded, Failed, + /// Deliberately dropped from the upload plan to honor [`MAX_SNAPSHOT_FILES_PER_RUN`]. + /// This is a policy decision rather than a failure, so a checkpoint attempt may still + /// commit the kept subset. Skipped, + /// The server returned no presigned target for this blob — a contract violation of + /// `upload-snapshot`'s positional alignment (see the length-mismatch warning in + /// [`upload_gathered_snapshot`]). + /// + /// Deliberately distinct from [`EntryStatus::Skipped`]: nothing intentional happened + /// here, so a checkpoint attempt that hits this must be withheld rather than committing + /// a silently smaller object set over a previously complete selected checkpoint. + NoTarget, GatherFailed, ReadFailed, } @@ -595,6 +606,7 @@ impl EntryStatus { Self::Uploaded => "uploaded", Self::Failed => "failed", Self::Skipped => "skipped", + Self::NoTarget => "no_target", Self::GatherFailed => "gather_failed", Self::ReadFailed => "read_failed", } @@ -613,6 +625,7 @@ struct SnapshotSummary { uploaded: usize, failed: usize, skipped: usize, + no_target: usize, gather_failed: usize, read_failed: usize, total: usize, @@ -625,6 +638,7 @@ impl SnapshotSummary { uploaded: 0, failed: 0, skipped: 0, + no_target: 0, gather_failed: 0, read_failed: 0, total: entries.len(), @@ -635,6 +649,7 @@ impl SnapshotSummary { EntryStatus::Uploaded => s.uploaded += 1, EntryStatus::Failed => s.failed += 1, EntryStatus::Skipped => s.skipped += 1, + EntryStatus::NoTarget => s.no_target += 1, EntryStatus::GatherFailed => s.gather_failed += 1, EntryStatus::ReadFailed => s.read_failed += 1, } @@ -671,10 +686,15 @@ pub(super) enum CheckpointResult { /// beyond reading local state were made. Skipped, /// A required upload (a non-cap-skipped blob, or the manifest), the upload-target - /// allocation, or the commit call itself failed. `generation` is `None` only when the - /// attempt was cut off before a generation was even minted (e.g. an external timeout - /// wrapping the whole attempt). Any minted generation's objects (if uploaded) are left - /// as uncommitted debris in storage; the server's existing marker (if any) is untouched. + /// allocation, or the commit call itself failed. Any minted generation's objects (if + /// uploaded) are left as uncommitted debris in storage; the server's existing marker + /// (if any) is untouched. + /// + /// `generation` is `None` when the attempt never reported one back to the caller. That + /// covers both "cut off before a generation was minted" and "cut off by an external + /// timeout wrapping the whole attempt" — in the latter case a generation may well have + /// been minted and objects uploaded, so `None` must not be read as "nothing landed in + /// storage". Failed { generation: Option, reason: String, @@ -731,6 +751,10 @@ pub(super) fn mint_generation() -> CheckpointGeneration { /// field, so no client-side renaming is needed before that point. #[allow(dead_code)] fn storage_name(generation: &CheckpointGeneration, logical: &str) -> String { + debug_assert!( + !logical.contains("__"), + "logical snapshot filename must not contain the reserved `__` separator: {logical}" + ); format!("checkpoint_{}__{logical}", generation.as_str()) } @@ -1036,10 +1060,6 @@ async fn run_checkpoint_pipeline( pre_upload_entries, ) .await; - log::info!( - "Checkpoint attempt generation={generation} pending commit", - generation = generation.as_str() - ); let Some(outcome) = outcome else { return CheckpointResult::Failed { generation: Some(generation), @@ -1047,6 +1067,10 @@ async fn run_checkpoint_pipeline( }; }; log_snapshot_outcome(&outcome); + log::info!( + "Checkpoint attempt generation={generation} pending commit", + generation = generation.as_str() + ); if !outcome.manifest_uploaded { return CheckpointResult::Failed { @@ -1054,14 +1078,19 @@ async fn run_checkpoint_pipeline( reason: "manifest failed to upload".to_string(), }; } + // `NoTarget` is fatal alongside `Failed`: the server owes us a presigned target for + // every requested filename, so a missing one means this attempt would otherwise commit + // a silently smaller object set and make it the selected checkpoint, discarding a + // previously complete one. Only `Skipped` (the deliberate per-run cap) is tolerated. if outcome .entries .iter() - .any(|e| e.status == EntryStatus::Failed) + .any(|e| matches!(e.status, EntryStatus::Failed | EntryStatus::NoTarget)) { return CheckpointResult::Failed { generation: Some(generation), - reason: "one or more required blobs failed to upload".to_string(), + reason: "one or more required blobs failed to upload or had no upload target" + .to_string(), }; } @@ -1388,10 +1417,16 @@ async fn gather_file( let path = Path::new(file_path); match tokio::fs::read(path).await { Ok(content) => { - let preferred = path - .file_name() - .map(|n| n.to_string_lossy().to_string()) - .unwrap_or_else(|| file_path.to_string()); + // Sanitize before uniquifying: the basename comes from an agent-created file, and + // it ends up inside the `checkpoint___` storage name + // that the exact-set commit has to reproduce byte for byte. + let preferred = sanitize_name_component( + &path + .file_name() + .map(|n| n.to_string_lossy().to_string()) + .unwrap_or_else(|| file_path.to_string()), + "snapshot_artifact", + ); let filename = unique_filename(&preferred, used_filenames); let mime = mime_guess::from_path(path) .first_or_octet_stream() @@ -1430,18 +1465,21 @@ async fn gather_file( } /// Upload a single prepared file through the retry helper. -/// Produces an [`EntryResult`] labelled with the file's filename, or marked `skipped` if the -/// server did not return a target for it. +/// Produces an [`EntryResult`] labelled with the file's filename, or marked +/// [`EntryStatus::NoTarget`] if the server did not return a target for it. async fn upload_entry( http: &http_client::Client, file: &SnapshotUploadFile, target_map: &HashMap, ) -> EntryResult { let Some(target) = target_map.get(&file.filename) else { - log::warn!("No upload target for file '{}', skipping", file.filename); + log::warn!( + "No upload target returned by the server for file '{}'; it will not be uploaded", + file.filename + ); return EntryResult { label: file.filename.clone(), - status: EntryStatus::Skipped, + status: EntryStatus::NoTarget, error: Some("no upload target returned by server".to_string()), }; }; @@ -1488,7 +1526,11 @@ fn fold_upload_results( repo_entry.status = "failed"; repo_entry.error = entry.error.clone(); } - EntryStatus::Skipped => { + // Both surface in the manifest as `skipped` so downstream rehydration + // consumers keep seeing a stable status vocabulary; the distinguishing + // detail lives in `error` (and in the checkpoint gate, which treats + // `NoTarget` as fatal). + EntryStatus::Skipped | EntryStatus::NoTarget => { repo_entry.uploaded = Some(false); repo_entry.status = "skipped"; repo_entry.error = entry.error.clone(); @@ -1514,7 +1556,7 @@ fn fold_upload_results( file_entry.status = "failed"; file_entry.error = entry.error.clone(); } - EntryStatus::Skipped => { + EntryStatus::Skipped | EntryStatus::NoTarget => { file_entry.uploaded = Some(false); file_entry.status = "skipped"; file_entry.error = entry.error.clone(); @@ -1615,11 +1657,13 @@ fn log_snapshot_outcome(outcome: &SnapshotOutcome) { "manifest: failed" }; let header = format!( - "Snapshot upload: {}/{} uploaded (failed: {}, skipped: {}, gather_failed: {}, read_failed: {}; {manifest_bit})", + "Snapshot upload: {}/{} uploaded (failed: {}, skipped: {}, no_target: {}, \ + gather_failed: {}, read_failed: {}; {manifest_bit})", summary.uploaded, summary.total, summary.failed, summary.skipped, + summary.no_target, summary.gather_failed, summary.read_failed, ); @@ -1715,25 +1759,44 @@ async fn git_output_string(repo_dir: &Path, args: &[&str]) -> Option { if value.is_empty() { None } else { Some(value) } } -fn sanitize_filename_component(value: &str) -> String { - let sanitized = value - .chars() - .map(|c| { - if c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | '_') { - c - } else { - '_' - } - }) - .collect::(); +/// Collapse `value` into the server's `[A-Za-z0-9._-]` charset, squashing runs of `_` so the +/// result can never contain `__`. +/// +/// `__` is reserved as the separator in `checkpoint___` storage +/// object names (see [`storage_name`] and [`CheckpointGeneration`]), and the logical name half +/// is derived from **agent-controlled** input: workspace file basenames and repo directory +/// names. Leaving it unsanitized lets an agent-created file such as `a__b.txt` (or, worse, +/// `checkpoint_1700000000000-0__x.txt`) produce an ambiguous storage name, which either fails +/// the server's existence check at commit time — losing the whole checkpoint — or lands under +/// a different generation than intended. +/// +/// Returns `fallback` when nothing usable survives sanitization. +fn sanitize_name_component(value: &str, fallback: &str) -> String { + let mut sanitized = String::with_capacity(value.len()); + for c in value.chars() { + let c = if c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | '_') { + c + } else { + '_' + }; + // Squash runs so `__` can never appear in the output. + if c == '_' && sanitized.ends_with('_') { + continue; + } + sanitized.push(c); + } let trimmed = sanitized.trim_matches('_'); if trimmed.is_empty() { - "repo".to_string() + fallback.to_string() } else { trimmed.to_string() } } +fn sanitize_filename_component(value: &str) -> String { + sanitize_name_component(value, "repo") +} + fn unique_filename(preferred: &str, used: &mut HashSet) -> String { let preferred = Path::new(preferred) .file_name() diff --git a/app/src/ai/agent_sdk/driver/snapshot_tests.rs b/app/src/ai/agent_sdk/driver/snapshot_tests.rs index a6791000016..d859553b167 100644 --- a/app/src/ai/agent_sdk/driver/snapshot_tests.rs +++ b/app/src/ai/agent_sdk/driver/snapshot_tests.rs @@ -41,8 +41,15 @@ struct TestClient { /// Number of trailing response entries to drop, simulating a server that returns fewer /// targets than the request contained (contract violation). Under the positional /// alignment contract, the trailing files in the request end up with no target and are - /// marked `skipped` downstream. + /// marked [`EntryStatus::NoTarget`] downstream. drop_trailing_targets: usize, + /// Restrict `drop_trailing_targets` to the first `get_snapshot_upload_targets` call. + /// + /// `upload_gathered_snapshot` always appends the manifest last, so truncating *every* + /// chunk always costs the manifest its target. Truncating only the first chunk (which + /// requires more than [`UPLOAD_BATCH_SIZE`] files, so there are at least two) isolates + /// the blob-level `NoTarget` path with the manifest still uploading cleanly. + drop_trailing_first_call_only: bool, /// Whether `commit_snapshot` should return an error, simulating a server-side commit /// rejection (e.g. a missing required object). fail_commit: bool, @@ -60,6 +67,7 @@ impl TestClient { http: build_test_http_client(), fail_get_targets: false, drop_trailing_targets: 0, + drop_trailing_first_call_only: false, fail_commit: false, upload_requests: Arc::new(StdMutex::new(Vec::new())), commit_requests: Arc::new(StdMutex::new(Vec::new())), @@ -72,6 +80,7 @@ impl TestClient { http: build_test_http_client(), fail_get_targets: true, drop_trailing_targets: 0, + drop_trailing_first_call_only: false, fail_commit: false, upload_requests: Arc::new(StdMutex::new(Vec::new())), commit_requests: Arc::new(StdMutex::new(Vec::new())), @@ -84,6 +93,25 @@ impl TestClient { http: build_test_http_client(), fail_get_targets: false, drop_trailing_targets: drop_trailing, + drop_trailing_first_call_only: false, + fail_commit: false, + upload_requests: Arc::new(StdMutex::new(Vec::new())), + commit_requests: Arc::new(StdMutex::new(Vec::new())), + }) + } + + /// Drops `drop_trailing` targets from the *first* upload-targets call only, leaving + /// later chunks (and therefore the manifest) intact. + fn new_dropping_trailing_on_first_call( + server_base_url: String, + drop_trailing: usize, + ) -> Arc { + Arc::new(Self { + server_base_url, + http: build_test_http_client(), + fail_get_targets: false, + drop_trailing_targets: drop_trailing, + drop_trailing_first_call_only: true, fail_commit: false, upload_requests: Arc::new(StdMutex::new(Vec::new())), commit_requests: Arc::new(StdMutex::new(Vec::new())), @@ -96,6 +124,7 @@ impl TestClient { http: build_test_http_client(), fail_get_targets: false, drop_trailing_targets: 0, + drop_trailing_first_call_only: false, fail_commit: true, upload_requests: Arc::new(StdMutex::new(Vec::new())), commit_requests: Arc::new(StdMutex::new(Vec::new())), @@ -170,7 +199,11 @@ impl HarnessSupportClient for TestClient { &self, request: &SnapshotUploadRequest, ) -> Result> { - self.upload_requests.lock().unwrap().push(request.clone()); + let call_index = { + let mut recorded = self.upload_requests.lock().unwrap(); + recorded.push(request.clone()); + recorded.len() - 1 + }; if self.fail_get_targets { anyhow::bail!("simulated get_snapshot_upload_targets failure"); } @@ -189,7 +222,12 @@ impl HarnessSupportClient for TestClient { fields: Vec::new(), }) .collect(); - let keep = targets.len().saturating_sub(self.drop_trailing_targets); + let drop_count = if self.drop_trailing_first_call_only && call_index > 0 { + 0 + } else { + self.drop_trailing_targets + }; + let keep = targets.len().saturating_sub(drop_count); targets.truncate(keep); Ok(targets) } @@ -1872,6 +1910,125 @@ fn checkpoint_commit_failure_reports_failed_result() { manifest_mock.assert(); } +#[test] +fn checkpoint_withholds_commit_when_the_server_omits_a_blob_upload_target() { + // Regression: a short `upload-snapshot` response used to mark the target-less blob + // `skipped`, which the checkpoint gate tolerated (it only rejected `failed`). The attempt + // then committed a silently smaller object set and made it the selected checkpoint, + // discarding a previously complete one. `EntryStatus::NoTarget` must now withhold commit. + // + // Declare more than UPLOAD_BATCH_SIZE files so the request is chunked and the truncation + // lands on a blob rather than on the always-last manifest entry. + let tempdir = snaptest_tempdir(); + let decl_dir = snaptest_tempdir(); + let declared_count = UPLOAD_BATCH_SIZE + 5; + let file_paths: Vec = (0..declared_count) + .map(|i| { + let path = tempdir.path().join(format!("file_{i:03}.txt")); + fs::write(&path, format!("content-{i}").as_bytes()).unwrap(); + path + }) + .collect(); + let file_refs: Vec<&Path> = file_paths.iter().map(|p| p.as_path()).collect(); + let declarations_path = write_declarations(decl_dir.path(), &[], &file_refs); + + let mut server = Server::new(); + let upload_mock = server + .mock("PUT", upload_path(r".+")) + .with_status(200) + .create(); + + let client = TestClient::new_dropping_trailing_on_first_call(server.url(), 1); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + + let CheckpointResult::Failed { reason, .. } = result else { + panic!("a blob without an upload target must fail the attempt, got {result:?}"); + }; + assert!( + reason.contains("upload target"), + "expected the missing-target reason specifically, got: {reason}" + ); + assert!( + client.commit_requests().is_empty(), + "commit must be withheld when the server omitted a blob's upload target" + ); + drop(upload_mock); +} + +#[test] +fn sanitize_name_component_never_yields_the_reserved_double_underscore() { + // The logical filename half of `checkpoint___` is derived from + // agent-controlled basenames, so it must never reintroduce the reserved separator. + for raw in [ + "a__b.txt", + "a b.txt", + "checkpoint_1700000000000-0__evil.txt", + "weird name?!.txt", + "___", + "ünïcödé.txt", + ] { + let sanitized = sanitize_name_component(raw, "snapshot_artifact"); + assert!( + !sanitized.contains("__"), + "sanitized {raw:?} still contains `__`: {sanitized}" + ); + assert!( + sanitized + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_')), + "sanitized {raw:?} left characters outside the server charset: {sanitized}" + ); + assert!(!sanitized.is_empty(), "sanitized {raw:?} became empty"); + } +} + +#[test] +fn checkpoint_storage_names_stay_unambiguous_for_hostile_filenames() { + // End-to-end guard for the same invariant: a file whose basename contains `__` must not + // produce a storage name with a second `__`, which would make the server's + // `checkpoint___` split ambiguous. + let tempdir = snaptest_tempdir(); + let hostile = tempdir.path().join("checkpoint_1700000000000-0__evil.txt"); + fs::write(&hostile, b"hostile").unwrap(); + let decl_dir = snaptest_tempdir(); + let declarations_path = write_declarations(decl_dir.path(), &[], &[&hostile]); + + let mut server = Server::new(); + let upload_mock = server + .mock("PUT", upload_path(r".+")) + .with_status(200) + .create(); + + let client = TestClient::new(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let CheckpointResult::Committed { generation } = result else { + panic!("expected Committed, got {result:?}"); + }; + + let commit = &client.commit_requests()[0]; + let prefix = format!("checkpoint_{}__", generation.as_str()); + for object in &commit.objects { + let suffix = object + .strip_prefix(&prefix) + .unwrap_or_else(|| panic!("object {object} is not under the generation prefix")); + assert!( + !suffix.contains("__"), + "storage name {object} has an ambiguous second `__` separator" + ); + } + drop(upload_mock); +} + #[test] fn checkpoint_new_gather_mints_a_fresh_generation_each_time() { // Two independent checkpoint attempts (each a fresh gather) must never reuse a generation. diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index 70156f7b213..442036f07f5 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -131,12 +131,33 @@ impl CheckpointGeneration { Self(value.into()) } + /// True when `value` satisfies the server's `[A-Za-z0-9._-]{1,128}` format and does not + /// contain the reserved `__` separator used by + /// `checkpoint___` storage object names. + fn is_valid(value: &str) -> bool { + !value.is_empty() + && value.len() <= 128 + && !value.contains("__") + && value + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_')) + } + /// Construct from a pre-validated string. Crate-visible so `driver::snapshot` can /// mint generations without duplicating this type. + /// + /// The debug assertion keeps the type's documented invariant honest: previously nothing + /// enforced it, so "validated" was aspirational. `snapshot::mint_generation` is the only + /// production caller and provably satisfies it, hence a `debug_assert!` rather than a + /// fallible constructor. // Only called by `snapshot::mint_generation`, which is itself unused until the // periodic checkpoint coordinator (a follow-up, stacked PR) lands. #[allow(dead_code)] pub(crate) fn from_validated(value: String) -> Self { + debug_assert!( + Self::is_valid(&value), + "checkpoint generation must match [A-Za-z0-9._-]{{1,128}} and exclude `__`: {value}" + ); Self(value) } From d23c0ae548614622d927a654d7abeb4419b00ab9 Mon Sep 17 00:00:00 2001 From: Joey Wang Date: Fri, 31 Jul 2026 15:18:24 -0400 Subject: [PATCH 3/4] Make minted snapshot filenames satisfy the server contract; retry commit Review follow-ups on the checkpoint upload/commit pipeline, plus a pass to cut the added comments down to what is not already obvious from the code. 1. Mint only filenames the server will accept. `sanitize_name_component` guaranteed the charset and the absence of the reserved `__` separator, but not the rest of the server's logical-name contract: at most 255 bytes, no leading `-`, not `.`/`..`, and -- on the legacy path -- nothing in the reserved `checkpoint_` namespace or the commit marker's own name. Those names come from agent-created workspace files, and the server rejects the *entire* upload-targets request when any one of them fails validation. So a single file called `-v.txt`, or one with a 300-character basename, permanently failed every subsequent checkpoint attempt (with a Sentry report each cycle), and a file called `checkpoint_notes.txt` took the end-of-run snapshot down with it. Sanitization now covers the whole contract, escaping reserved names rather than dropping the file. 2. Stop `unique_filename` from undoing that sanitization. Two files named `a_.txt` produced stem `a_` and therefore `a__2.txt`, reintroducing the separator and tripping `storage_name`'s debug assertion -- a panic in debug builds, from a function documented never to panic. The stem's trailing `_` is now trimmed, which makes the invariant hold at the single place logical names are minted; the assertion at the consumer is redundant and goes away. 3. Retry the commit call. Every blob and the manifest upload through `with_bounded_retry`, but the commit -- the one call where all the work is already done and paid for -- had none, so a single transient failure discarded a complete checkpoint and left its objects as debris. Committing the same generation twice is idempotent server-side, so the same bounded retry applies. Tests: the sanitizer and `unique_filename` are checked against a mirror of the server's validation, an end-to-end attempt with hostile basenames asserts every name we send and commit passes it, the commit-failure test now pins the retry count, and new cases pin the JSON for legacy uploads (unchanged wire format), checkpoint uploads, and commits, none of which the in-process test client exercises. Co-Authored-By: Oz --- app/src/ai/agent_sdk/driver/snapshot.rs | 213 +++++++------- app/src/ai/agent_sdk/driver/snapshot_tests.rs | 277 +++++++++++++----- app/src/server/server_api/harness_support.rs | 81 ++--- 3 files changed, 336 insertions(+), 235 deletions(-) diff --git a/app/src/ai/agent_sdk/driver/snapshot.rs b/app/src/ai/agent_sdk/driver/snapshot.rs index 50ac72bf2ca..7c4cdd768a5 100644 --- a/app/src/ai/agent_sdk/driver/snapshot.rs +++ b/app/src/ai/agent_sdk/driver/snapshot.rs @@ -584,17 +584,12 @@ struct SnapshotUploadFile { enum EntryStatus { Uploaded, Failed, - /// Deliberately dropped from the upload plan to honor [`MAX_SNAPSHOT_FILES_PER_RUN`]. - /// This is a policy decision rather than a failure, so a checkpoint attempt may still - /// commit the kept subset. + /// Deliberately dropped to honor [`MAX_SNAPSHOT_FILES_PER_RUN`]. A policy decision, not a + /// failure, so a checkpoint attempt may still commit the kept subset. Skipped, - /// The server returned no presigned target for this blob — a contract violation of - /// `upload-snapshot`'s positional alignment (see the length-mismatch warning in - /// [`upload_gathered_snapshot`]). - /// - /// Deliberately distinct from [`EntryStatus::Skipped`]: nothing intentional happened - /// here, so a checkpoint attempt that hits this must be withheld rather than committing - /// a silently smaller object set over a previously complete selected checkpoint. + /// The server returned no presigned target for this blob, violating `upload-snapshot`'s + /// positional alignment. Distinct from [`EntryStatus::Skipped`] because nothing + /// intentional happened: committing here would silently shrink the object set. NoTarget, GatherFailed, ReadFailed, @@ -668,69 +663,50 @@ struct SnapshotOutcome { manifest_uploaded: bool, } -/// Outcome of one checkpoint attempt, as opposed to [`SnapshotOutcome`] which only tracks -/// per-entry upload results within a single attempt. -// The whole checkpoint pipeline below (through `run_checkpoint_pipeline`) has no -// production caller yet -- the periodic checkpoint coordinator that drives it lands -// in a follow-up, stacked PR. `#[allow(dead_code)]` is temporary and should be -// removable once that PR is merged on top of this one. +/// 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 { - /// Every required object (blobs plus manifest) for `generation` uploaded successfully - /// and the exact-set commit call succeeded; `generation` is now the server's selected - /// checkpoint. + /// Every required object uploaded and the exact-set commit succeeded, so `generation` is + /// now the server's selected checkpoint. Committed { generation: CheckpointGeneration }, - /// There were no usable declarations to checkpoint (declarations file missing, empty, - /// or containing no valid entries). No generation was minted and no network calls - /// beyond reading local state were made. + /// Nothing to checkpoint: the declarations file was missing, empty, or had no valid + /// entries. No generation was minted and no network calls were made. Skipped, - /// A required upload (a non-cap-skipped blob, or the manifest), the upload-target - /// allocation, or the commit call itself failed. Any minted generation's objects (if - /// uploaded) are left as uncommitted debris in storage; the server's existing marker - /// (if any) is untouched. + /// A required upload, the upload-target allocation, or the commit failed. Uploaded + /// objects are left as uncommitted debris; the server's existing marker is untouched. /// - /// `generation` is `None` when the attempt never reported one back to the caller. That - /// covers both "cut off before a generation was minted" and "cut off by an external - /// timeout wrapping the whole attempt" — in the latter case a generation may well have - /// been minted and objects uploaded, so `None` must not be read as "nothing landed in - /// storage". + /// `generation` is `None` whenever the attempt never reported one back, which includes + /// being cut off by an external timeout after uploading — so `None` does not mean + /// "nothing landed in storage". Failed { generation: Option, reason: String, }, } -/// Selects which upload-accounting path the shared gather/upload pipeline uses for a given -/// attempt. See `SnapshotUploadMode` (`crate::server::server_api::harness_support`) for the -/// server-side semantics. +/// Which upload-accounting path the shared gather/upload pipeline uses. See +/// [`SnapshotUploadMode`] for the server-side semantics. enum PipelineMode { - /// One-shot end-of-run upload: unprefixed object names, counted against the - /// execution's cumulative lifetime attachment quota. Legacy, - /// Periodic or finalization checkpoint attempt: the server stores each requested file - /// as `checkpoint___` and does not charge the cumulative quota. - // Not constructed until the coordinator PR (see the allow(dead_code) note above). + // Not constructed until the coordinator PR. #[allow(dead_code)] Checkpoint(CheckpointGeneration), } -/// Monotonic disambiguator for [`mint_generation`] so two attempts minted within the same -/// millisecond (e.g. in tests, or on a very fast retry) never collide. +/// 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 new checkpoint generation identifier. -/// -/// Must be called exactly once per checkpoint *attempt*, and only after that attempt's -/// payload has been gathered ("frozen") — retrying the same already-gathered payload (e.g. -/// after a transient upload failure) must reuse the previously minted generation rather than -/// calling this again; any newly gathered payload always mints a fresh one. Enforcing that -/// distinction is the caller's responsibility (see the coordinator in -/// `checkpoint_coordinator.rs`). +/// Mint a `-` generation identifier, which satisfies +/// [`CheckpointGeneration`]'s format by construction. /// -/// Format: `-`. This satisfies the server's -/// `[A-Za-z0-9._-]{1,128}` charset and never contains the reserved `__` separator. +/// 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() @@ -741,20 +717,14 @@ pub(super) fn mint_generation() -> CheckpointGeneration { CheckpointGeneration::from_validated(format!("{millis}-{counter}")) } -/// Compute the generation-prefixed storage object name for a logical filename (a blob or the -/// manifest), matching the server's `checkpoint___` convention. +/// Reproduce the server's `checkpoint___` storage name for a +/// logical filename. /// -/// Used only when assembling the exact-set [`CommitSnapshotRequest`] after upload — the -/// *logical* name (produced by [`unique_filename`]) is what flows through gather, manifest -/// building, and the upload-targets request; the server itself derives the storage name for -/// each presigned upload target from that logical filename plus the request's `generation` -/// field, so no client-side renaming is needed before that point. +/// Only needed to assemble the exact-set [`CommitSnapshotRequest`]: everything earlier in the +/// pipeline speaks logical names, and the server derives the storage name for each presigned +/// target itself. #[allow(dead_code)] fn storage_name(generation: &CheckpointGeneration, logical: &str) -> String { - debug_assert!( - !logical.contains("__"), - "logical snapshot filename must not contain the reserved `__` separator: {logical}" - ); format!("checkpoint_{}__{logical}", generation.as_str()) } @@ -1004,14 +974,13 @@ async fn run_pipeline( .await } -/// Run one checkpoint attempt from the declarations file at `path`: read declarations, gather -/// the payload, mint a generation for it, upload every blob plus the manifest in checkpoint -/// mode, and commit the exact set that landed. Never panics; all failure modes are reported -/// via the returned [`CheckpointResult`] (and, for unexpected failures, `report_error!`). +/// Run one checkpoint attempt from the declarations file at `path`: gather the payload, mint a +/// generation, upload it in checkpoint mode, and commit the exact set that landed. Never +/// panics; every failure mode comes back as a [`CheckpointResult`]. /// -/// Unlike [`upload_snapshot_from_declarations_file`], a missing/empty/unusable declarations -/// file is reported as [`CheckpointResult::Skipped`] rather than `None`, since the coordinator -/// needs to distinguish "nothing to do" from "tried and failed" to drive its state machine. +/// 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, @@ -1027,15 +996,13 @@ pub(super) async fn run_checkpoint_from_declarations_file( return CheckpointResult::Skipped; } let gathered = gather_snapshot_entries(declarations).await; - // The generation is minted here, once the gathered payload (blob contents, manifest - // stubs) is frozen for this attempt — see `mint_generation`'s contract. + // Mint only once the payload is frozen — see `mint_generation`'s contract. let generation = mint_generation(); run_checkpoint_pipeline(client, generation, gathered).await } -/// Upload and commit an already-gathered payload under `generation`. Split out from -/// [`run_checkpoint_from_declarations_file`] so a caller retrying the exact same attempt (as -/// opposed to gathering fresh) can reuse both the payload and the generation. +/// 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, @@ -1078,10 +1045,9 @@ async fn run_checkpoint_pipeline( reason: "manifest failed to upload".to_string(), }; } - // `NoTarget` is fatal alongside `Failed`: the server owes us a presigned target for - // every requested filename, so a missing one means this attempt would otherwise commit - // a silently smaller object set and make it the selected checkpoint, discarding a - // previously complete one. Only `Skipped` (the deliberate per-run cap) is tolerated. + // Committing while any entry is `NoTarget` would make a silently smaller object set the + // selected checkpoint, discarding a previously complete one. Only `Skipped` (the + // deliberate per-run cap) is tolerated. if outcome .entries .iter() @@ -1094,9 +1060,7 @@ async fn run_checkpoint_pipeline( }; } - // Exact-set commit: the manifest object plus every blob whose own upload actually - // succeeded (cap-skipped, gather-failed, and read-failed entries are never included, - // matching the server's exact-set contract). + // Exact-set commit: the manifest plus every blob that actually uploaded. let manifest_object = storage_name(&generation, &manifest_filename); let mut objects: Vec = outcome .entries @@ -1111,7 +1075,11 @@ async fn run_checkpoint_pipeline( manifest_object, objects, }; - match client.commit_snapshot(&commit_request).await { + // Every object is already in storage by this point, so a transient failure here would + // throw away the whole attempt. Committing the same generation twice is idempotent + // server-side, which makes retrying safe. + let operation = format!("checkpoint commit '{}'", generation.as_str()); + match with_bounded_retry(&operation, || client.commit_snapshot(&commit_request)).await { Ok(response) => { log::info!("Checkpoint committed: generation={}", response.generation); CheckpointResult::Committed { generation } @@ -1417,15 +1385,14 @@ async fn gather_file( let path = Path::new(file_path); match tokio::fs::read(path).await { Ok(content) => { - // Sanitize before uniquifying: the basename comes from an agent-created file, and - // it ends up inside the `checkpoint___` storage name - // that the exact-set commit has to reproduce byte for byte. + // Sanitize before uniquifying: the basename is agent-controlled and ends up in the + // storage name the exact-set commit has to reproduce byte for byte. let preferred = sanitize_name_component( &path .file_name() .map(|n| n.to_string_lossy().to_string()) .unwrap_or_else(|| file_path.to_string()), - "snapshot_artifact", + FALLBACK_SNAPSHOT_FILENAME, ); let filename = unique_filename(&preferred, used_filenames); let mime = mime_guess::from_path(path) @@ -1526,10 +1493,8 @@ fn fold_upload_results( repo_entry.status = "failed"; repo_entry.error = entry.error.clone(); } - // Both surface in the manifest as `skipped` so downstream rehydration - // consumers keep seeing a stable status vocabulary; the distinguishing - // detail lives in `error` (and in the checkpoint gate, which treats - // `NoTarget` as fatal). + // Both surface as `skipped` to keep the manifest's status vocabulary stable + // for rehydration consumers; the distinguishing detail lives in `error`. EntryStatus::Skipped | EntryStatus::NoTarget => { repo_entry.uploaded = Some(false); repo_entry.status = "skipped"; @@ -1759,18 +1724,25 @@ async fn git_output_string(repo_dir: &Path, args: &[&str]) -> Option { if value.is_empty() { None } else { Some(value) } } -/// Collapse `value` into the server's `[A-Za-z0-9._-]` charset, squashing runs of `_` so the -/// result can never contain `__`. -/// -/// `__` is reserved as the separator in `checkpoint___` storage -/// object names (see [`storage_name`] and [`CheckpointGeneration`]), and the logical name half -/// is derived from **agent-controlled** input: workspace file basenames and repo directory -/// names. Leaving it unsanitized lets an agent-created file such as `a__b.txt` (or, worse, -/// `checkpoint_1700000000000-0__x.txt`) produce an ambiguous storage name, which either fails -/// the server's existence check at commit time — losing the whole checkpoint — or lands under -/// a different generation than intended. +/// Fallback used wherever a name sanitizes down to nothing usable. +const FALLBACK_SNAPSHOT_FILENAME: &str = "snapshot_artifact"; + +/// Prepended to names that would otherwise collide with the server's reserved namespace. +const RESERVED_NAME_ESCAPE: &str = "snapshot-"; + +/// Longest logical filename we will mint. The server rejects names over 255 bytes; the +/// remainder is headroom for the `_` de-duplication suffix [`unique_filename`] may append. +const MAX_SNAPSHOT_FILENAME_LEN: usize = 240; + +/// Reshape `value` into a logical snapshot filename the server will accept, falling back to +/// `fallback` when nothing usable survives. /// -/// Returns `fallback` when nothing usable survives sanitization. +/// Logical names are agent-controlled and the server rejects the *entire* upload-targets +/// request if one is malformed, so a single awkward basename would otherwise cost the whole +/// snapshot. Its rules: `[A-Za-z0-9._-]` only, at most 255 bytes, not `.` or `..`, no leading +/// `-`, and — on the legacy path — nothing in the reserved `checkpoint_` namespace. Runs of +/// `_` are squashed on top of that so the `checkpoint___` separator +/// stays unambiguous. fn sanitize_name_component(value: &str, fallback: &str) -> String { let mut sanitized = String::with_capacity(value.len()); for c in value.chars() { @@ -1779,18 +1751,31 @@ fn sanitize_name_component(value: &str, fallback: &str) -> String { } else { '_' }; - // Squash runs so `__` can never appear in the output. if c == '_' && sanitized.ends_with('_') { continue; } sanitized.push(c); } - let trimmed = sanitized.trim_matches('_'); - if trimmed.is_empty() { - fallback.to_string() - } else { - trimmed.to_string() + + let trimmed = sanitized + .trim_start_matches(['_', '-']) + .trim_end_matches('_'); + let mut name = match trimmed { + "" | "." | ".." => fallback.to_string(), + other => other.to_string(), + }; + if is_reserved_snapshot_name(&name) { + name.insert_str(0, RESERVED_NAME_ESCAPE); } + // Sanitized names are pure ASCII, so this always lands on a char boundary. + name.truncate(MAX_SNAPSHOT_FILENAME_LEN); + name +} + +/// Names the server refuses to hand out presigned legacy upload targets for, because they +/// belong to the checkpoint protocol's own object namespace. +fn is_reserved_snapshot_name(name: &str) -> bool { + name.starts_with("checkpoint_") || name == "latest-checkpoint.json" } fn sanitize_filename_component(value: &str) -> String { @@ -1801,23 +1786,21 @@ fn unique_filename(preferred: &str, used: &mut HashSet) -> String { let preferred = Path::new(preferred) .file_name() .map(|n| n.to_string_lossy().to_string()) - .unwrap_or_else(|| "snapshot_artifact".to_string()); - let preferred = if preferred.is_empty() { - "snapshot_artifact".to_string() - } else { - preferred - }; + .filter(|n| !n.is_empty()) + .unwrap_or_else(|| FALLBACK_SNAPSHOT_FILENAME.to_string()); if used.insert(preferred.clone()) { return preferred; } let path = Path::new(&preferred); + // Trailing `_` is trimmed so `a_.txt` de-duplicates to `a_2.txt` rather than reintroducing + // the reserved `__` separator that sanitization just squashed out. let stem = path .file_stem() - .map(|s| s.to_string_lossy().to_string()) + .map(|s| s.to_string_lossy().trim_end_matches('_').to_string()) .filter(|s| !s.is_empty()) - .unwrap_or_else(|| "snapshot_artifact".to_string()); + .unwrap_or_else(|| FALLBACK_SNAPSHOT_FILENAME.to_string()); let extension = path.extension().map(|e| e.to_string_lossy().to_string()); for suffix in 2.. { diff --git a/app/src/ai/agent_sdk/driver/snapshot_tests.rs b/app/src/ai/agent_sdk/driver/snapshot_tests.rs index d859553b167..44bc639dbb9 100644 --- a/app/src/ai/agent_sdk/driver/snapshot_tests.rs +++ b/app/src/ai/agent_sdk/driver/snapshot_tests.rs @@ -11,6 +11,7 @@ use tokio::runtime::Runtime; use super::*; use crate::ai::agent::conversation::AIConversationId; +use crate::ai::agent_sdk::retry::MAX_ATTEMPTS; use crate::ai::agent_sdk::test_support::build_test_http_client; use crate::ai::artifacts::Artifact; use crate::server::server_api::harness_support::{ @@ -39,30 +40,24 @@ struct TestClient { http: http_client::Client, fail_get_targets: bool, /// Number of trailing response entries to drop, simulating a server that returns fewer - /// targets than the request contained (contract violation). Under the positional - /// alignment contract, the trailing files in the request end up with no target and are - /// marked [`EntryStatus::NoTarget`] downstream. + /// targets than the request contained. Under positional alignment those files end up + /// with no target and are marked [`EntryStatus::NoTarget`] downstream. drop_trailing_targets: usize, /// Restrict `drop_trailing_targets` to the first `get_snapshot_upload_targets` call. - /// - /// `upload_gathered_snapshot` always appends the manifest last, so truncating *every* - /// chunk always costs the manifest its target. Truncating only the first chunk (which - /// requires more than [`UPLOAD_BATCH_SIZE`] files, so there are at least two) isolates - /// the blob-level `NoTarget` path with the manifest still uploading cleanly. + /// The manifest is always the last entry of the last chunk, so truncating every chunk + /// would always cost the manifest its target instead of a blob. drop_trailing_first_call_only: bool, - /// Whether `commit_snapshot` should return an error, simulating a server-side commit - /// rejection (e.g. a missing required object). + /// Whether `commit_snapshot` should return an error. fail_commit: bool, - /// Every `get_snapshot_upload_targets` request received, in order, for checkpoint-mode - /// assertions (mode/generation used, plain logical filenames). + /// Every request received, in order, for wire-shape and exact-set assertions. upload_requests: Arc>>, - /// Every `commit_snapshot` request received, in order, for exact-set assertions. commit_requests: Arc>>, } impl TestClient { - fn new(server_base_url: String) -> Arc { - Arc::new(Self { + /// Happy-path client; the `new_*` constructors below flip one failure mode each. + fn base(server_base_url: String) -> Self { + Self { server_base_url, http: build_test_http_client(), fail_get_targets: false, @@ -71,32 +66,24 @@ impl TestClient { fail_commit: false, upload_requests: Arc::new(StdMutex::new(Vec::new())), commit_requests: Arc::new(StdMutex::new(Vec::new())), - }) + } + } + + fn new(server_base_url: String) -> Arc { + Arc::new(Self::base(server_base_url)) } fn new_failing_get_targets(server_base_url: String) -> Arc { Arc::new(Self { - server_base_url, - http: build_test_http_client(), fail_get_targets: true, - drop_trailing_targets: 0, - drop_trailing_first_call_only: false, - fail_commit: false, - upload_requests: Arc::new(StdMutex::new(Vec::new())), - commit_requests: Arc::new(StdMutex::new(Vec::new())), + ..Self::base(server_base_url) }) } fn new_dropping_trailing(server_base_url: String, drop_trailing: usize) -> Arc { Arc::new(Self { - server_base_url, - http: build_test_http_client(), - fail_get_targets: false, drop_trailing_targets: drop_trailing, - drop_trailing_first_call_only: false, - fail_commit: false, - upload_requests: Arc::new(StdMutex::new(Vec::new())), - commit_requests: Arc::new(StdMutex::new(Vec::new())), + ..Self::base(server_base_url) }) } @@ -107,27 +94,16 @@ impl TestClient { drop_trailing: usize, ) -> Arc { Arc::new(Self { - server_base_url, - http: build_test_http_client(), - fail_get_targets: false, drop_trailing_targets: drop_trailing, drop_trailing_first_call_only: true, - fail_commit: false, - upload_requests: Arc::new(StdMutex::new(Vec::new())), - commit_requests: Arc::new(StdMutex::new(Vec::new())), + ..Self::base(server_base_url) }) } fn new_failing_commit(server_base_url: String) -> Arc { Arc::new(Self { - server_base_url, - http: build_test_http_client(), - fail_get_targets: false, - drop_trailing_targets: 0, - drop_trailing_first_call_only: false, fail_commit: true, - upload_requests: Arc::new(StdMutex::new(Vec::new())), - commit_requests: Arc::new(StdMutex::new(Vec::new())), + ..Self::base(server_base_url) }) } @@ -1652,12 +1628,9 @@ fn checkpoint_commits_generation_prefixed_storage_names_while_upload_targets_use #[test] fn checkpoint_withholds_commit_when_a_required_blob_fails() { - // A blob upload fails permanently (404 is a non-retryable status). The exact-set contract - // requires withholding the commit entirely rather than committing a partial set. The - // manifest mock must return success here so this test isolates the blob-failure branch: - // without it, an unmocked manifest PUT would also fail against mockito's default - // (non-2xx) response for unmatched routes, and this test would pass for the wrong reason - // (falling into the "manifest failed to upload" branch instead of the intended one). + // 404 is non-retryable. The manifest mock must succeed so this isolates the blob-failure + // branch: an unmocked manifest PUT would fail too and the test would pass for the wrong + // reason. let tempdir = snaptest_tempdir(); let file_path = tempdir.path().join("bad.txt"); fs::write(&file_path, b"will-fail").unwrap(); @@ -1699,9 +1672,7 @@ fn checkpoint_withholds_commit_when_a_required_blob_fails() { #[test] fn checkpoint_manifest_upload_failure_withholds_commit() { - // The manifest itself fails to upload while the only blob succeeds. Losing the manifest - // means there is no rehydration catalogue even if blobs landed, so commit must still be - // withheld -- this exercises a distinct branch from the blob-failure test above. + // Without the manifest there is no rehydration catalogue, even though the blob landed. let tempdir = snaptest_tempdir(); let file_path = tempdir.path().join("ok.txt"); fs::write(&file_path, b"fine").unwrap(); @@ -1713,8 +1684,7 @@ fn checkpoint_manifest_upload_failure_withholds_commit() { .mock("PUT", upload_path("ok\\.txt")) .with_status(200) .create(); - // The upload path retries transient-looking failures with bounded retry before giving - // up, so a persistently failing manifest upload is attempted more than once. + // A persistent 5xx is retried, so the manifest PUT lands more than once. let manifest_mock = server .mock("PUT", upload_path("snapshot_state\\.json")) .with_status(500) @@ -1902,23 +1872,21 @@ fn checkpoint_commit_failure_reports_failed_result() { matches!(result, CheckpointResult::Failed { .. }), "expected Failed, got {result:?}" ); + // Everything is already uploaded by commit time, so a commit failure is worth retrying + // before the attempt is abandoned. assert_eq!( client.commit_requests().len(), - 1, - "commit should still be attempted exactly once" + MAX_ATTEMPTS, + "commit should exhaust its bounded retries before failing the attempt" ); manifest_mock.assert(); } #[test] fn checkpoint_withholds_commit_when_the_server_omits_a_blob_upload_target() { - // Regression: a short `upload-snapshot` response used to mark the target-less blob - // `skipped`, which the checkpoint gate tolerated (it only rejected `failed`). The attempt - // then committed a silently smaller object set and made it the selected checkpoint, - // discarding a previously complete one. `EntryStatus::NoTarget` must now withhold commit. - // - // Declare more than UPLOAD_BATCH_SIZE files so the request is chunked and the truncation - // lands on a blob rather than on the always-last manifest entry. + // Committing here would make a smaller object set the selected checkpoint, discarding a + // previously complete one. Declare more than UPLOAD_BATCH_SIZE files so the request is + // chunked and the truncation lands on a blob rather than the always-last manifest. let tempdir = snaptest_tempdir(); let decl_dir = snaptest_tempdir(); let declared_count = UPLOAD_BATCH_SIZE + 5; @@ -1962,8 +1930,6 @@ fn checkpoint_withholds_commit_when_the_server_omits_a_blob_upload_target() { #[test] fn sanitize_name_component_never_yields_the_reserved_double_underscore() { - // The logical filename half of `checkpoint___` is derived from - // agent-controlled basenames, so it must never reintroduce the reserved separator. for raw in [ "a__b.txt", "a b.txt", @@ -1972,7 +1938,7 @@ fn sanitize_name_component_never_yields_the_reserved_double_underscore() { "___", "ünïcödé.txt", ] { - let sanitized = sanitize_name_component(raw, "snapshot_artifact"); + let sanitized = sanitize_name_component(raw, FALLBACK_SNAPSHOT_FILENAME); assert!( !sanitized.contains("__"), "sanitized {raw:?} still contains `__`: {sanitized}" @@ -1989,9 +1955,8 @@ fn sanitize_name_component_never_yields_the_reserved_double_underscore() { #[test] fn checkpoint_storage_names_stay_unambiguous_for_hostile_filenames() { - // End-to-end guard for the same invariant: a file whose basename contains `__` must not - // produce a storage name with a second `__`, which would make the server's - // `checkpoint___` split ambiguous. + // End-to-end guard: a basename containing `__` must not produce a storage name with a + // second `__`, which would make the server's split ambiguous. let tempdir = snaptest_tempdir(); let hostile = tempdir.path().join("checkpoint_1700000000000-0__evil.txt"); fs::write(&hostile, b"hostile").unwrap(); @@ -2072,3 +2037,177 @@ fn checkpoint_new_gather_mints_a_fresh_generation_each_time() { ); manifest_mock.assert(); } + +/// Mirror of the server's logical-name validation. The server rejects the *whole* +/// upload-targets request when any name fails these, so every name we mint must pass. +fn assert_server_accepts_logical_name(name: &str) { + assert!(!name.is_empty(), "name is empty"); + assert!(name.len() <= 255, "name exceeds 255 bytes: {name}"); + assert!(name != "." && name != "..", "name is a directory alias"); + assert!(!name.starts_with('-'), "name parses as a flag: {name}"); + assert!( + !name.starts_with("checkpoint_") && name != "latest-checkpoint.json", + "name collides with the reserved checkpoint namespace: {name}" + ); + assert!( + name.bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_')), + "name leaves the server charset: {name}" + ); +} + +#[test] +fn sanitize_name_component_satisfies_the_server_name_contract() { + let long_name = format!("{}.txt", "a".repeat(400)); + for raw in [ + "-rf.txt", + "--force", + "-", + ".", + "..", + "checkpoint_1700000000000-0__evil.txt", + "latest-checkpoint.json", + "a__b.txt", + "weird name?!.txt", + "___", + "ünïcödé.txt", + &long_name, + ] { + let sanitized = sanitize_name_component(raw, FALLBACK_SNAPSHOT_FILENAME); + assert_server_accepts_logical_name(&sanitized); + assert!( + !sanitized.contains("__"), + "sanitized {raw:?} contains the reserved separator: {sanitized}" + ); + } +} + +#[test] +fn unique_filename_keeps_deduplicated_names_within_the_contract() { + // `a_.txt` used to de-duplicate to `a__2.txt`, reintroducing the reserved separator that + // sanitization had just squashed out. + let mut used = HashSet::new(); + let names: Vec = (0..3) + .map(|_| unique_filename(&sanitize_name_component("a_.txt", "repo"), &mut used)) + .collect(); + for name in &names { + assert_server_accepts_logical_name(name); + assert!( + !name.contains("__"), + "de-duplicated name is ambiguous: {name}" + ); + } + assert_eq!(names.len(), used.len(), "names must stay unique: {names:?}"); +} + +#[test] +fn checkpoint_commits_server_valid_names_for_hostile_basenames() { + // A single name the server would reject fails the entire upload-targets request, so an + // awkward basename must not be able to cost the whole checkpoint. + let tempdir = snaptest_tempdir(); + let decl_dir = snaptest_tempdir(); + let hostile_names = ["-rf.txt", "latest-checkpoint.json", "spaced name!.txt"]; + let file_paths: Vec = hostile_names + .iter() + .map(|name| { + let path = tempdir.path().join(name); + fs::write(&path, b"content").unwrap(); + path + }) + .collect(); + let file_refs: Vec<&Path> = file_paths.iter().map(|p| p.as_path()).collect(); + let declarations_path = write_declarations(decl_dir.path(), &[], &file_refs); + + let mut server = Server::new(); + let upload_mock = server + .mock("PUT", upload_path(r".+")) + .with_status(200) + .create(); + + let client = TestClient::new(server.url()); + let result = Runtime::new() + .unwrap() + .block_on(run_checkpoint_from_declarations_file( + &declarations_path, + client.clone(), + )); + let CheckpointResult::Committed { generation } = result else { + panic!("expected Committed, got {result:?}"); + }; + + for request in client.upload_requests() { + for file in &request.files { + assert_server_accepts_logical_name(&file.filename); + } + } + let prefix = format!("checkpoint_{}__", generation.as_str()); + for object in &client.commit_requests()[0].objects { + let logical = object + .strip_prefix(&prefix) + .unwrap_or_else(|| panic!("object {object} is not under the generation prefix")); + assert_server_accepts_logical_name(logical); + } + drop(upload_mock); +} + +// ------------------------------------------------------------------------------------------------ +// Wire format. These pin the JSON the server actually parses; the in-process `TestClient` +// never exercises serde. +// ------------------------------------------------------------------------------------------------ + +fn test_file_info() -> SnapshotFileInfo { + SnapshotFileInfo { + filename: "note.txt".to_string(), + mime_type: "text/plain".to_string(), + } +} + +#[test] +fn legacy_upload_request_omits_the_checkpoint_fields() { + // The end-of-run path must keep emitting exactly the pre-checkpoint payload. + assert_eq!( + serde_json::to_value(SnapshotUploadRequest::legacy(vec![test_file_info()])).unwrap(), + serde_json::json!({ + "files": [{"filename": "note.txt", "mime_type": "text/plain"}], + }) + ); +} + +#[test] +fn checkpoint_upload_request_sends_mode_and_generation() { + let request = SnapshotUploadRequest::checkpoint( + CheckpointGeneration::new_for_test("1700000000000-0"), + vec![test_file_info()], + ); + assert_eq!( + serde_json::to_value(request).unwrap(), + serde_json::json!({ + "mode": "checkpoint", + "generation": "1700000000000-0", + "files": [{"filename": "note.txt", "mime_type": "text/plain"}], + }) + ); +} + +#[test] +fn commit_snapshot_request_matches_the_server_schema() { + let request = CommitSnapshotRequest { + generation: "1700000000000-0".to_string(), + manifest_object: "checkpoint_1700000000000-0__snapshot_state.json".to_string(), + objects: vec![ + "checkpoint_1700000000000-0__note.txt".to_string(), + "checkpoint_1700000000000-0__snapshot_state.json".to_string(), + ], + }; + assert_eq!( + serde_json::to_value(request).unwrap(), + serde_json::json!({ + "generation": "1700000000000-0", + "manifest_object": "checkpoint_1700000000000-0__snapshot_state.json", + "objects": [ + "checkpoint_1700000000000-0__note.txt", + "checkpoint_1700000000000-0__snapshot_state.json", + ], + }) + ); +} diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index 442036f07f5..2f5e4b557cb 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -53,13 +53,10 @@ pub enum UploadFieldValue { ContentData, } -/// Selects how the server accounts for a `SnapshotUploadRequest`'s uploads. +/// Selects how the server names and accounts for a [`SnapshotUploadRequest`]'s uploads. /// -/// `Legacy` (the default) uses unprefixed object names and counts uploads against -/// the execution's cumulative lifetime attachment quota, matching today's one-shot -/// end-of-run snapshot. `Checkpoint` signs generation-prefixed object names and does -/// not consume that cumulative quota; the server enforces per-attempt limits instead -/// when the generation is committed via [`HarnessSupportClient::commit_snapshot`]. +/// `Legacy` uses unprefixed names and charges the execution's cumulative attachment quota. +/// `Checkpoint` signs generation-prefixed names and is charged per attempt at commit time. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize)] #[serde(rename_all = "snake_case")] pub enum SnapshotUploadMode { @@ -71,13 +68,11 @@ pub enum SnapshotUploadMode { /// Request body for upload-snapshot upload targets. #[derive(Debug, Clone, serde::Serialize)] pub struct SnapshotUploadRequest { - /// Upload accounting mode. Omitted (default) is equivalent to `legacy` on the - /// server; see [`SnapshotUploadMode`]. + /// Omitted when legacy, which the server treats as the default. #[serde(skip_serializing_if = "is_default_mode")] pub mode: SnapshotUploadMode, - /// Required when `mode` is [`SnapshotUploadMode::Checkpoint`]. Identifies the - /// checkpoint attempt; every requested file is uploaded by the server as - /// `checkpoint___`. Ignored for `legacy` mode. + /// Required in checkpoint mode; the server uploads each file as + /// `checkpoint___`. #[serde(skip_serializing_if = "Option::is_none")] pub generation: Option, pub files: Vec, @@ -88,7 +83,6 @@ fn is_default_mode(mode: &SnapshotUploadMode) -> bool { } impl SnapshotUploadRequest { - /// Build a legacy-mode request, matching today's one-shot end-of-run upload. pub fn legacy(files: Vec) -> Self { Self { mode: SnapshotUploadMode::Legacy, @@ -97,7 +91,6 @@ impl SnapshotUploadRequest { } } - /// Build a checkpoint-mode request for the given generation. pub fn checkpoint(generation: CheckpointGeneration, files: Vec) -> Self { Self { mode: SnapshotUploadMode::Checkpoint, @@ -107,33 +100,26 @@ impl SnapshotUploadRequest { } } -/// A checkpoint generation identifier minted by the client for one checkpoint attempt. +/// Client-minted identifier for one checkpoint attempt, used to key that attempt's storage +/// objects as `checkpoint___`. /// -/// Must match the server's `[A-Za-z0-9._-]{1,128}` format and must not contain the -/// reserved `__` separator (validated by -/// [`crate::ai::agent_sdk::driver::snapshot::mint_generation`], the only production -/// constructor). Storage object basenames are `checkpoint___`; -/// the generation is a GCS keying detail only and must never leak into agent-visible -/// paths or restore commands. +/// A generation is a storage-keying detail and must never leak into agent-visible paths or +/// restore commands. #[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize)] #[serde(transparent)] pub struct CheckpointGeneration(String); impl CheckpointGeneration { - /// Wrap an already-validated generation string. Exposed for tests; production - /// code should go through `snapshot::mint_generation` instead. - /// - /// The only caller (`driver::snapshot`'s test module) is itself excluded on - /// Windows (snapshot upload is cloud-agent-only and Linux-only), so this must - /// be gated the same way or it is dead code under Windows clippy/test builds. + /// Test-only escape hatch; production code mints generations via + /// `snapshot::mint_generation`. Gated to match `driver::snapshot`'s test module, which + /// does not build on Windows. #[cfg(all(test, not(windows)))] pub(crate) fn new_for_test(value: impl Into) -> Self { Self(value.into()) } - /// True when `value` satisfies the server's `[A-Za-z0-9._-]{1,128}` format and does not - /// contain the reserved `__` separator used by - /// `checkpoint___` storage object names. + /// Mirrors the server's `[A-Za-z0-9._-]{1,128}` format check, including the reserved `__` + /// separator that would make `checkpoint___` ambiguous. fn is_valid(value: &str) -> bool { !value.is_empty() && value.len() <= 128 @@ -143,15 +129,10 @@ impl CheckpointGeneration { .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_')) } - /// Construct from a pre-validated string. Crate-visible so `driver::snapshot` can - /// mint generations without duplicating this type. - /// - /// The debug assertion keeps the type's documented invariant honest: previously nothing - /// enforced it, so "validated" was aspirational. `snapshot::mint_generation` is the only - /// production caller and provably satisfies it, hence a `debug_assert!` rather than a - /// fallible constructor. - // Only called by `snapshot::mint_generation`, which is itself unused until the - // periodic checkpoint coordinator (a follow-up, stacked PR) lands. + /// 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. + // Unused until the periodic checkpoint coordinator (a follow-up, stacked PR) lands. #[allow(dead_code)] pub(crate) fn from_validated(value: String) -> Self { debug_assert!( @@ -177,12 +158,11 @@ impl std::fmt::Display for CheckpointGeneration { } } -/// Request body for committing a fully uploaded checkpoint generation. Exact-set: the -/// server persists `objects` verbatim as the commit marker and later selection returns -/// exactly that set, never every object sharing the generation prefix. See -/// `docs/remote-2111-checkpoint-spec.md` (warp-server) for the full protocol. -// Not constructed until the periodic checkpoint coordinator (a follow-up, stacked PR) -// starts calling `HarnessSupportClient::commit_snapshot`. +/// Request body for committing a fully uploaded checkpoint generation. +/// +/// Exact-set: the server persists `objects` verbatim as the commit marker and selection +/// later returns exactly that set, not everything sharing the generation prefix. +// Not constructed until the periodic checkpoint coordinator (a follow-up, stacked PR) lands. #[allow(dead_code)] #[derive(Debug, Clone, serde::Serialize)] pub struct CommitSnapshotRequest { @@ -362,13 +342,12 @@ pub trait HarnessSupportClient: 'static + Send + Sync { request: &SnapshotUploadRequest, ) -> Result>; - /// Commit a fully uploaded checkpoint generation for the active execution's exact - /// object set. Must only be called after every object named in `request.objects` - /// (including `request.manifest_object`) has itself uploaded successfully; the - /// server verifies existence and per-attempt size limits before this becomes the - /// selected checkpoint. - // Not called until the periodic checkpoint coordinator (a follow-up, stacked PR) - // lands. + /// Make a fully uploaded checkpoint generation the selected checkpoint. + /// + /// 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. + // Not called until the periodic checkpoint coordinator (a follow-up, stacked PR) lands. #[allow(dead_code)] async fn commit_snapshot( &self, From b88ff2504dd43c18e6b6351cd5f0371339882a4b Mon Sep 17 00:00:00 2001 From: Joey Wang Date: Fri, 31 Jul 2026 16:07:49 -0400 Subject: [PATCH 4/4] Trim redundant comments from the checkpoint mechanics Follow-up comment pass. `harness_support.rs` explained the same `allow(dead_code)` reason three times; that is now one note covering the file, matching what `snapshot.rs` already does. Dropped the doc comments that only restated `FALLBACK_SNAPSHOT_FILENAME` and `RESERVED_NAME_ESCAPE`, and shortened the rest where the code already carried the meaning. Kept the ones that encode contracts the code cannot show on its own: the server's filename rules, why the length cap is 240 rather than 255, why `String::truncate` cannot split a char, why `unique_filename` trims a trailing underscore, and why re-committing a generation is safe. Co-Authored-By: Oz --- app/src/ai/agent_sdk/driver/snapshot.rs | 36 +++++++------------- app/src/server/server_api/harness_support.rs | 5 ++- 2 files changed, 15 insertions(+), 26 deletions(-) diff --git a/app/src/ai/agent_sdk/driver/snapshot.rs b/app/src/ai/agent_sdk/driver/snapshot.rs index 7c4cdd768a5..6c8bbd8c0c6 100644 --- a/app/src/ai/agent_sdk/driver/snapshot.rs +++ b/app/src/ai/agent_sdk/driver/snapshot.rs @@ -670,8 +670,7 @@ struct SnapshotOutcome { #[allow(dead_code)] #[derive(Debug)] pub(super) enum CheckpointResult { - /// Every required object uploaded and the exact-set commit succeeded, so `generation` is - /// now the server's selected checkpoint. + /// `generation` is now the server's selected checkpoint. Committed { generation: CheckpointGeneration }, /// Nothing to checkpoint: the declarations file was missing, empty, or had no valid /// entries. No generation was minted and no network calls were made. @@ -692,7 +691,6 @@ pub(super) enum CheckpointResult { /// [`SnapshotUploadMode`] for the server-side semantics. enum PipelineMode { Legacy, - // Not constructed until the coordinator PR. #[allow(dead_code)] Checkpoint(CheckpointGeneration), } @@ -701,8 +699,8 @@ enum PipelineMode { #[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 construction. +/// Mint a generation identifier, which satisfies [`CheckpointGeneration`]'s format by +/// construction. /// /// 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 @@ -717,12 +715,10 @@ pub(super) fn mint_generation() -> CheckpointGeneration { CheckpointGeneration::from_validated(format!("{millis}-{counter}")) } -/// Reproduce the server's `checkpoint___` storage name for a -/// logical filename. +/// Reproduce the server's `checkpoint___` storage name. /// -/// Only needed to assemble the exact-set [`CommitSnapshotRequest`]: everything earlier in the -/// pipeline speaks logical names, and the server derives the storage name for each presigned -/// target itself. +/// 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()) @@ -1045,9 +1041,8 @@ async fn run_checkpoint_pipeline( reason: "manifest failed to upload".to_string(), }; } - // Committing while any entry is `NoTarget` would make a silently smaller object set the - // selected checkpoint, discarding a previously complete one. Only `Skipped` (the - // deliberate per-run cap) is tolerated. + // Committing with a `NoTarget` entry would make a silently smaller object set the selected + // checkpoint, discarding a previously complete one. if outcome .entries .iter() @@ -1075,9 +1070,8 @@ async fn run_checkpoint_pipeline( manifest_object, objects, }; - // Every object is already in storage by this point, so a transient failure here would - // throw away the whole attempt. Committing the same generation twice is idempotent - // server-side, which makes retrying safe. + // Every object is already in storage, so a transient failure here would throw away the + // whole attempt. Re-committing the same generation is idempotent server-side. let operation = format!("checkpoint commit '{}'", generation.as_str()); match with_bounded_retry(&operation, || client.commit_snapshot(&commit_request)).await { Ok(response) => { @@ -1385,8 +1379,8 @@ async fn gather_file( let path = Path::new(file_path); match tokio::fs::read(path).await { Ok(content) => { - // Sanitize before uniquifying: the basename is agent-controlled and ends up in the - // storage name the exact-set commit has to reproduce byte for byte. + // Sanitize before uniquifying so the de-duplication suffix cannot break the + // invariants; see `sanitize_name_component`. let preferred = sanitize_name_component( &path .file_name() @@ -1724,10 +1718,7 @@ async fn git_output_string(repo_dir: &Path, args: &[&str]) -> Option { if value.is_empty() { None } else { Some(value) } } -/// Fallback used wherever a name sanitizes down to nothing usable. const FALLBACK_SNAPSHOT_FILENAME: &str = "snapshot_artifact"; - -/// Prepended to names that would otherwise collide with the server's reserved namespace. const RESERVED_NAME_ESCAPE: &str = "snapshot-"; /// Longest logical filename we will mint. The server rejects names over 255 bytes; the @@ -1772,8 +1763,7 @@ fn sanitize_name_component(value: &str, fallback: &str) -> String { name } -/// Names the server refuses to hand out presigned legacy upload targets for, because they -/// belong to the checkpoint protocol's own object namespace. +/// Names owned by the checkpoint protocol, which the server refuses to sign legacy uploads for. fn is_reserved_snapshot_name(name: &str) -> bool { name.starts_with("checkpoint_") || name == "latest-checkpoint.json" } diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index 2f5e4b557cb..e1bbac47323 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -132,7 +132,8 @@ 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. - // Unused until the periodic checkpoint coordinator (a follow-up, stacked PR) lands. + // 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!( @@ -162,7 +163,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. -// Not constructed until the periodic checkpoint coordinator (a follow-up, stacked PR) lands. #[allow(dead_code)] #[derive(Debug, Clone, serde::Serialize)] pub struct CommitSnapshotRequest { @@ -347,7 +347,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. - // Not called until the periodic checkpoint coordinator (a follow-up, stacked PR) lands. #[allow(dead_code)] async fn commit_snapshot( &self,