diff --git a/crates/guest-agent/src/main.rs b/crates/guest-agent/src/main.rs index 5a545179966..b1a4ccb5a0f 100644 --- a/crates/guest-agent/src/main.rs +++ b/crates/guest-agent/src/main.rs @@ -62,33 +62,65 @@ fn helper_exit_code_from_args() -> Option { let mut args = std::env::args_os(); let _program = args.next()?; let command = args.next()?; - if command != "verify-session-history-identity" { - return None; - } - let metadata_path = args - .next() - .unwrap_or_else(final_session_history_identity_path_from_process_env); - let remaining = args.collect::>(); - let expected = match parse_session_history_identity_expectation(&remaining) { - Ok(expected) => expected, - Err(()) => return Some(SESSION_HISTORY_IDENTITY_VERIFY_EXIT_INVALID_ARGS), - }; - Some( - match session_history_identity::verify_final_session_history_identity_file( - metadata_path, - expected.as_ref(), - ) { - Ok(()) => SESSION_HISTORY_IDENTITY_VERIFY_EXIT_SUCCESS, - Err(error) => { - let exit_code = error.helper_exit_code(); - if exit_code == SESSION_HISTORY_IDENTITY_VERIFY_EXIT_SUCCESS { - SESSION_HISTORY_IDENTITY_VERIFY_EXIT_FAILURE - } else { - exit_code - } + match command.to_str()? { + "verify-session-history-identity" => { + let metadata_path = args + .next() + .unwrap_or_else(final_session_history_identity_path_from_process_env); + let remaining = args.collect::>(); + let expected = match parse_session_history_identity_expectation(&remaining) { + Ok(expected) => expected, + Err(()) => return Some(SESSION_HISTORY_IDENTITY_VERIFY_EXIT_INVALID_ARGS), + }; + Some( + match session_history_identity::verify_final_session_history_identity_file( + metadata_path, + expected.as_ref(), + ) { + Ok(()) => SESSION_HISTORY_IDENTITY_VERIFY_EXIT_SUCCESS, + Err(error) => session_history_identity_helper_exit_code(&error), + }, + ) + } + "export-session-history-sidecar" => { + let Some(metadata_path) = args.next() else { + return Some(SESSION_HISTORY_IDENTITY_VERIFY_EXIT_INVALID_ARGS); + }; + let Some(export_path) = args.next() else { + return Some(SESSION_HISTORY_IDENTITY_VERIFY_EXIT_INVALID_ARGS); + }; + if args.next().is_some() { + return Some(SESSION_HISTORY_IDENTITY_VERIFY_EXIT_INVALID_ARGS); } - }, - ) + Some( + match session_history_identity::export_final_session_history_sidecar_file( + metadata_path, + export_path, + ) { + Ok(metadata) => match serde_json::to_string(&metadata) { + Ok(json) => { + println!("{json}"); + SESSION_HISTORY_IDENTITY_VERIFY_EXIT_SUCCESS + } + Err(_) => SESSION_HISTORY_IDENTITY_VERIFY_EXIT_FAILURE, + }, + Err(error) => session_history_identity_helper_exit_code(&error), + }, + ) + } + _ => None, + } +} + +fn session_history_identity_helper_exit_code( + error: &session_history_identity::FinalSessionHistoryIdentityVerifyError, +) -> i32 { + let exit_code = error.helper_exit_code(); + if exit_code == SESSION_HISTORY_IDENTITY_VERIFY_EXIT_SUCCESS { + SESSION_HISTORY_IDENTITY_VERIFY_EXIT_FAILURE + } else { + exit_code + } } #[allow(clippy::panic)] diff --git a/crates/guest-agent/src/session_history_identity.rs b/crates/guest-agent/src/session_history_identity.rs index 2bf77c03ba3..3190bafa398 100644 --- a/crates/guest-agent/src/session_history_identity.rs +++ b/crates/guest-agent/src/session_history_identity.rs @@ -15,7 +15,8 @@ use guest_contracts::session_history_identity::{ SESSION_HISTORY_IDENTITY_VERIFY_EXIT_HISTORY_READ, SESSION_HISTORY_IDENTITY_VERIFY_EXIT_HISTORY_TOO_LARGE, SESSION_HISTORY_IDENTITY_VERIFY_EXIT_INVALID_METADATA, - SESSION_HISTORY_IDENTITY_VERIFY_EXIT_METADATA_READ, + SESSION_HISTORY_IDENTITY_VERIFY_EXIT_METADATA_READ, SESSION_HISTORY_SIDECAR_MAX_BYTES, + SessionHistorySidecarExportMetadata, SessionHistorySidecarRepresentation, }; use sha2::{Digest, Sha256}; use std::fmt; @@ -73,6 +74,34 @@ pub fn verify_final_session_history_identity_file( verify_final_session_history_identity(&identity) } +pub fn export_final_session_history_sidecar_file( + metadata_path: impl AsRef, + export_path: impl AsRef, +) -> Result { + let identity = read_final_session_history_identity(metadata_path)?; + verify_final_session_history_identity(&identity)?; + let source = session_history::read_session_history_checkpoint_source_from_payload_bounded( + &identity.history_marker_payload, + SESSION_HISTORY_SIDECAR_MAX_BYTES, + ) + .map_err(FinalSessionHistoryIdentityVerifyError::HistoryRead)?; + let (representation, bytes) = match source { + session_history::SessionHistoryCheckpointSource::Decoded(bytes) => { + (SessionHistorySidecarRepresentation::Raw, bytes) + } + session_history::SessionHistoryCheckpointSource::CodexZstd { encoded } => { + (SessionHistorySidecarRepresentation::CodexZstd, encoded) + } + }; + crate::paths::write_private(export_path.as_ref(), &bytes) + .map_err(AgentError::from) + .map_err(FinalSessionHistoryIdentityVerifyError::HistoryRead)?; + Ok(SessionHistorySidecarExportMetadata { + representation, + encoded_size: bytes.len() as u64, + }) +} + fn read_final_session_history_identity( metadata_path: impl AsRef, ) -> Result { @@ -261,6 +290,35 @@ mod tests { verify_final_session_history_identity_file(metadata_path, None).unwrap(); } + #[test] + fn exports_claude_literal_history_sidecar() { + let dir = tempfile::tempdir().unwrap(); + let history = br#"{"type":"system"}"#; + let history_path = dir.path().join("history.jsonl"); + std::fs::write(&history_path, history).unwrap(); + let identity = FinalSessionHistoryIdentity::new( + FinalSessionHistoryFramework::ClaudeCode, + "a".repeat(64), + FinalSessionHistoryRefKind::Blob, + hex::encode(Sha256::digest(history)), + history.len() as u64, + history_path.to_string_lossy(), + ) + .unwrap(); + let metadata_path = write_metadata(&dir, &identity); + let export_path = dir.path().join("sidecar.blob"); + + let export_metadata = + export_final_session_history_sidecar_file(metadata_path, &export_path).unwrap(); + + assert_eq!( + export_metadata.representation, + SessionHistorySidecarRepresentation::Raw + ); + assert_eq!(export_metadata.encoded_size, history.len() as u64); + assert_eq!(std::fs::read(export_path).unwrap(), history); + } + #[test] fn verifies_codex_marker_history() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/guest-contracts/src/runtime_paths.rs b/crates/guest-contracts/src/runtime_paths.rs index ce676496c38..b95baddabf9 100644 --- a/crates/guest-contracts/src/runtime_paths.rs +++ b/crates/guest-contracts/src/runtime_paths.rs @@ -159,6 +159,11 @@ pub fn final_session_history_identity_file(run_dir: impl AsRef) -> PathBuf file(run_dir, "final-session-history-identity.json") } +/// Return the run-root session-history sidecar export file. +pub fn session_history_sidecar_export_file(run_dir: impl AsRef) -> PathBuf { + file(run_dir, "session-history-sidecar") +} + /// Return the run-root `failure-diagnostic.json` file. pub fn failure_diagnostic_file(run_dir: impl AsRef) -> PathBuf { file(run_dir, "failure-diagnostic.json") diff --git a/crates/guest-contracts/src/session_history_identity.rs b/crates/guest-contracts/src/session_history_identity.rs index bf11d5350ab..fa411bf19b9 100644 --- a/crates/guest-contracts/src/session_history_identity.rs +++ b/crates/guest-contracts/src/session_history_identity.rs @@ -34,6 +34,9 @@ pub const SESSION_HISTORY_IDENTITY_VERIFY_EXIT_HISTORY_MISMATCH: i32 = 8; /// Guest helper exit code for local history exceeding the guest verification budget. pub const SESSION_HISTORY_IDENTITY_VERIFY_EXIT_HISTORY_TOO_LARGE: i32 = 9; +/// Maximum native sidecar bytes runner can copy from guest to host. +pub const SESSION_HISTORY_SIDECAR_MAX_BYTES: u64 = 64 * 1024 * 1024; + /// Framework that owns a final session-history file. #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(rename_all = "kebab-case")] @@ -52,6 +55,26 @@ pub enum FinalSessionHistoryRefKind { Blob, } +/// Native on-disk representation used for cached session-history sidecars. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "kebab-case")] +pub enum SessionHistorySidecarRepresentation { + /// Uncompressed session-history bytes. + Raw, + /// Codex zstd session-history bytes. + CodexZstd, +} + +/// Metadata printed by the guest sidecar export helper. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionHistorySidecarExportMetadata { + /// Representation written to the exported sidecar file. + pub representation: SessionHistorySidecarRepresentation, + /// Exact byte length of the exported sidecar file. + pub encoded_size: u64, +} + /// Run-private final session-history identity metadata. #[derive(Clone, Deserialize, Eq, PartialEq, Serialize)] #[serde(rename_all = "camelCase")] diff --git a/crates/runner/src/cmd/start/job_discovery.rs b/crates/runner/src/cmd/start/job_discovery.rs index df911236c0a..2f41880a0ef 100644 --- a/crates/runner/src/cmd/start/job_discovery.rs +++ b/crates/runner/src/cmd/start/job_discovery.rs @@ -287,6 +287,10 @@ fn build_session_history_restore_plan( | SandboxReuseResult::UnparkFailed => Some(SessionHistoryRestoreFallback::NonReuse), }; + if reuse.result != SandboxReuseResult::Reused { + return SessionHistoryRestorePlan::DeferredHashBacked { fallback }; + } + let started_at = Instant::now(); let materializer = SessionHistoryMaterializer::start_cancellable( http, @@ -1215,6 +1219,34 @@ mod tests { } } + #[test] + fn restore_plan_defers_hash_backed_history_for_non_reuse() { + let http = test_http_client(); + let context = context_with_history_ref("history-hash-a"); + let cancel = RunCancellationHandle::new(); + let mut timing = RunnerPreSpawnTiming::start_after_claim(); + + let plan = build_session_history_restore_plan( + &http, + &context, + true, + &cancel, + SessionHistoryRestoreReuse { + entry: None, + result: SandboxReuseResult::PoolMiss, + }, + &mut timing, + None, + ); + + match plan { + SessionHistoryRestorePlan::DeferredHashBacked { fallback } => { + assert_eq!(fallback, Some(SessionHistoryRestoreFallback::NonReuse)); + } + _ => panic!("non-reuse hash-backed history should defer materialization"), + } + } + #[tokio::test] async fn restore_plan_classifies_session_identity_mismatch() { let http = test_http_client(); diff --git a/crates/runner/src/cmd/start/sandbox_finalization.rs b/crates/runner/src/cmd/start/sandbox_finalization.rs index 1b95998479a..28f682de4af 100644 --- a/crates/runner/src/cmd/start/sandbox_finalization.rs +++ b/crates/runner/src/cmd/start/sandbox_finalization.rs @@ -144,6 +144,7 @@ pub(super) async fn finalize_sandbox_for_completion( run_id, sandbox_id, cli_agent_session_id_override: resolved_cli_agent_session_id, + restored_session_identity: restored_session_identity.as_ref(), terminal_status, completed_at: completed_at.clone(), storage_fingerprints: storage_fingerprints.clone(), @@ -726,6 +727,7 @@ mod tests { run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status, completed_at: local_completed_at(), storage_fingerprints, diff --git a/crates/runner/src/cmd/start/tests/support/idle_pool.rs b/crates/runner/src/cmd/start/tests/support/idle_pool.rs index ace899553dc..e93251f34e0 100644 --- a/crates/runner/src/cmd/start/tests/support/idle_pool.rs +++ b/crates/runner/src/cmd/start/tests/support/idle_pool.rs @@ -138,6 +138,7 @@ pub(in super::super) async fn seed_idle_pool_with_workspace_promotion( run_id, sandbox_id, cli_agent_session_id_override: Some(spec.session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: TEST_SESSION_LAST_COMPLETED_AT.into(), storage_fingerprints: StorageFingerprints::default(), diff --git a/crates/runner/src/executor/agent_run.rs b/crates/runner/src/executor/agent_run.rs index 87835c19abf..e0762bb6327 100644 --- a/crates/runner/src/executor/agent_run.rs +++ b/crates/runner/src/executor/agent_run.rs @@ -17,6 +17,7 @@ use sandbox::{ Sandbox, StartProcessRequest, }; use shell_quote::quote_shell_arg; +use tokio::io::AsyncReadExt; use tokio_util::sync::CancellationToken; use tracing::{debug, info, warn}; @@ -39,6 +40,7 @@ use super::guest_state::{restore_guest_state, sync_guest_timezone}; use super::session_history_download::{ SessionHistoryDownloadPhaseTiming, SessionHistoryDownloadTimings, SessionHistoryMaterialization, SessionHistoryMaterializer, + verify_codex_zstd_session_history_bytes, verify_identity_session_history_bytes, }; use super::session_restore::{MaterializedResumeSession, restore_session}; use super::storage::{apply_storage_fingerprint_reuse, download_storages, guest_download_has_work}; @@ -59,6 +61,9 @@ use crate::restored_session_identity::{ }; use crate::telemetry::{JobTelemetry, SessionHistoryTelemetryMetadata}; use crate::types::{ExecutionContext, GuestDownloadManifest}; +use crate::workspace_image_cache::{ + WorkspaceSessionHistorySidecar, WorkspaceSessionHistorySidecarRepresentation, +}; const AGENT_WRAPPER_STDERR_CAPTURE_LIMIT_BYTES: u32 = 64 * 1024; const SESSION_HISTORY_DOWNLOAD_TELEMETRY_ERROR: &str = "session history download failed"; @@ -294,15 +299,96 @@ fn record_session_history_materializer_state( } } +fn record_session_history_restore_fallback( + telemetry: &mut JobTelemetry, + fallback: Option, +) { + if let Some(fallback) = fallback { + telemetry.record(fallback.action_type(), Duration::ZERO, true, None); + if let Some(reason) = fallback.identity_mismatch_reason() { + record_session_history_identity_mismatch_reason(telemetry, reason); + } + if matches!(fallback, SessionHistoryRestoreFallback::MissingIdleIdentity) { + telemetry.record( + "session_history_identity_reuse_missing", + Duration::ZERO, + true, + None, + ); + record_session_history_identity_reason( + telemetry, + SessionHistoryIdentityReason::ReuseMissingNoIdleIdentity, + ); + } + } +} + +async fn materialize_session_history_sidecar( + context: &ExecutionContext, + sidecar: &WorkspaceSessionHistorySidecar, +) -> RunnerResult> { + let resume_session = context.resume_session.as_ref().ok_or_else(|| { + RunnerError::Internal("resume session missing for sidecar restore".into()) + })?; + let history_ref = resume_session.history_ref().ok_or_else(|| { + RunnerError::Internal("resume session history ref missing for sidecar restore".into()) + })?; + let bytes = read_session_history_sidecar_bytes(sidecar).await?; + match sidecar.representation { + WorkspaceSessionHistorySidecarRepresentation::Raw => { + verify_identity_session_history_bytes(&bytes, history_ref.raw_size, &history_ref.hash)?; + Ok(MaterializedResumeSession::new( + resume_session.cli_agent_session_id.clone(), + bytes, + )) + } + WorkspaceSessionHistorySidecarRepresentation::CodexZstd => { + let timestamp = verify_codex_zstd_session_history_bytes( + &bytes, + history_ref.raw_size, + &history_ref.hash, + )?; + Ok(MaterializedResumeSession::new_codex_zstd( + resume_session.cli_agent_session_id.clone(), + bytes, + timestamp, + )) + } + } +} + +async fn read_session_history_sidecar_bytes( + sidecar: &WorkspaceSessionHistorySidecar, +) -> RunnerResult> { + let file = tokio::fs::File::open(&sidecar.path).await?; + let mut bytes = Vec::with_capacity(sidecar.encoded_size.min(1024 * 1024) as usize); + file.take(sidecar.encoded_size.saturating_add(1)) + .read_to_end(&mut bytes) + .await?; + if bytes.len() as u64 != sidecar.encoded_size { + return Err(RunnerError::Internal( + "workspace session history sidecar size mismatch".into(), + )); + } + Ok(bytes) +} + #[derive(Default)] #[must_use = "restore plans decide whether resume history download can be skipped"] pub(crate) enum SessionHistoryRestorePlan { #[default] Default, + DeferredHashBacked { + fallback: Option, + }, Prestarted { materializer: SessionHistoryMaterializer, fallback: Option, }, + LocalSidecar { + sidecar: WorkspaceSessionHistorySidecar, + fallback: Option, + }, SkipVerified(RestoredSessionIdentity), } @@ -924,7 +1010,8 @@ pub(super) async fn run_in_sandbox_with_process_cancel_timeouts( let mut session_restore_diagnostics = None; let mut restored_session_identity = None; - let session_history_materializer = match session_history_restore_plan { + let mut local_session_history_sidecar = None; + let mut session_history_materializer = match session_history_restore_plan { SessionHistoryRestorePlan::SkipVerified(identity) => { match verify_restored_session_identity_for_reuse(sandbox, context, identity).await { Ok(identity) => { @@ -956,6 +1043,16 @@ pub(super) async fn run_in_sandbox_with_process_cancel_timeouts( } } } + SessionHistoryRestorePlan::DeferredHashBacked { fallback } => { + record_session_history_restore_fallback(telemetry, fallback); + Some(SessionHistoryMaterializer::start_cancellable( + &config.http, + context.resume_session.as_ref(), + effective_cli_framework(&context.cli_agent_type), + cancel.clone(), + Some(&config.session_history_probe), + )) + } SessionHistoryRestorePlan::Default => Some(SessionHistoryMaterializer::start_cancellable( &config.http, context.resume_session.as_ref(), @@ -967,27 +1064,85 @@ pub(super) async fn run_in_sandbox_with_process_cancel_timeouts( materializer, fallback, } => { - if let Some(fallback) = fallback { - telemetry.record(fallback.action_type(), Duration::ZERO, true, None); - if let Some(reason) = fallback.identity_mismatch_reason() { - record_session_history_identity_mismatch_reason(telemetry, reason); + record_session_history_restore_fallback(telemetry, fallback); + Some(materializer) + } + SessionHistoryRestorePlan::LocalSidecar { sidecar, fallback } => { + record_session_history_restore_fallback(telemetry, fallback); + local_session_history_sidecar = Some(sidecar); + None + } + }; + if let Some(sidecar) = local_session_history_sidecar { + let restore_started = Instant::now(); + match materialize_session_history_sidecar(context, &sidecar).await { + Ok(session) => match restore_session(sandbox, context, &session).await { + Ok(diagnostics) => { + telemetry.record( + "session_history_workspace_cache_restore", + restore_started.elapsed(), + true, + None, + ); + telemetry.record("session_restore", restore_started.elapsed(), true, None); + session_restore_diagnostics = Some(diagnostics); } - if matches!(fallback, SessionHistoryRestoreFallback::MissingIdleIdentity) { + Err(error) => { + telemetry.record( + "session_history_workspace_cache_restore", + restore_started.elapsed(), + false, + Some("restore_error"), + ); telemetry.record( - "session_history_identity_reuse_missing", + "session_history_workspace_cache_miss", Duration::ZERO, true, - None, + Some("restore_error"), ); - record_session_history_identity_reason( - telemetry, - SessionHistoryIdentityReason::ReuseMissingNoIdleIdentity, + warn!( + run_id = %context.run_id, + error = %error, + "workspace session history sidecar restore failed; falling back to remote history" ); + session_history_materializer = + Some(SessionHistoryMaterializer::start_cancellable( + &config.http, + context.resume_session.as_ref(), + effective_cli_framework(&context.cli_agent_type), + cancel.clone(), + Some(&config.session_history_probe), + )); } + }, + Err(error) => { + telemetry.record( + "session_history_workspace_cache_restore", + restore_started.elapsed(), + false, + Some("materialize_error"), + ); + telemetry.record( + "session_history_workspace_cache_miss", + Duration::ZERO, + true, + Some("materialize_error"), + ); + warn!( + run_id = %context.run_id, + error = %error, + "workspace session history sidecar materialization failed; falling back to remote history" + ); + session_history_materializer = Some(SessionHistoryMaterializer::start_cancellable( + &config.http, + context.resume_session.as_ref(), + effective_cli_framework(&context.cli_agent_type), + cancel.clone(), + Some(&config.session_history_probe), + )); } - Some(materializer) } - }; + } if let Some(session_history_materializer) = session_history_materializer { // 4. Restore session history. Hash-backed history downloads can start // before sandbox preparation, then materialize here right before restore. diff --git a/crates/runner/src/executor/sandbox_run.rs b/crates/runner/src/executor/sandbox_run.rs index 3bf44206bda..45981dbf66d 100644 --- a/crates/runner/src/executor/sandbox_run.rs +++ b/crates/runner/src/executor/sandbox_run.rs @@ -8,9 +8,12 @@ use futures_util::FutureExt; use sandbox::{ Sandbox, SandboxConfig, SandboxCreateObserver, SandboxCreateStage, SandboxFactory, SandboxId, }; +use tokio_util::sync::CancellationToken; use tracing::{info, warn}; -use super::agent_run::{AgentExecutionResult, RunControls, RunStart, run_in_sandbox}; +use super::agent_run::{ + AgentExecutionResult, RunControls, RunStart, SessionHistoryRestorePlan, run_in_sandbox, +}; use super::cli_framework::{ EffectiveCliFramework, effective_cli_framework, normalized_cli_agent_type, }; @@ -24,13 +27,14 @@ use super::session_id::{ use super::telemetry::record_workspace_cache_result; use super::{ ExecuteOutcome, ExecutionFailure, ExecutorConfig, JobParams, NewSandboxDispatch, RunnerError, - RunnerResult, SandboxPreparedNotifier, SandboxReuseResult, + RunnerResult, SandboxPreparedNotifier, SandboxReuseResult, SessionHistoryMaterializer, }; use crate::duration::duration_ms; use crate::ids::RunId; use crate::network_log_manager::NetworkLogSession; use crate::provider::NetworkPolicyRefreshRegistration; use crate::proxy; +use crate::restored_session_identity::RestoredSessionIdentity; use crate::telemetry::JobTelemetry; use crate::types::{ExecutionContext, FirewallEntry}; use crate::workspace_image_cache::{ @@ -154,7 +158,7 @@ pub(super) async fn execute_new_sandbox_with_prepared_notifier( reuse_result, } = dispatch; let NewSandboxHooks { - controls, + mut controls, sandbox_prepared, } = hooks; let prepare_started = Instant::now(); @@ -167,6 +171,15 @@ pub(super) async fn execute_new_sandbox_with_prepared_notifier( telemetry, ) .await; + controls.session_history_restore_plan = resolve_fresh_session_history_restore_plan( + std::mem::take(&mut controls.session_history_restore_plan), + workspace_image.as_ref(), + context, + config, + controls.cancel.clone(), + telemetry, + ) + .await; let prepared = match create_started_sandbox( factory, context, @@ -208,6 +221,14 @@ pub(super) async fn execute_new_sandbox_with_prepared_notifier( None, ); workspace_image = None; + controls.session_history_restore_plan = + discard_local_sidecar_restore_plan_for_workspace_retry( + std::mem::take(&mut controls.session_history_restore_plan), + context, + config, + controls.cancel.clone(), + telemetry, + ); match create_started_sandbox( factory, context, @@ -343,6 +364,119 @@ pub(super) async fn prepare_workspace_image( Some(lease) } +async fn resolve_fresh_session_history_restore_plan( + plan: SessionHistoryRestorePlan, + workspace_image: Option<&WorkspaceImageLease>, + context: &ExecutionContext, + config: &ExecutorConfig, + cancel: CancellationToken, + telemetry: &mut JobTelemetry, +) -> SessionHistoryRestorePlan { + let fallback = match plan { + SessionHistoryRestorePlan::DeferredHashBacked { fallback } => fallback, + other => return other, + }; + let Some(workspace_image) = workspace_image else { + return start_fresh_session_history_materializer( + context, config, cancel, telemetry, fallback, + ); + }; + if !workspace_image.is_cache_hit() { + return start_fresh_session_history_materializer( + context, config, cancel, telemetry, fallback, + ); + } + let Some(expected) = RestoredSessionIdentity::from_context(context) else { + telemetry.record( + "session_history_workspace_cache_miss", + Duration::ZERO, + true, + Some("identity_mismatch"), + ); + return start_fresh_session_history_materializer( + context, config, cancel, telemetry, fallback, + ); + }; + let probe_started = Instant::now(); + telemetry.record( + "session_history_workspace_cache_probe", + Duration::ZERO, + true, + None, + ); + match workspace_image + .probe_session_history_sidecar(&expected) + .await + { + Ok(sidecar) => { + telemetry.record( + "session_history_workspace_cache_hit", + probe_started.elapsed(), + true, + None, + ); + SessionHistoryRestorePlan::LocalSidecar { sidecar, fallback } + } + Err(reason) => { + telemetry.record( + "session_history_workspace_cache_miss", + probe_started.elapsed(), + true, + Some(reason.as_str()), + ); + start_fresh_session_history_materializer(context, config, cancel, telemetry, fallback) + } + } +} + +fn discard_local_sidecar_restore_plan_for_workspace_retry( + plan: SessionHistoryRestorePlan, + context: &ExecutionContext, + config: &ExecutorConfig, + cancel: CancellationToken, + telemetry: &mut JobTelemetry, +) -> SessionHistoryRestorePlan { + match plan { + SessionHistoryRestorePlan::LocalSidecar { fallback, .. } => { + telemetry.record( + "session_history_workspace_cache_miss", + Duration::ZERO, + true, + Some("sandbox_retry_without_workspace_image"), + ); + start_fresh_session_history_materializer(context, config, cancel, telemetry, fallback) + } + other => other, + } +} + +fn start_fresh_session_history_materializer( + context: &ExecutionContext, + config: &ExecutorConfig, + cancel: CancellationToken, + telemetry: &mut JobTelemetry, + fallback: Option, +) -> SessionHistoryRestorePlan { + let started = Instant::now(); + let materializer = SessionHistoryMaterializer::start_cancellable( + &config.http, + context.resume_session.as_ref(), + effective_cli_framework(&context.cli_agent_type), + cancel, + Some(&config.session_history_probe), + ); + telemetry.record( + "runner_fresh_session_history_materializer_start", + started.elapsed(), + true, + None, + ); + SessionHistoryRestorePlan::Prestarted { + materializer, + fallback, + } +} + fn workspace_image_prepare_error(result: WorkspaceCacheCheckoutResult) -> Option<&'static str> { match result { WorkspaceCacheCheckoutResult::Hit diff --git a/crates/runner/src/executor/session_history_download.rs b/crates/runner/src/executor/session_history_download.rs index 5dab10aed86..f4571195e6e 100644 --- a/crates/runner/src/executor/session_history_download.rs +++ b/crates/runner/src/executor/session_history_download.rs @@ -780,6 +780,36 @@ fn validate_identity_body_size( Ok(()) } +pub(super) fn verify_identity_session_history_bytes( + bytes: &[u8], + expected_raw_size: u64, + expected_hash: &str, +) -> RunnerResult<()> { + if expected_raw_size == 0 { + return Err(RunnerError::Internal( + "identity session history rawSize must be positive".into(), + )); + } + if expected_raw_size > RESUME_SESSION_HISTORY_MAX_BYTES { + return Err(RunnerError::Internal(format!( + "session history is too large: {expected_raw_size} bytes exceeds {RESUME_SESSION_HISTORY_MAX_BYTES} bytes" + ))); + } + if bytes.len() as u64 != expected_raw_size { + return Err(RunnerError::Internal(format!( + "session history size mismatch: expected {expected_raw_size} bytes, got {} bytes", + bytes.len() + ))); + } + let actual_hash = hex::encode(Sha256::digest(bytes)); + if actual_hash != expected_hash { + return Err(RunnerError::Internal( + "session history hash mismatch".into(), + )); + } + Ok(()) +} + fn validate_compressed_ref( encoding: &str, history_ref: &ResumeSessionHistoryRef, @@ -831,6 +861,20 @@ fn validate_decompressed_raw_size( Ok(()) } +pub(super) fn verify_codex_zstd_session_history_bytes( + encoded_bytes: &[u8], + expected_raw_size: u64, + expected_hash: &str, +) -> RunnerResult>> { + let mut timings = SessionHistoryDownloadTimings::default(); + verify_codex_zstd_session_history( + encoded_bytes, + expected_raw_size, + expected_hash, + &mut timings, + ) +} + fn verify_codex_zstd_session_history( encoded_bytes: &[u8], expected_raw_size: u64, diff --git a/crates/runner/src/executor/tests/agent_run_tests.rs b/crates/runner/src/executor/tests/agent_run_tests.rs index 9405fc5b4ff..d623832d368 100644 --- a/crates/runner/src/executor/tests/agent_run_tests.rs +++ b/crates/runner/src/executor/tests/agent_run_tests.rs @@ -52,6 +52,9 @@ use crate::types::{ ResumeSessionHistoryEncoding, ResumeSessionHistoryRef, ResumeSessionHistoryRefKind, SandboxReuseResult, }; +use crate::workspace_image_cache::{ + WorkspaceSessionHistorySidecar, WorkspaceSessionHistorySidecarRepresentation, +}; const LARGE_SESSION_HISTORY_SIZE_BYTES: usize = 1024 * 1024 + 1; @@ -1180,6 +1183,225 @@ async fn run_in_sandbox_uses_prestarted_session_history_materializer() { ); } +#[tokio::test] +async fn run_in_sandbox_restores_session_history_from_workspace_sidecar() { + let dir = tempfile::tempdir().unwrap(); + let config = test_executor_config(dir.path()).await; + let sandbox = sandbox_mock::MockSandbox::new("test"); + let history = br#"{"type":"init"}"#; + let sidecar_path = dir.path().join("session-history.blob"); + tokio::fs::write(&sidecar_path, history).await.unwrap(); + let server = MockServer::start_async().await; + let history_mock = server + .mock_async(|when, then| { + when.method(GET).path("/history.blob"); + then.status(200).body(history); + }) + .await; + let mut ctx = minimal_context(); + ctx.resume_session = Some(ResumeSession { + cli_agent_session_id: "sess-sidecar-123".into(), + history: ResumeSessionHistory::Ref { + history_ref: ResumeSessionHistoryRef { + kind: ResumeSessionHistoryRefKind::Blob, + hash: hex::encode(Sha256::digest(history)), + url: server.url("/history.blob?token=secret"), + encoding: None, + raw_size: history.len() as u64, + encoded_size: history.len() as u64, + download_source: None, + }, + }, + }); + + let mut telemetry = test_telemetry(&config, &ctx); + let result = run_in_sandbox( + &sandbox, + &ctx, + &config, + RunStart { + restore_guest_state: false, + reuse_result: SandboxReuseResult::PoolMiss, + prev_storage: None, + }, + &mut telemetry, + RunControls::new(tokio_util::sync::CancellationToken::new(), None) + .with_session_history_restore_plan(SessionHistoryRestorePlan::LocalSidecar { + sidecar: WorkspaceSessionHistorySidecar { + path: sidecar_path, + representation: WorkspaceSessionHistorySidecarRepresentation::Raw, + encoded_size: history.len() as u64, + }, + fallback: None, + }), + ) + .await + .unwrap(); + + assert!(result.failure.is_none()); + let writes = sandbox.write_file_calls(); + assert_eq!(writes.len(), 1); + assert_eq!( + writes[0].path, + "/home/user/.claude/projects/-home-user-workspace/sess-sidecar-123.jsonl" + ); + assert_eq!(writes[0].content, history); + history_mock.assert_calls_async(0).await; + let ops = telemetry.pending_ops_snapshot(); + assert_successful_action_once(&ops, "session_history_workspace_cache_restore"); + assert_successful_action_once(&ops, "session_restore"); + assert_no_action(&ops, "session_history_download"); +} + +#[tokio::test] +async fn run_in_sandbox_falls_back_when_workspace_sidecar_hash_mismatches() { + let dir = tempfile::tempdir().unwrap(); + let config = test_executor_config(dir.path()).await; + let sandbox = sandbox_mock::MockSandbox::new("test"); + let history = br#"{"type":"init"}"#; + let corrupt_history = vec![b'x'; history.len()]; + let sidecar_path = dir.path().join("session-history.blob"); + tokio::fs::write(&sidecar_path, corrupt_history) + .await + .unwrap(); + let server = MockServer::start_async().await; + let history_mock = server + .mock_async(|when, then| { + when.method(GET).path("/history.blob"); + then.status(200).body(history); + }) + .await; + let mut ctx = minimal_context(); + ctx.resume_session = Some(ResumeSession { + cli_agent_session_id: "sess-sidecar-fallback-123".into(), + history: ResumeSessionHistory::Ref { + history_ref: ResumeSessionHistoryRef { + kind: ResumeSessionHistoryRefKind::Blob, + hash: hex::encode(Sha256::digest(history)), + url: server.url("/history.blob?token=secret"), + encoding: None, + raw_size: history.len() as u64, + encoded_size: history.len() as u64, + download_source: None, + }, + }, + }); + + let mut telemetry = test_telemetry(&config, &ctx); + let result = run_in_sandbox( + &sandbox, + &ctx, + &config, + RunStart { + restore_guest_state: false, + reuse_result: SandboxReuseResult::PoolMiss, + prev_storage: None, + }, + &mut telemetry, + RunControls::new(tokio_util::sync::CancellationToken::new(), None) + .with_session_history_restore_plan(SessionHistoryRestorePlan::LocalSidecar { + sidecar: WorkspaceSessionHistorySidecar { + path: sidecar_path, + representation: WorkspaceSessionHistorySidecarRepresentation::Raw, + encoded_size: history.len() as u64, + }, + fallback: None, + }), + ) + .await + .unwrap(); + + assert!(result.failure.is_none()); + let writes = sandbox.write_file_calls(); + assert_eq!(writes.len(), 1); + assert_eq!( + writes[0].path, + "/home/user/.claude/projects/-home-user-workspace/sess-sidecar-fallback-123.jsonl" + ); + assert_eq!(writes[0].content, history); + history_mock.assert_calls_async(1).await; + let ops = telemetry.pending_ops_snapshot(); + assert_failed_action_error_once( + &ops, + "session_history_workspace_cache_restore", + "materialize_error", + ); + assert_successful_action_once(&ops, "session_history_download"); +} + +#[tokio::test] +async fn run_in_sandbox_restores_codex_zstd_sidecar_with_session_timestamp() { + let dir = tempfile::tempdir().unwrap(); + let config = test_executor_config(dir.path()).await; + let sandbox = sandbox_mock::MockSandbox::new("test"); + let session_id = "019e9154-c304-70f0-adde-36efb1be1701"; + let history = br#"{"type":"session_meta","payload":{"timestamp":"2026-06-04T07:18:08Z"}}"#; + let mut history = history.to_vec(); + history.push(b'\n'); + let compressed_history = zstd_bytes(&history); + let sidecar_path = dir.path().join("session-history.blob"); + tokio::fs::write(&sidecar_path, &compressed_history) + .await + .unwrap(); + let server = MockServer::start_async().await; + let history_mock = server + .mock_async(|when, then| { + when.method(GET).path("/history.blob"); + then.status(200).body(&compressed_history); + }) + .await; + let mut ctx = minimal_context(); + ctx.cli_agent_type = "codex".into(); + ctx.resume_session = Some(ResumeSession { + cli_agent_session_id: session_id.into(), + history: ResumeSessionHistory::Ref { + history_ref: ResumeSessionHistoryRef { + kind: ResumeSessionHistoryRefKind::Blob, + hash: hex::encode(Sha256::digest(&history)), + url: server.url("/history.blob?token=secret"), + encoding: Some(ResumeSessionHistoryEncoding::Zstd), + raw_size: history.len() as u64, + encoded_size: compressed_history.len() as u64, + download_source: None, + }, + }, + }); + + let mut telemetry = test_telemetry(&config, &ctx); + let result = run_in_sandbox( + &sandbox, + &ctx, + &config, + RunStart { + restore_guest_state: false, + reuse_result: SandboxReuseResult::PoolMiss, + prev_storage: None, + }, + &mut telemetry, + RunControls::new(tokio_util::sync::CancellationToken::new(), None) + .with_session_history_restore_plan(SessionHistoryRestorePlan::LocalSidecar { + sidecar: WorkspaceSessionHistorySidecar { + path: sidecar_path, + representation: WorkspaceSessionHistorySidecarRepresentation::CodexZstd, + encoded_size: compressed_history.len() as u64, + }, + fallback: None, + }), + ) + .await + .unwrap(); + + assert!(result.failure.is_none()); + let writes = sandbox.write_file_calls(); + assert_eq!(writes.len(), 1); + assert_eq!( + writes[0].path, + "/home/user/.codex/sessions/2026/06/04/rollout-2026-06-04T07-18-08-019e9154-c304-70f0-adde-36efb1be1701.jsonl.zst" + ); + assert_eq!(writes[0].content, compressed_history); + history_mock.assert_calls_async(0).await; +} + #[tokio::test] async fn run_in_sandbox_records_completed_prestarted_materializer_failure() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/runner/src/executor/tests/sandbox_run_tests/workspace_cache.rs b/crates/runner/src/executor/tests/sandbox_run_tests/workspace_cache.rs index 56665a73afa..122031c0696 100644 --- a/crates/runner/src/executor/tests/sandbox_run_tests/workspace_cache.rs +++ b/crates/runner/src/executor/tests/sandbox_run_tests/workspace_cache.rs @@ -748,6 +748,7 @@ async fn reusable_idle_sandbox_with_workspace_promotion( run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-06-01T00:00:01.000Z".into(), storage_fingerprints: StorageFingerprints::default(), @@ -854,6 +855,7 @@ async fn reusable_idle_sandbox_with_unlocked_workspace_promotion( run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-06-01T00:00:01.000Z".into(), storage_fingerprints: StorageFingerprints::default(), diff --git a/crates/runner/src/restored_session_identity.rs b/crates/runner/src/restored_session_identity.rs index 6a240f21172..4afdf0573db 100644 --- a/crates/runner/src/restored_session_identity.rs +++ b/crates/runner/src/restored_session_identity.rs @@ -84,6 +84,14 @@ pub(crate) struct RestoredSessionFinalMetadataVerification<'a> { pub(crate) history_size_bytes: u64, } +pub(crate) struct RestoredSessionIdentityFields<'a> { + pub(crate) framework: FinalSessionHistoryFramework, + pub(crate) session_id_hash: &'a str, + pub(crate) history_ref_kind: FinalSessionHistoryRefKind, + pub(crate) history_hash: &'a str, + pub(crate) history_size_bytes: u64, +} + impl RestoredSessionIdentity { pub(crate) fn new( framework: RestoredSessionFramework, @@ -186,6 +194,16 @@ impl RestoredSessionIdentity { } } + pub(crate) fn cache_fields(&self) -> Option> { + Some(RestoredSessionIdentityFields { + framework: final_session_history_framework(self.framework), + session_id_hash: &self.session_id_hash, + history_ref_kind: final_session_history_ref_kind(self.history_ref_kind), + history_hash: &self.history_hash, + history_size_bytes: self.history_size_bytes?, + }) + } + pub(crate) fn has_final_metadata_verification(&self) -> bool { self.final_metadata_verification().is_some() } diff --git a/crates/runner/src/workspace_image_cache/entry.rs b/crates/runner/src/workspace_image_cache/entry.rs index ac8740705aa..41ab18560f9 100644 --- a/crates/runner/src/workspace_image_cache/entry.rs +++ b/crates/runner/src/workspace_image_cache/entry.rs @@ -14,6 +14,8 @@ pub(super) struct CacheEntryPaths { entry_dir: PathBuf, metadata: PathBuf, current_image: PathBuf, + session_history_sidecar: PathBuf, + session_history_sidecar_metadata: PathBuf, } impl CacheEntryPaths { @@ -22,6 +24,8 @@ impl CacheEntryPaths { Self { metadata: entry_dir.join("metadata.json"), current_image: entry_dir.join("current.ext4"), + session_history_sidecar: entry_dir.join("session-history.blob"), + session_history_sidecar_metadata: entry_dir.join("session-history.metadata.json"), entry_dir, } } @@ -42,9 +46,27 @@ impl CacheEntryPaths { &self.current_image } + pub(super) fn session_history_sidecar(&self) -> &Path { + &self.session_history_sidecar + } + + pub(super) fn session_history_sidecar_metadata(&self) -> &Path { + &self.session_history_sidecar_metadata + } + pub(super) fn tmp_image(&self, run_id: RunId) -> PathBuf { self.entry_dir.join(format!("current.ext4.tmp.{run_id}")) } + + pub(super) fn tmp_session_history_sidecar(&self, run_id: RunId) -> PathBuf { + self.entry_dir + .join(format!("session-history.blob.tmp.{run_id}")) + } + + pub(super) fn tmp_session_history_sidecar_metadata(&self, run_id: RunId) -> PathBuf { + self.entry_dir + .join(format!("session-history.metadata.json.tmp.{run_id}")) + } } impl SessionWorkspaceCache { @@ -76,6 +98,24 @@ impl SessionWorkspaceCache { self.entry_paths(cache_key).tmp_image(run_id) } + pub(super) fn session_workspace_cache_tmp_sidecar( + &self, + cache_key: &str, + run_id: RunId, + ) -> PathBuf { + self.entry_paths(cache_key) + .tmp_session_history_sidecar(run_id) + } + + pub(super) fn session_workspace_cache_tmp_sidecar_metadata( + &self, + cache_key: &str, + run_id: RunId, + ) -> PathBuf { + self.entry_paths(cache_key) + .tmp_session_history_sidecar_metadata(run_id) + } + pub(super) fn scoped_cache_key( &self, profile_name: &str, diff --git a/crates/runner/src/workspace_image_cache/fs.rs b/crates/runner/src/workspace_image_cache/fs.rs index 747e7714ba2..8db9b22e738 100644 --- a/crates/runner/src/workspace_image_cache/fs.rs +++ b/crates/runner/src/workspace_image_cache/fs.rs @@ -33,7 +33,10 @@ impl SessionWorkspaceCache { } pub(super) fn is_workspace_tmp_path_name(name: &str) -> bool { - name.starts_with("current.ext4.tmp.") || name.starts_with("metadata.json.tmp.") + name.starts_with("current.ext4.tmp.") + || name.starts_with("metadata.json.tmp.") + || name.starts_with("session-history.blob.tmp.") + || name.starts_with("session-history.metadata.json.tmp.") } pub(super) fn allocated_bytes(metadata: &std::fs::Metadata) -> u64 { diff --git a/crates/runner/src/workspace_image_cache/gc.rs b/crates/runner/src/workspace_image_cache/gc.rs index cdf2637d924..14bafd6ba36 100644 --- a/crates/runner/src/workspace_image_cache/gc.rs +++ b/crates/runner/src/workspace_image_cache/gc.rs @@ -485,9 +485,12 @@ impl SessionWorkspaceCache { Err(_) if self.inner.cache_scope.is_empty() => String::new(), Err(_) => return None, }; + let sidecar_allocated = self + .session_history_sidecar_allocated_bytes(&cache_key) + .await; Some(GcCandidate { cache_key, - allocated_bytes: allocated_bytes(&file_metadata), + allocated_bytes: allocated_bytes(&file_metadata).saturating_add(sidecar_allocated), file_dev: file_metadata.dev(), file_ino: file_metadata.ino(), last_used_at, diff --git a/crates/runner/src/workspace_image_cache/inspection.rs b/crates/runner/src/workspace_image_cache/inspection.rs index 2458a00b76c..190f3068d55 100644 --- a/crates/runner/src/workspace_image_cache/inspection.rs +++ b/crates/runner/src/workspace_image_cache/inspection.rs @@ -145,6 +145,10 @@ impl SessionWorkspaceCache { Some(metadata) => allocated_bytes(metadata), None => 0, }; + let persistent_allocated_bytes = current_allocated_bytes.saturating_add( + self.session_history_sidecar_allocated_bytes(&cache_key) + .await, + ); let entry = match (current_metadata, metadata) { (None, metadata) => { @@ -164,7 +168,7 @@ impl SessionWorkspaceCache { Some(reason.into()), metadata.as_ref(), None, - 0, + persistent_allocated_bytes, temporary, ) } @@ -178,7 +182,7 @@ impl SessionWorkspaceCache { Some(reason), None, Some(¤t_metadata), - current_allocated_bytes, + persistent_allocated_bytes, temporary, ) } @@ -197,7 +201,7 @@ impl SessionWorkspaceCache { reason, Some(&metadata), Some(¤t_metadata), - current_allocated_bytes, + persistent_allocated_bytes, temporary, ) } diff --git a/crates/runner/src/workspace_image_cache/lifecycle.rs b/crates/runner/src/workspace_image_cache/lifecycle.rs index 3b609abb117..7473c56a798 100644 --- a/crates/runner/src/workspace_image_cache/lifecycle.rs +++ b/crates/runner/src/workspace_image_cache/lifecycle.rs @@ -29,6 +29,8 @@ use super::types::{ WorkspaceImageActiveLeaseRequest, WorkspaceImageLeaseIdentity, WorkspaceImagePrepareRequest, WorkspaceImagePromotionIdentity, WorkspaceImagePromotionIdentityMismatch, WorkspaceImagePromotionIdentityRequest, WorkspaceImagePromotionRequest, + WorkspaceSessionHistorySidecar, WorkspaceSessionHistorySidecarMiss, + WorkspaceSessionHistorySidecarPromotionSource, }; use super::{ CACHE_FORMAT_VERSION, CACHE_KEY_VERSION, SessionWorkspaceCache, WORKSPACE_DRIVE_LAYOUT, @@ -65,6 +67,7 @@ pub(crate) struct WorkspaceImagePromotionContext { terminal_status: WorkspaceCacheTerminalStatus, completed_at: String, storage_fingerprints: StorageFingerprints, + restored_session_identity: Option, } pub(crate) struct WorkspaceImagePromotionIdentityFailure { @@ -90,6 +93,7 @@ struct WorkspaceImagePromotionInput<'a> { terminal_status: WorkspaceCacheTerminalStatus, completed_at: &'a str, storage_fingerprints: &'a StorageFingerprints, + session_history_sidecar: Option<&'a WorkspaceSessionHistorySidecarPromotionSource>, } struct WorkspaceImagePromotionTarget { @@ -859,6 +863,21 @@ impl SessionWorkspaceCache { current.display() ))); } + if let Err(e) = self + .publish_session_history_sidecar( + input.cache_key, + input.run_id, + input.session_history_sidecar, + ) + .await + { + warn!( + run_id = %input.run_id, + cache_key = input.cache_key, + error = %e, + "workspace image cache session history sidecar publish failed" + ); + } let allocated = allocated_bytes(¤t_metadata); let metadata = WorkspaceCacheMetadata { format_version: CACHE_FORMAT_VERSION, @@ -886,6 +905,7 @@ impl SessionWorkspaceCache { .await { let _ = remove_workspace_cache_path_if_exists(¤t).await; + let _ = self.prune_session_history_sidecar(input.cache_key).await; return Err(e); } info!( @@ -919,6 +939,22 @@ impl WorkspaceImageLease { self.result == WorkspaceCacheCheckoutResult::Hit } + pub(crate) async fn probe_session_history_sidecar( + &self, + expected: &crate::restored_session_identity::RestoredSessionIdentity, + ) -> Result { + if !self.is_cache_hit() { + return Err(WorkspaceSessionHistorySidecarMiss::NoCacheHit); + } + let cache_key = self + .cache_key + .as_deref() + .ok_or(WorkspaceSessionHistorySidecarMiss::Missing)?; + self.cache + .probe_session_history_sidecar(cache_key, expected) + .await + } + pub(crate) fn previous_storage(&self) -> Option<&StorageFingerprints> { self.previous_storage.as_ref() } @@ -966,6 +1002,7 @@ impl WorkspaceImageLease { run_id, sandbox_id: sandbox::SandboxId::new_v4(), cli_agent_session_id_override, + restored_session_identity: None, terminal_status, completed_at, storage_fingerprints: storage_fingerprints.clone(), @@ -1039,6 +1076,7 @@ impl WorkspaceImageLease { terminal_status: request.terminal_status, completed_at: request.completed_at, storage_fingerprints: request.storage_fingerprints, + restored_session_identity: request.restored_session_identity.cloned(), }) } } @@ -1060,6 +1098,40 @@ impl WorkspaceImagePromotionContext { &self.cli_agent_session_id } + pub(crate) fn restored_session_identity( + &self, + ) -> Option<&crate::restored_session_identity::RestoredSessionIdentity> { + self.restored_session_identity.as_ref() + } + + pub(crate) fn session_history_sidecar_tmp_path(&self) -> PathBuf { + self.cache + .session_workspace_cache_tmp_sidecar(&self.cache_key, self.run_id) + } + + pub(crate) fn session_history_sidecar_source( + &self, + tmp_path: PathBuf, + representation: super::types::WorkspaceSessionHistorySidecarRepresentation, + encoded_size: u64, + ) -> Option { + Some(WorkspaceSessionHistorySidecarPromotionSource { + tmp_path, + representation, + encoded_size, + restored_session_identity: self.restored_session_identity.clone()?, + }) + } + + pub(crate) async fn discard_session_history_sidecar_source( + &self, + source: &WorkspaceSessionHistorySidecarPromotionSource, + ) { + self.cache + .discard_session_history_sidecar_source(source) + .await; + } + pub(crate) fn validate_identity( &self, expected: &WorkspaceImagePromotionIdentity, @@ -1104,7 +1176,15 @@ impl WorkspaceImagePromotionContext { self.validate_expected_identity(&self.cache, request) } + #[cfg(test)] pub(crate) async fn promote(&self) -> RunnerResult { + self.promote_with_session_history_sidecar(None).await + } + + pub(crate) async fn promote_with_session_history_sidecar( + &self, + session_history_sidecar: Option<&WorkspaceSessionHistorySidecarPromotionSource>, + ) -> RunnerResult { let tainted_storage_fingerprints; let promotion_storage_fingerprints = match self.terminal_status { WorkspaceCacheTerminalStatus::Success => &self.storage_fingerprints, @@ -1144,6 +1224,7 @@ impl WorkspaceImagePromotionContext { terminal_status: self.terminal_status, completed_at: &self.completed_at, storage_fingerprints: promotion_storage_fingerprints, + session_history_sidecar, }) .await } @@ -1261,6 +1342,7 @@ impl WorkspaceImagePromotionContext { terminal_status: _, completed_at: _, storage_fingerprints: _, + restored_session_identity: _, } = self; let base = WorkspaceImageLeaseBase { cache, diff --git a/crates/runner/src/workspace_image_cache/mod.rs b/crates/runner/src/workspace_image_cache/mod.rs index fe541c7f567..579365db55c 100644 --- a/crates/runner/src/workspace_image_cache/mod.rs +++ b/crates/runner/src/workspace_image_cache/mod.rs @@ -39,6 +39,7 @@ mod inspection; mod lifecycle; mod metadata; mod path_safety; +mod sidecar; mod types; #[cfg(test)] @@ -55,6 +56,8 @@ pub(crate) use types::{ WorkspaceImageCacheInspectionSummary, WorkspaceImageLeaseIdentity, WorkspaceImagePrepareRequest, WorkspaceImagePromotionIdentityMismatch, WorkspaceImagePromotionIdentityRequest, WorkspaceImagePromotionRequest, + WorkspaceSessionHistorySidecar, WorkspaceSessionHistorySidecarPromotionSource, + WorkspaceSessionHistorySidecarRepresentation, }; const CACHE_FORMAT_VERSION: u32 = 1; diff --git a/crates/runner/src/workspace_image_cache/sidecar.rs b/crates/runner/src/workspace_image_cache/sidecar.rs new file mode 100644 index 00000000000..2c82653522b --- /dev/null +++ b/crates/runner/src/workspace_image_cache/sidecar.rs @@ -0,0 +1,263 @@ +use std::path::Path; + +use guest_contracts::session_history_identity::{ + FinalSessionHistoryFramework, FinalSessionHistoryRefKind, SESSION_HISTORY_SIDECAR_MAX_BYTES, +}; +use serde::{Deserialize, Serialize}; +use tokio::fs; + +use crate::error::{RunnerError, RunnerResult}; +use crate::restored_session_identity::{RestoredSessionIdentity, RestoredSessionIdentityFields}; + +use super::fs::{ + allocated_bytes, ensure_workspace_cache_entry_dir, remove_workspace_cache_path_if_exists, + workspace_cache_existing_path_allocated_bytes, +}; +use super::metadata::WorkspaceImageFileIdentity; +use super::types::{ + WorkspaceSessionHistorySidecar, WorkspaceSessionHistorySidecarMiss, + WorkspaceSessionHistorySidecarPromotionSource, WorkspaceSessionHistorySidecarRepresentation, +}; +use super::{SessionWorkspaceCache, entry::CacheEntryPaths}; + +const SESSION_HISTORY_SIDECAR_FORMAT_VERSION: u8 = 1; + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +struct WorkspaceSessionHistorySidecarMetadata { + version: u8, + framework: FinalSessionHistoryFramework, + session_id_hash: String, + history_ref_kind: FinalSessionHistoryRefKind, + history_hash: String, + history_size_bytes: u64, + representation: WorkspaceSessionHistorySidecarRepresentation, + encoded_size: u64, + body_file: WorkspaceImageFileIdentity, + allocated_bytes: u64, +} + +impl WorkspaceSessionHistorySidecarMetadata { + fn from_source( + source: &WorkspaceSessionHistorySidecarPromotionSource, + body_metadata: &std::fs::Metadata, + ) -> Option { + let RestoredSessionIdentityFields { + framework, + session_id_hash, + history_ref_kind, + history_hash, + history_size_bytes, + } = source.restored_session_identity.cache_fields()?; + Some(Self { + version: SESSION_HISTORY_SIDECAR_FORMAT_VERSION, + framework, + session_id_hash: session_id_hash.to_owned(), + history_ref_kind, + history_hash: history_hash.to_owned(), + history_size_bytes, + representation: source.representation, + encoded_size: source.encoded_size, + body_file: WorkspaceImageFileIdentity::from_metadata(body_metadata), + allocated_bytes: allocated_bytes(body_metadata), + }) + } + + fn validate_for_request( + &self, + expected: &RestoredSessionIdentity, + ) -> Result<(), WorkspaceSessionHistorySidecarMiss> { + if self.version != SESSION_HISTORY_SIDECAR_FORMAT_VERSION { + return Err(WorkspaceSessionHistorySidecarMiss::InvalidMetadata); + } + if self.encoded_size == 0 || self.encoded_size > SESSION_HISTORY_SIDECAR_MAX_BYTES { + return Err(WorkspaceSessionHistorySidecarMiss::InvalidMetadata); + } + let fields = expected + .cache_fields() + .ok_or(WorkspaceSessionHistorySidecarMiss::IdentityMismatch)?; + if self.framework != fields.framework + || self.session_id_hash != fields.session_id_hash + || self.history_ref_kind != fields.history_ref_kind + || self.history_hash != fields.history_hash + || self.history_size_bytes != fields.history_size_bytes + { + return Err(WorkspaceSessionHistorySidecarMiss::IdentityMismatch); + } + match (self.framework, self.representation) { + (_, WorkspaceSessionHistorySidecarRepresentation::Raw) + if self.encoded_size == self.history_size_bytes => + { + Ok(()) + } + ( + FinalSessionHistoryFramework::Codex, + WorkspaceSessionHistorySidecarRepresentation::CodexZstd, + ) => Ok(()), + _ => Err(WorkspaceSessionHistorySidecarMiss::UnsupportedFormat), + } + } +} + +impl SessionWorkspaceCache { + pub(super) async fn probe_session_history_sidecar( + &self, + cache_key: &str, + expected: &RestoredSessionIdentity, + ) -> Result { + let paths = self.entry_paths(cache_key); + let metadata = self + .read_session_history_sidecar_metadata(paths.session_history_sidecar_metadata()) + .await?; + metadata.validate_for_request(expected)?; + let body_metadata = fs::symlink_metadata(paths.session_history_sidecar()) + .await + .map_err(|e| { + if e.kind() == std::io::ErrorKind::NotFound { + WorkspaceSessionHistorySidecarMiss::BodyMissing + } else { + WorkspaceSessionHistorySidecarMiss::FileIdentityMismatch + } + })?; + if !body_metadata.is_file() + || WorkspaceImageFileIdentity::from_metadata(&body_metadata) != metadata.body_file + || body_metadata.len() != metadata.encoded_size + { + return Err(WorkspaceSessionHistorySidecarMiss::FileIdentityMismatch); + } + Ok(WorkspaceSessionHistorySidecar { + path: paths.session_history_sidecar().to_path_buf(), + representation: metadata.representation, + encoded_size: metadata.encoded_size, + }) + } + + pub(super) async fn publish_session_history_sidecar( + &self, + cache_key: &str, + run_id: crate::ids::RunId, + source: Option<&WorkspaceSessionHistorySidecarPromotionSource>, + ) -> RunnerResult<()> { + let paths = self.entry_paths(cache_key); + match source { + Some(source) => { + if let Err(error) = self + .publish_session_history_sidecar_source(cache_key, run_id, &paths, source) + .await + { + let _ = self.prune_session_history_sidecar(cache_key).await; + return Err(error); + } + } + None => self.prune_session_history_sidecar(cache_key).await?, + } + Ok(()) + } + + pub(super) async fn discard_session_history_sidecar_source( + &self, + source: &WorkspaceSessionHistorySidecarPromotionSource, + ) { + let _ = remove_workspace_cache_path_if_exists(&source.tmp_path).await; + } + + pub(super) async fn session_history_sidecar_allocated_bytes(&self, cache_key: &str) -> u64 { + let paths = self.entry_paths(cache_key); + let body = workspace_cache_existing_path_allocated_bytes(paths.session_history_sidecar()) + .await + .ok() + .flatten() + .unwrap_or(0); + let metadata = + workspace_cache_existing_path_allocated_bytes(paths.session_history_sidecar_metadata()) + .await + .ok() + .flatten() + .unwrap_or(0); + body.saturating_add(metadata) + } + + async fn publish_session_history_sidecar_source( + &self, + cache_key: &str, + run_id: crate::ids::RunId, + paths: &CacheEntryPaths, + source: &WorkspaceSessionHistorySidecarPromotionSource, + ) -> RunnerResult<()> { + let tmp_metadata = fs::symlink_metadata(&source.tmp_path).await?; + if !tmp_metadata.is_file() + || tmp_metadata.len() != source.encoded_size + || source.encoded_size == 0 + || source.encoded_size > SESSION_HISTORY_SIDECAR_MAX_BYTES + { + let _ = remove_workspace_cache_path_if_exists(&source.tmp_path).await; + self.prune_session_history_sidecar(cache_key).await?; + return Ok(()); + } + let Some(sidecar_metadata) = + WorkspaceSessionHistorySidecarMetadata::from_source(source, &tmp_metadata) + else { + let _ = remove_workspace_cache_path_if_exists(&source.tmp_path).await; + self.prune_session_history_sidecar(cache_key).await?; + return Ok(()); + }; + ensure_workspace_cache_entry_dir(paths.entry_dir()).await?; + let tmp_metadata_path = + self.session_workspace_cache_tmp_sidecar_metadata(cache_key, run_id); + let sidecar_metadata_path = paths.session_history_sidecar_metadata(); + let sidecar_body_path = paths.session_history_sidecar(); + let _ = remove_workspace_cache_path_if_exists(&tmp_metadata_path).await; + let bytes = serde_json::to_vec_pretty(&sidecar_metadata).map_err(|e| { + RunnerError::Internal(format!("serialize workspace session history sidecar: {e}")) + })?; + if let Err(e) = fs::write(&tmp_metadata_path, bytes).await { + let _ = remove_workspace_cache_path_if_exists(&tmp_metadata_path).await; + let _ = remove_workspace_cache_path_if_exists(&source.tmp_path).await; + return Err(e.into()); + } + let _ = remove_workspace_cache_path_if_exists(sidecar_metadata_path).await; + let _ = remove_workspace_cache_path_if_exists(sidecar_body_path).await; + if let Err(e) = fs::rename(&source.tmp_path, sidecar_body_path).await { + let _ = remove_workspace_cache_path_if_exists(&tmp_metadata_path).await; + let _ = remove_workspace_cache_path_if_exists(&source.tmp_path).await; + return Err(e.into()); + } + if let Err(e) = fs::rename(&tmp_metadata_path, sidecar_metadata_path).await { + let _ = remove_workspace_cache_path_if_exists(&tmp_metadata_path).await; + let _ = remove_workspace_cache_path_if_exists(sidecar_body_path).await; + return Err(e.into()); + } + Ok(()) + } + + pub(super) async fn prune_session_history_sidecar(&self, cache_key: &str) -> RunnerResult<()> { + let paths = self.entry_paths(cache_key); + let metadata_result = + remove_workspace_cache_path_if_exists(paths.session_history_sidecar_metadata()).await; + let body_result = + remove_workspace_cache_path_if_exists(paths.session_history_sidecar()).await; + metadata_result?; + body_result?; + Ok(()) + } + + async fn read_session_history_sidecar_metadata( + &self, + path: &Path, + ) -> Result { + let bytes = crate::state_file::read_to_bytes_required( + path, + crate::state_file::WORKSPACE_METADATA_MAX_BYTES, + crate::state_file::OwnerCheck::None, + ) + .await + .map_err(|error| match error { + RunnerError::Io(e) if e.kind() == std::io::ErrorKind::NotFound => { + WorkspaceSessionHistorySidecarMiss::Missing + } + _ => WorkspaceSessionHistorySidecarMiss::InvalidMetadata, + })?; + serde_json::from_slice(&bytes) + .map_err(|_| WorkspaceSessionHistorySidecarMiss::InvalidMetadata) + } +} diff --git a/crates/runner/src/workspace_image_cache/tests/mod.rs b/crates/runner/src/workspace_image_cache/tests/mod.rs index 1786f690b11..c9231f645f3 100644 --- a/crates/runner/src/workspace_image_cache/tests/mod.rs +++ b/crates/runner/src/workspace_image_cache/tests/mod.rs @@ -2,6 +2,8 @@ use std::collections::HashMap; use std::os::unix::fs::MetadataExt; use std::path::Path; +use api_contracts::generated::constants::runners::paths::CANONICAL_WORKING_DIR; + use super::fs::{ allocated_bytes, has_copy_headroom, local_timestamp, sparse_copy_with_timeout, workspace_cache_path_allocated_bytes, @@ -15,13 +17,16 @@ use super::path_safety::{ filter_storage_fingerprints_for_working_dir, is_safe_guest_working_dir, normalize_safe_guest_working_dir, }; +use super::types::WorkspaceSessionHistorySidecarMiss; use super::*; use crate::error::RunnerError; use crate::ids::RunId; use crate::paths::{RunnerPaths, scoped_session_workspace_cache_key, session_workspace_cache_key}; +use crate::restored_session_identity::{RestoredSessionFramework, RestoredSessionIdentity}; use crate::storage_fingerprints::StorageFingerprint; use crate::storage_fingerprints::StorageFingerprints; -use crate::types::{HeldSessionState, MAX_HELD_SESSION_STATES}; +use crate::types::{HeldSessionState, MAX_HELD_SESSION_STATES, ResumeSessionHistoryRefKind}; +use sha2::{Digest, Sha256}; use tokio::fs; const TEST_PROFILE_NAME: &str = "vm0/default"; @@ -91,6 +96,39 @@ async fn write_current_cache_entry( key } +fn test_restored_session_identity(session_id: &str, history: &[u8]) -> RestoredSessionIdentity { + RestoredSessionIdentity::new( + RestoredSessionFramework::ClaudeCode, + session_id, + ResumeSessionHistoryRefKind::Blob, + hex::encode(Sha256::digest(history)), + Some(history.len() as u64), + ) +} + +async fn publish_test_session_history_sidecar( + cache: &SessionWorkspaceCache, + cache_key: &str, + run_id: RunId, + session_id: &str, + history: &[u8], +) -> RestoredSessionIdentity { + let tmp_path = cache.session_workspace_cache_tmp_sidecar(cache_key, run_id); + fs::write(&tmp_path, history).await.unwrap(); + let identity = test_restored_session_identity(session_id, history); + let source = WorkspaceSessionHistorySidecarPromotionSource { + tmp_path, + representation: WorkspaceSessionHistorySidecarRepresentation::Raw, + encoded_size: history.len() as u64, + restored_session_identity: identity.clone(), + }; + cache + .publish_session_history_sidecar(cache_key, run_id, Some(&source)) + .await + .unwrap(); + identity +} + async fn promote_current_cache_entry( cache: &SessionWorkspaceCache, paths: &RunnerPaths, @@ -1383,6 +1421,7 @@ async fn abandoned_cache_hit_promotion_context_invalidates_consumed_entry() { run_id, sandbox_id, cli_agent_session_id_override: None, + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-05-02T00:00:00.000Z".into(), storage_fingerprints: StorageFingerprints::default(), @@ -1455,6 +1494,7 @@ async fn promotion_context_preserves_existing_newer_cache_entry() { run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-05-01T00:00:00.000Z".into(), storage_fingerprints: StorageFingerprints::default(), @@ -1477,6 +1517,201 @@ async fn promotion_context_preserves_existing_newer_cache_entry() { ); } +#[tokio::test] +async fn session_history_sidecar_publish_and_probe_hit() { + let dir = tempfile::tempdir().unwrap(); + let paths = RunnerPaths::new(dir.path().join("runner")); + tokio::fs::create_dir_all(paths.base_dir()).await.unwrap(); + let cache = SessionWorkspaceCache::new(paths.clone()); + let run_id = RunId::new_v4(); + let session_id = "sess-sidecar-hit"; + let history = br#"{"type":"message","content":"hello"}"#; + let cache_key = write_current_cache_entry( + &cache, + run_id, + session_id, + CANONICAL_WORKING_DIR, + "2026-05-01T00:00:00.000Z", + "2026-05-01T00:00:00.000Z", + ) + .await; + + let identity = + publish_test_session_history_sidecar(&cache, &cache_key, run_id, session_id, history).await; + let sidecar = cache + .probe_session_history_sidecar(&cache_key, &identity) + .await + .unwrap(); + + assert_eq!( + sidecar.representation, + WorkspaceSessionHistorySidecarRepresentation::Raw + ); + assert_eq!(sidecar.encoded_size, history.len() as u64); + assert_eq!(fs::read(sidecar.path).await.unwrap(), history); +} + +#[tokio::test] +async fn session_history_sidecar_publish_none_prunes_existing_sidecar() { + let dir = tempfile::tempdir().unwrap(); + let paths = RunnerPaths::new(dir.path().join("runner")); + tokio::fs::create_dir_all(paths.base_dir()).await.unwrap(); + let cache = SessionWorkspaceCache::new(paths.clone()); + let run_id = RunId::new_v4(); + let session_id = "sess-sidecar-prune"; + let history = br#"{"type":"message","content":"old"}"#; + let cache_key = write_current_cache_entry( + &cache, + run_id, + session_id, + CANONICAL_WORKING_DIR, + "2026-05-01T00:00:00.000Z", + "2026-05-01T00:00:00.000Z", + ) + .await; + let identity = + publish_test_session_history_sidecar(&cache, &cache_key, run_id, session_id, history).await; + + cache + .publish_session_history_sidecar(&cache_key, RunId::new_v4(), None) + .await + .unwrap(); + + assert_eq!( + cache + .probe_session_history_sidecar(&cache_key, &identity) + .await + .unwrap_err(), + WorkspaceSessionHistorySidecarMiss::Missing + ); +} + +#[tokio::test] +async fn session_history_sidecar_invalid_source_prunes_existing_sidecar() { + let dir = tempfile::tempdir().unwrap(); + let paths = RunnerPaths::new(dir.path().join("runner")); + tokio::fs::create_dir_all(paths.base_dir()).await.unwrap(); + let cache = SessionWorkspaceCache::new(paths.clone()); + let run_id = RunId::new_v4(); + let session_id = "sess-sidecar-invalid-source-prune"; + let history = br#"{"type":"message","content":"old"}"#; + let cache_key = write_current_cache_entry( + &cache, + run_id, + session_id, + CANONICAL_WORKING_DIR, + "2026-05-01T00:00:00.000Z", + "2026-05-01T00:00:00.000Z", + ) + .await; + let identity = + publish_test_session_history_sidecar(&cache, &cache_key, run_id, session_id, history).await; + let invalid_tmp_path = cache.session_workspace_cache_tmp_sidecar(&cache_key, RunId::new_v4()); + fs::write(&invalid_tmp_path, b"new").await.unwrap(); + let invalid_source = WorkspaceSessionHistorySidecarPromotionSource { + tmp_path: invalid_tmp_path, + representation: WorkspaceSessionHistorySidecarRepresentation::Raw, + encoded_size: 0, + restored_session_identity: identity.clone(), + }; + + cache + .publish_session_history_sidecar(&cache_key, RunId::new_v4(), Some(&invalid_source)) + .await + .unwrap(); + + assert_eq!( + cache + .probe_session_history_sidecar(&cache_key, &identity) + .await + .unwrap_err(), + WorkspaceSessionHistorySidecarMiss::Missing + ); +} + +#[tokio::test] +async fn session_history_sidecar_probe_rejects_mismatched_body_identity() { + let dir = tempfile::tempdir().unwrap(); + let paths = RunnerPaths::new(dir.path().join("runner")); + tokio::fs::create_dir_all(paths.base_dir()).await.unwrap(); + let cache = SessionWorkspaceCache::new(paths.clone()); + let run_id = RunId::new_v4(); + let session_id = "sess-sidecar-mismatch"; + let history = br#"{"type":"message","content":"stable"}"#; + let cache_key = write_current_cache_entry( + &cache, + run_id, + session_id, + CANONICAL_WORKING_DIR, + "2026-05-01T00:00:00.000Z", + "2026-05-01T00:00:00.000Z", + ) + .await; + let identity = + publish_test_session_history_sidecar(&cache, &cache_key, run_id, session_id, history).await; + fs::write( + paths + .session_workspace_cache_entry_dir(&cache_key) + .join("session-history.blob"), + b"changed", + ) + .await + .unwrap(); + + assert_eq!( + cache + .probe_session_history_sidecar(&cache_key, &identity) + .await + .unwrap_err(), + WorkspaceSessionHistorySidecarMiss::FileIdentityMismatch + ); +} + +#[tokio::test] +async fn session_history_sidecar_counts_toward_gc_candidate_and_inspection() { + let dir = tempfile::tempdir().unwrap(); + let paths = RunnerPaths::new(dir.path().join("runner")); + tokio::fs::create_dir_all(paths.base_dir()).await.unwrap(); + let cache = SessionWorkspaceCache::new(paths.clone()); + let run_id = RunId::new_v4(); + let session_id = "sess-sidecar-accounting"; + let history = br#"{"type":"message","content":"accounting"}"#; + let cache_key = write_current_cache_entry( + &cache, + run_id, + session_id, + CANONICAL_WORKING_DIR, + "2026-05-01T00:00:00.000Z", + "2026-05-01T00:00:00.000Z", + ) + .await; + publish_test_session_history_sidecar(&cache, &cache_key, run_id, session_id, history).await; + let current_allocated = workspace_cache_path_allocated_bytes( + &paths.session_workspace_cache_current_image(&cache_key), + ) + .await; + let sidecar_allocated = cache + .session_history_sidecar_allocated_bytes(&cache_key) + .await; + + let candidate = cache.gc_candidate(cache_key.clone()).await.unwrap(); + assert_eq!( + candidate.allocated_bytes, + current_allocated.saturating_add(sidecar_allocated) + ); + + let inspection = cache.inspect().await.unwrap(); + let entry = inspection + .entries + .iter() + .find(|entry| entry.cache_key == cache_key) + .unwrap(); + assert_eq!( + entry.allocated_bytes, + current_allocated.saturating_add(sidecar_allocated) + ); +} + #[tokio::test] async fn no_lock_promotion_context_abandonment_preserves_existing_entry() { let dir = tempfile::tempdir().unwrap(); @@ -1515,6 +1750,7 @@ async fn no_lock_promotion_context_abandonment_preserves_existing_entry() { run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-05-01T00:00:00.000Z".into(), storage_fingerprints: StorageFingerprints::default(), @@ -2920,6 +3156,7 @@ async fn promotion_context_keeps_entry_locked_until_reused_active_lease_drops() run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: local_timestamp(), storage_fingerprints: StorageFingerprints::default(), @@ -3019,6 +3256,7 @@ async fn promotion_context_validates_expected_identity() { run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: local_timestamp(), storage_fingerprints: StorageFingerprints::default(), @@ -4246,6 +4484,7 @@ async fn no_session_checkout_without_late_cli_agent_session_id_has_no_promotion_ run_id, sandbox_id, cli_agent_session_id_override: None, + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-05-01T00:00:00.000Z".into(), storage_fingerprints: StorageFingerprints::default(), @@ -4345,6 +4584,7 @@ async fn late_session_promotion_skips_when_entry_lock_is_busy() { run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-05-01T00:00:00.000Z".into(), storage_fingerprints: StorageFingerprints::default(), @@ -4395,6 +4635,7 @@ async fn no_lock_promotion_context_survives_reuse_active_lease() { run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-05-01T00:00:00.000Z".into(), storage_fingerprints: StorageFingerprints::default(), @@ -4419,6 +4660,7 @@ async fn no_lock_promotion_context_survives_reuse_active_lease() { run_id: RunId::new_v4(), sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-05-01T00:00:01.000Z".into(), storage_fingerprints: StorageFingerprints::default(), diff --git a/crates/runner/src/workspace_image_cache/types.rs b/crates/runner/src/workspace_image_cache/types.rs index 9813c296bdb..f411359d632 100644 --- a/crates/runner/src/workspace_image_cache/types.rs +++ b/crates/runner/src/workspace_image_cache/types.rs @@ -105,11 +105,56 @@ pub(crate) struct WorkspaceImagePromotionRequest<'a> { pub(crate) run_id: RunId, pub(crate) sandbox_id: sandbox::SandboxId, pub(crate) cli_agent_session_id_override: Option<&'a str>, + pub(crate) restored_session_identity: + Option<&'a crate::restored_session_identity::RestoredSessionIdentity>, pub(crate) terminal_status: WorkspaceCacheTerminalStatus, pub(crate) completed_at: String, pub(crate) storage_fingerprints: StorageFingerprints, } +pub(crate) type WorkspaceSessionHistorySidecarRepresentation = + guest_contracts::session_history_identity::SessionHistorySidecarRepresentation; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct WorkspaceSessionHistorySidecar { + pub(crate) path: PathBuf, + pub(crate) representation: WorkspaceSessionHistorySidecarRepresentation, + pub(crate) encoded_size: u64, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct WorkspaceSessionHistorySidecarPromotionSource { + pub(crate) tmp_path: PathBuf, + pub(crate) representation: WorkspaceSessionHistorySidecarRepresentation, + pub(crate) encoded_size: u64, + pub(crate) restored_session_identity: crate::restored_session_identity::RestoredSessionIdentity, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum WorkspaceSessionHistorySidecarMiss { + NoCacheHit, + Missing, + InvalidMetadata, + IdentityMismatch, + UnsupportedFormat, + BodyMissing, + FileIdentityMismatch, +} + +impl WorkspaceSessionHistorySidecarMiss { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::NoCacheHit => "no_cache_hit", + Self::Missing => "missing", + Self::InvalidMetadata => "invalid_metadata", + Self::IdentityMismatch => "identity_mismatch", + Self::UnsupportedFormat => "unsupported_format", + Self::BodyMissing => "body_missing", + Self::FileIdentityMismatch => "file_identity_mismatch", + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] #[serde(rename_all = "camelCase")] pub(crate) struct FsStats { diff --git a/crates/runner/src/workspace_promotion.rs b/crates/runner/src/workspace_promotion.rs index 3241f071b40..5f4a68af8dd 100644 --- a/crates/runner/src/workspace_promotion.rs +++ b/crates/runner/src/workspace_promotion.rs @@ -1,20 +1,65 @@ use std::panic::AssertUnwindSafe; +use std::path::PathBuf; +use std::time::Duration; use futures_util::FutureExt; -use sandbox::Sandbox; +use guest_contracts::session_history_identity::{ + SESSION_HISTORY_SIDECAR_MAX_BYTES, SessionHistorySidecarExportMetadata, +}; +use sandbox::{CopyFileOptions, EXEC_OUTPUT_LIMIT_64_KIB, ExecRequest, Sandbox}; +use shell_quote::quote_shell_arg; +use tokio::fs; use tracing::warn; +use crate::helper_exec::{format_helper_exec_failure, helper_exec_succeeded}; +use crate::paths::guest; use crate::workspace_image_cache::{ WorkspaceImagePromotionContext, WorkspaceImagePromotionOutcome, + WorkspaceSessionHistorySidecarPromotionSource, }; use crate::workspace_mount::flush_and_unmount_workspace_drive; +const SESSION_HISTORY_SIDECAR_EXPORT_TIMEOUT: Duration = Duration::from_secs(10); +const SESSION_HISTORY_SIDECAR_COPY_TIMEOUT: Duration = Duration::from_secs(30); +const SESSION_HISTORY_SIDECAR_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5); + enum WorkspacePromotionAction { Promoted, PreservedExisting, AbandonUnpublished, } +struct SessionHistorySidecarSourceGuard { + source: Option, +} + +impl SessionHistorySidecarSourceGuard { + fn new(source: Option) -> Self { + Self { source } + } + + fn as_ref(&self) -> Option<&WorkspaceSessionHistorySidecarPromotionSource> { + self.source.as_ref() + } + + async fn discard(&mut self, promotion: &WorkspaceImagePromotionContext) { + if let Some(source) = self.source.as_ref() { + promotion + .discard_session_history_sidecar_source(source) + .await; + } + self.source = None; + } +} + +impl Drop for SessionHistorySidecarSourceGuard { + fn drop(&mut self) { + if let Some(source) = self.source.take() { + let _ = std::fs::remove_file(source.tmp_path); + } + } +} + pub(crate) async fn promote_workspace_image_from_active_sandbox( sandbox: &dyn Sandbox, promotion: Option, @@ -56,9 +101,13 @@ async fn promote_workspace_image_from_active_sandbox_inner( promotion: &WorkspaceImagePromotionContext, reason: &'static str, ) -> WorkspacePromotionAction { + let mut sidecar_source = SessionHistorySidecarSourceGuard::new( + export_session_history_sidecar(sandbox, promotion, reason).await, + ); match flush_and_unmount_workspace_drive(sandbox, promotion.run_id()).await { Ok(()) => {} Err(e) => { + sidecar_source.discard(promotion).await; warn!( run_id = %promotion.run_id(), sandbox_id = %promotion.sandbox_id(), @@ -72,7 +121,13 @@ async fn promote_workspace_image_from_active_sandbox_inner( } } - match promotion.promote().await { + let outcome = promotion + .promote_with_session_history_sidecar(sidecar_source.as_ref()) + .await; + if !matches!(outcome, Ok(WorkspaceImagePromotionOutcome::Promoted)) { + sidecar_source.discard(promotion).await; + } + match outcome { Ok(WorkspaceImagePromotionOutcome::Promoted) => WorkspacePromotionAction::Promoted, Ok(WorkspaceImagePromotionOutcome::PreservedExisting) => { WorkspacePromotionAction::PreservedExisting @@ -95,6 +150,188 @@ async fn promote_workspace_image_from_active_sandbox_inner( } } +async fn export_session_history_sidecar( + sandbox: &dyn Sandbox, + promotion: &WorkspaceImagePromotionContext, + reason: &'static str, +) -> Option { + let verification = promotion + .restored_session_identity()? + .final_metadata_verification()?; + let export_path = guest_contracts::runtime_paths::session_history_sidecar_export_file( + PathBuf::from(verification.runtime_dir), + ); + let export_path = export_path.to_string_lossy().into_owned(); + let command = [ + quote_shell_arg(guest::RUN_AGENT), + "export-session-history-sidecar".to_string(), + quote_shell_arg(verification.metadata_path), + quote_shell_arg(&export_path), + ] + .join(" "); + let env = [( + guest_contracts::runtime_paths::GUEST_RUNTIME_DIR_ENV, + verification.runtime_dir, + )]; + let request = ExecRequest { + cmd: &command, + timeout: SESSION_HISTORY_SIDECAR_EXPORT_TIMEOUT, + env: &env, + sudo: false, + stdin_bytes: None, + output_limits: EXEC_OUTPUT_LIMIT_64_KIB, + }; + let result = match sandbox + .exec_with_diagnostic_label(&request, "session-history-sidecar-export") + .await + { + Ok(result) => result, + Err(e) => { + warn!( + run_id = %promotion.run_id(), + sandbox_id = %promotion.sandbox_id(), + profile_name = promotion.profile_name(), + session_id = %promotion.cli_agent_session_id(), + reason, + error = %e, + "workspace image cache session history sidecar export errored" + ); + cleanup_guest_session_history_sidecar_export(sandbox, promotion, &export_path, reason) + .await; + return None; + } + }; + if !helper_exec_succeeded(&result) { + warn!( + run_id = %promotion.run_id(), + sandbox_id = %promotion.sandbox_id(), + profile_name = promotion.profile_name(), + session_id = %promotion.cli_agent_session_id(), + reason, + error = %format_helper_exec_failure("session history sidecar export", &result), + "workspace image cache session history sidecar export failed" + ); + cleanup_guest_session_history_sidecar_export(sandbox, promotion, &export_path, reason) + .await; + return None; + } + let metadata = match serde_json::from_slice::( + result.stdout.as_slice(), + ) { + Ok(metadata) + if metadata.encoded_size > 0 + && metadata.encoded_size <= SESSION_HISTORY_SIDECAR_MAX_BYTES => + { + metadata + } + Ok(_) | Err(_) => { + warn!( + run_id = %promotion.run_id(), + sandbox_id = %promotion.sandbox_id(), + profile_name = promotion.profile_name(), + session_id = %promotion.cli_agent_session_id(), + reason, + "workspace image cache session history sidecar export returned invalid metadata" + ); + cleanup_guest_session_history_sidecar_export(sandbox, promotion, &export_path, reason) + .await; + return None; + } + }; + let tmp_path = promotion.session_history_sidecar_tmp_path(); + let _ = fs::remove_file(&tmp_path).await; + let copied = match sandbox + .copy_file( + &export_path, + &tmp_path, + CopyFileOptions { + max_bytes: SESSION_HISTORY_SIDECAR_MAX_BYTES, + timeout: SESSION_HISTORY_SIDECAR_COPY_TIMEOUT, + missing_ok: false, + }, + ) + .await + { + Ok(result) => result, + Err(e) => { + cleanup_guest_session_history_sidecar_export(sandbox, promotion, &export_path, reason) + .await; + let _ = fs::remove_file(&tmp_path).await; + warn!( + run_id = %promotion.run_id(), + sandbox_id = %promotion.sandbox_id(), + profile_name = promotion.profile_name(), + session_id = %promotion.cli_agent_session_id(), + reason, + error = %e, + "workspace image cache session history sidecar copy failed" + ); + return None; + } + }; + cleanup_guest_session_history_sidecar_export(sandbox, promotion, &export_path, reason).await; + if copied.bytes_copied != metadata.encoded_size { + let _ = fs::remove_file(&tmp_path).await; + warn!( + run_id = %promotion.run_id(), + sandbox_id = %promotion.sandbox_id(), + profile_name = promotion.profile_name(), + session_id = %promotion.cli_agent_session_id(), + reason, + copied_bytes = copied.bytes_copied, + encoded_size = metadata.encoded_size, + "workspace image cache session history sidecar copy size mismatch" + ); + return None; + } + promotion.session_history_sidecar_source( + tmp_path, + metadata.representation, + metadata.encoded_size, + ) +} + +async fn cleanup_guest_session_history_sidecar_export( + sandbox: &dyn Sandbox, + promotion: &WorkspaceImagePromotionContext, + export_path: &str, + reason: &'static str, +) { + let command = ["rm -f --".to_string(), quote_shell_arg(export_path)].join(" "); + let request = ExecRequest { + cmd: &command, + timeout: SESSION_HISTORY_SIDECAR_CLEANUP_TIMEOUT, + env: &[], + sudo: false, + stdin_bytes: None, + output_limits: EXEC_OUTPUT_LIMIT_64_KIB, + }; + match sandbox + .exec_with_diagnostic_label(&request, "session-history-sidecar-cleanup") + .await + { + Ok(result) if helper_exec_succeeded(&result) => {} + Ok(result) => warn!( + run_id = %promotion.run_id(), + sandbox_id = %promotion.sandbox_id(), + profile_name = promotion.profile_name(), + session_id = %promotion.cli_agent_session_id(), + reason, + error = %format_helper_exec_failure("session history sidecar cleanup", &result), + "workspace image cache session history sidecar cleanup failed" + ), + Err(e) => warn!( + run_id = %promotion.run_id(), + sandbox_id = %promotion.sandbox_id(), + profile_name = promotion.profile_name(), + session_id = %promotion.cli_agent_session_id(), + reason, + error = %e, + "workspace image cache session history sidecar cleanup errored" + ), + } +} + pub(crate) async fn promote_workspace_image_from_parked_sandbox( sandbox: &mut dyn Sandbox, promotion: Option, diff --git a/crates/runner/src/workspace_promotion/test_support.rs b/crates/runner/src/workspace_promotion/test_support.rs index 9ef3efa0e37..64c7fa876ed 100644 --- a/crates/runner/src/workspace_promotion/test_support.rs +++ b/crates/runner/src/workspace_promotion/test_support.rs @@ -3,6 +3,7 @@ use sandbox::SandboxId; use crate::ids::RunId; use crate::paths::RunnerPaths; +use crate::restored_session_identity::RestoredSessionIdentity; use crate::storage_fingerprints::StorageFingerprints; use crate::workspace_image_cache::{ SessionWorkspaceCache, WorkspaceCacheCheckoutResult, WorkspaceCacheTerminalStatus, @@ -12,6 +13,7 @@ use crate::workspace_image_cache::{ pub(crate) const TEST_COMPLETED_AT: &str = "2026-06-03T00:00:00.000Z"; const TEST_WORKSPACE_IMAGE: &[u8] = b"workspace image"; +pub(crate) const TEST_WORKSPACE_IMAGE_SIZE_BYTES: u64 = TEST_WORKSPACE_IMAGE.len() as u64; pub(crate) struct WorkspacePromotionFixture { pub(crate) _dir: tempfile::TempDir, @@ -23,6 +25,13 @@ pub(crate) struct WorkspacePromotionFixture { impl WorkspacePromotionFixture { pub(crate) async fn new(session_id: &str) -> Self { + Self::new_with_restored_session_identity(session_id, None).await + } + + pub(crate) async fn new_with_restored_session_identity( + session_id: &str, + restored_session_identity: Option<&RestoredSessionIdentity>, + ) -> Self { let dir = tempfile::tempdir().unwrap(); let paths = RunnerPaths::new(dir.path().join("runner")); tokio::fs::create_dir_all(paths.base_dir()).await.unwrap(); @@ -37,7 +46,7 @@ impl WorkspacePromotionFixture { profile_name: "vm0/default", cli_agent_session_id: Some(session_id), working_dir: CANONICAL_WORKING_DIR, - image_size_bytes: TEST_WORKSPACE_IMAGE.len() as u64, + image_size_bytes: TEST_WORKSPACE_IMAGE_SIZE_BYTES, }, workspace_drive_required: true, }) @@ -56,6 +65,7 @@ impl WorkspacePromotionFixture { run_id, sandbox_id, cli_agent_session_id_override: Some(session_id), + restored_session_identity, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: TEST_COMPLETED_AT.into(), storage_fingerprints: StorageFingerprints::default(), @@ -96,7 +106,7 @@ impl WorkspacePromotionFixture { profile_name: "vm0/default", cli_agent_session_id: Some(&session_id), working_dir: CANONICAL_WORKING_DIR, - image_size_bytes: TEST_WORKSPACE_IMAGE.len() as u64, + image_size_bytes: TEST_WORKSPACE_IMAGE_SIZE_BYTES, }, workspace_drive_required: true, }) @@ -107,6 +117,7 @@ impl WorkspacePromotionFixture { run_id, sandbox_id, cli_agent_session_id_override: Some(&session_id), + restored_session_identity: None, terminal_status: WorkspaceCacheTerminalStatus::Success, completed_at: "2026-06-04T00:00:00.000Z".into(), storage_fingerprints: StorageFingerprints::default(), @@ -134,7 +145,7 @@ impl WorkspacePromotionFixture { profile_name: "vm0/default", cli_agent_session_id: Some(session_id), working_dir: CANONICAL_WORKING_DIR, - image_size_bytes: TEST_WORKSPACE_IMAGE.len() as u64, + image_size_bytes: TEST_WORKSPACE_IMAGE_SIZE_BYTES, }, workspace_drive_required: true, }) diff --git a/crates/runner/src/workspace_promotion/tests.rs b/crates/runner/src/workspace_promotion/tests.rs index e3401b38046..a0128e3ea36 100644 --- a/crates/runner/src/workspace_promotion/tests.rs +++ b/crates/runner/src/workspace_promotion/tests.rs @@ -1,19 +1,29 @@ -use super::test_support::WorkspacePromotionFixture; +use super::test_support::{TEST_WORKSPACE_IMAGE_SIZE_BYTES, WorkspacePromotionFixture}; use super::*; use std::sync::Arc; use std::time::Duration; +use api_contracts::generated::constants::runners::paths::CANONICAL_WORKING_DIR; use async_trait::async_trait; +use guest_contracts::session_history_identity::{ + FinalSessionHistoryFramework, FinalSessionHistoryIdentity, FinalSessionHistoryRefKind, + SessionHistorySidecarExportMetadata, SessionHistorySidecarRepresentation, +}; use sandbox::{ CopyFileOptions, CopyFileResult, ExecRequest, ExecResult, GuestProcessHandle, ProcessExit, Sandbox, SandboxFactory, SandboxId, StartProcessRequest, }; use sandbox_mock::{ExecMatcher, MockSandboxFactory, MockSandboxOverrides}; +use sha2::{Digest, Sha256}; use tracing_subscriber::prelude::*; use tracing_test_support::{CapturedEvent, CapturedEvents}; -use crate::workspace_image_cache::WorkspaceCacheCheckoutResult; +use crate::ids::RunId; +use crate::restored_session_identity::RestoredSessionIdentity; +use crate::workspace_image_cache::{ + WorkspaceCacheCheckoutResult, WorkspaceImageLeaseIdentity, WorkspaceImagePrepareRequest, +}; async fn mock_sandbox_with_overrides( sandbox_id: SandboxId, @@ -139,6 +149,92 @@ async fn parked_workspace_promotion_unparks_unmounts_and_promotes_cache_entry() assert_eq!(states[0].session_id, fixture.session_id); } +#[tokio::test] +async fn active_workspace_promotion_exports_session_history_sidecar() { + let session_id = "sess-active-sidecar-promote"; + let history = br#"{"type":"message","content":"cached"}"#; + let metadata = FinalSessionHistoryIdentity::new( + FinalSessionHistoryFramework::ClaudeCode, + hex::encode(Sha256::digest(session_id.as_bytes())), + FinalSessionHistoryRefKind::Blob, + hex::encode(Sha256::digest(history)), + history.len() as u64, + format!("/home/user/.claude/projects/-home-user-workspace/{session_id}.jsonl"), + ) + .unwrap(); + let restored_identity = RestoredSessionIdentity::from_final_metadata( + metadata, + "/home/user/.vm0/guest-agent/runs/run-1/final-session-history-identity.json", + "/home/user/.vm0/guest-agent/runs/run-1", + ) + .unwrap(); + let fixture = WorkspacePromotionFixture::new_with_restored_session_identity( + session_id, + Some(&restored_identity), + ) + .await; + assert!(fixture.promotion.restored_session_identity().is_some()); + let sandbox = sandbox_mock::MockSandbox::new(fixture.sandbox_id.to_string()); + let export_metadata = SessionHistorySidecarExportMetadata { + representation: SessionHistorySidecarRepresentation::Raw, + encoded_size: history.len() as u64, + }; + sandbox.push_exec_result(Ok(ExecResult::new( + 0, + serde_json::to_vec(&export_metadata).unwrap(), + Vec::new(), + ))); + sandbox.push_copy_file_result(Ok(history.to_vec())); + + let (promoted, events) = capture_promotion_events(promote_workspace_image_from_active_sandbox( + &sandbox, + Some(fixture.promotion), + "test", + )) + .await; + + assert!(promoted); + let exec_calls = sandbox.exec_calls(); + assert_eq!(exec_calls.len(), 3); + assert!(exec_calls[0].cmd.contains("export-session-history-sidecar")); + assert!(exec_calls[1].cmd.contains("rm -f --")); + assert!(exec_calls[1].cmd.contains("/session-history-sidecar")); + assert!(exec_calls[2].sudo); + let copy_calls = sandbox.copy_file_calls(); + assert_eq!(copy_calls.len(), 1); + assert!(copy_calls[0].path.ends_with("/session-history-sidecar")); + let sidecar_entry_dir = copy_calls[0].host_path.parent().unwrap(); + let sidecar_metadata_path = sidecar_entry_dir.join("session-history.metadata.json"); + if !sidecar_metadata_path.is_file() { + let entries = std::fs::read_dir(sidecar_entry_dir) + .unwrap() + .map(|entry| entry.unwrap().file_name()) + .collect::>(); + panic!("missing sidecar metadata; entries={entries:?}; events={events:#?}"); + } + + let lease = fixture + .cache + .prepare(WorkspaceImagePrepareRequest { + identity: WorkspaceImageLeaseIdentity { + run_id: RunId::new_v4(), + sandbox_id: SandboxId::new_v4(), + profile_name: "vm0/default", + cli_agent_session_id: Some(session_id), + working_dir: CANONICAL_WORKING_DIR, + image_size_bytes: TEST_WORKSPACE_IMAGE_SIZE_BYTES, + }, + workspace_drive_required: true, + }) + .await; + assert_eq!(lease.result(), WorkspaceCacheCheckoutResult::Hit); + let sidecar = lease + .probe_session_history_sidecar(&restored_identity) + .await + .unwrap(); + assert_eq!(tokio::fs::read(sidecar.path).await.unwrap(), history); +} + #[tokio::test] async fn parked_workspace_promotion_unpark_error_skips_cache() { let fixture = WorkspacePromotionFixture::new("sess-parked-unpark-error").await;