1use std::collections::BTreeSet;
4use std::ffi::OsStr;
5use std::path::{Path, PathBuf};
6use std::time::{Duration, Instant};
7
8use anyhow::{Context, Result, bail, ensure};
9
10use crate::checkpoint_transfer::{
11 CheckpointTransfer, capture_stdin_command, export_command, export_stdin_command,
12 pack_stdin_command,
13};
14use crate::session_manager::{
15 ManagedSessionHandle, ManagedSessionLease, SessionManagerControl, StandaloneSession,
16 new_command_id, worker_connect_needs_restart,
17};
18use crate::worker_client::RelayRejected;
19use mj_checkpoint::archive::{
20 BundleManifest, CanonicalSessionSnapshot, SessionManifest, TargetManifest,
21 verify_archive_streaming,
22};
23use mj_checkpoint::checkpoint::{
24 CHECKPOINT_EXPORT_PROTOCOL_VERSION, CHECKPOINT_STAGING_PROTOCOL_VERSION, CapturedCheckpoint,
25 CheckpointCaptureSpec, CheckpointExportSpec, CheckpointPackSpec, CheckpointRepositoryCapture,
26 CheckpointRepositorySpec, canonical_session_contains_prompt, checkpoint_sha256,
27};
28use mj_core::config::{HarnessKind, sessions_dir};
29use mj_core::state::{
30 CheckpointMetadata, ManagedSessionSnapshot, SessionRecord, SessionState, State,
31};
32use mj_transcript::projection::canonical_session_from_materialized;
33
34use crate::targets::{
35 self, CommandExecutor, CommandOutput, CommandSpec, ProcessExecutor, ProvisionStage,
36 ProvisionStageGuard,
37};
38use mj_core::relay::{RelayCommand, RelayCursor, RelayExecutionState};
39
40use super::backend::backend_locator;
41use super::readiness::wait_for_native_session_in_stage;
42use super::worker_restart::{InstalledWorkerRestart, RESTART_FOR_CHECKPOINT};
43use super::{
44 Controller, execute_checked, now, persist_session_record_transition_or_restore,
45 scp_command_spec, ssh_command_spec, target_kind, target_profile_home,
46};
47
48#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct SessionExportLayout {
55 pub backend: targets::TargetLocator,
57 pub workspace_root: String,
59 pub primary_repository: String,
62 pub repositories: Vec<CheckpointRepositorySpec>,
63 pub managed_worktree: Option<mj_core::state::ManagedWorktree>,
67}
68
69const CHECKPOINT_BARRIER_TIMEOUT: Duration = Duration::from_secs(30);
74const CHECKPOINT_CANCEL_TIMEOUT: Duration = Duration::from_secs(30);
78const CHECKPOINT_BARRIER_TIMEOUT_AFTER_RESTART: Duration = Duration::from_secs(300);
81
82pub fn reconcile_managed_checkpoint_archives() -> Result<usize> {
86 let mut state = crate::database::load_state_migrating()?;
87 for operation in crate::database::load_move_operations()? {
90 if operation.retains_checkpoint()
91 && let Some(checkpoint) = operation.checkpoint
92 && let Some(mut session) = state.sessions.get(&operation.selection.session_id).cloned()
93 {
94 session.checkpoint = Some(checkpoint);
95 state
96 .sessions
97 .insert(format!("move:{}", operation.operation_id), session);
98 }
99 }
100 reconcile_managed_checkpoint_archives_in(&sessions_dir(), &state)
101}
102
103fn reconcile_managed_checkpoint_archives_in(directory: &Path, state: &State) -> Result<usize> {
104 if !directory.exists() {
105 return Ok(0);
106 }
107 let referenced_names = state
108 .sessions
109 .values()
110 .filter_map(|session| session.checkpoint.as_ref())
111 .filter_map(|checkpoint| checkpoint.archive_path.file_name())
112 .map(ToOwned::to_owned)
113 .collect::<BTreeSet<_>>();
114 let mut removed = 0;
115 for entry in std::fs::read_dir(directory)
116 .with_context(|| format!("scan checkpoint directory {}", directory.display()))?
117 {
118 let entry = entry?;
119 let file_type = entry.file_type()?;
120 if !file_type.is_file()
121 || !is_managed_checkpoint_archive_name(&entry.file_name())
122 || referenced_names.contains(&entry.file_name())
123 {
124 continue;
125 }
126 std::fs::remove_file(entry.path()).with_context(|| {
127 format!(
128 "remove unreferenced managed checkpoint {}",
129 entry.path().display()
130 )
131 })?;
132 removed += 1;
133 }
134 Ok(removed)
135}
136
137fn is_managed_checkpoint_archive_name(name: &OsStr) -> bool {
138 let Some(stem) = name.to_str().and_then(|name| name.strip_suffix(".hel.zip")) else {
139 return false;
140 };
141 let Some((frontier_prefix, nonce)) = stem.rsplit_once("-archive-") else {
142 return false;
143 };
144 if nonce.len() != 32
145 || !nonce
146 .bytes()
147 .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase())
148 {
149 return false;
150 }
151 let Some((session_id, frontier)) = frontier_prefix.rsplit_once('-') else {
152 return false;
153 };
154 !session_id.is_empty()
155 && frontier.parse::<u64>().is_ok()
156 && mj_core::config::validate_id("session", session_id).is_ok()
157}
158
159#[derive(Debug, Clone)]
160pub struct CheckpointArtifact {
161 pub metadata: CheckpointMetadata,
162 pub native_session_id: String,
163 pub event_frontier_digest: String,
165}
166
167pub(super) enum ControllerRelayLease {
175 Managed {
176 handle: ManagedSessionHandle,
177 lease: Option<ManagedSessionLease>,
178 },
179 Standalone(StandaloneSession),
180}
181
182impl ControllerRelayLease {
183 pub(super) fn connection_mut(&mut self) -> &mut StandaloneSession {
186 match self {
187 Self::Managed { lease, .. } => lease
188 .as_mut()
189 .expect("checkpoint latch has already returned its connection")
190 .connection_mut(),
191 Self::Standalone(connection) => connection,
192 }
193 }
194
195 async fn submit(&mut self, command_id: String, command: RelayCommand) -> Result<u64> {
196 match self {
197 Self::Managed {
198 lease: Some(lease), ..
199 } => lease.connection_mut().submit(command_id, command).await,
200 Self::Managed { handle, .. } => handle.submit(command_id, command).await,
201 Self::Standalone(connection) => connection.submit(command_id, command).await,
202 }
203 }
204
205 async fn sync_snapshot(&mut self) -> Result<ManagedSessionSnapshot> {
206 match self {
207 Self::Managed {
208 lease: Some(lease), ..
209 } => lease.connection_mut().sync().await,
210 Self::Managed { handle, .. } => {
211 handle.sync_now().await?;
212 handle
213 .view()
214 .snapshot
215 .context("managed session has no snapshot")
216 }
217 Self::Standalone(connection) => connection.sync().await,
218 }
219 }
220
221 fn replace_connection(&mut self, connection: StandaloneSession) {
223 match self {
224 Self::Managed {
225 lease: Some(lease), ..
226 } => lease.replace_connection(connection),
227 Self::Standalone(existing) => *existing = connection,
228 Self::Managed { lease: None, .. } => {
229 *self = Self::Standalone(connection);
230 }
231 }
232 }
233
234 fn end_latch(&mut self) {
238 if let Self::Managed { lease, .. } = self
239 && let Some(lease) = lease.take()
240 {
241 lease.release();
242 }
243 }
244
245 async fn cancel_abandoned_barrier(&mut self) -> Result<()> {
253 let Self::Managed { handle, lease } = self else {
254 return Ok(());
257 };
258 match lease.take() {
259 Some(lease) => drop(lease),
260 None => drop(handle.lease_connection().await?),
261 }
262 Ok(())
263 }
264
265 pub(super) fn release(self) {
266 if let Self::Managed {
267 lease: Some(lease), ..
268 } = self
269 {
270 lease.release();
271 }
272 }
273}
274
275#[derive(Debug, Clone, Copy, PartialEq, Eq)]
277pub(super) enum LatchExclusivity {
278 ReleaseAfterLatch,
283 HoldThroughClose,
286}
287
288#[derive(Debug, Clone, Copy, PartialEq, Eq)]
290pub(super) enum CheckpointExportPolicy {
291 Always,
293 ReuseUnchangedArchive,
297}
298
299#[derive(Debug, Clone, Copy, PartialEq, Eq)]
301pub(super) enum CheckpointCompletion {
302 HeldBarrier,
306 ReleasedAfterCapture,
309}
310
311pub(super) struct LatchedCheckpoint {
312 pub(super) artifact: CheckpointArtifact,
313 pub(super) relay: ControllerRelayLease,
314 pub(super) barrier_command_id: String,
315 pub(super) cursor: RelayCursor,
316 pub(super) completion: CheckpointCompletion,
317}
318
319impl LatchedCheckpoint {
326 async fn complete(mut self) -> Result<()> {
328 let (prefix, command) = match self.completion {
329 CheckpointCompletion::HeldBarrier => (
330 "checkpoint-complete",
331 RelayCommand::CompleteCheckpoint {
332 barrier_command_id: self.barrier_command_id.clone(),
333 },
334 ),
335 CheckpointCompletion::ReleasedAfterCapture => (
338 "checkpoint-floor",
339 RelayCommand::AdvanceRecoveryFloor {
340 through: self.cursor.clone(),
341 },
342 ),
343 };
344 let command_id = new_command_id(prefix)?;
345 self.relay.submit(command_id, command).await.map(|_| ())
346 }
347
348 async fn abandon(mut self, session_id: &str) {
354 if self.completion == CheckpointCompletion::ReleasedAfterCapture {
355 return;
359 }
360 if let Err(error) = self.relay.cancel_abandoned_barrier().await {
361 tracing::warn!(
362 session_id,
363 "abandoned checkpoint could not cancel its relay barrier: {error:#}"
364 );
365 }
366 }
367}
368
369impl Controller {
370 pub(super) fn persist_checkpoint_transition_or_restore(
371 &mut self,
372 session_id: &str,
373 previous: &SessionRecord,
374 context: &'static str,
375 ) -> Result<()> {
376 persist_session_record_transition_or_restore(
377 &mut self.state,
378 session_id,
379 previous,
380 context,
381 &crate::database::save_checkpointed_session,
382 )
383 }
384
385 pub(super) fn persist_failed_checkpoint_state_or_restore(
386 &mut self,
387 session_id: &str,
388 previous: &SessionRecord,
389 primary: anyhow::Error,
390 ) -> anyhow::Error {
391 match self.persist_session_state(session_id) {
392 Ok(()) => primary,
393 Err(error) => self.restore_prior_session_after_persistence_failure(
394 session_id,
395 previous,
396 primary.context(format!(
397 "failed to persist the checkpoint rollback state: {error:#}"
398 )),
399 ),
400 }
401 }
402
403 pub async fn checkpoint_session(&mut self, session_id: &str) -> Result<CheckpointMetadata> {
407 self.checkpoint_session_controlled(session_id, &ProcessExecutor)
408 .await
409 }
410
411 pub async fn checkpoint_session_controlled(
412 &mut self,
413 session_id: &str,
414 executor: &(impl CommandExecutor + Sync),
415 ) -> Result<CheckpointMetadata> {
416 self.checkpoint_session_controlled_with_manager(session_id, executor, None)
417 .await
418 }
419
420 async fn checkpoint_session_controlled_with_manager(
421 &mut self,
422 session_id: &str,
423 executor: &(impl CommandExecutor + Sync),
424 manager: Option<&SessionManagerControl>,
425 ) -> Result<CheckpointMetadata> {
426 let previous = self
427 .state
428 .sessions
429 .get(session_id)
430 .with_context(|| format!("unknown session {session_id}"))?
431 .clone();
432 previous.validate_configuration(&self.config)?;
433 ensure!(
434 !matches!(
435 previous.state,
436 SessionState::Closing | SessionState::Destroying
437 ),
438 "session {session_id} is already closing; resume that close instead of starting an ordinary checkpoint"
439 );
440 let record = self.state.sessions.get_mut(session_id).unwrap();
441 record.state = SessionState::Checkpointing;
442 record.updated_at = now();
443 record.last_checkpoint_error = None;
444 self.persist_session_transition_or_restore(
445 session_id,
446 &previous,
447 "persist checkpointing state before creating a checkpoint",
448 )?;
449
450 match self
451 .checkpoint_session_latched(
452 session_id,
453 executor,
454 manager,
455 LatchExclusivity::ReleaseAfterLatch,
456 CheckpointExportPolicy::Always,
457 )
458 .await
459 {
460 Ok(latched) => {
461 let artifact = latched.artifact.clone();
462 if let Err(error) = mj_core::test_hooks::reach_test_hook(
463 "checkpoint_archive_before_database_publication",
464 ) {
465 latched.abandon(session_id).await;
466 return Err(remove_uninstalled_checkpoint(
467 &artifact.metadata.archive_path,
468 error,
469 ));
470 }
471 {
472 let record = self.state.sessions.get_mut(session_id).unwrap();
473 record.state = SessionState::Running;
474 record.native_session_id = Some(artifact.native_session_id.clone());
475 record.checkpoint = Some(artifact.metadata.clone());
476 record.updated_at = now();
477 record.last_error = None;
478 record.last_checkpoint_error = None;
479 }
480 let persist_started = Instant::now();
481 if let Err(error) = self.persist_checkpoint_transition_or_restore(
482 session_id,
483 &previous,
484 "persist verified checkpoint before releasing relay history",
485 ) {
486 latched.abandon(session_id).await;
487 return Err(error);
488 }
489 tracing::info!(
490 session_id,
491 persist_ms = persist_started.elapsed().as_millis() as u64,
492 "checkpoint metadata persisted"
493 );
494 prune_replaced_checkpoint(previous.checkpoint.as_ref(), &artifact.metadata);
495 release_projection_behind_checkpoint(session_id, &artifact.metadata);
496 if let Err(error) = latched.complete().await {
497 tracing::warn!(
503 session_id,
504 "verified checkpoint was saved, but the relay could not be told to release the history it covers: {error:#}"
505 );
506 }
507 Ok(artifact.metadata)
508 }
509 Err(error) => {
510 let deferred = checkpoint_was_deferred(&error);
515 if let Some(record) = self.state.sessions.get_mut(session_id) {
516 record.state = if previous.state == SessionState::Checkpointing {
517 SessionState::Running
518 } else {
519 previous.state
520 };
521 record.updated_at = now();
522 if !deferred {
523 record.last_checkpoint_error = Some(format!("{error:#}"));
524 }
525 }
526 Err(self.persist_failed_checkpoint_state_or_restore(session_id, &previous, error))
527 }
528 }
529 }
530
531 pub async fn create_recovery_checkpoint_managed_controlled(
534 &self,
535 session_id: &str,
536 manager: &SessionManagerControl,
537 executor: &(impl CommandExecutor + Sync),
538 ) -> Result<CheckpointArtifact> {
539 self.create_recovery_checkpoint_with_manager(session_id, Some(manager), executor)
540 .await
541 }
542
543 async fn create_recovery_checkpoint_with_manager(
544 &self,
545 session_id: &str,
546 manager: Option<&SessionManagerControl>,
547 executor: &(impl CommandExecutor + Sync),
548 ) -> Result<CheckpointArtifact> {
549 let previous_checkpoint = self
550 .state
551 .sessions
552 .get(session_id)
553 .with_context(|| format!("unknown session {session_id}"))?
554 .checkpoint
555 .clone();
556 let latched = self
557 .checkpoint_session_latched_with_recovery_stage(
558 session_id,
559 executor,
560 manager,
561 LatchExclusivity::ReleaseAfterLatch,
562 CheckpointExportPolicy::Always,
563 true,
564 )
565 .await?;
566 let artifact = latched.artifact.clone();
567 let verification = {
568 let _verifying = ProvisionStageGuard::new(executor, ProvisionStage::Verifying);
569 verify_checkpoint_artifact(session_id, &artifact)
570 };
571 if let Err(error) = verification {
572 latched.abandon(session_id).await;
573 return Err(remove_uninstalled_checkpoint(
574 &artifact.metadata.archive_path,
575 error.context("final recovery checkpoint verification"),
576 ));
577 }
578 if let Err(error) =
579 mj_core::test_hooks::reach_test_hook("checkpoint_archive_before_database_publication")
580 {
581 latched.abandon(session_id).await;
582 return Err(remove_uninstalled_checkpoint(
583 &artifact.metadata.archive_path,
584 error,
585 ));
586 }
587 let persist_started = Instant::now();
588 if let Err(error) = crate::database::record_recovery_success(
589 session_id,
590 &artifact.native_session_id,
591 &artifact.metadata,
592 ) {
593 latched.abandon(session_id).await;
594 return Err(error
595 .context("persist verified recovery checkpoint before releasing relay history"));
596 }
597 tracing::info!(
598 session_id,
599 persist_ms = persist_started.elapsed().as_millis() as u64,
600 "recovery checkpoint metadata persisted"
601 );
602 if let Err(error) = latched.complete().await {
603 tracing::warn!(
608 session_id,
609 "recovery checkpoint was saved, but the relay could not be told to release the history it covers: {error:#}"
610 );
611 }
612 prune_replaced_checkpoint(previous_checkpoint.as_ref(), &artifact.metadata);
613 release_projection_behind_checkpoint(session_id, &artifact.metadata);
614 Ok(artifact)
615 }
616
617 pub fn session_export_layout(
625 &self,
626 session_id: &str,
627 executor: &(impl CommandExecutor + Sync),
628 ) -> Result<SessionExportLayout> {
629 let session = self
630 .state
631 .sessions
632 .get(session_id)
633 .with_context(|| format!("unknown session {session_id}"))?
634 .clone();
635 let locator = session
636 .target
637 .as_ref()
638 .context("session has no live target")?;
639 let backend = backend_locator(locator, &session, &self.config)?;
640 let (workspace_root, primary_repository, repositories) =
641 if let Some(project_directory) = &session.project_directory {
642 let parent = project_directory
643 .parent()
644 .context("bare project directory has no parent")?;
645 let destination = project_directory
646 .file_name()
647 .context("bare project directory cannot be the filesystem root")?;
648 (
649 parent.to_string_lossy().into_owned(),
650 "project".to_owned(),
651 vec![CheckpointRepositorySpec {
652 id: "project".into(),
653 relative_destination: PathBuf::from(destination),
654 capture: if session.managed_worktree.is_some() {
660 CheckpointRepositoryCapture::DeltaFrom {
661 base_commit: super::worktree::raw_checkout_position(
662 &session,
663 &self.config,
664 project_directory,
665 executor,
666 )?
667 .head_commit,
668 }
669 } else {
670 CheckpointRepositoryCapture::MetadataOnly
671 },
672 origin_override: None,
673 }],
674 )
675 } else {
676 let bundle = self
677 .config
678 .bundles
679 .get(&session.bundle_id)
680 .context("session bundle is missing")?;
681 let workspace_root = match &backend {
682 targets::TargetLocator::LocalPodman { .. }
683 | targets::TargetLocator::LocalDocker { .. }
684 | targets::TargetLocator::AppleContainer { .. }
685 | targets::TargetLocator::SshPodman { .. }
686 | targets::TargetLocator::SshDocker { .. } => "/workspace".to_string(),
687 targets::TargetLocator::AwsEc2 { workspace, .. }
688 | targets::TargetLocator::SshBare { workspace, .. } => workspace.clone(),
689 targets::TargetLocator::LocalBare { worker_root } => worker_root.clone(),
690 };
691 let repositories = bundle
692 .repositories
693 .iter()
694 .map(|repository| CheckpointRepositorySpec {
695 id: repository.id.clone(),
696 relative_destination: repository.destination.clone(),
697 capture: CheckpointRepositoryCapture::RemoteWorkspace,
698 origin_override: None,
699 })
700 .collect();
701 (workspace_root, bundle.primary_repo.clone(), repositories)
702 };
703 Ok(SessionExportLayout {
704 backend,
705 workspace_root,
706 primary_repository,
707 repositories,
708 managed_worktree: session.managed_worktree,
709 })
710 }
711
712 pub(super) async fn checkpoint_session_latched(
713 &self,
714 session_id: &str,
715 executor: &(impl CommandExecutor + Sync),
716 manager: Option<&SessionManagerControl>,
717 exclusivity: LatchExclusivity,
718 export_policy: CheckpointExportPolicy,
719 ) -> Result<LatchedCheckpoint> {
720 self.checkpoint_session_latched_with_recovery_stage(
721 session_id,
722 executor,
723 manager,
724 exclusivity,
725 export_policy,
726 exclusivity == LatchExclusivity::HoldThroughClose,
727 )
728 .await
729 }
730
731 async fn checkpoint_session_latched_with_recovery_stage(
732 &self,
733 session_id: &str,
734 executor: &(impl CommandExecutor + Sync),
735 manager: Option<&SessionManagerControl>,
736 exclusivity: LatchExclusivity,
737 export_policy: CheckpointExportPolicy,
738 recovery_copy: bool,
739 ) -> Result<LatchedCheckpoint> {
740 if let Some(operation) = crate::database::load_move_operation(session_id)?
741 && operation.queue_admission_started
742 && !operation.queue_admission_finished
743 {
744 bail!(
747 "move queue admission is incomplete; retry Move before checkpointing this destination"
748 );
749 }
750 let session = self
751 .state
752 .sessions
753 .get(session_id)
754 .with_context(|| format!("unknown session {session_id}"))?
755 .clone();
756 session.validate_configuration(&self.config)?;
757 let layout = self.session_export_layout(session_id, executor)?;
758 let backend = layout.backend.clone();
759 let profile = self
760 .config
761 .profiles
762 .get(&session.last_profile)
763 .context("session profile is missing")?;
764 let reconnect = targets::reconnect_plan(&backend, session_id)?
765 .commands
766 .into_iter()
767 .next()
768 .context("reconnect plan is empty")?;
769 let worker_root = targets::worker_root(&backend, session_id)?;
770 let harness_home = target_profile_home(&backend, session_id, profile);
771 let SessionExportLayout {
772 workspace_root,
773 primary_repository,
774 repositories,
775 ..
776 } = layout;
777 let target_path = |path: &str| match &backend {
778 targets::TargetLocator::AwsEc2 { .. } | targets::TargetLocator::SshBare { .. }
779 if !path.starts_with('/') =>
780 {
781 PathBuf::from(format!("~/{path}"))
782 }
783 _ => PathBuf::from(path),
784 };
785 let operation_id = new_command_id("checkpoint")?;
788 let remote_spec = format!("{worker_root}/{operation_id}-spec.json");
789 let remote_archive = format!("{worker_root}/{operation_id}.hel.zip");
790 let remote_stage = format!("{worker_root}/{operation_id}-stage");
791 let checkpointed_at = now();
792 let target_manifest = TargetManifest {
793 template_id: session.target_template_id.clone(),
794 target_kind: target_kind(&backend).into(),
795 details: Default::default(),
796 };
797 let bundle_manifest = BundleManifest {
798 id: session.bundle_id.clone(),
799 primary_repository,
800 };
801 let session_manifest = |native_session_id: &str| SessionManifest {
802 id: session.id.clone(),
803 title: session.title.clone(),
804 harness_kind: session.harness_kind,
805 profile_id: session.last_profile.clone(),
806 native_session_id: native_session_id.to_owned(),
807 created_at: session.created_at.clone(),
808 checkpointed_at: checkpointed_at.clone(),
809 hel_version: env!("CARGO_PKG_VERSION").into(),
810 relay_version: env!("CARGO_PKG_VERSION").into(),
811 adapter_version: "acp-v1".into(),
812 };
813 let releases_after_capture = exclusivity == LatchExclusivity::ReleaseAfterLatch;
814 if releases_after_capture
815 && let Some(native_session_id) = session.native_session_id.as_deref()
816 {
817 let prestage = CheckpointCaptureSpec {
818 protocol_version: CHECKPOINT_STAGING_PROTOCOL_VERSION,
819 session: session_manifest(native_session_id),
820 target: target_manifest.clone(),
821 bundle: bundle_manifest.clone(),
822 relay_root: target_path(&worker_root),
823 harness_home: target_path(&harness_home),
824 workspace_root: target_path(&workspace_root),
825 repositories: repositories.clone(),
826 allow_empty_native: false,
827 stage_path: target_path(&remote_stage),
828 refresh_existing: false,
829 };
830 let prestage_started = Instant::now();
831 let prestaged = {
832 let _recovery_copy = recovery_copy
833 .then(|| ProvisionStageGuard::new(executor, ProvisionStage::RecoveryCopy));
834 run_checkpoint_staging_command(
835 executor,
836 &backend,
837 session_id,
838 &prestage,
839 capture_stdin_command,
840 "prestage target checkpoint",
841 )
842 };
843 match prestaged {
844 Ok(output) => match serde_json::from_slice::<CapturedCheckpoint>(&output.stdout) {
845 Ok(captured) => tracing::info!(
846 session_id,
847 prestage_ms = prestage_started.elapsed().as_millis() as u64,
848 native_bytes = captured.native_bytes,
849 repository_bytes = captured.repository_bytes,
850 reused_native = captured.reused_native,
851 "checkpoint target state prestaged while ACP dispatch remained active"
852 ),
853 Err(error) => tracing::warn!(
854 session_id,
855 error = format!("{error:#}"),
856 "checkpoint prestage returned an invalid result; barrier capture will replace it"
857 ),
858 },
859 Err(error) => {
860 if executor.cancellation_requested() {
861 return Err(error.context("checkpoint prestage was cancelled"));
862 }
863 tracing::warn!(
864 session_id,
865 error = format!("{error:#}"),
866 "checkpoint prestage failed; barrier capture will collect a fresh generation"
867 );
868 }
869 }
870 }
871 let (mut relay, mut restarted_worker) = self
872 .open_checkpoint_relay(
873 session_id,
874 executor,
875 manager,
876 InstalledWorkerRestart {
877 backend: &backend,
878 worker_root: &worker_root,
879 reconnect: &reconnect,
880 launch: None,
881 messages: &RESTART_FOR_CHECKPOINT,
882 },
883 exclusivity == LatchExclusivity::HoldThroughClose
884 || session.harness_kind != HarnessKind::Kimi,
885 )
886 .await?;
887 let (barrier, barrier_command_id) = loop {
888 let checkpoint_only = relay
892 .connection_mut()
893 .sync()
894 .await?
895 .operational
896 .checkpoint_only;
897 if !checkpoint_only {
898 wait_for_native_session_in_stage(
899 relay.connection_mut(),
900 executor,
901 targets::ProvisionStage::Starting,
902 )
903 .await?;
904 }
905 if exclusivity == LatchExclusivity::HoldThroughClose {
906 let snapshot = relay.connection_mut().sync().await?;
907 if !checkpoint_only && snapshot.operational.capacity_retry.is_some() {
908 relay
911 .connection_mut()
912 .submit(
913 new_command_id("cancel-capacity-retry")?,
914 RelayCommand::CancelTurn,
915 )
916 .await?;
917 }
918 }
919 if exclusivity == LatchExclusivity::ReleaseAfterLatch {
920 let snapshot = relay.connection_mut().sync().await?;
921 if snapshot.operational.execution == RelayExecutionState::Running {
922 relay.release();
925 return Err(CheckpointDeferred::harness_busy().into());
926 }
927 if !snapshot
928 .operational
929 .safe_for_checkpoint(session.harness_kind)
930 {
931 relay.release();
936 return Err(CheckpointDeferred::background_snapshot(
937 &snapshot.operational,
938 session.harness_kind,
939 )
940 .into());
941 }
942 }
943 let barrier_command_id = new_command_id("checkpoint")?;
944 let timeout = if restarted_worker {
945 CHECKPOINT_BARRIER_TIMEOUT_AFTER_RESTART
946 } else {
947 CHECKPOINT_BARRIER_TIMEOUT
948 };
949 let result = {
950 let connection = relay.connection_mut();
951 connection
952 .submit(
953 barrier_command_id.clone(),
954 RelayCommand::BeginCheckpoint {
955 reason: Some("controller archive checkpoint".into()),
956 },
957 )
958 .await?;
959 wait_for_checkpoint_barrier(
960 connection,
961 session_id,
962 &barrier_command_id,
963 timeout,
964 BarrierBusyPolicy::of(exclusivity),
965 session.harness_kind,
966 )
967 .await
968 };
969 match result {
970 Ok(barrier) => break (barrier, barrier_command_id),
971 Err(error)
972 if !restarted_worker && checkpoint_barrier_needs_worker_restart(&error) =>
973 {
974 if exclusivity == LatchExclusivity::ReleaseAfterLatch
975 && matches!(session.harness_kind, HarnessKind::Kimi | HarnessKind::Codex)
976 {
977 let safe_to_restart =
978 relay.connection_mut().sync().await.is_ok_and(|snapshot| {
979 snapshot.operational.safe_to_replace(session.harness_kind)
980 });
981 if !safe_to_restart {
982 return Err(error.context(CheckpointDeferred::background_work()));
983 }
984 }
985 tracing::warn!(
986 session_id,
987 "checkpoint requires a worker restart; restarting and retrying: {error:#}"
988 );
989 let connection = self
990 .restart_worker_for_checkpoint(
991 session_id,
992 executor,
993 &backend,
994 &worker_root,
995 &reconnect,
996 )
997 .await?;
998 relay.replace_connection(connection);
999 restarted_worker = true;
1000 }
1001 Err(error) => return Err(error),
1002 }
1003 };
1004 let barrier_ready_at = Instant::now();
1005 relay
1009 .connection_mut()
1010 .sync_project_memory()
1011 .await
1012 .context("synchronize project memory for checkpoint")?;
1013 let cursor = barrier
1014 .operational
1015 .checkpoint_ready
1016 .clone()
1017 .context("relay reported a checkpoint barrier without its ready cursor")?;
1018 let materialized = barrier.materialized;
1019 let expected_ordinal = materialized.applied_event_ordinal;
1020 let expected_digest = materialized.applied_event_digest.clone();
1021 ensure!(
1022 expected_ordinal == barrier.operational.latest_ordinal,
1023 "checkpoint projection frontier {expected_ordinal} does not match relay frontier {}",
1024 barrier.operational.latest_ordinal
1025 );
1026 ensure!(
1027 expected_digest == barrier.operational.latest_digest,
1028 "checkpoint projection digest does not match the relay frontier digest"
1029 );
1030 ensure_exact_checkpoint_cut(&cursor, expected_ordinal, &expected_digest)?;
1031 let canonical_session = canonical_session_from_materialized(&materialized)?;
1032 let native_session_id = barrier
1033 .operational
1034 .native_session_id
1035 .or_else(|| session.native_session_id.clone())
1036 .context("harness did not report its native session ID")?;
1037
1038 if exclusivity == LatchExclusivity::ReleaseAfterLatch {
1043 relay.end_latch();
1044 }
1045
1046 if export_policy == CheckpointExportPolicy::ReuseUnchangedArchive
1052 && session.managed_worktree.is_none()
1056 && let Some(artifact) = reusable_installed_checkpoint(
1057 session_id,
1058 session.checkpoint.as_ref(),
1059 &native_session_id,
1060 cursor.ordinal,
1061 &canonical_session,
1062 )
1063 {
1064 return Ok(LatchedCheckpoint {
1065 artifact,
1066 relay,
1067 barrier_command_id,
1068 cursor,
1069 completion: CheckpointCompletion::HeldBarrier,
1070 });
1071 }
1072
1073 let mut completion = CheckpointCompletion::HeldBarrier;
1078
1079 let exported: Result<CheckpointArtifact> = async {
1080 let spec = CheckpointExportSpec {
1081 protocol_version: CHECKPOINT_EXPORT_PROTOCOL_VERSION,
1082 session: session_manifest(&native_session_id),
1083 target: target_manifest,
1084 bundle: bundle_manifest,
1085 relay_root: target_path(&worker_root),
1086 harness_home: target_path(&harness_home),
1087 workspace_root: target_path(&workspace_root),
1088 repositories,
1089 canonical_session,
1090 output_path: target_path(&remote_archive),
1091 };
1092 let mut export_ms: Option<u64> = None;
1095 let exported = if releases_after_capture {
1096 let capture_spec = CheckpointCaptureSpec {
1097 protocol_version: CHECKPOINT_STAGING_PROTOCOL_VERSION,
1098 session: spec.session.clone(),
1099 target: spec.target.clone(),
1100 bundle: spec.bundle.clone(),
1101 relay_root: spec.relay_root.clone(),
1102 harness_home: spec.harness_home.clone(),
1103 workspace_root: spec.workspace_root.clone(),
1104 repositories: spec.repositories.clone(),
1105 allow_empty_native: !canonical_session_contains_prompt(&spec.canonical_session),
1106 stage_path: target_path(&remote_stage),
1107 refresh_existing: true,
1108 };
1109 let capture_started = Instant::now();
1110 let captured = {
1111 let _recovery_copy = recovery_copy.then(|| {
1112 ProvisionStageGuard::new(executor, ProvisionStage::RecoveryCopy)
1113 });
1114 run_checkpoint_staging_command(
1115 executor,
1116 &backend,
1117 session_id,
1118 &capture_spec,
1119 capture_stdin_command,
1120 "capture target checkpoint",
1121 )?
1122 };
1123 let captured: CapturedCheckpoint = serde_json::from_slice(&captured.stdout)
1124 .context("decode captured checkpoint result")?;
1125 tracing::info!(
1126 session_id,
1127 capture_ms = capture_started.elapsed().as_millis() as u64,
1128 barrier_held_ms = barrier_ready_at.elapsed().as_millis() as u64,
1129 native_bytes = captured.native_bytes,
1130 repository_bytes = captured.repository_bytes,
1131 reused_native = captured.reused_native,
1132 "checkpoint target state captured; releasing ACP dispatch"
1133 );
1134 completion = release_checkpoint_after_capture(
1135 &mut relay,
1136 session_id,
1137 &barrier_command_id,
1138 &cursor,
1139 session.harness_kind,
1140 )
1141 .await?;
1142 let pack_spec = CheckpointPackSpec {
1143 protocol_version: CHECKPOINT_STAGING_PROTOCOL_VERSION,
1144 relay_root: spec.relay_root.clone(),
1145 stage_path: target_path(&remote_stage),
1146 canonical_session: spec.canonical_session.clone(),
1147 output_path: spec.output_path.clone(),
1148 };
1149 let pack_started = Instant::now();
1150 let output = {
1151 let _recovery_copy = recovery_copy.then(|| {
1152 ProvisionStageGuard::new(executor, ProvisionStage::RecoveryCopy)
1153 });
1154 run_checkpoint_staging_command(
1155 executor,
1156 &backend,
1157 session_id,
1158 &pack_spec,
1159 pack_stdin_command,
1160 "pack target checkpoint",
1161 )?
1162 };
1163 tracing::info!(
1164 session_id,
1165 pack_ms = pack_started.elapsed().as_millis() as u64,
1166 "checkpoint archive packaged after ACP dispatch resumed"
1167 );
1168 output
1169 } else {
1170 let export_started = Instant::now();
1171 let output = {
1172 let _recovery_copy = recovery_copy.then(|| {
1173 ProvisionStageGuard::new(executor, ProvisionStage::RecoveryCopy)
1174 });
1175 export_target_checkpoint(
1176 executor,
1177 &backend,
1178 session_id,
1179 &spec,
1180 &remote_spec,
1181 )?
1182 };
1183 export_ms = Some(export_started.elapsed().as_millis() as u64);
1184 output
1185 };
1186 let target_checkpoint: mj_checkpoint::checkpoint::TargetCheckpoint =
1187 serde_json::from_slice(&exported.stdout)
1188 .context("decode target checkpoint result")?;
1189 if let Some(export_ms) = export_ms {
1190 let timings = target_checkpoint.timings.unwrap_or_default();
1193 tracing::info!(
1194 session_id,
1195 export_ms,
1196 timings_reported = target_checkpoint.timings.is_some(),
1197 native_ms = timings.native_ms,
1198 repositories_ms = timings.repositories_ms,
1199 archive_ms = timings.archive_ms,
1200 worker_total_ms = timings.total_ms,
1201 "checkpoint archive exported on the target"
1202 );
1203 }
1204 if target_checkpoint.event_frontier != expected_ordinal {
1205 bail!(
1206 "target checkpoint event frontier changed: expected {expected_ordinal}, found {}",
1207 target_checkpoint.event_frontier
1208 );
1209 }
1210 if target_checkpoint.event_frontier_digest != expected_digest {
1211 bail!("target checkpoint event frontier digest changed");
1212 }
1213
1214 let archive_id = new_command_id("archive")?;
1219 let destination = sessions_dir().join(format!(
1220 "{session_id}-{}-{archive_id}.hel.zip",
1221 target_checkpoint.event_frontier
1222 ));
1223 let transfer = CheckpointTransfer {
1224 locator: &backend,
1225 session_id,
1226 operation_id: &operation_id,
1227 remote_archive: &remote_archive,
1228 destination: &destination,
1229 expected_sha256: &target_checkpoint.sha256,
1230 expected_event_frontier: target_checkpoint.event_frontier,
1231 expected_event_frontier_digest: &target_checkpoint.event_frontier_digest,
1232 };
1233 let metadata = {
1234 let _verifying = ProvisionStageGuard::new(executor, ProvisionStage::Verifying);
1235 let transfer_started = Instant::now();
1236 let verified = transfer.execute(executor)?;
1237 tracing::info!(
1238 session_id,
1239 transfer_and_checksum_ms = transfer_started.elapsed().as_millis() as u64,
1240 "checkpoint archive transferred and checksum-verified"
1241 );
1242 let installed_archive = verified.archive_path().to_path_buf();
1243 let validate_transferred = || -> Result<()> {
1244 ensure!(
1245 verified.sha256() == target_checkpoint.sha256,
1246 "target and controller checkpoint checksums differ"
1247 );
1248 ensure!(
1249 verified.event_frontier_digest() == expected_digest,
1250 "verified checkpoint event frontier digest changed"
1251 );
1252 Ok(())
1253 };
1254 if let Err(error) = validate_transferred() {
1255 return Err(remove_uninstalled_checkpoint(&installed_archive, error));
1256 }
1257 if completion == CheckpointCompletion::HeldBarrier {
1261 let revalidated = relay.sync_snapshot().await.and_then(|snapshot| {
1262 if releases_after_capture {
1263 validate_automatic_checkpoint_barrier_snapshot(
1264 &snapshot,
1265 &barrier_command_id,
1266 &cursor,
1267 session.harness_kind,
1268 )
1269 } else {
1270 validate_checkpoint_barrier_snapshot(
1271 &snapshot,
1272 &barrier_command_id,
1273 &cursor,
1274 )
1275 }
1276 });
1277 if let Err(error) = revalidated {
1278 return Err(remove_uninstalled_checkpoint(
1279 &installed_archive,
1280 error.context(
1281 "checkpoint barrier changed while transferring its archive",
1282 ),
1283 ));
1284 }
1285 }
1286 if let Err(error) = transfer
1287 .cleanup_plan(&verified)
1288 .and_then(|plan| plan.execute(executor).map(|_| ()))
1289 {
1290 return Err(remove_uninstalled_checkpoint(
1291 &installed_archive,
1292 error.context("clean target checkpoint staging"),
1293 ));
1294 }
1295 CheckpointMetadata {
1296 archive_path: verified.archive_path().to_path_buf(),
1297 sha256: verified.sha256().to_string(),
1298 created_at: checkpointed_at.clone(),
1299 event_frontier: verified.event_frontier(),
1300 }
1301 };
1302 Ok(CheckpointArtifact {
1303 metadata,
1304 native_session_id,
1305 event_frontier_digest: expected_digest,
1306 })
1307 }
1308 .await;
1309
1310 let artifact = match exported {
1311 Ok(artifact) => artifact,
1312 Err(error) => {
1313 if completion == CheckpointCompletion::HeldBarrier
1319 && let Err(cancel_error) = relay.cancel_abandoned_barrier().await
1320 {
1321 tracing::warn!(
1322 session_id,
1323 "failed checkpoint could not cancel its relay barrier: {cancel_error:#}"
1324 );
1325 }
1326 return Err(error);
1327 }
1328 };
1329 Ok(LatchedCheckpoint {
1330 artifact,
1331 relay,
1332 barrier_command_id,
1333 cursor,
1334 completion,
1335 })
1336 }
1337
1338 pub(super) async fn prepare_move_source_checkpoint(
1339 &self,
1340 session_id: &str,
1341 executor: &(impl CommandExecutor + Sync),
1342 manager: &SessionManagerControl,
1343 operation: &mut mj_core::state::MoveOperation,
1344 ) -> Result<()> {
1345 let snapshot = super::move_session::refresh_move_source(manager, session_id).await?;
1346 if snapshot
1347 .as_ref()
1348 .is_some_and(|snapshot| snapshot.operational.checkpoint_only)
1349 {
1350 operation.source_checkpoint_only = true;
1351 crate::database::save_move_operation(operation)?;
1352 return Ok(());
1353 }
1354 if snapshot.as_ref().is_some_and(|snapshot| {
1355 matches!(
1356 snapshot.operational.execution,
1357 RelayExecutionState::Closing | RelayExecutionState::Closed
1358 )
1359 }) {
1360 return Ok(());
1361 }
1362 if !operation.source_checkpoint_only
1363 && snapshot
1364 .as_ref()
1365 .is_some_and(|snapshot| snapshot.operational.native_session_is_ready())
1366 {
1367 return Ok(());
1368 }
1369 ensure!(
1370 !executor.cancellation_requested() && !operation.cancellation_requested,
1371 "Move cancelled before source recovery"
1372 );
1373 operation.source_checkpoint_only = true;
1374 crate::database::save_move_operation(operation)?;
1375 executor.notify_notice("Recovering source data without starting its old harness");
1376 let (backend, worker_root) = self.worker_placement(session_id)?;
1377 let reconnect = targets::reconnect_plan(&backend, session_id)?
1378 .commands
1379 .into_iter()
1380 .next()
1381 .context("reconnect plan is empty")?;
1382 let launch = self.current_worker_launch_config(session_id, &backend)?;
1383 let connection = self
1384 .restart_worker_with_installed_binary(
1385 session_id,
1386 executor,
1387 InstalledWorkerRestart {
1388 backend: &backend,
1389 worker_root: &worker_root,
1390 reconnect: &reconnect,
1391 launch: Some(&launch),
1392 messages: &RESTART_FOR_CHECKPOINT,
1393 },
1394 )
1395 .await?;
1396 adopt_restarted_checkpoint_relay(session_id, Some(manager), connection)
1397 .await?
1398 .release();
1399 Ok(())
1400 }
1401
1402 async fn open_checkpoint_relay(
1406 &self,
1407 session_id: &str,
1408 executor: &(impl CommandExecutor + Sync),
1409 manager: Option<&SessionManagerControl>,
1410 target: InstalledWorkerRestart<'_>,
1411 restart_if_unreachable: bool,
1412 ) -> Result<(ControllerRelayLease, bool)> {
1413 let project_memory = match self.project_memory_sync_target(session_id) {
1414 Ok(target) => Some(target),
1415 Err(error) => {
1416 tracing::warn!(
1417 session_id,
1418 error = format!("{error:#}"),
1419 "project memory will not be synchronized during checkpoint reconnect"
1420 );
1421 None
1422 }
1423 };
1424 match connect_checkpoint_relay(
1425 session_id,
1426 manager,
1427 target.reconnect,
1428 project_memory.clone(),
1429 )
1430 .await
1431 {
1432 Ok(relay) => Ok((relay, false)),
1433 Err(error) if worker_connect_needs_restart(&error) && restart_if_unreachable => {
1434 tracing::warn!(
1435 session_id,
1436 "checkpoint could not reach the worker; restarting it: {error:#}"
1437 );
1438 let mut connection = self
1439 .restart_worker_for_checkpoint(
1440 session_id,
1441 executor,
1442 target.backend,
1443 target.worker_root,
1444 target.reconnect,
1445 )
1446 .await?;
1447 connection.set_project_memory_target(project_memory);
1448 let relay =
1449 adopt_restarted_checkpoint_relay(session_id, manager, connection).await?;
1450 Ok((relay, true))
1451 }
1452 Err(error) if worker_connect_needs_restart(&error) => {
1453 Err(error.context(CheckpointDeferred::background_work()))
1454 }
1455 Err(error) => Err(error).context("connect to the session worker for checkpoint"),
1456 }
1457 }
1458
1459 async fn restart_worker_for_checkpoint(
1463 &self,
1464 session_id: &str,
1465 executor: &(impl CommandExecutor + Sync),
1466 backend: &targets::TargetLocator,
1467 worker_root: &str,
1468 reconnect: &targets::CommandSpec,
1469 ) -> Result<StandaloneSession> {
1470 self.restart_worker_with_installed_binary(
1471 session_id,
1472 executor,
1473 InstalledWorkerRestart {
1474 backend,
1475 worker_root,
1476 reconnect,
1477 launch: None,
1478 messages: &RESTART_FOR_CHECKPOINT,
1479 },
1480 )
1481 .await
1482 }
1483}
1484
1485async fn connect_checkpoint_relay(
1486 session_id: &str,
1487 manager: Option<&SessionManagerControl>,
1488 reconnect: &targets::CommandSpec,
1489 project_memory: Option<crate::session_manager::ProjectMemorySyncTarget>,
1490) -> Result<ControllerRelayLease> {
1491 if let Some(manager) = manager {
1492 let handle = manager
1493 .wait_for_session(session_id, Duration::from_secs(5))
1494 .await?;
1495 let mut lease = handle.lease_connection().await?;
1496 lease
1497 .connection_mut()
1498 .set_project_memory_target(project_memory);
1499 Ok(ControllerRelayLease::Managed {
1500 handle,
1501 lease: Some(lease),
1502 })
1503 } else {
1504 let target = crate::session_manager::RelaySessionTarget {
1505 session_id: session_id.to_owned(),
1506 spec: reconnect.clone(),
1507 worker_recovery: None,
1508 project_memory,
1509 };
1510 Ok(ControllerRelayLease::Standalone(
1511 StandaloneSession::connect(&target).await?,
1512 ))
1513 }
1514}
1515
1516async fn adopt_restarted_checkpoint_relay(
1517 session_id: &str,
1518 manager: Option<&SessionManagerControl>,
1519 connection: StandaloneSession,
1520) -> Result<ControllerRelayLease> {
1521 let Some(manager) = manager else {
1522 return Ok(ControllerRelayLease::Standalone(connection));
1523 };
1524 let handle = manager
1525 .wait_for_session(session_id, Duration::from_secs(5))
1526 .await?;
1527 match handle.lease_connection().await {
1528 Ok(mut lease) => {
1529 lease.replace_connection(connection);
1530 Ok(ControllerRelayLease::Managed {
1531 handle,
1532 lease: Some(lease),
1533 })
1534 }
1535 Err(error) => {
1536 tracing::warn!(
1537 session_id,
1538 "session actor could not lease after worker restart; using the restarted proxy: {error:#}"
1539 );
1540 Ok(ControllerRelayLease::Standalone(connection))
1541 }
1542 }
1543}
1544
1545#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1547enum BarrierBusyPolicy {
1548 DeferWhileRunning,
1554 InterruptWhileRunning,
1558}
1559
1560impl BarrierBusyPolicy {
1561 fn of(exclusivity: LatchExclusivity) -> Self {
1562 match exclusivity {
1563 LatchExclusivity::ReleaseAfterLatch => Self::DeferWhileRunning,
1564 LatchExclusivity::HoldThroughClose => Self::InterruptWhileRunning,
1565 }
1566 }
1567}
1568
1569async fn wait_for_checkpoint_barrier(
1570 relay: &mut StandaloneSession,
1571 session_id: &str,
1572 command_id: &str,
1573 timeout: Duration,
1574 busy: BarrierBusyPolicy,
1575 harness: HarnessKind,
1576) -> Result<ManagedSessionSnapshot> {
1577 let deadline = tokio::time::Instant::now() + timeout;
1578 let mut cancel_submitted = false;
1579 let mut cancel_deadline = None;
1580 let mut cancel_started_at: Option<Instant> = None;
1581 loop {
1582 let snapshot = relay.sync().await?;
1583 if busy == BarrierBusyPolicy::DeferWhileRunning
1584 && !snapshot.operational.safe_for_checkpoint(harness)
1585 {
1586 return Err(
1591 CheckpointDeferred::background_snapshot(&snapshot.operational, harness).into(),
1592 );
1593 }
1594 if checkpoint_barrier_is_ready(&snapshot, command_id) {
1595 if let Some(started_at) = cancel_started_at {
1596 tracing::info!(
1597 session_id,
1598 barrier_command_id = command_id,
1599 cancellation_ms = started_at.elapsed().as_millis() as u64,
1600 "active turn cancellation settled before checkpoint barrier"
1601 );
1602 }
1603 return Ok(snapshot);
1604 }
1605 if busy == BarrierBusyPolicy::InterruptWhileRunning
1606 && snapshot.operational.execution == RelayExecutionState::Running
1607 && !cancel_submitted
1608 {
1609 let cancel_turn = RelayCommand::CancelTurn;
1610 if relay.protocol_version() < cancel_turn.minimum_protocol() {
1611 return Err(CheckpointBarrierUnreachable::cancel_turn_unavailable(
1612 command_id,
1613 relay.protocol_version(),
1614 )
1615 .into());
1616 }
1617 let cancel_command_id = new_command_id("checkpoint-cancel-turn")?;
1618 match relay.submit(cancel_command_id, cancel_turn).await {
1619 Ok(_) => {
1620 cancel_submitted = true;
1621 cancel_started_at = Some(Instant::now());
1622 cancel_deadline = Some(tokio::time::Instant::now() + CHECKPOINT_CANCEL_TIMEOUT);
1623 tracing::info!(
1624 session_id,
1625 barrier_command_id = command_id,
1626 "requested active turn cancellation before checkpoint barrier"
1627 );
1628 }
1629 Err(error) if checkpoint_cancel_turn_needs_worker_restart(&error) => {
1630 return Err(error.context(
1631 CheckpointBarrierUnreachable::cancel_turn_unavailable(
1632 command_id,
1633 relay.protocol_version(),
1634 ),
1635 ));
1636 }
1637 Err(error) if worker_connect_needs_restart(&error) => {
1638 return Err(error.context(
1639 CheckpointBarrierUnreachable::cancel_turn_unreachable(command_id),
1640 ));
1641 }
1642 Err(error) => {
1643 if let Ok(snapshot) = relay.sync().await
1647 && checkpoint_barrier_is_ready(&snapshot, command_id)
1648 {
1649 tracing::info!(
1650 session_id,
1651 barrier_command_id = command_id,
1652 "active turn settled while submitting checkpoint cancellation"
1653 );
1654 return Ok(snapshot);
1655 }
1656 return Err(error.context("cancel active ACP turn before checkpoint barrier"));
1657 }
1658 }
1659 continue;
1660 }
1661 let out_of_time = tokio::time::Instant::now() >= cancel_deadline.unwrap_or(deadline);
1662 if let Some(error) = checkpoint_barrier_wait_ended(
1663 &snapshot,
1664 command_id,
1665 busy,
1666 out_of_time,
1667 cancel_submitted,
1668 ) {
1669 return Err(error);
1670 }
1671 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
1672 }
1673}
1674
1675fn checkpoint_barrier_wait_ended(
1682 snapshot: &ManagedSessionSnapshot,
1683 command_id: &str,
1684 busy: BarrierBusyPolicy,
1685 out_of_time: bool,
1686 cancel_submitted: bool,
1687) -> Option<anyhow::Error> {
1688 if snapshot.operational.execution == RelayExecutionState::Closed {
1689 return Some(CheckpointBarrierUnreachable::runtime_stopped().into());
1690 }
1691 if snapshot.operational.execution == RelayExecutionState::Running {
1692 return Some(match busy {
1693 BarrierBusyPolicy::DeferWhileRunning => CheckpointDeferred::harness_busy().into(),
1694 BarrierBusyPolicy::InterruptWhileRunning if out_of_time && cancel_submitted => {
1695 CheckpointBarrierUnreachable::cancel_timed_out(command_id).into()
1696 }
1697 BarrierBusyPolicy::InterruptWhileRunning => return None,
1698 });
1699 }
1700 out_of_time.then(|| CheckpointBarrierUnreachable::not_admitted(command_id).into())
1701}
1702
1703#[derive(Debug)]
1710struct CheckpointBarrierUnreachable(String);
1711
1712impl CheckpointBarrierUnreachable {
1713 fn runtime_stopped() -> Self {
1714 Self("ACP runtime stopped before reaching the checkpoint barrier".to_owned())
1715 }
1716
1717 fn not_admitted(command_id: &str) -> Self {
1718 Self(format!(
1719 "ACP relay did not reach checkpoint barrier {command_id}"
1720 ))
1721 }
1722
1723 fn cancel_timed_out(command_id: &str) -> Self {
1724 Self(format!(
1725 "active ACP turn did not settle after cancellation before checkpoint barrier {command_id}"
1726 ))
1727 }
1728
1729 fn cancel_turn_unavailable(command_id: &str, protocol_version: u32) -> Self {
1730 Self(format!(
1731 "worker protocol {protocol_version} cannot cancel the active ACP turn before checkpoint barrier {command_id} (requires protocol {})",
1732 RelayCommand::CancelTurn.minimum_protocol(),
1733 ))
1734 }
1735
1736 fn cancel_turn_unreachable(command_id: &str) -> Self {
1737 Self(format!(
1738 "worker transport became unavailable while cancelling the active ACP turn before checkpoint barrier {command_id}"
1739 ))
1740 }
1741}
1742
1743impl std::fmt::Display for CheckpointBarrierUnreachable {
1744 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1745 formatter.write_str(&self.0)
1746 }
1747}
1748
1749impl std::error::Error for CheckpointBarrierUnreachable {}
1750
1751fn checkpoint_barrier_needs_worker_restart(error: &anyhow::Error) -> bool {
1752 error
1753 .downcast_ref::<CheckpointBarrierUnreachable>()
1754 .is_some()
1755}
1756
1757fn checkpoint_cancel_turn_needs_worker_restart(error: &anyhow::Error) -> bool {
1762 error.chain().any(|cause| {
1763 let Some(rejected) = cause.downcast_ref::<RelayRejected>() else {
1764 return false;
1765 };
1766 rejected.0.code == mj_core::relay::RelayErrorCode::IncompatibleProtocol
1767 })
1768}
1769
1770#[derive(Debug)]
1780pub struct CheckpointDeferred(String);
1781
1782impl CheckpointDeferred {
1783 pub fn harness_busy() -> Self {
1784 Self("the agent is working; try again when it is idle".to_owned())
1785 }
1786
1787 fn background_work() -> Self {
1788 Self("Kimi background-agent state could not be synchronized; checkpoint requires a synchronized empty task list".into())
1789 }
1790
1791 fn background_snapshot(
1792 state: &mj_core::relay::RelayOperationalState,
1793 harness: HarnessKind,
1794 ) -> Self {
1795 Self(
1796 state
1797 .checkpoint_background_blocker(harness)
1798 .unwrap_or("background state changed during checkpoint")
1799 .into(),
1800 )
1801 }
1802
1803 fn frontier_moved() -> Self {
1804 Self(
1805 "the session moved past the checkpoint-ready cursor before the barrier latched, so this checkpoint was deferred"
1806 .to_owned(),
1807 )
1808 }
1809
1810 fn harness_turn_during_capture() -> Self {
1811 Self(
1812 "the agent started a turn of its own while target state was captured, so this checkpoint was deferred"
1813 .to_owned(),
1814 )
1815 }
1816}
1817
1818impl std::fmt::Display for CheckpointDeferred {
1819 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1820 formatter.write_str(&self.0)
1821 }
1822}
1823
1824impl std::error::Error for CheckpointDeferred {}
1825
1826pub fn checkpoint_was_deferred(error: &anyhow::Error) -> bool {
1831 error
1832 .chain()
1833 .any(|cause| cause.downcast_ref::<CheckpointDeferred>().is_some())
1834}
1835
1836pub struct IdleWorkspaceLease {
1840 lease: ManagedSessionLease,
1841 command_id: String,
1842 harness: HarnessKind,
1843}
1844
1845impl IdleWorkspaceLease {
1846 pub async fn acquire(handle: &ManagedSessionHandle, harness: HarnessKind) -> Result<Self> {
1847 tokio::time::timeout(Duration::from_secs(30), async {
1848 let mut lease = handle.lease_connection().await?;
1849 let snapshot = lease.connection_mut().sync().await?;
1850 ensure!(
1851 snapshot.operational.safe_to_replace(harness),
1852 "session must be live and idle with no queued or background work"
1853 );
1854 let command_id = new_command_id("workspace-write")?;
1855 lease
1856 .connection_mut()
1857 .submit(
1858 command_id.clone(),
1859 RelayCommand::BeginCheckpoint {
1860 reason: Some("API workspace file write".into()),
1861 },
1862 )
1863 .await?;
1864 loop {
1865 let snapshot = lease.connection_mut().sync().await?;
1866 if checkpoint_barrier_is_ready(&snapshot, &command_id) {
1867 let mut operation = Self {
1868 lease,
1869 command_id,
1870 harness,
1871 };
1872 operation.verify().await?;
1873 return Ok(operation);
1874 }
1875 ensure!(
1876 snapshot.operational.execution != RelayExecutionState::Running,
1877 "session started work before the file barrier was ready"
1878 );
1879 tokio::time::sleep(Duration::from_millis(25)).await;
1880 }
1881 })
1882 .await
1883 .context("session did not become available for a file write within 30 seconds")?
1884 }
1885
1886 pub async fn verify(&mut self) -> Result<()> {
1887 let mut snapshot = self.lease.connection_mut().sync().await?;
1888 ensure!(
1889 checkpoint_barrier_is_ready(&snapshot, &self.command_id),
1890 "file write lost its workspace barrier"
1891 );
1892 snapshot.operational.checkpoint_barrier = None;
1893 snapshot.operational.queued_prompts.clear();
1896 ensure!(
1897 snapshot.operational.safe_to_replace(self.harness),
1898 "session is no longer idle for the file write"
1899 );
1900 Ok(())
1901 }
1902
1903 pub async fn release(mut self) -> Result<()> {
1904 tokio::time::timeout(Duration::from_secs(30), async {
1905 self.lease
1906 .connection_mut()
1907 .submit(
1908 new_command_id("workspace-release")?,
1909 RelayCommand::ReleaseCheckpoint {
1910 barrier_command_id: self.command_id.clone(),
1911 },
1912 )
1913 .await?;
1914 loop {
1915 let snapshot = self.lease.connection_mut().sync().await?;
1916 if snapshot.operational.checkpoint_barrier.as_deref() != Some(&self.command_id) {
1917 return Ok::<_, anyhow::Error>(());
1918 }
1919 tokio::time::sleep(Duration::from_millis(25)).await;
1920 }
1921 })
1922 .await
1923 .context("release file write barrier timed out")??;
1924 self.lease.release();
1925 Ok(())
1926 }
1927}
1928
1929fn checkpoint_barrier_is_ready(snapshot: &ManagedSessionSnapshot, command_id: &str) -> bool {
1930 snapshot.operational.checkpoint_barrier.as_deref() == Some(command_id)
1931 && snapshot.operational.checkpoint_ready.is_some()
1932}
1933
1934fn ensure_exact_checkpoint_cut(
1942 cursor: &RelayCursor,
1943 expected_ordinal: u64,
1944 expected_digest: &str,
1945) -> Result<()> {
1946 if cursor.ordinal != expected_ordinal || cursor.digest != expected_digest {
1947 bail!(CheckpointDeferred::frontier_moved());
1948 }
1949 Ok(())
1950}
1951
1952fn validate_checkpoint_barrier_snapshot(
1967 snapshot: &ManagedSessionSnapshot,
1968 command_id: &str,
1969 expected: &RelayCursor,
1970) -> Result<()> {
1971 ensure!(
1972 snapshot.operational.checkpoint_barrier.as_deref() == Some(command_id),
1973 "checkpoint barrier {command_id} is no longer active"
1974 );
1975 ensure!(
1976 snapshot.operational.checkpoint_ready.as_ref() == Some(expected),
1977 "checkpoint barrier {command_id} has a different ready cursor"
1978 );
1979 if snapshot
1980 .operational
1981 .last_harness_turn_started_ordinal
1982 .is_some_and(|ordinal| ordinal > expected.ordinal)
1983 {
1984 bail!(CheckpointDeferred::harness_turn_during_capture());
1985 }
1986 Ok(())
1987}
1988
1989fn validate_automatic_checkpoint_barrier_snapshot(
1993 snapshot: &ManagedSessionSnapshot,
1994 command_id: &str,
1995 expected: &RelayCursor,
1996 harness: HarnessKind,
1997) -> Result<()> {
1998 validate_checkpoint_barrier_snapshot(snapshot, command_id, expected)?;
1999 ensure!(
2000 snapshot.operational.safe_for_checkpoint(harness),
2001 CheckpointDeferred::background_snapshot(&snapshot.operational, harness)
2002 );
2003 Ok(())
2004}
2005
2006fn remove_uninstalled_checkpoint(path: &Path, error: anyhow::Error) -> anyhow::Error {
2007 match std::fs::remove_file(path) {
2008 Ok(()) => error,
2009 Err(remove_error) if remove_error.kind() == std::io::ErrorKind::NotFound => error,
2010 Err(remove_error) => error.context(format!(
2011 "also failed to remove uninstalled checkpoint {}: {remove_error}",
2012 path.display()
2013 )),
2014 }
2015}
2016
2017pub(super) async fn wait_for_relay_closed(relay: &mut StandaloneSession) -> Result<()> {
2018 let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(30);
2019 loop {
2020 if relay.sync().await?.operational.execution == RelayExecutionState::Closed {
2021 return Ok(());
2022 }
2023 if tokio::time::Instant::now() >= deadline {
2024 bail!("ACP runtime did not close within 30 seconds");
2025 }
2026 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
2027 }
2028}
2029
2030async fn release_checkpoint_after_capture(
2042 relay: &mut ControllerRelayLease,
2043 session_id: &str,
2044 barrier_command_id: &str,
2045 cursor: &RelayCursor,
2046 harness: HarnessKind,
2047) -> Result<CheckpointCompletion> {
2048 relay
2049 .sync_snapshot()
2050 .await
2051 .and_then(|snapshot| {
2052 validate_automatic_checkpoint_barrier_snapshot(
2053 &snapshot,
2054 barrier_command_id,
2055 cursor,
2056 harness,
2057 )
2058 })
2059 .context("checkpoint barrier changed while capturing target state")?;
2060 match relay
2061 .submit(
2062 new_command_id("checkpoint-release")?,
2063 RelayCommand::ReleaseCheckpoint {
2064 barrier_command_id: barrier_command_id.to_owned(),
2065 },
2066 )
2067 .await
2068 {
2069 Ok(_) => Ok(CheckpointCompletion::ReleasedAfterCapture),
2070 Err(error) => {
2071 tracing::debug!(
2072 session_id,
2073 "relay kept the checkpoint barrier through the transfer: {error:#}"
2074 );
2075 Ok(CheckpointCompletion::HeldBarrier)
2076 }
2077 }
2078}
2079
2080fn run_checkpoint_staging_command<T: serde::Serialize>(
2081 executor: &impl CommandExecutor,
2082 locator: &targets::TargetLocator,
2083 session_id: &str,
2084 spec: &T,
2085 command: fn(&targets::TargetLocator, &str) -> Result<CommandSpec>,
2086 operation: &str,
2087) -> Result<CommandOutput> {
2088 let body = serde_json::to_vec(spec).with_context(|| format!("serialize {operation} spec"))?;
2089 let mut replaced_worker = false;
2090 loop {
2091 let command = command(locator, session_id)?;
2092 let output = executor.execute_with_stdin(&command, &mut body.as_slice())?;
2093 if output.status == 0 {
2094 return Ok(output);
2095 }
2096 let failure = String::from_utf8_lossy(&output.stderr).into_owned();
2097 if staging_protocol_unsupported(&failure)
2098 && replace_stale_export_worker(
2099 executor,
2100 locator,
2101 session_id,
2102 None,
2103 &failure,
2104 &mut replaced_worker,
2105 )?
2106 {
2107 continue;
2108 }
2109 bail!(
2110 "{operation} failed with status {}: {failure}",
2111 output.status
2112 );
2113 }
2114}
2115
2116fn export_target_checkpoint(
2121 executor: &impl CommandExecutor,
2122 locator: &targets::TargetLocator,
2123 session_id: &str,
2124 spec: &CheckpointExportSpec,
2125 remote_spec: &str,
2126) -> Result<CommandOutput> {
2127 export_target_checkpoint_with_worker(executor, locator, session_id, spec, remote_spec, None)
2128}
2129
2130fn export_target_checkpoint_with_worker(
2131 executor: &impl CommandExecutor,
2132 locator: &targets::TargetLocator,
2133 session_id: &str,
2134 spec: &CheckpointExportSpec,
2135 remote_spec: &str,
2136 worker_binary: Option<&Path>,
2137) -> Result<CommandOutput> {
2138 let body = serde_json::to_vec(spec).context("serialize checkpoint export spec")?;
2139 let mut replaced_worker = false;
2140 loop {
2141 let streamed = export_stdin_command(locator, session_id)?;
2142 let output = executor.execute_with_stdin(&streamed, &mut body.as_slice())?;
2143 if output.status == 0 {
2144 return Ok(output);
2145 }
2146 let failure = String::from_utf8_lossy(&output.stderr).into_owned();
2147 if export_spec_stdin_unsupported(&failure) {
2148 tracing::debug!(
2149 session_id,
2150 "target worker predates streamed checkpoint specs; uploading the spec file instead"
2151 );
2152 let output = export_uploaded_spec(executor, locator, session_id, spec, remote_spec)?;
2153 if output.status == 0 {
2154 return Ok(output);
2155 }
2156 let failure = String::from_utf8_lossy(&output.stderr).into_owned();
2157 if replace_stale_export_worker(
2158 executor,
2159 locator,
2160 session_id,
2161 worker_binary,
2162 &failure,
2163 &mut replaced_worker,
2164 )? {
2165 continue;
2166 }
2167 bail!(
2168 "export target checkpoint failed with status {}: {failure}",
2169 output.status
2170 );
2171 }
2172 if replace_stale_export_worker(
2173 executor,
2174 locator,
2175 session_id,
2176 worker_binary,
2177 &failure,
2178 &mut replaced_worker,
2179 )? {
2180 continue;
2181 }
2182 bail!(
2183 "{} failed with status {}: {failure}",
2184 streamed.purpose,
2185 output.status
2186 );
2187 }
2188}
2189
2190fn export_uploaded_spec(
2191 executor: &impl CommandExecutor,
2192 locator: &targets::TargetLocator,
2193 session_id: &str,
2194 spec: &CheckpointExportSpec,
2195 remote_spec: &str,
2196) -> Result<CommandOutput> {
2197 let staging = tempfile::tempdir().context("create checkpoint staging")?;
2198 let local_spec = staging.path().join("checkpoint-spec.json");
2199 spec.write(&local_spec)?;
2200 upload_checkpoint_spec(executor, locator, session_id, &local_spec, remote_spec)?;
2201 let output = executor.execute(&export_command(locator, session_id, remote_spec)?)?;
2202 if output.status == 0 {
2203 execute_checked(
2204 executor,
2205 targets::command_on_locator(
2206 locator,
2207 session_id,
2208 ["rm", "-f", "--", remote_spec].map(str::to_owned).to_vec(),
2209 "remove uploaded checkpoint specification",
2210 )?,
2211 )
2212 .context("clean successful checkpoint export specification")?;
2213 }
2214 Ok(output)
2215}
2216
2217fn replace_stale_export_worker(
2222 executor: &impl CommandExecutor,
2223 locator: &targets::TargetLocator,
2224 session_id: &str,
2225 worker_binary: Option<&Path>,
2226 failure: &str,
2227 replaced_worker: &mut bool,
2228) -> Result<bool> {
2229 if *replaced_worker || !staging_protocol_unsupported(failure) {
2230 return Ok(false);
2231 }
2232 tracing::debug!(
2233 session_id,
2234 "target worker does not support this checkpoint export protocol; replacing the installed Mjolnir binary and retrying"
2235 );
2236 let owned_binary;
2237 let binary = if let Some(path) = worker_binary {
2238 path
2239 } else {
2240 owned_binary = super::worker_binary::worker_binary_for(locator, executor)?;
2241 owned_binary.as_path()
2242 };
2243 super::worker_binary::replace_installed_worker_binary(executor, locator, session_id, binary)?;
2244 *replaced_worker = true;
2245 Ok(true)
2246}
2247
2248fn export_spec_stdin_unsupported(failure: &str) -> bool {
2256 failure.contains("read checkpoint export spec -")
2257 || failure.contains("unexpected argument")
2258 || failure.contains("invalid value")
2259}
2260
2261fn export_spec_schema_unsupported(failure: &str) -> bool {
2266 failure.contains("parse checkpoint")
2267 && (failure.contains("unknown field") || failure.contains("unknown variant"))
2268}
2269
2270fn export_protocol_unsupported(failure: &str) -> bool {
2271 export_spec_schema_unsupported(failure)
2272 || failure.contains("unsupported checkpoint export protocol version")
2273}
2274
2275fn staging_protocol_unsupported(failure: &str) -> bool {
2276 export_protocol_unsupported(failure)
2277 || failure.contains("unsupported checkpoint staging protocol version")
2278 || failure.contains("unrecognized subcommand")
2279 || failure.contains("unexpected argument")
2280}
2281
2282pub(super) fn upload_checkpoint_spec(
2283 executor: &impl CommandExecutor,
2284 locator: &targets::TargetLocator,
2285 session_id: &str,
2286 local: &Path,
2287 remote: &str,
2288) -> Result<()> {
2289 match locator {
2290 targets::TargetLocator::LocalBare { .. } => {
2291 std::fs::copy(local, remote)
2292 .with_context(|| format!("copy checkpoint specification to {remote}"))?;
2293 Ok(())
2294 }
2295 targets::TargetLocator::LocalPodman { container_id, .. } => execute_checked(
2296 executor,
2297 CommandSpec::new(
2298 "podman",
2299 [
2300 "cp".into(),
2301 local.to_string_lossy().into_owned(),
2302 format!("{container_id}:{remote}"),
2303 ],
2304 )
2305 .purpose("upload checkpoint specification"),
2306 )
2307 .map(|_| ()),
2308 targets::TargetLocator::LocalDocker { container_id } => execute_checked(
2309 executor,
2310 CommandSpec::new(
2311 "docker",
2312 [
2313 "cp".into(),
2314 local.to_string_lossy().into_owned(),
2315 format!("{container_id}:{remote}"),
2316 ],
2317 )
2318 .purpose("upload checkpoint specification"),
2319 )
2320 .map(|_| ()),
2321 targets::TargetLocator::AppleContainer { container_id } => execute_checked(
2322 executor,
2323 CommandSpec::new(
2324 "container",
2325 [
2326 "cp".into(),
2327 local.to_string_lossy().into_owned(),
2328 format!("{container_id}:{remote}"),
2329 ],
2330 )
2331 .purpose("upload checkpoint specification"),
2332 )
2333 .map(|_| ()),
2334 targets::TargetLocator::AwsEc2 { ssh, .. }
2335 | targets::TargetLocator::SshBare { ssh, .. } => execute_checked(
2336 executor,
2337 scp_command_spec(ssh, local, remote, false).purpose("upload checkpoint specification"),
2338 )
2339 .map(|_| ()),
2340 targets::TargetLocator::SshPodman {
2341 ssh, container_id, ..
2342 }
2343 | targets::TargetLocator::SshDocker { ssh, container_id } => {
2344 let engine = match locator {
2345 targets::TargetLocator::SshPodman { .. } => "podman",
2346 targets::TargetLocator::SshDocker { .. } => "docker",
2347 _ => unreachable!("matched remote container target"),
2348 };
2349 let staging = format!(".local/share/hel/uploads/{session_id}-checkpoint.json");
2350 execute_checked(
2351 executor,
2352 ssh_command_spec(ssh, ["mkdir", "-p", ".local/share/hel/uploads"])
2353 .purpose("create remote checkpoint staging"),
2354 )?;
2355 execute_checked(
2356 executor,
2357 scp_command_spec(ssh, local, &staging, false)
2358 .purpose("upload remote container checkpoint specification"),
2359 )?;
2360 execute_checked(
2361 executor,
2362 ssh_command_spec(
2363 ssh,
2364 [engine, "cp", &staging, &format!("{container_id}:{remote}")],
2365 )
2366 .purpose("install remote container checkpoint specification"),
2367 )?;
2368 execute_checked(
2369 executor,
2370 ssh_command_spec(ssh, ["rm", "-f", "--", &staging])
2371 .purpose("remove remote checkpoint staging"),
2372 )?;
2373 Ok(())
2374 }
2375 }?;
2376 Ok(())
2377}
2378
2379fn reusable_installed_checkpoint(
2387 session_id: &str,
2388 installed: Option<&CheckpointMetadata>,
2389 native_session_id: &str,
2390 latched_ordinal: u64,
2391 latched_session: &CanonicalSessionSnapshot,
2392) -> Option<CheckpointArtifact> {
2393 let installed = installed?;
2394 if installed.event_frontier > latched_ordinal {
2395 tracing::warn!(
2396 session_id,
2397 installed_frontier = installed.event_frontier,
2398 latched_ordinal,
2399 "installed checkpoint is ahead of the latched cursor; exporting a fresh archive"
2400 );
2401 return None;
2402 }
2403 let verified = match verify_archive_streaming(&installed.archive_path) {
2404 Ok(verified) => verified,
2405 Err(error) => {
2406 tracing::warn!(
2407 session_id,
2408 path = %installed.archive_path.display(),
2409 "installed checkpoint could not be verified for reuse: {error:#}"
2410 );
2411 return None;
2412 }
2413 };
2414 if verified.archive_sha256 != installed.sha256
2415 || verified.manifest.session.id != session_id
2416 || verified.canonical_session.event_frontier != installed.event_frontier
2417 {
2418 tracing::warn!(
2419 session_id,
2420 path = %installed.archive_path.display(),
2421 "installed checkpoint no longer matches its controller metadata; exporting a fresh archive"
2422 );
2423 return None;
2424 }
2425 if !verified.canonical_session.content_matches(latched_session) {
2426 tracing::info!(
2427 session_id,
2428 archive_frontier = verified.canonical_session.event_frontier,
2429 latched_ordinal,
2430 "session content changed since the installed checkpoint; exporting a fresh archive"
2431 );
2432 return None;
2433 }
2434 tracing::info!(
2435 session_id,
2436 archive_frontier = verified.canonical_session.event_frontier,
2437 latched_ordinal,
2438 "reusing the installed checkpoint archive; only relay bookkeeping moved"
2439 );
2440 Some(CheckpointArtifact {
2441 metadata: installed.clone(),
2442 native_session_id: native_session_id.to_owned(),
2443 event_frontier_digest: verified.canonical_session.event_frontier_digest,
2444 })
2445}
2446
2447pub(super) fn verify_installed_checkpoint_gate(
2448 session_id: &str,
2449 checkpoint: &CheckpointMetadata,
2450) -> Result<()> {
2451 let sha256 = checkpoint_sha256(&checkpoint.archive_path).with_context(|| {
2452 format!(
2453 "hash installed checkpoint {} before target cleanup",
2454 checkpoint.archive_path.display()
2455 )
2456 })?;
2457 ensure!(
2458 sha256 == checkpoint.sha256,
2459 "refusing target cleanup for session {session_id}: installed checkpoint SHA changed"
2460 );
2461 Ok(())
2462}
2463
2464fn verify_checkpoint_artifact(session_id: &str, artifact: &CheckpointArtifact) -> Result<()> {
2465 let sha256 = checkpoint_sha256(&artifact.metadata.archive_path).with_context(|| {
2466 format!(
2467 "hash completed checkpoint {}",
2468 artifact.metadata.archive_path.display()
2469 )
2470 })?;
2471 ensure!(
2472 sha256 == artifact.metadata.sha256,
2473 "completed checkpoint SHA changed before persistence for session {session_id}"
2474 );
2475 Ok(())
2476}
2477
2478pub(super) fn release_projection_behind_checkpoint(session_id: &str, current: &CheckpointMetadata) {
2486 match crate::database::compact_materialized_transcript_through(
2487 session_id,
2488 current.event_frontier,
2489 ) {
2490 Ok(retention) if retention.items == 0 => {}
2491 Ok(retention) => tracing::info!(
2492 session_id,
2493 items = retention.items,
2494 bytes = retention.bytes,
2495 remaining = retention.remaining,
2496 event_frontier = current.event_frontier,
2497 "released projection history the checkpoint covers"
2498 ),
2499 Err(error) => tracing::warn!(
2500 session_id,
2501 "checkpoint was saved, but the projection history it covers could not be released: {error:#}"
2502 ),
2503 }
2504}
2505
2506pub(super) fn prune_replaced_checkpoint(
2507 previous: Option<&CheckpointMetadata>,
2508 current: &CheckpointMetadata,
2509) {
2510 let Some(previous) = previous.filter(|old| old.archive_path != current.archive_path) else {
2511 return;
2512 };
2513 match crate::database::move_checkpoint_is_retained(&previous.archive_path) {
2514 Ok(true) => return,
2515 Ok(false) => {}
2516 Err(error) => {
2517 tracing::warn!(%error, "could not check move retention; keeping superseded checkpoint");
2518 return;
2519 }
2520 }
2521 if let Err(error) = std::fs::remove_file(&previous.archive_path)
2522 && error.kind() != std::io::ErrorKind::NotFound
2523 {
2524 tracing::warn!(
2525 path = %previous.archive_path.display(),
2526 "could not remove superseded recovery copy: {error}"
2527 );
2528 }
2529}
2530
2531#[cfg(test)]
2532mod tests {
2533 use std::cell::{Cell, RefCell};
2534 use std::collections::BTreeMap;
2535 use std::fs::OpenOptions;
2536 use std::path::{Path, PathBuf};
2537 #[cfg(unix)]
2538 use std::process::Command;
2539 #[cfg(unix)]
2540 use std::time::Duration;
2541
2542 #[cfg(unix)]
2543 use agent_client_protocol::schema::v1::{ContentBlock, TextContent};
2544 use anyhow::Result;
2545
2546 #[cfg(unix)]
2547 use crate::controller::now;
2548 use crate::controller::restore_session_after_persistence_failure;
2549 use crate::controller::test_support::{checkpoint_test_session, write_checkpoint_gate_archive};
2550 #[cfg(unix)]
2551 use crate::session_manager::{ManagedSessionHandle, new_command_id};
2552 use crate::worker_client::RelayTransportDead;
2553 use mj_checkpoint::archive::{
2554 BundleManifest, CanonicalTranscriptBody, CanonicalTranscriptItem, TargetManifest,
2555 };
2556 use mj_checkpoint::checkpoint::CheckpointExportSpec;
2557 #[cfg(unix)]
2558 use mj_core::config::{
2559 Config, HarnessProfile, ProjectBundle, ProjectRepository, TargetTemplate,
2560 };
2561 #[cfg(unix)]
2562 use mj_core::state::TargetLocator;
2563 use mj_core::state::{
2564 CheckpointMetadata, ManagedSessionSnapshot, MaterializedSession, SessionState, State,
2565 };
2566 use mj_transcript::projection::canonical_session_from_materialized;
2567
2568 #[cfg(unix)]
2569 use crate::targets::ProvisionStage;
2570 use crate::targets::{self, CommandExecutor, CommandOutput, CommandSpec};
2571 #[cfg(unix)]
2572 use mj_core::relay::RelayCommandOutcome;
2573 use mj_core::relay::{RelayCommand, RelayCursor, RelayExecutionState};
2574
2575 use super::*;
2576
2577 struct UnusedExecutor;
2581
2582 impl CommandExecutor for UnusedExecutor {
2583 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
2584 panic!("the export layout ran {command:?}");
2585 }
2586 }
2587
2588 #[test]
2589 fn the_export_layout_places_each_session_kind_in_its_workspace() {
2590 let session_id = "1123456789abcdef0123456789abcdef";
2591 let mut config = crate::controller::test_support::resume_compatibility_config();
2592 config.bundles.insert(
2593 "app-bundle".into(),
2594 mj_core::config::ProjectBundle {
2595 primary_repo: "app".into(),
2596 repositories: vec![mj_core::config::ProjectRepository {
2597 id: "app".into(),
2598 github: None,
2599 local: None,
2600 destination: PathBuf::from("app"),
2601 git_ref: None,
2602 }],
2603 },
2604 );
2605
2606 let mut session = checkpoint_test_session(session_id);
2609 session.bundle_id = "app-bundle".into();
2610 session.target = Some(mj_core::state::TargetLocator::LocalPodman {
2611 container_id: "hel-session".into(),
2612 workspace_storage: Default::default(),
2613 });
2614 let mut state = State::default();
2615 state.sessions.insert(session_id.into(), session.clone());
2616 let controller = Controller {
2617 config: config.clone(),
2618 state,
2619 };
2620
2621 let layout = controller
2622 .session_export_layout(session_id, &UnusedExecutor)
2623 .unwrap();
2624 assert_eq!(layout.workspace_root, "/workspace");
2625 assert_eq!(layout.primary_repository, "app");
2626 assert_eq!(
2627 layout
2628 .repositories
2629 .iter()
2630 .map(|repository| (
2631 repository.id.clone(),
2632 repository.relative_destination.clone()
2633 ))
2634 .collect::<Vec<_>>(),
2635 [("app".to_owned(), PathBuf::from("app"))]
2636 );
2637 assert!(matches!(
2638 layout.repositories[0].capture,
2639 CheckpointRepositoryCapture::RemoteWorkspace
2640 ));
2641 assert!(layout.managed_worktree.is_none());
2642
2643 let mut raw = session;
2646 raw.target_template_id = "local-bare".into();
2647 raw.project_directory = Some(PathBuf::from("/home/dev/project"));
2648 raw.target = Some(mj_core::state::TargetLocator::LocalBare {
2649 worker_root: PathBuf::from("/home/dev/.local/share/hel/workers/session"),
2650 });
2651 let mut state = State::default();
2652 state.sessions.insert(session_id.into(), raw);
2653 let controller = Controller { config, state };
2654
2655 let layout = controller
2656 .session_export_layout(session_id, &UnusedExecutor)
2657 .unwrap();
2658 assert_eq!(layout.workspace_root, "/home/dev");
2659 assert_eq!(layout.primary_repository, "project");
2660 assert_eq!(
2661 layout.repositories[0].relative_destination,
2662 PathBuf::from("project")
2663 );
2664 assert!(matches!(
2665 layout.repositories[0].capture,
2666 CheckpointRepositoryCapture::MetadataOnly
2667 ));
2668 }
2669
2670 #[test]
2671 fn startup_reconciliation_only_removes_unreferenced_controller_checkpoints() {
2672 let directory = tempfile::tempdir().unwrap();
2673 let session_id = "1123456789abcdef0123456789abcdef";
2674 let referenced_name =
2675 format!("{session_id}-7-archive-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.hel.zip");
2676 let orphan_name =
2677 format!("{session_id}-8-archive-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb.hel.zip");
2678 let imported_name = format!("{session_id}.hel.zip");
2679 for name in [
2680 &referenced_name,
2681 &orphan_name,
2682 &imported_name,
2683 "notes.hel.zip",
2684 ] {
2685 std::fs::write(directory.path().join(name), b"test").unwrap();
2686 }
2687 let mut state = State::default();
2688 let mut session = checkpoint_test_session(session_id);
2689 session.checkpoint = Some(CheckpointMetadata {
2690 archive_path: directory.path().join(&referenced_name),
2691 sha256: "c".repeat(64),
2692 created_at: "2026-08-12T00:00:00Z".into(),
2693 event_frontier: 7,
2694 });
2695 state.sessions.insert(session_id.into(), session);
2696
2697 assert_eq!(
2698 reconcile_managed_checkpoint_archives_in(directory.path(), &state).unwrap(),
2699 1
2700 );
2701 assert!(directory.path().join(referenced_name).exists());
2702 assert!(!directory.path().join(orphan_name).exists());
2703 assert!(directory.path().join(imported_name).exists());
2704 assert!(directory.path().join("notes.hel.zip").exists());
2705 }
2706 #[test]
2707 fn recovery_artifact_final_verification_checks_the_archive_digest() {
2708 let directory = tempfile::tempdir().unwrap();
2709 let session_id = "1123456789abcdef0123456789abcdef";
2710 let metadata = write_checkpoint_gate_archive(directory.path(), session_id, 7);
2711 let mut artifact = CheckpointArtifact {
2712 metadata,
2713 native_session_id: "native-session".into(),
2714 event_frontier_digest: "a".repeat(64),
2715 };
2716
2717 verify_checkpoint_artifact(session_id, &artifact).unwrap();
2718 artifact.metadata.sha256 = "b".repeat(64);
2719 assert!(
2720 verify_checkpoint_artifact(session_id, &artifact)
2721 .unwrap_err()
2722 .to_string()
2723 .contains("checkpoint SHA changed")
2724 );
2725 }
2726 fn checkpoint_barrier_snapshot(cursor: &RelayCursor) -> ManagedSessionSnapshot {
2729 let mut materialized = MaterializedSession::empty("session-1");
2730 materialized.applied_event_ordinal = cursor.ordinal;
2731 materialized.applied_event_digest = cursor.digest.clone();
2732 ManagedSessionSnapshot {
2733 subagent_requests: Vec::new(),
2734 subagent_results: Vec::new(),
2735 window: mj_core::state::ProjectionWindow::of(&materialized),
2736 materialized,
2737 latest_credential_sync_signal: None,
2738 worker_build: None,
2739 operational: mj_core::relay::RelayOperationalState {
2740 goal: serde_json::from_value(
2741 serde_json::json!({"known":true,"execution":{"version":1,"status":"idle"}}),
2742 )
2743 .unwrap(),
2744 capacity_retry: None,
2745 activity_turn_started_at_ms: None,
2746 checkpoint_only: false,
2747 acp_ready: None,
2748 store_id: None,
2749 idle_since_ms: None,
2750 session_id: "session-1".into(),
2751 execution: RelayExecutionState::Idle,
2752 latest_ordinal: cursor.ordinal,
2753 latest_digest: cursor.digest.clone(),
2754 acknowledged_through: cursor.ordinal,
2755 acknowledged_digest: cursor.digest.clone(),
2756 recovery_floor_ordinal: 0,
2757 recovery_floor_digest: mj_core::relay::RELAY_EVENT_GENESIS_DIGEST.into(),
2758 native_session_id: Some("native-session".into()),
2759 agent_capabilities: None,
2760 agent_info: None,
2761 steering_supported: None,
2762 config_options: Vec::new(),
2763 modes: None,
2764 available_commands: Vec::new(),
2765 config: BTreeMap::new(),
2766 active_prompt: None,
2767 queued_prompts: Vec::new(),
2768 active_user_shells: Vec::new(),
2769 active_agent_terminals: Vec::new(),
2770 checkpoint_barrier: Some("checkpoint-1".into()),
2771 checkpoint_ready: None,
2772 last_acp_activity_at_ms: None,
2773 current_step_started_at_ms: None,
2774 foreground_tool_started_at_ms: None,
2775 harness_turn: None,
2776 last_harness_turn_started_ordinal: None,
2777 background_commands: Vec::new(),
2778 background_work_known: None,
2779 },
2780 }
2781 }
2782 #[test]
2783 fn checkpoint_barrier_is_not_reached_until_its_ready_cursor_is_projected() {
2784 let cursor = RelayCursor {
2785 ordinal: 7,
2786 digest: "a".repeat(64),
2787 };
2788 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2789
2790 assert!(!checkpoint_barrier_is_ready(&snapshot, "checkpoint-1"));
2791 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2792 assert!(checkpoint_barrier_is_ready(&snapshot, "checkpoint-1"));
2793 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).unwrap();
2794 }
2795 #[test]
2796 fn checkpoint_revalidation_accepts_a_frontier_that_moved_past_the_ready_cursor() {
2797 let cursor = RelayCursor {
2798 ordinal: 7,
2799 digest: "a".repeat(64),
2800 };
2801 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2802 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2803
2804 snapshot.operational.latest_ordinal = cursor.ordinal + 2;
2808 snapshot.operational.latest_digest = "b".repeat(64);
2809 snapshot.materialized.applied_event_ordinal = cursor.ordinal + 2;
2810 snapshot.materialized.applied_event_digest = "b".repeat(64);
2811 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).unwrap();
2812
2813 snapshot.operational.checkpoint_ready = Some(RelayCursor {
2815 ordinal: cursor.ordinal + 1,
2816 digest: "c".repeat(64),
2817 });
2818 assert!(validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).is_err());
2819 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2820 snapshot.operational.checkpoint_barrier = None;
2821 assert!(validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).is_err());
2822 }
2823
2824 #[test]
2825 fn routine_kimi_checkpoint_defers_when_background_liveness_is_not_safe() {
2826 let cursor = RelayCursor {
2827 ordinal: 7,
2828 digest: "a".repeat(64),
2829 };
2830 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2831 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2832
2833 for (known, has_task, label) in [
2834 (Some(false), false, "tracker reported a failure"),
2835 (Some(true), true, "a native task is still active"),
2836 (None, false, "an older worker omitted the tracker field"),
2837 ] {
2838 snapshot.operational.background_work_known = known;
2839 snapshot.operational.background_commands = has_task
2840 .then(|| mj_core::relay::BackgroundCommand {
2841 id: "kimi:agent-1".into(),
2842 started_at_ms: 1,
2843 command: "background agent".into(),
2844 can_stop: false,
2845 })
2846 .into_iter()
2847 .collect();
2848 let error = validate_automatic_checkpoint_barrier_snapshot(
2849 &snapshot,
2850 "checkpoint-1",
2851 &cursor,
2852 HarnessKind::Kimi,
2853 )
2854 .expect_err(label);
2855 assert!(checkpoint_was_deferred(&error), "{label}: {error:#}");
2856 assert!(!checkpoint_barrier_needs_worker_restart(&error));
2857 }
2858
2859 snapshot.operational.background_work_known = None;
2863 snapshot.operational.background_commands = vec![mj_core::relay::BackgroundCommand {
2864 id: "legacy-task".into(),
2865 started_at_ms: 1,
2866 command: "legacy background work".into(),
2867 can_stop: false,
2868 }];
2869 validate_automatic_checkpoint_barrier_snapshot(
2870 &snapshot,
2871 "checkpoint-1",
2872 &cursor,
2873 HarnessKind::Codex,
2874 )
2875 .expect("non-Kimi checkpoint compatibility");
2876 }
2877 fn exported_checkpoint_json() -> Vec<u8> {
2879 serde_json::to_vec(&mj_checkpoint::checkpoint::TargetCheckpoint {
2880 path: PathBuf::from("/var/lib/hel/workers/session/checkpoint.hel.zip"),
2881 sha256: "c".repeat(64),
2882 event_frontier: 7,
2883 event_frontier_digest: "d".repeat(64),
2884 timings: None,
2885 })
2886 .unwrap()
2887 }
2888 fn export_spec_fixture() -> CheckpointExportSpec {
2889 CheckpointExportSpec {
2890 protocol_version: CHECKPOINT_EXPORT_PROTOCOL_VERSION,
2891 session: mj_checkpoint::archive::SessionManifest {
2892 id: LATCH_RELAY_SESSION.into(),
2893 title: "streamed spec".into(),
2894 harness_kind: mj_core::config::HarnessKind::Codex,
2895 profile_id: "codex".into(),
2896 native_session_id: "native-session".into(),
2897 created_at: "2026-08-12T00:00:00Z".into(),
2898 checkpointed_at: "2026-08-16T00:00:00Z".into(),
2899 hel_version: "test".into(),
2900 relay_version: "test".into(),
2901 adapter_version: "acp-v1".into(),
2902 },
2903 target: TargetManifest {
2904 template_id: "podman".into(),
2905 target_kind: "local-podman".into(),
2906 details: BTreeMap::new(),
2907 },
2908 bundle: BundleManifest {
2909 id: "project".into(),
2910 primary_repository: "app".into(),
2911 },
2912 relay_root: PathBuf::from("/var/lib/hel/workers/session"),
2913 harness_home: PathBuf::from("/var/lib/hel/profiles/codex"),
2914 workspace_root: PathBuf::from("/workspace"),
2915 repositories: Vec::new(),
2916 canonical_session: canonical_session_from_materialized(&MaterializedSession::empty(
2917 LATCH_RELAY_SESSION.to_owned(),
2918 ))
2919 .unwrap(),
2920 output_path: PathBuf::from("/var/lib/hel/workers/session/checkpoint.hel.zip"),
2921 }
2922 }
2923 struct ExportExecutor {
2926 streamed_status: i32,
2927 streamed_stderr: String,
2928 retry_stdin_after_failure: bool,
2929 stdin_calls: Cell<usize>,
2930 purposes: RefCell<Vec<String>>,
2931 streamed_spec: RefCell<Vec<u8>>,
2932 }
2933 impl ExportExecutor {
2934 fn new(streamed_status: i32, streamed_stderr: &str) -> Self {
2935 Self {
2936 streamed_status,
2937 streamed_stderr: streamed_stderr.to_owned(),
2938 retry_stdin_after_failure: false,
2939 stdin_calls: Cell::new(0),
2940 purposes: RefCell::new(Vec::new()),
2941 streamed_spec: RefCell::new(Vec::new()),
2942 }
2943 }
2944
2945 fn retry_stdin_after_failure(mut self) -> Self {
2946 self.retry_stdin_after_failure = true;
2947 self
2948 }
2949 }
2950 impl CommandExecutor for ExportExecutor {
2951 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
2952 self.purposes.borrow_mut().push(command.purpose.clone());
2953 Ok(CommandOutput {
2954 status: 0,
2955 stdout: exported_checkpoint_json(),
2956 stderr: Vec::new(),
2957 })
2958 }
2959
2960 fn execute_with_stdin(
2961 &self,
2962 command: &CommandSpec,
2963 input: &mut (dyn std::io::Read + Send),
2964 ) -> Result<CommandOutput> {
2965 self.purposes.borrow_mut().push(command.purpose.clone());
2966 let mut spec = Vec::new();
2967 input.read_to_end(&mut spec)?;
2968 *self.streamed_spec.borrow_mut() = spec;
2969 let attempt = self.stdin_calls.get();
2970 self.stdin_calls.set(attempt + 1);
2971 let failed =
2972 self.streamed_status != 0 && (attempt == 0 || !self.retry_stdin_after_failure);
2973 Ok(CommandOutput {
2974 status: if failed { self.streamed_status } else { 0 },
2975 stdout: if failed {
2976 Vec::new()
2977 } else {
2978 exported_checkpoint_json()
2979 },
2980 stderr: if failed {
2981 self.streamed_stderr.clone().into_bytes()
2982 } else {
2983 Vec::new()
2984 },
2985 })
2986 }
2987 }
2988 #[test]
2989 fn docker_checkpoint_fallback_upload_uses_docker_cp() {
2990 struct RecordingExecutor {
2991 commands: RefCell<Vec<CommandSpec>>,
2992 }
2993 impl CommandExecutor for RecordingExecutor {
2994 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
2995 self.commands.borrow_mut().push(command.clone());
2996 Ok(CommandOutput {
2997 status: 0,
2998 stdout: Vec::new(),
2999 stderr: Vec::new(),
3000 })
3001 }
3002 }
3003
3004 let executor = RecordingExecutor {
3005 commands: RefCell::new(Vec::new()),
3006 };
3007 let locator = targets::TargetLocator::LocalDocker {
3008 container_id: "hel-session-12345678".to_owned(),
3009 };
3010 upload_checkpoint_spec(
3011 &executor,
3012 &locator,
3013 LATCH_RELAY_SESSION,
3014 Path::new("checkpoint-spec.json"),
3015 "/var/lib/hel/workers/session/checkpoint-spec.json",
3016 )
3017 .unwrap();
3018
3019 let commands = executor.commands.borrow();
3020 assert_eq!(commands.len(), 1);
3021 assert_eq!(commands[0].program, "docker");
3022 assert_eq!(
3023 commands[0].args,
3024 [
3025 "cp",
3026 "checkpoint-spec.json",
3027 "hel-session-12345678:/var/lib/hel/workers/session/checkpoint-spec.json"
3028 ]
3029 );
3030 assert_eq!(commands[0].purpose, "upload checkpoint specification");
3031 }
3032 #[test]
3033 fn checkpoint_export_streams_its_spec_instead_of_uploading_it() {
3034 let locator = targets::TargetLocator::LocalPodman {
3035 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3036 workspace_storage: Default::default(),
3037 };
3038 let spec = export_spec_fixture();
3039 let executor = ExportExecutor::new(0, "");
3040
3041 let output = export_target_checkpoint(
3042 &executor,
3043 &locator,
3044 LATCH_RELAY_SESSION,
3045 &spec,
3046 "/var/lib/hel/workers/session/checkpoint-spec.json",
3047 )
3048 .unwrap();
3049
3050 assert_eq!(output.stdout, exported_checkpoint_json());
3051 assert_eq!(
3052 serde_json::from_slice::<CheckpointExportSpec>(&executor.streamed_spec.borrow())
3053 .unwrap(),
3054 spec
3055 );
3056 assert_eq!(
3057 executor.purposes.into_inner(),
3058 vec!["export target checkpoint".to_owned()]
3059 );
3060 }
3061 #[test]
3064 fn an_export_that_cannot_read_stdin_falls_back_to_uploading_the_spec() {
3065 let locator = targets::TargetLocator::LocalPodman {
3066 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3067 workspace_storage: Default::default(),
3068 };
3069 let executor = ExportExecutor::new(
3070 1,
3071 "Error: read checkpoint export spec -\n\nCaused by:\n \
3072 No such file or directory (os error 2)\n",
3073 );
3074
3075 let output = export_target_checkpoint(
3076 &executor,
3077 &locator,
3078 LATCH_RELAY_SESSION,
3079 &export_spec_fixture(),
3080 "/var/lib/hel/workers/session/checkpoint-spec.json",
3081 )
3082 .unwrap();
3083
3084 assert_eq!(output.stdout, exported_checkpoint_json());
3085 assert_eq!(
3086 executor.purposes.into_inner(),
3087 vec![
3088 "export target checkpoint".to_owned(),
3089 "upload checkpoint specification".to_owned(),
3090 "export target checkpoint".to_owned(),
3091 "remove uploaded checkpoint specification".to_owned(),
3092 ]
3093 );
3094 }
3095 #[test]
3096 fn a_failing_export_is_not_retried_as_an_old_worker() {
3097 let locator = targets::TargetLocator::LocalPodman {
3098 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3099 workspace_storage: Default::default(),
3100 };
3101 let executor = ExportExecutor::new(1, "Error: repository 'app' is missing\n");
3102
3103 let error = export_target_checkpoint(
3104 &executor,
3105 &locator,
3106 LATCH_RELAY_SESSION,
3107 &export_spec_fixture(),
3108 "/var/lib/hel/workers/session/checkpoint-spec.json",
3109 )
3110 .unwrap_err();
3111
3112 assert!(
3113 format!("{error:#}").contains("repository 'app' is missing"),
3114 "{error:#}"
3115 );
3116 assert_eq!(
3117 executor.purposes.into_inner(),
3118 vec!["export target checkpoint".to_owned()]
3119 );
3120 }
3121 #[test]
3124 fn a_legacy_export_worker_is_replaced_before_it_runs_obsolete_behavior() {
3125 let locator = targets::TargetLocator::LocalPodman {
3126 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3127 workspace_storage: Default::default(),
3128 };
3129 let spec = export_spec_fixture();
3130 let executor = ExportExecutor::new(
3131 1,
3132 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3133 unknown field `protocol_version`, expected `session` at line 1 column 20\n",
3134 )
3135 .retry_stdin_after_failure();
3136 let worker_binary = Path::new("/hel-test-worker");
3137
3138 let output = export_target_checkpoint_with_worker(
3139 &executor,
3140 &locator,
3141 LATCH_RELAY_SESSION,
3142 &spec,
3143 "/var/lib/hel/workers/session/checkpoint-spec.json",
3144 Some(worker_binary),
3145 )
3146 .unwrap();
3147
3148 assert_eq!(output.stdout, exported_checkpoint_json());
3149 assert_eq!(
3150 serde_json::from_slice::<CheckpointExportSpec>(&executor.streamed_spec.borrow())
3151 .unwrap(),
3152 spec
3153 );
3154 assert_eq!(
3155 executor.purposes.into_inner(),
3156 vec![
3157 "export target checkpoint".to_owned(),
3158 "stage replacement Mjolnir worker".to_owned(),
3159 "assign replacement worker to the worker user".to_owned(),
3160 "replace installed Mjolnir worker".to_owned(),
3161 "make replaced Mjolnir worker executable".to_owned(),
3162 "export target checkpoint".to_owned(),
3163 ]
3164 );
3165 }
3166 #[test]
3167 fn a_schema_mismatch_after_uploading_the_spec_still_replaces_the_worker_binary() {
3168 let locator = targets::TargetLocator::LocalPodman {
3169 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3170 workspace_storage: Default::default(),
3171 };
3172 struct FileThenRefreshExecutor {
3173 purposes: RefCell<Vec<String>>,
3174 file_export_calls: Cell<usize>,
3175 }
3176 impl CommandExecutor for FileThenRefreshExecutor {
3177 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
3178 self.purposes.borrow_mut().push(command.purpose.clone());
3179 if command.purpose == "export target checkpoint" {
3180 let attempt = self.file_export_calls.get();
3181 self.file_export_calls.set(attempt + 1);
3182 if attempt == 0 {
3183 return Ok(CommandOutput {
3184 status: 1,
3185 stdout: Vec::new(),
3186 stderr: b"Error: parse checkpoint export spec /spec.json\n\nCaused by:\n unknown variant `terminal_output`, expected one of `user`, `agent`, `thought`, `tool`, `plan`, `system`\n".to_vec(),
3187 });
3188 }
3189 }
3190 Ok(CommandOutput {
3191 status: 0,
3192 stdout: exported_checkpoint_json(),
3193 stderr: Vec::new(),
3194 })
3195 }
3196
3197 fn execute_with_stdin(
3198 &self,
3199 command: &CommandSpec,
3200 input: &mut (dyn std::io::Read + Send),
3201 ) -> Result<CommandOutput> {
3202 self.purposes.borrow_mut().push(command.purpose.clone());
3203 let mut discarded = Vec::new();
3204 input.read_to_end(&mut discarded)?;
3205 let stdin_calls = self
3206 .purposes
3207 .borrow()
3208 .iter()
3209 .filter(|purpose| *purpose == "export target checkpoint")
3210 .count();
3211 if stdin_calls == 1 {
3212 return Ok(CommandOutput {
3213 status: 1,
3214 stdout: Vec::new(),
3215 stderr: b"Error: read checkpoint export spec -\n\nCaused by:\n No such file or directory (os error 2)\n".to_vec(),
3216 });
3217 }
3218 Ok(CommandOutput {
3219 status: 0,
3220 stdout: exported_checkpoint_json(),
3221 stderr: Vec::new(),
3222 })
3223 }
3224 }
3225
3226 let executor = FileThenRefreshExecutor {
3227 purposes: RefCell::new(Vec::new()),
3228 file_export_calls: Cell::new(0),
3229 };
3230 let output = export_target_checkpoint_with_worker(
3231 &executor,
3232 &locator,
3233 LATCH_RELAY_SESSION,
3234 &export_spec_fixture(),
3235 "/var/lib/hel/workers/session/checkpoint-spec.json",
3236 Some(Path::new("/hel-test-worker")),
3237 )
3238 .unwrap();
3239
3240 assert_eq!(output.stdout, exported_checkpoint_json());
3241 assert_eq!(
3242 executor.purposes.into_inner(),
3243 vec![
3244 "export target checkpoint".to_owned(),
3245 "upload checkpoint specification".to_owned(),
3246 "export target checkpoint".to_owned(),
3247 "stage replacement Mjolnir worker".to_owned(),
3248 "assign replacement worker to the worker user".to_owned(),
3249 "replace installed Mjolnir worker".to_owned(),
3250 "make replaced Mjolnir worker executable".to_owned(),
3251 "export target checkpoint".to_owned(),
3252 ]
3253 );
3254 }
3255 #[test]
3259 fn a_working_session_defers_but_close_waits_for_cancellation_before_recovery() {
3260 let cursor = RelayCursor {
3261 ordinal: 7,
3262 digest: "a".repeat(64),
3263 };
3264 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
3265 snapshot.operational.execution = RelayExecutionState::Running;
3266
3267 let deferred = checkpoint_barrier_wait_ended(
3268 &snapshot,
3269 "checkpoint-1",
3270 BarrierBusyPolicy::DeferWhileRunning,
3271 false,
3272 false,
3273 )
3274 .expect("a working session ends the wait at once");
3275 assert!(checkpoint_was_deferred(&deferred), "{deferred:#}");
3276 assert!(
3277 !checkpoint_barrier_needs_worker_restart(&deferred),
3278 "a deferred copy must never restart the worker: {deferred:#}"
3279 );
3280 assert_eq!(
3281 BarrierBusyPolicy::of(LatchExclusivity::HoldThroughClose),
3282 BarrierBusyPolicy::InterruptWhileRunning
3283 );
3284
3285 assert!(
3288 checkpoint_barrier_wait_ended(
3289 &snapshot,
3290 "checkpoint-1",
3291 BarrierBusyPolicy::InterruptWhileRunning,
3292 false,
3293 false,
3294 )
3295 .is_none()
3296 );
3297 let interrupted = checkpoint_barrier_wait_ended(
3298 &snapshot,
3299 "checkpoint-1",
3300 BarrierBusyPolicy::InterruptWhileRunning,
3301 true,
3302 true,
3303 )
3304 .expect("an unresponsive cancellation ends the wait at the deadline");
3305 assert!(
3306 checkpoint_barrier_needs_worker_restart(&interrupted),
3307 "{interrupted:#}"
3308 );
3309 assert!(!checkpoint_was_deferred(&interrupted), "{interrupted:#}");
3310
3311 snapshot.operational.execution = RelayExecutionState::Idle;
3314 let wedged = checkpoint_barrier_wait_ended(
3315 &snapshot,
3316 "checkpoint-1",
3317 BarrierBusyPolicy::DeferWhileRunning,
3318 true,
3319 false,
3320 )
3321 .expect("the deadline ends the wait");
3322 assert!(
3323 checkpoint_barrier_needs_worker_restart(&wedged),
3324 "{wedged:#}"
3325 );
3326 assert!(!checkpoint_was_deferred(&wedged), "{wedged:#}");
3327 }
3328
3329 #[test]
3332 fn a_frontier_that_moved_before_the_latch_defers_the_checkpoint() {
3333 let cursor = RelayCursor {
3334 ordinal: 220,
3335 digest: "a".repeat(64),
3336 };
3337 ensure_exact_checkpoint_cut(&cursor, cursor.ordinal, &cursor.digest)
3338 .expect("a projection latched at the ready cursor is an exact cut");
3339
3340 for (ordinal, digest) in [(223, "a".repeat(64)), (220, "b".repeat(64))] {
3341 let error = ensure_exact_checkpoint_cut(&cursor, ordinal, &digest)
3342 .expect_err("a projection past the ready cursor is not an exact cut");
3343 assert!(checkpoint_was_deferred(&error), "{error:#}");
3344 assert!(
3345 !checkpoint_barrier_needs_worker_restart(&error),
3346 "{error:#}"
3347 );
3348 }
3349 }
3350
3351 #[test]
3355 fn a_harness_turn_started_during_capture_abandons_the_archive() {
3356 let cursor = RelayCursor {
3357 ordinal: 220,
3358 digest: "a".repeat(64),
3359 };
3360 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
3361 snapshot.operational.checkpoint_ready = Some(cursor.clone());
3362
3363 snapshot.operational.last_harness_turn_started_ordinal = Some(cursor.ordinal);
3364 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor)
3365 .expect("a turn that started at or before the cursor is covered by the archive");
3366
3367 snapshot.operational.last_harness_turn_started_ordinal = Some(cursor.ordinal + 1);
3368 let error = validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor)
3369 .expect_err("a turn that started after the cursor invalidates the capture");
3370 assert!(checkpoint_was_deferred(&error), "{error:#}");
3371 }
3372
3373 #[test]
3374 fn a_stuck_checkpoint_barrier_is_retried_by_restarting_the_worker() {
3375 for failure in [
3378 CheckpointBarrierUnreachable::not_admitted(
3379 "checkpoint-976f6746887c5ccd93b9d8bbe120ef06",
3380 ),
3381 CheckpointBarrierUnreachable::runtime_stopped(),
3382 ] {
3383 let error = anyhow::Error::new(failure).context("latch a session checkpoint");
3384 assert!(checkpoint_barrier_needs_worker_restart(&error), "{error:#}");
3385 }
3386 assert!(!checkpoint_barrier_needs_worker_restart(&anyhow::anyhow!(
3387 "export target checkpoint failed with status 1"
3388 )));
3389 assert!(!checkpoint_barrier_needs_worker_restart(&anyhow::anyhow!(
3392 "ACP relay did not reach checkpoint barrier checkpoint-1"
3393 )));
3394 }
3395
3396 #[test]
3397 fn an_incompatible_cancel_turn_requests_worker_recovery() {
3398 let error = anyhow::Error::new(RelayRejected(mj_core::relay::RelayProtocolError {
3399 code: mj_core::relay::RelayErrorCode::IncompatibleProtocol,
3400 message: "request uses protocol 6".into(),
3401 retryable: false,
3402 detail: None,
3403 }))
3404 .context("cancel active ACP turn before checkpoint barrier");
3405 assert!(
3406 checkpoint_cancel_turn_needs_worker_restart(&error),
3407 "{error:#}"
3408 );
3409 assert!(checkpoint_barrier_needs_worker_restart(&error.context(
3410 CheckpointBarrierUnreachable::cancel_turn_unavailable("checkpoint-1", 6,)
3411 )));
3412 }
3413 #[test]
3414 fn a_dead_worker_hello_failure_is_retried_by_restarting_the_worker() {
3415 let dead = anyhow::Error::new(RelayTransportDead::new("the proxy is gone"))
3416 .context("connect to the session worker for checkpoint");
3417 assert!(worker_connect_needs_restart(&dead), "{dead:#}");
3418 assert!(!worker_connect_needs_restart(&anyhow::anyhow!(
3419 "unknown session"
3420 )));
3421 }
3422 #[cfg(unix)]
3423 #[tokio::test]
3424 async fn checkpoint_restart_stop_failure_names_mjolnir() {
3425 struct FailingStop;
3426
3427 impl CommandExecutor for FailingStop {
3428 fn execute(&self, _command: &CommandSpec) -> Result<CommandOutput> {
3429 Ok(CommandOutput {
3430 status: 1,
3431 stdout: Vec::new(),
3432 stderr: b"permission denied".to_vec(),
3433 })
3434 }
3435 }
3436
3437 let session_id = "0123456789abcdef0123456789abcdef";
3438 let worker_root = format!("/tmp/mjolnir-checkpoint-test/{session_id}");
3439 let backend = targets::TargetLocator::LocalBare {
3440 worker_root: worker_root.clone(),
3441 };
3442 let controller = Controller {
3443 config: Config::default(),
3444 state: State::default(),
3445 };
3446 let reconnect = CommandSpec::new("unused", std::iter::empty::<&str>());
3447
3448 let result = controller
3449 .restart_worker_for_checkpoint(
3450 session_id,
3451 &FailingStop,
3452 &backend,
3453 &worker_root,
3454 &reconnect,
3455 )
3456 .await;
3457 let error = match result {
3458 Ok(_) => panic!("a failed worker stop unexpectedly restarted the checkpoint worker"),
3459 Err(error) => error,
3460 };
3461 let detail = format!("{error:#}");
3462 assert!(
3463 detail.starts_with("stop wedged Mjolnir worker before retrying checkpoint"),
3464 "{detail}"
3465 );
3466 assert!(detail.contains("permission denied"), "{detail}");
3467 }
3468 #[test]
3469 fn export_spec_schema_mismatch_is_detected_from_the_parse_error() {
3470 assert!(export_spec_schema_unsupported(
3471 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3472 unknown field `terminal_refs`, expected `call` at line 1 column 7276552\n"
3473 ));
3474 assert!(export_spec_schema_unsupported(
3475 "Error: parse checkpoint export spec /spec.json\n\nCaused by:\n \
3476 unknown variant `terminal_output`, expected one of `user`, `agent`\n"
3477 ));
3478 assert!(!export_spec_schema_unsupported(
3479 "Error: repository 'app' is missing\n"
3480 ));
3481 assert!(!export_spec_schema_unsupported(
3482 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3483 missing field `relay_root`\n"
3484 ));
3485 assert!(export_protocol_unsupported(
3486 "Error: unsupported checkpoint export protocol version 3; worker supports 2\n"
3487 ));
3488 }
3489 const LATCH_RELAY_ROOT: &str = "MJ_TEST_LATCH_RELAY_ROOT";
3490 const LATCH_RELAY_STARTS: &str = "MJ_TEST_LATCH_RELAY_STARTS";
3491 const LATCH_RELAY_REJECT_RELEASE: &str = "MJ_TEST_LATCH_REJECT_RELEASE";
3492 #[cfg(unix)]
3493 const LATCH_RELAY_RUNNING: &str = "MJ_TEST_LATCH_RELAY_RUNNING";
3494 #[cfg(unix)]
3495 const LATCH_TEST_CHILD: &str = "MJ_TEST_LATCH_CHILD";
3496 #[cfg(unix)]
3497 const ABANDON_TEST_CHILD: &str = "MJ_TEST_ABANDON_LATCH_CHILD";
3498 #[cfg(unix)]
3499 const RELEASE_TEST_CHILD: &str = "MJ_TEST_RELEASE_LATCH_CHILD";
3500 #[cfg(unix)]
3501 const LEGACY_RELEASE_TEST_CHILD: &str = "MJ_TEST_LEGACY_RELEASE_LATCH_CHILD";
3502 #[cfg(unix)]
3503 const REUSE_TEST_CHILD: &str = "MJ_TEST_REUSE_LATCH_CHILD";
3504 const LATCH_CHECKPOINT_ONLY: &str = "MJ_TEST_LATCH_CHECKPOINT_ONLY";
3505 const LATCH_RELAY_STARTUP_DELAY_MS: &str = "MJ_TEST_LATCH_STARTUP_DELAY_MS";
3506 const LATCH_RELAY_SESSION: &str = "018f9dd2-a3b4-7c8d-9000-0123456789ab";
3507 #[cfg(unix)]
3509 #[derive(Clone, Copy, PartialEq, Eq)]
3510 enum ReleaseSupport {
3511 Supported,
3512 Rejected,
3515 }
3516 #[test]
3523 fn latch_relay_child_serves_stdio() {
3524 let Some(root) = std::env::var_os(LATCH_RELAY_ROOT) else {
3525 return;
3526 };
3527 println!();
3531 if let Some(starts) = std::env::var_os(LATCH_RELAY_STARTS) {
3534 use std::io::Write;
3535 let mut log = OpenOptions::new()
3536 .create(true)
3537 .append(true)
3538 .open(starts)
3539 .expect("open the relay start log");
3540 writeln!(log, "{}", std::process::id()).expect("record this relay start");
3541 }
3542 let checkpoint_only = std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some();
3543 let mut relay = if checkpoint_only {
3544 mj_worker::relay::DurableRelay::open_for_checkpoint(
3545 Path::new(&root),
3546 LATCH_RELAY_SESSION,
3547 "1.0.0",
3548 )
3549 } else {
3550 mj_worker::relay::DurableRelay::open(Path::new(&root), LATCH_RELAY_SESSION, "1.0.0")
3551 }
3552 .expect("open the test relay journal");
3553 if relay.operational_state().native_session_id.is_none() {
3554 relay
3555 .record_observation(mj_core::relay::RelayObservation::SessionOpened {
3556 native_session_id: "native-session".into(),
3557 resumed: true,
3558 })
3559 .unwrap();
3560 }
3561 if !relay.operational_state().goal.synchronized() && !checkpoint_only {
3562 relay.record_session_update(serde_json::from_value(serde_json::json!({
3563 "sessionUpdate":"session_info_update", "_meta":{"goal":null,"execution":{"version":1,"status":"idle"}}
3564 })).unwrap()).unwrap();
3565 }
3566 let ready_at = Instant::now()
3567 + Duration::from_millis(
3568 std::env::var(LATCH_RELAY_STARTUP_DELAY_MS)
3569 .ok()
3570 .map(|value| value.parse::<u64>().unwrap())
3571 .unwrap_or(0),
3572 );
3573 let reject_release = std::env::var_os(LATCH_RELAY_REJECT_RELEASE).is_some();
3574 #[cfg(unix)]
3575 let running = std::env::var_os(LATCH_RELAY_RUNNING).is_some();
3576 #[cfg(unix)]
3577 if running && relay.operational_state().active_prompt.is_none() {
3578 let response = relay.handle(mj_core::relay::RelayRequestEnvelope {
3579 request_id: "seed-running-request".into(),
3580 protocol_version: mj_core::relay::RELAY_PROTOCOL_VERSION,
3581 request: mj_core::relay::RelayRequest::Submit {
3582 command_id: "seed-running-prompt".into(),
3583 command: RelayCommand::Prompt {
3584 prompt: vec![ContentBlock::Text(TextContent::new("running"))],
3585 },
3586 },
3587 });
3588 assert!(matches!(
3589 response.body,
3590 mj_core::relay::RelayResponseBody::Ok {
3591 payload: mj_core::relay::RelayResponsePayload::Accepted { .. }
3592 }
3593 ));
3594 let claimed = relay
3595 .claim_pending_commands(true)
3596 .expect("seed the running prompt");
3597 assert_eq!(claimed.len(), 1);
3598 assert_eq!(claimed[0].command_id, "seed-running-prompt");
3599 }
3600 let mut reader = std::io::stdin().lock();
3601 let mut writer = std::io::stdout().lock();
3602 let mut configured = false;
3603 while let Some(request) =
3604 mj_core::relay::read_relay_frame(&mut reader).expect("read a relay request")
3605 {
3606 if !checkpoint_only && !configured && Instant::now() >= ready_at {
3607 relay
3608 .record_observation(mj_core::relay::RelayObservation::SessionConfigured {
3609 config_options: Vec::new(),
3610 })
3611 .unwrap();
3612 configured = true;
3613 }
3614 if matches!(
3615 &request.request,
3616 mj_core::relay::RelayRequest::Submit {
3617 command: RelayCommand::BeginCheckpoint { .. },
3618 ..
3619 }
3620 ) {
3621 assert!(
3622 checkpoint_only || relay.operational_state().native_session_is_ready(),
3623 "checkpoint submitted before current ACP startup finished"
3624 );
3625 }
3626 let response = if reject_release && requests_checkpoint_release(&request) {
3627 unparseable_request_response(&request)
3628 } else {
3629 relay.handle(request)
3630 };
3631 mj_core::relay::write_relay_frame(&mut writer, &response)
3632 .expect("answer a relay request");
3633 if checkpoint_only {
3634 relay.dispatch_checkpoint_only().unwrap();
3635 }
3636 for claimed in relay
3637 .claim_pending_commands(true)
3638 .expect("claim relay commands")
3639 {
3640 match claimed.command {
3641 RelayCommand::BeginCheckpoint { .. } => {
3642 relay
3643 .record_checkpoint_ready(&claimed.command_id)
3644 .expect("report the checkpoint barrier ready");
3645 }
3646 #[cfg(unix)]
3647 RelayCommand::CancelTurn => {
3648 let prompt_id = relay
3649 .operational_state()
3650 .active_prompt
3651 .as_ref()
3652 .map(|prompt| prompt.command_id.clone())
3653 .expect("a prompt to cancel");
3654 relay
3655 .record_command_completed(
3656 &claimed.command_id,
3657 RelayCommandOutcome::Cancelled,
3658 )
3659 .expect("complete the cancellation");
3660 relay
3661 .record_command_completed(
3662 &prompt_id,
3663 RelayCommandOutcome::Prompt {
3664 diagnostic: None,
3665 stop_reason: "cancelled".into(),
3666 usage: None,
3667 },
3668 )
3669 .expect("complete the cancelled prompt");
3670 }
3671 _ => {}
3672 }
3673 }
3674 }
3675 }
3676 fn requests_checkpoint_release(request: &mj_core::relay::RelayRequestEnvelope) -> bool {
3677 matches!(
3678 &request.request,
3679 mj_core::relay::RelayRequest::Submit {
3680 command: RelayCommand::ReleaseCheckpoint { .. },
3681 ..
3682 }
3683 )
3684 }
3685 fn unparseable_request_response(
3689 request: &mj_core::relay::RelayRequestEnvelope,
3690 ) -> mj_core::relay::RelayResponseEnvelope {
3691 mj_core::relay::RelayResponseEnvelope {
3692 request_id: request.request_id.clone(),
3693 protocol_version: request.protocol_version,
3694 body: mj_core::relay::RelayResponseBody::Error {
3695 error: mj_core::relay::RelayProtocolError {
3696 code: mj_core::relay::RelayErrorCode::InvalidRequest,
3697 message: "unknown variant `release_checkpoint`".into(),
3698 retryable: false,
3699 detail: None,
3700 },
3701 },
3702 }
3703 }
3704 #[cfg(unix)]
3707 fn latch_relay_target(
3708 relay_root: &Path,
3709 starts: Option<&Path>,
3710 release: ReleaseSupport,
3711 running: bool,
3712 ) -> crate::session_manager::RelaySessionTarget {
3713 let script = format!(
3716 "\"$0\" --exact {}::latch_relay_child_serves_stdio --nocapture | \
3717 grep --line-buffered '^{{'",
3718 module_path!()
3719 .strip_prefix("mj_controller::")
3720 .unwrap_or(module_path!())
3721 );
3722 let mut spec = CommandSpec::new(
3723 "sh",
3724 [
3725 "-c".to_owned(),
3726 script,
3727 std::env::current_exe()
3728 .unwrap()
3729 .to_string_lossy()
3730 .into_owned(),
3731 ],
3732 )
3733 .purpose("test latch relay");
3734 spec.env.insert(
3735 LATCH_RELAY_ROOT.to_owned(),
3736 relay_root.to_string_lossy().into_owned(),
3737 );
3738 if let Some(starts) = starts {
3739 spec.env.insert(
3740 LATCH_RELAY_STARTS.to_owned(),
3741 starts.to_string_lossy().into_owned(),
3742 );
3743 }
3744 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
3745 spec.env.insert(LATCH_CHECKPOINT_ONLY.into(), "1".into());
3746 }
3747 if release == ReleaseSupport::Rejected {
3748 spec.env
3749 .insert(LATCH_RELAY_REJECT_RELEASE.to_owned(), "1".to_owned());
3750 }
3751 if running {
3752 spec.env
3753 .insert(LATCH_RELAY_RUNNING.to_owned(), "1".to_owned());
3754 }
3755 crate::session_manager::RelaySessionTarget {
3756 session_id: LATCH_RELAY_SESSION.to_owned(),
3757 spec,
3758 worker_recovery: None,
3759 project_memory: None,
3760 }
3761 }
3762 #[cfg(unix)]
3765 async fn latch_a_live_checkpoint(
3766 relay_root: &Path,
3767 starts: Option<&Path>,
3768 release: ReleaseSupport,
3769 running: bool,
3770 ) -> (
3771 crate::session_manager::SessionManagerChannels,
3772 ManagedSessionHandle,
3773 ControllerRelayLease,
3774 String,
3775 RelayCursor,
3776 ) {
3777 crate::database::save_session(&checkpoint_test_session(LATCH_RELAY_SESSION)).unwrap();
3780 let channels = crate::session_manager::spawn_session_manager().unwrap();
3781 channels
3782 .targets
3783 .send(vec![latch_relay_target(
3784 relay_root, starts, release, running,
3785 )])
3786 .unwrap();
3787 let handle = channels
3788 .control
3789 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
3790 .await
3791 .unwrap();
3792
3793 let lease = handle.lease_connection().await.unwrap();
3794 let mut relay = ControllerRelayLease::Managed {
3795 handle: handle.clone(),
3796 lease: Some(lease),
3797 };
3798 let barrier_command_id = new_command_id("checkpoint").unwrap();
3799 let connection = relay.connection_mut();
3800 connection
3801 .submit(
3802 barrier_command_id.clone(),
3803 RelayCommand::BeginCheckpoint { reason: None },
3804 )
3805 .await
3806 .unwrap();
3807 let barrier = wait_for_checkpoint_barrier(
3808 connection,
3809 LATCH_RELAY_SESSION,
3810 &barrier_command_id,
3811 CHECKPOINT_BARRIER_TIMEOUT,
3812 BarrierBusyPolicy::InterruptWhileRunning,
3813 HarnessKind::Codex,
3814 )
3815 .await
3816 .unwrap();
3817 assert_eq!(
3818 barrier.materialized.applied_event_ordinal,
3819 barrier.operational.latest_ordinal
3820 );
3821 let cursor = barrier.operational.checkpoint_ready.clone().unwrap();
3822 (channels, handle, relay, barrier_command_id, cursor)
3823 }
3824
3825 #[cfg(unix)]
3830 #[tokio::test]
3831 async fn a_close_checkpoint_cancels_a_running_turn_without_restarting_the_worker() {
3832 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
3835 let directory = tempfile::tempdir().unwrap();
3836 let test_name = format!(
3837 "{}::a_close_checkpoint_cancels_a_running_turn_without_restarting_the_worker",
3838 module_path!()
3839 .strip_prefix("mj_controller::")
3840 .unwrap_or(module_path!())
3841 );
3842 let output = Command::new(std::env::current_exe().unwrap())
3843 .args(["--exact", &test_name, "--nocapture"])
3844 .env(LATCH_TEST_CHILD, "1")
3845 .env("MJ_DATA_DIR", directory.path())
3846 .output()
3847 .unwrap();
3848 assert!(
3849 output.status.success(),
3850 "isolated cancellation checkpoint test failed\nstdout:\n{}\nstderr:\n{}",
3851 String::from_utf8_lossy(&output.stdout),
3852 String::from_utf8_lossy(&output.stderr)
3853 );
3854 return;
3855 }
3856 let _writer = crate::database::install_isolated_test_writer();
3857 let relay_root = tempfile::tempdir().unwrap();
3858 let start_log_directory = tempfile::tempdir().unwrap();
3859 let start_log = start_log_directory.path().join("relay-starts");
3860 let (_channels, _handle, mut relay, _barrier_command_id, _cursor) =
3861 latch_a_live_checkpoint(
3862 relay_root.path(),
3863 Some(&start_log),
3864 ReleaseSupport::Supported,
3865 true,
3866 )
3867 .await;
3868 let snapshot = relay.sync_snapshot().await.unwrap();
3869 assert_eq!(
3870 snapshot.operational.execution,
3871 RelayExecutionState::Idle,
3872 "the close wait returned before the cancelled turn became idle"
3873 );
3874 assert!(
3875 snapshot.operational.active_prompt.is_none(),
3876 "the close wait returned before the cancelled prompt settled"
3877 );
3878 assert_eq!(
3879 relay_starts(&start_log),
3880 1,
3881 "responsive cancellation restarted worker"
3882 );
3883 }
3884 #[cfg(unix)]
3887 async fn wait_until_the_actor_serves_again(handle: &ManagedSessionHandle) {
3888 for attempt in 0.. {
3889 if handle.sync_now().await.is_ok() {
3890 return;
3891 }
3892 assert!(attempt < 200, "the actor never took its connection back");
3893 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
3894 }
3895 }
3896 #[cfg(unix)]
3900 #[tokio::test]
3901 async fn ending_the_checkpoint_latch_returns_the_connection_to_its_actor() {
3902 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
3905 let directory = tempfile::tempdir().unwrap();
3906 let test_name = format!(
3907 "{}::ending_the_checkpoint_latch_returns_the_connection_to_its_actor",
3908 module_path!()
3909 .strip_prefix("mj_controller::")
3910 .unwrap_or(module_path!())
3911 );
3912 let output = Command::new(std::env::current_exe().unwrap())
3913 .args(["--exact", &test_name, "--nocapture"])
3914 .env(LATCH_TEST_CHILD, "1")
3915 .env("MJ_DATA_DIR", directory.path())
3916 .output()
3917 .unwrap();
3918 assert!(
3919 output.status.success(),
3920 "isolated checkpoint latch test failed\nstdout:\n{}\nstderr:\n{}",
3921 String::from_utf8_lossy(&output.stdout),
3922 String::from_utf8_lossy(&output.stderr)
3923 );
3924 return;
3925 }
3926 let _writer = crate::database::install_isolated_test_writer();
3928
3929 std::thread::spawn(|| {
3932 std::thread::sleep(std::time::Duration::from_secs(120));
3933 eprintln!("the checkpoint latch never returned its connection");
3934 std::process::exit(101);
3935 });
3936
3937 let relay_root = tempfile::tempdir().unwrap();
3938 let (_channels, handle, mut relay, barrier_command_id, cursor) =
3939 latch_a_live_checkpoint(relay_root.path(), None, ReleaseSupport::Supported, false)
3940 .await;
3941
3942 assert!(
3945 handle.sync_now().await.is_err(),
3946 "a latched projection must not be advanced by its own actor"
3947 );
3948
3949 relay.end_latch();
3950 wait_until_the_actor_serves_again(&handle).await;
3951
3952 let latched = relay.sync_snapshot().await.unwrap();
3956 validate_checkpoint_barrier_snapshot(&latched, &barrier_command_id, &cursor).unwrap();
3957
3958 let prompt_ordinal = relay
3961 .submit(
3962 new_command_id("prompt").unwrap(),
3963 RelayCommand::Prompt {
3964 prompt: vec![ContentBlock::Text(TextContent::new("hello"))],
3965 },
3966 )
3967 .await
3968 .unwrap();
3969 assert!(prompt_ordinal > cursor.ordinal);
3970 let snapshot = relay.sync_snapshot().await.unwrap();
3971 assert!(snapshot.operational.latest_ordinal > cursor.ordinal);
3972 validate_checkpoint_barrier_snapshot(&snapshot, &barrier_command_id, &cursor).unwrap();
3973
3974 latched_checkpoint(
3975 relay,
3976 barrier_command_id,
3977 cursor,
3978 CheckpointCompletion::HeldBarrier,
3979 )
3980 .complete()
3981 .await
3982 .unwrap();
3983 handle.sync_now().await.unwrap();
3984 assert_eq!(
3985 handle
3986 .view()
3987 .snapshot
3988 .expect("the actor published the completed barrier")
3989 .operational
3990 .checkpoint_barrier,
3991 None
3992 );
3993 }
3994 #[cfg(unix)]
3998 #[tokio::test]
3999 async fn releasing_a_checkpoint_after_capture_defers_only_the_recovery_floor() {
4000 if std::env::var_os(RELEASE_TEST_CHILD).is_none() {
4003 let directory = tempfile::tempdir().unwrap();
4004 let test_name = format!(
4005 "{}::releasing_a_checkpoint_after_capture_defers_only_the_recovery_floor",
4006 module_path!()
4007 .strip_prefix("mj_controller::")
4008 .unwrap_or(module_path!())
4009 );
4010 let output = Command::new(std::env::current_exe().unwrap())
4011 .args(["--exact", &test_name, "--nocapture"])
4012 .env(RELEASE_TEST_CHILD, "1")
4013 .env("MJ_DATA_DIR", directory.path())
4014 .output()
4015 .unwrap();
4016 assert!(
4017 output.status.success(),
4018 "isolated checkpoint release test failed\nstdout:\n{}\nstderr:\n{}",
4019 String::from_utf8_lossy(&output.stdout),
4020 String::from_utf8_lossy(&output.stderr)
4021 );
4022 return;
4023 }
4024 let _writer = crate::database::install_isolated_test_writer();
4026
4027 std::thread::spawn(|| {
4030 std::thread::sleep(std::time::Duration::from_secs(120));
4031 eprintln!("the captured checkpoint never released its barrier");
4032 std::process::exit(101);
4033 });
4034
4035 let relay_root = tempfile::tempdir().unwrap();
4036 let (_channels, handle, mut relay, barrier_command_id, cursor) =
4037 latch_a_live_checkpoint(relay_root.path(), None, ReleaseSupport::Supported, false)
4038 .await;
4039 relay.end_latch();
4040 wait_until_the_actor_serves_again(&handle).await;
4041
4042 let completion = release_checkpoint_after_capture(
4045 &mut relay,
4046 LATCH_RELAY_SESSION,
4047 &barrier_command_id,
4048 &cursor,
4049 HarnessKind::Codex,
4050 )
4051 .await
4052 .unwrap();
4053 assert_eq!(completion, CheckpointCompletion::ReleasedAfterCapture);
4054 let released = relay.sync_snapshot().await.unwrap();
4055 assert_eq!(released.operational.checkpoint_barrier, None);
4056 assert_eq!(released.operational.checkpoint_ready, None);
4057 assert_eq!(
4058 released.operational.recovery_floor_ordinal, 0,
4059 "an exported archive that is not installed may not release journal history"
4060 );
4061
4062 relay
4065 .submit(
4066 new_command_id("prompt").unwrap(),
4067 RelayCommand::Prompt {
4068 prompt: vec![ContentBlock::Text(TextContent::new("during transfer"))],
4069 },
4070 )
4071 .await
4072 .unwrap();
4073 let mut dispatched = None;
4074 for attempt in 0.. {
4075 let snapshot = relay.sync_snapshot().await.unwrap();
4076 if let Some(active) = snapshot.operational.active_prompt {
4077 dispatched = Some(active);
4078 break;
4079 }
4080 assert!(attempt < 200, "a released barrier still froze ACP dispatch");
4081 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
4082 }
4083 assert!(dispatched.is_some());
4084
4085 latched_checkpoint(
4088 relay,
4089 barrier_command_id,
4090 cursor.clone(),
4091 CheckpointCompletion::ReleasedAfterCapture,
4092 )
4093 .complete()
4094 .await
4095 .unwrap();
4096 handle.sync_now().await.unwrap();
4097 let installed = handle
4098 .view()
4099 .snapshot
4100 .expect("the actor published the advanced recovery floor");
4101 assert_eq!(installed.operational.recovery_floor_ordinal, cursor.ordinal);
4102 assert_eq!(installed.operational.recovery_floor_digest, cursor.digest);
4103 }
4104 #[cfg(unix)]
4107 #[tokio::test]
4108 async fn a_worker_that_rejects_the_release_keeps_its_barrier_through_the_transfer() {
4109 if std::env::var_os(LEGACY_RELEASE_TEST_CHILD).is_none() {
4112 let directory = tempfile::tempdir().unwrap();
4113 let test_name = format!(
4114 "{}::a_worker_that_rejects_the_release_keeps_its_barrier_through_the_transfer",
4115 module_path!()
4116 .strip_prefix("mj_controller::")
4117 .unwrap_or(module_path!())
4118 );
4119 let output = Command::new(std::env::current_exe().unwrap())
4120 .args(["--exact", &test_name, "--nocapture"])
4121 .env(LEGACY_RELEASE_TEST_CHILD, "1")
4122 .env("MJ_DATA_DIR", directory.path())
4123 .output()
4124 .unwrap();
4125 assert!(
4126 output.status.success(),
4127 "isolated legacy checkpoint release test failed\nstdout:\n{}\nstderr:\n{}",
4128 String::from_utf8_lossy(&output.stdout),
4129 String::from_utf8_lossy(&output.stderr)
4130 );
4131 return;
4132 }
4133 let _writer = crate::database::install_isolated_test_writer();
4135
4136 std::thread::spawn(|| {
4139 std::thread::sleep(std::time::Duration::from_secs(120));
4140 eprintln!("the rejected release never finished its checkpoint");
4141 std::process::exit(101);
4142 });
4143
4144 let relay_root = tempfile::tempdir().unwrap();
4145 let start_log = tempfile::tempdir().unwrap();
4146 let start_log = start_log.path().join("relay-starts");
4147 let (_channels, handle, mut relay, barrier_command_id, cursor) = latch_a_live_checkpoint(
4148 relay_root.path(),
4149 Some(&start_log),
4150 ReleaseSupport::Rejected,
4151 false,
4152 )
4153 .await;
4154 relay.end_latch();
4155 wait_until_the_actor_serves_again(&handle).await;
4156
4157 let completion = release_checkpoint_after_capture(
4158 &mut relay,
4159 LATCH_RELAY_SESSION,
4160 &barrier_command_id,
4161 &cursor,
4162 HarnessKind::Codex,
4163 )
4164 .await
4165 .unwrap();
4166 assert_eq!(completion, CheckpointCompletion::HeldBarrier);
4167 assert_eq!(relay_starts(&start_log), 1);
4170
4171 let transferring = relay.sync_snapshot().await.unwrap();
4175 validate_checkpoint_barrier_snapshot(&transferring, &barrier_command_id, &cursor).unwrap();
4176 latched_checkpoint(relay, barrier_command_id, cursor.clone(), completion)
4177 .complete()
4178 .await
4179 .unwrap();
4180 handle.sync_now().await.unwrap();
4181 let completed = handle
4182 .view()
4183 .snapshot
4184 .expect("the actor published the completed barrier");
4185 assert_eq!(completed.operational.checkpoint_barrier, None);
4186 assert_eq!(completed.operational.recovery_floor_ordinal, cursor.ordinal);
4187 }
4188 #[cfg(unix)]
4193 #[tokio::test]
4194 async fn abandoning_a_latched_checkpoint_drops_the_connection_that_opened_its_barrier() {
4195 if std::env::var_os(ABANDON_TEST_CHILD).is_none() {
4198 let directory = tempfile::tempdir().unwrap();
4199 let test_name = format!(
4200 "{}::abandoning_a_latched_checkpoint_drops_the_connection_that_opened_its_barrier",
4201 module_path!()
4202 .strip_prefix("mj_controller::")
4203 .unwrap_or(module_path!())
4204 );
4205 let output = Command::new(std::env::current_exe().unwrap())
4206 .args(["--exact", &test_name, "--nocapture"])
4207 .env(ABANDON_TEST_CHILD, "1")
4208 .env("MJ_DATA_DIR", directory.path())
4209 .output()
4210 .unwrap();
4211 assert!(
4212 output.status.success(),
4213 "isolated abandoned checkpoint test failed\nstdout:\n{}\nstderr:\n{}",
4214 String::from_utf8_lossy(&output.stdout),
4215 String::from_utf8_lossy(&output.stderr)
4216 );
4217 return;
4218 }
4219 let _writer = crate::database::install_isolated_test_writer();
4221
4222 std::thread::spawn(|| {
4225 std::thread::sleep(std::time::Duration::from_secs(120));
4226 eprintln!("an abandoned checkpoint never released its relay connection");
4227 std::process::exit(101);
4228 });
4229
4230 let relay_root = tempfile::tempdir().unwrap();
4231 let start_log = tempfile::tempdir().unwrap();
4232 let start_log = start_log.path().join("relay-starts");
4233 let (_channels, handle, mut relay, barrier_command_id, cursor) = latch_a_live_checkpoint(
4234 relay_root.path(),
4235 Some(&start_log),
4236 ReleaseSupport::Supported,
4237 false,
4238 )
4239 .await;
4240 relay.end_latch();
4241 wait_until_the_actor_serves_again(&handle).await;
4242 assert_eq!(relay_starts(&start_log), 1);
4243
4244 latched_checkpoint(
4245 relay,
4246 barrier_command_id,
4247 cursor,
4248 CheckpointCompletion::HeldBarrier,
4249 )
4250 .abandon(LATCH_RELAY_SESSION)
4251 .await;
4252
4253 wait_until_the_actor_serves_again(&handle).await;
4258 assert_eq!(relay_starts(&start_log), 2);
4259 }
4260 #[cfg(unix)]
4265 #[test]
4266 fn a_move_checkpoint_can_verify_its_archive_without_source_harness_readiness() {
4267 let directory = tempfile::tempdir().unwrap();
4268 let name = format!(
4269 "{}::a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content",
4270 module_path!()
4271 .strip_prefix("mj_controller::")
4272 .unwrap_or(module_path!())
4273 );
4274 let output = Command::new(std::env::current_exe().unwrap())
4275 .args(["--exact", &name, "--nocapture"])
4276 .env(REUSE_TEST_CHILD, "1")
4277 .env(LATCH_CHECKPOINT_ONLY, "1")
4278 .env("MJ_DATA_DIR", directory.path())
4279 .output()
4280 .unwrap();
4281 assert!(
4282 output.status.success(),
4283 "checkpoint-only capture failed: {}\n{}",
4284 String::from_utf8_lossy(&output.stdout),
4285 String::from_utf8_lossy(&output.stderr)
4286 );
4287 }
4288
4289 #[cfg(unix)]
4290 #[tokio::test]
4291 async fn a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content() {
4292 if std::env::var_os(REUSE_TEST_CHILD).is_none() {
4295 let directory = tempfile::tempdir().unwrap();
4296 let test_name = format!(
4297 "{}::a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content",
4298 module_path!()
4299 .strip_prefix("mj_controller::")
4300 .unwrap_or(module_path!())
4301 );
4302 let output = Command::new(std::env::current_exe().unwrap())
4303 .args(["--exact", &test_name, "--nocapture"])
4304 .env(REUSE_TEST_CHILD, "1")
4305 .env(LATCH_RELAY_STARTUP_DELAY_MS, "31000")
4308 .env("MJ_DATA_DIR", directory.path())
4309 .output()
4310 .unwrap();
4311 assert!(
4312 output.status.success(),
4313 "isolated checkpoint reuse test failed\nstdout:\n{}\nstderr:\n{}",
4314 String::from_utf8_lossy(&output.stdout),
4315 String::from_utf8_lossy(&output.stderr)
4316 );
4317 return;
4318 }
4319 let _writer = crate::database::install_isolated_test_writer();
4321
4322 std::thread::spawn(|| {
4325 std::thread::sleep(std::time::Duration::from_secs(120));
4326 eprintln!("the reuse checkpoint never finished its latch");
4327 std::process::exit(101);
4328 });
4329
4330 #[derive(Default)]
4331 struct RecordingExecutor {
4332 purposes: std::sync::Mutex<Vec<String>>,
4333 active_stages: std::sync::Mutex<Vec<ProvisionStage>>,
4334 stage_events: std::sync::Mutex<Vec<(ProvisionStage, bool)>>,
4335 observed_stages: std::sync::Mutex<Vec<(String, Vec<ProvisionStage>)>>,
4336 }
4337
4338 impl RecordingExecutor {
4339 fn refused(&self, command: &CommandSpec) -> Result<CommandOutput> {
4340 self.purposes.lock().unwrap().push(command.purpose.clone());
4341 self.observed_stages.lock().unwrap().push((
4342 command.purpose.clone(),
4343 self.active_stages.lock().unwrap().clone(),
4344 ));
4345 Ok(CommandOutput {
4346 status: 1,
4347 stdout: Vec::new(),
4348 stderr: b"no target is provisioned for this test".to_vec(),
4349 })
4350 }
4351
4352 fn purposes(&self) -> Vec<String> {
4353 self.purposes.lock().unwrap().clone()
4354 }
4355
4356 fn observed_stages(&self) -> Vec<(String, Vec<ProvisionStage>)> {
4357 self.observed_stages.lock().unwrap().clone()
4358 }
4359
4360 fn stage_events(&self) -> Vec<(ProvisionStage, bool)> {
4361 self.stage_events.lock().unwrap().clone()
4362 }
4363 }
4364
4365 impl CommandExecutor for RecordingExecutor {
4366 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
4367 self.refused(command)
4368 }
4369
4370 fn execute_with_stdin(
4371 &self,
4372 command: &CommandSpec,
4373 _input: &mut (dyn std::io::Read + Send),
4374 ) -> Result<CommandOutput> {
4375 self.refused(command)
4376 }
4377
4378 fn stage_started(&self, stage: ProvisionStage) {
4379 self.active_stages.lock().unwrap().push(stage);
4380 self.stage_events.lock().unwrap().push((stage, true));
4381 }
4382
4383 fn stage_finished(&self, stage: ProvisionStage) {
4384 let mut active = self.active_stages.lock().unwrap();
4385 let position = active
4386 .iter()
4387 .position(|active_stage| *active_stage == stage)
4388 .expect("stage finished without a matching start");
4389 active.remove(position);
4390 self.stage_events.lock().unwrap().push((stage, false));
4391 }
4392 }
4393
4394 let data_directory = PathBuf::from(std::env::var_os("MJ_DATA_DIR").unwrap());
4395 let relay_root = data_directory.join("relay");
4396 let profile_home = data_directory.join("profile");
4397 let archive_directory = data_directory.join("archives");
4398 for directory in [&relay_root, &profile_home, &archive_directory] {
4399 std::fs::create_dir_all(directory).unwrap();
4400 }
4401 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
4402 let mut seed =
4403 mj_worker::relay::DurableRelay::open(&relay_root, LATCH_RELAY_SESSION, "1.0.0")
4404 .unwrap();
4405 seed.record_observation(mj_core::relay::RelayObservation::SessionOpened {
4406 native_session_id: "native-session".into(),
4407 resumed: true,
4408 })
4409 .unwrap();
4410 seed.record_observation(mj_core::relay::RelayObservation::SessionConfigured {
4411 config_options: Vec::new(),
4412 })
4413 .unwrap();
4414 }
4415 let checkpoint = write_checkpoint_gate_archive(&archive_directory, LATCH_RELAY_SESSION, 2);
4418
4419 let mut session = checkpoint_test_session(LATCH_RELAY_SESSION);
4420 session.target_template_id = "local".into();
4421 session.target = Some(TargetLocator::LocalBare {
4422 worker_root: data_directory.join("workers").join(LATCH_RELAY_SESSION),
4423 });
4424 session.checkpoint = Some(checkpoint.clone());
4425 crate::database::save_session(&session).unwrap();
4426
4427 let mut config = Config::default();
4428 config.profiles.insert(
4429 "codex".into(),
4430 HarnessProfile {
4431 enabled: true,
4432 kind: mj_core::config::HarnessKind::Codex,
4433 home: profile_home,
4434 environment: BTreeMap::new(),
4435 context_window_bytes: None,
4436 },
4437 );
4438 config
4439 .targets
4440 .insert("local".into(), TargetTemplate::LocalBare);
4441 config.bundles.insert(
4442 "project".into(),
4443 ProjectBundle {
4444 primary_repo: "project".into(),
4445 repositories: vec![ProjectRepository {
4446 id: "project".into(),
4447 github: Some("example/project".into()),
4448 local: None,
4449 destination: "project".into(),
4450 git_ref: None,
4451 }],
4452 },
4453 );
4454 let controller = Controller {
4455 config,
4456 state: State {
4457 sessions: BTreeMap::from([(LATCH_RELAY_SESSION.into(), session)]),
4458 ..State::default()
4459 },
4460 };
4461
4462 let channels = crate::session_manager::spawn_session_manager().unwrap();
4463 channels
4464 .targets
4465 .send(vec![latch_relay_target(
4466 &relay_root,
4467 None,
4468 ReleaseSupport::Supported,
4469 false,
4470 )])
4471 .unwrap();
4472 let handle = channels
4473 .control
4474 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
4475 .await
4476 .unwrap();
4477
4478 let executor = RecordingExecutor::default();
4479 let latched = controller
4480 .checkpoint_session_latched(
4481 LATCH_RELAY_SESSION,
4482 &executor,
4483 Some(&channels.control),
4484 LatchExclusivity::HoldThroughClose,
4485 CheckpointExportPolicy::ReuseUnchangedArchive,
4486 )
4487 .await
4488 .unwrap();
4489
4490 assert!(
4491 executor.purposes().is_empty(),
4492 "an unchanged session exported an archive anyway: {:?}",
4493 executor.purposes()
4494 );
4495 assert_eq!(latched.artifact.metadata, checkpoint);
4496 assert!(checkpoint.archive_path.exists());
4497
4498 assert!(latched.cursor.ordinal > checkpoint.event_frontier);
4501 let cursor = latched.cursor.clone();
4502 latched.complete().await.unwrap();
4503 wait_until_the_actor_serves_again(&handle).await;
4504
4505 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
4506 let snapshot = handle.view().snapshot.unwrap();
4507 assert!(snapshot.operational.checkpoint_only);
4508 assert!(!snapshot.operational.native_session_is_ready());
4509 assert_eq!(
4510 verify_archive_streaming(&checkpoint.archive_path)
4511 .unwrap()
4512 .manifest
4513 .session
4514 .native_session_id,
4515 "native-session"
4516 );
4517 channels.shutdown.shutdown().await.unwrap();
4518 return;
4519 }
4520
4521 handle
4525 .submit(
4526 new_command_id("busy-prompt").unwrap(),
4527 RelayCommand::Prompt {
4528 prompt: vec![ContentBlock::Text(TextContent::new("keep working"))],
4529 },
4530 )
4531 .await
4532 .unwrap();
4533 let mut connection = handle.lease_connection().await.unwrap();
4534 let before = connection.connection_mut().sync().await.unwrap();
4535 assert_eq!(before.operational.execution, RelayExecutionState::Running);
4536 connection.release();
4537 let deferred = controller
4538 .checkpoint_session_latched(
4539 LATCH_RELAY_SESSION,
4540 &executor,
4541 Some(&channels.control),
4542 LatchExclusivity::ReleaseAfterLatch,
4543 CheckpointExportPolicy::ReuseUnchangedArchive,
4544 )
4545 .await;
4546 assert!(
4547 matches!(deferred, Err(ref error) if error.downcast_ref::<CheckpointDeferred>().is_some())
4548 );
4549 wait_until_the_actor_serves_again(&handle).await;
4550 let mut connection = handle.lease_connection().await.unwrap();
4551 let after = connection.connection_mut().sync().await.unwrap();
4552 assert_eq!(after.operational.execution, RelayExecutionState::Running);
4553 assert!(after.operational.checkpoint_barrier.is_none());
4554 let journal =
4555 std::fs::read_to_string(relay_root.join("relay-journal/active.jsonl")).unwrap();
4556 for line in journal.lines() {
4557 let event: mj_core::relay::RelayEvent = serde_json::from_str(line).unwrap();
4558 if event.ordinal > before.operational.latest_ordinal {
4559 assert!(
4560 !matches!(
4561 event.observation,
4562 mj_core::relay::RelayObservation::CommandQueued {
4563 command: RelayCommand::BeginCheckpoint { .. },
4564 ..
4565 } | mj_core::relay::RelayObservation::CommandInterrupted {
4566 command: mj_core::relay::RelayCommandKind::BeginCheckpoint,
4567 ..
4568 }
4569 ),
4570 "busy deferral journaled checkpoint activity: {event:?}"
4571 );
4572 }
4573 }
4574 connection.release();
4575 handle
4576 .submit(
4577 new_command_id("finish-busy-prompt").unwrap(),
4578 RelayCommand::CancelTurn,
4579 )
4580 .await
4581 .unwrap();
4582 handle.sync_now().await.unwrap();
4583
4584 handle
4586 .submit(
4587 new_command_id("resume-notice").unwrap(),
4588 RelayCommand::RecordNotice {
4589 text: "the session changed".into(),
4590 },
4591 )
4592 .await
4593 .unwrap();
4594 for attempt in 0.. {
4595 handle.sync_now().await.unwrap();
4596 let materialized = handle.view().snapshot.map(|snapshot| snapshot.materialized);
4597 if materialized.is_some_and(|materialized| {
4598 materialized.applied_event_ordinal > cursor.ordinal
4599 && !materialized.transcript.is_empty()
4600 }) {
4601 break;
4602 }
4603 assert!(attempt < 200, "the notice never reached the projection");
4604 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
4605 }
4606
4607 let changed = controller
4608 .checkpoint_session_latched(
4609 LATCH_RELAY_SESSION,
4610 &executor,
4611 Some(&channels.control),
4612 LatchExclusivity::HoldThroughClose,
4613 CheckpointExportPolicy::ReuseUnchangedArchive,
4614 )
4615 .await;
4616 let Err(error) = changed else {
4617 panic!("a changed session reused its installed archive");
4618 };
4619
4620 assert!(
4621 executor
4622 .purposes()
4623 .contains(&"export target checkpoint".to_owned()),
4624 "a changed session skipped its export: {:?}",
4625 executor.purposes()
4626 );
4627 assert!(
4628 format!("{error:#}").contains("no target is provisioned for this test"),
4629 "{error:#}"
4630 );
4631 assert!(
4632 executor.observed_stages().iter().any(|(purpose, stages)| {
4633 purpose == "export target checkpoint"
4634 && stages.contains(&ProvisionStage::RecoveryCopy)
4635 }),
4636 "close checkpoint export did not run inside RecoveryCopy: {:?}",
4637 executor.observed_stages()
4638 );
4639 assert_eq!(
4640 executor
4641 .stage_events()
4642 .into_iter()
4643 .filter(|(stage, _)| *stage == ProvisionStage::RecoveryCopy)
4644 .collect::<Vec<_>>(),
4645 vec![
4646 (ProvisionStage::RecoveryCopy, true),
4647 (ProvisionStage::RecoveryCopy, false)
4648 ]
4649 );
4650 assert!(executor.active_stages.lock().unwrap().is_empty());
4651 assert!(checkpoint.archive_path.exists());
4652 }
4653 #[cfg(unix)]
4654 fn relay_starts(path: &Path) -> usize {
4655 std::fs::read_to_string(path)
4656 .unwrap_or_default()
4657 .lines()
4658 .count()
4659 }
4660 #[cfg(unix)]
4663 fn latched_checkpoint(
4664 relay: ControllerRelayLease,
4665 barrier_command_id: String,
4666 cursor: RelayCursor,
4667 completion: CheckpointCompletion,
4668 ) -> LatchedCheckpoint {
4669 LatchedCheckpoint {
4670 artifact: CheckpointArtifact {
4671 metadata: CheckpointMetadata {
4672 archive_path: PathBuf::from("checkpoint.hel.zip"),
4673 sha256: "a".repeat(64),
4674 created_at: now(),
4675 event_frontier: cursor.ordinal,
4676 },
4677 native_session_id: "native-session".into(),
4678 event_frontier_digest: cursor.digest.clone(),
4679 },
4680 relay,
4681 barrier_command_id,
4682 cursor,
4683 completion,
4684 }
4685 }
4686 #[test]
4687 fn checkpoint_persistence_rollback_restores_memory_and_reports_both_failures() {
4688 let session_id = "0123456789abcdef0123456789abcdef";
4689 let previous = checkpoint_test_session(session_id);
4690 let mut changed = previous.clone();
4691 changed.state = SessionState::Closing;
4692 changed.last_checkpoint_error = Some("partially installed checkpoint".into());
4693 let mut state = State::default();
4694 state.sessions.insert(session_id.into(), changed);
4695
4696 let error = restore_session_after_persistence_failure(
4697 &mut state,
4698 session_id,
4699 &previous,
4700 anyhow::anyhow!("verified checkpoint persistence failed"),
4701 |record| {
4702 assert_eq!(record, &previous);
4703 Err(anyhow::anyhow!("rollback database write failed"))
4704 },
4705 );
4706
4707 assert_eq!(state.sessions.get(session_id), Some(&previous));
4708 let detail = format!("{error:#}");
4709 assert!(detail.contains("verified checkpoint persistence failed"));
4710 assert!(detail.contains("rollback database write failed"));
4711 }
4712 #[test]
4713 fn installed_checkpoint_gate_reopens_and_checks_sha() {
4714 let directory = tempfile::tempdir().unwrap();
4715 let session_id = "0123456789abcdef0123456789abcdef";
4716 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4717 verify_installed_checkpoint_gate(session_id, &checkpoint).unwrap();
4718
4719 let mut wrong_sha = checkpoint.clone();
4720 wrong_sha.sha256 = "b".repeat(64);
4721 assert!(
4722 verify_installed_checkpoint_gate(session_id, &wrong_sha)
4723 .unwrap_err()
4724 .to_string()
4725 .contains("SHA changed")
4726 );
4727 std::fs::write(
4728 &checkpoint.archive_path,
4729 b"changed after first verification",
4730 )
4731 .unwrap();
4732 assert!(
4733 format!(
4734 "{:#}",
4735 verify_installed_checkpoint_gate(session_id, &checkpoint).unwrap_err()
4736 )
4737 .contains("installed checkpoint SHA changed")
4738 );
4739 }
4740 #[test]
4741 fn an_installed_archive_is_reused_when_only_relay_bookkeeping_moved() {
4742 let directory = tempfile::tempdir().unwrap();
4743 let session_id = "0123456789abcdef0123456789abcdef";
4744 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4745 let archived = verify_archive_streaming(&checkpoint.archive_path)
4746 .unwrap()
4747 .canonical_session;
4748
4749 let mut latched = archived.clone();
4752 latched.event_frontier += 6;
4753 latched.event_frontier_digest = "b".repeat(64);
4754 latched.session.last_activity_at_ms = Some(9_999);
4755
4756 let artifact = reusable_installed_checkpoint(
4757 session_id,
4758 Some(&checkpoint),
4759 "native-session",
4760 latched.event_frontier,
4761 &latched,
4762 )
4763 .expect("an unchanged session reuses its installed archive");
4764
4765 assert_eq!(artifact.metadata, checkpoint);
4766 assert_eq!(artifact.native_session_id, "native-session");
4767 assert_eq!(
4768 artifact.event_frontier_digest,
4769 archived.event_frontier_digest
4770 );
4771 verify_checkpoint_artifact(session_id, &artifact).unwrap();
4773 verify_installed_checkpoint_gate(session_id, &artifact.metadata).unwrap();
4774 }
4775 #[test]
4776 fn archive_reuse_falls_back_to_a_full_export_for_anything_but_bookkeeping() {
4777 let directory = tempfile::tempdir().unwrap();
4778 let session_id = "0123456789abcdef0123456789abcdef";
4779 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4780 let archived = verify_archive_streaming(&checkpoint.archive_path)
4781 .unwrap()
4782 .canonical_session;
4783 let mut latched = archived.clone();
4784 latched.event_frontier += 6;
4785 let reuse = |installed: Option<&CheckpointMetadata>,
4786 ordinal: u64,
4787 session: &CanonicalSessionSnapshot| {
4788 reusable_installed_checkpoint(session_id, installed, "native-session", ordinal, session)
4789 };
4790
4791 assert!(reuse(None, latched.event_frontier, &latched).is_none());
4792
4793 let mut with_new_content = latched.clone();
4794 with_new_content.transcript.push(CanonicalTranscriptItem {
4795 stable_id: "system:notice:notice-1".into(),
4796 position: latched.event_frontier,
4797 latest_content_event_ordinal: None,
4798 created_at_ms: 2_000,
4799 last_changed_at_ms: 2_000,
4800 body: CanonicalTranscriptBody::System {
4801 text: "resumed".into(),
4802 },
4803 });
4804 assert!(reuse(Some(&checkpoint), latched.event_frontier, &with_new_content).is_none());
4805
4806 assert!(reuse(Some(&checkpoint), checkpoint.event_frontier - 1, &latched).is_none());
4808
4809 let mut wrong_sha = checkpoint.clone();
4810 wrong_sha.sha256 = "b".repeat(64);
4811 assert!(reuse(Some(&wrong_sha), latched.event_frontier, &latched).is_none());
4812
4813 let another_session =
4814 write_checkpoint_gate_archive(directory.path(), "1123456789abcdef0123456789abcdef", 7);
4815 assert!(reuse(Some(&another_session), latched.event_frontier, &latched).is_none());
4816
4817 std::fs::write(&checkpoint.archive_path, b"not an archive any more").unwrap();
4818 assert!(reuse(Some(&checkpoint), latched.event_frontier, &latched).is_none());
4819 }
4820 #[cfg(unix)]
4821 #[tokio::test]
4822 async fn workspace_lease_blocks_prompts_and_releases_without_advancing_recovery() {
4823 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
4824 let directory = tempfile::tempdir().unwrap();
4825 let name = format!(
4826 "{}::workspace_lease_blocks_prompts_and_releases_without_advancing_recovery",
4827 module_path!()
4828 .strip_prefix("mj_controller::")
4829 .unwrap_or(module_path!())
4830 );
4831 let mut command = crate::targets::CommandSpec::new(
4832 std::env::current_exe().unwrap().to_string_lossy(),
4833 ["--exact", &name, "--nocapture"],
4834 );
4835 command.env.insert(LATCH_TEST_CHILD.into(), "1".into());
4836 command.env.insert(
4837 "MJ_DATA_DIR".into(),
4838 directory.path().to_string_lossy().into(),
4839 );
4840 let result =
4841 crate::targets::CancellableProcessExecutor::with_timeout(Duration::from_secs(60))
4842 .execute(&command)
4843 .unwrap();
4844 assert_eq!(
4845 result.status,
4846 0,
4847 "{}\n{}",
4848 String::from_utf8_lossy(&result.stdout),
4849 String::from_utf8_lossy(&result.stderr)
4850 );
4851 assert!(
4852 String::from_utf8_lossy(&result.stdout).contains("1 passed"),
4853 "child did not run its test"
4854 );
4855 return;
4856 }
4857 let _writer = crate::database::install_isolated_test_writer();
4858 let root = tempfile::tempdir().unwrap();
4859 let (_channels, handle, mut relay, barrier, _cursor) =
4860 latch_a_live_checkpoint(root.path(), None, ReleaseSupport::Supported, false).await;
4861 relay
4862 .connection_mut()
4863 .submit(
4864 new_command_id("release-initial").unwrap(),
4865 RelayCommand::ReleaseCheckpoint {
4866 barrier_command_id: barrier,
4867 },
4868 )
4869 .await
4870 .unwrap();
4871 relay.release();
4872 wait_until_the_actor_serves_again(&handle).await;
4873 let before = handle
4874 .view()
4875 .snapshot
4876 .unwrap()
4877 .operational
4878 .recovery_floor_ordinal;
4879 let mut workspace = IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
4880 .await
4881 .unwrap();
4882 workspace.verify().await.unwrap();
4883 drop(workspace);
4884 wait_until_the_actor_serves_again(&handle).await;
4885 assert!(
4886 handle
4887 .view()
4888 .snapshot
4889 .unwrap()
4890 .operational
4891 .checkpoint_barrier
4892 .is_none()
4893 );
4894 let mut workspace = IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
4895 .await
4896 .unwrap();
4897 let submitting = handle.clone();
4898 let mut prompt = tokio::spawn(async move {
4899 submitting
4900 .submit(
4901 new_command_id("after-write").unwrap(),
4902 RelayCommand::Prompt {
4903 prompt: vec![ContentBlock::Text(TextContent::new("go"))],
4904 },
4905 )
4906 .await
4907 });
4908 assert!(
4909 tokio::time::timeout(Duration::from_millis(50), &mut prompt)
4910 .await
4911 .is_err(),
4912 "prompt must wait for the workspace owner"
4913 );
4914 workspace.verify().await.unwrap();
4915 workspace.release().await.unwrap();
4916 tokio::time::timeout(Duration::from_secs(10), prompt)
4917 .await
4918 .unwrap()
4919 .unwrap()
4920 .unwrap();
4921 wait_until_the_actor_serves_again(&handle).await;
4922 let after = handle.view().snapshot.unwrap();
4923 assert!(after.operational.checkpoint_barrier.is_none());
4924 assert_eq!(after.operational.recovery_floor_ordinal, before);
4925 assert!(
4926 IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
4927 .await
4928 .is_err(),
4929 "a queued or running prompt must prevent file injection"
4930 );
4931 }
4932}