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: match &session.managed_worktree {
662 Some(worktree) => CheckpointRepositoryCapture::DeltaFrom {
663 base_commit: super::worktree::managed_worktree_base_commit(
664 worktree, executor,
665 )?,
666 },
667 None => CheckpointRepositoryCapture::MetadataOnly,
668 },
669 origin_override: None,
670 }],
671 )
672 } else {
673 let bundle = self
674 .config
675 .bundles
676 .get(&session.bundle_id)
677 .context("session bundle is missing")?;
678 let workspace_root = match &backend {
679 targets::TargetLocator::LocalPodman { .. }
680 | targets::TargetLocator::LocalDocker { .. }
681 | targets::TargetLocator::AppleContainer { .. }
682 | targets::TargetLocator::SshPodman { .. }
683 | targets::TargetLocator::SshDocker { .. } => "/workspace".to_string(),
684 targets::TargetLocator::AwsEc2 { workspace, .. }
685 | targets::TargetLocator::SshBare { workspace, .. } => workspace.clone(),
686 targets::TargetLocator::LocalBare { worker_root } => worker_root.clone(),
687 };
688 let repositories = bundle
689 .repositories
690 .iter()
691 .map(|repository| CheckpointRepositorySpec {
692 id: repository.id.clone(),
693 relative_destination: repository.destination.clone(),
694 capture: CheckpointRepositoryCapture::RemoteWorkspace,
695 origin_override: None,
696 })
697 .collect();
698 (workspace_root, bundle.primary_repo.clone(), repositories)
699 };
700 Ok(SessionExportLayout {
701 backend,
702 workspace_root,
703 primary_repository,
704 repositories,
705 managed_worktree: session.managed_worktree,
706 })
707 }
708
709 pub(super) async fn checkpoint_session_latched(
710 &self,
711 session_id: &str,
712 executor: &(impl CommandExecutor + Sync),
713 manager: Option<&SessionManagerControl>,
714 exclusivity: LatchExclusivity,
715 export_policy: CheckpointExportPolicy,
716 ) -> Result<LatchedCheckpoint> {
717 self.checkpoint_session_latched_with_recovery_stage(
718 session_id,
719 executor,
720 manager,
721 exclusivity,
722 export_policy,
723 exclusivity == LatchExclusivity::HoldThroughClose,
724 )
725 .await
726 }
727
728 async fn checkpoint_session_latched_with_recovery_stage(
729 &self,
730 session_id: &str,
731 executor: &(impl CommandExecutor + Sync),
732 manager: Option<&SessionManagerControl>,
733 exclusivity: LatchExclusivity,
734 export_policy: CheckpointExportPolicy,
735 recovery_copy: bool,
736 ) -> Result<LatchedCheckpoint> {
737 if let Some(operation) = crate::database::load_move_operation(session_id)?
738 && operation.queue_admission_started
739 && !operation.queue_admission_finished
740 {
741 bail!(
744 "move queue admission is incomplete; retry Move before checkpointing this destination"
745 );
746 }
747 let session = self
748 .state
749 .sessions
750 .get(session_id)
751 .with_context(|| format!("unknown session {session_id}"))?
752 .clone();
753 session.validate_configuration(&self.config)?;
754 let layout = self.session_export_layout(session_id, executor)?;
755 let backend = layout.backend.clone();
756 let profile = self
757 .config
758 .profiles
759 .get(&session.last_profile)
760 .context("session profile is missing")?;
761 let reconnect = targets::reconnect_plan(&backend, session_id)?
762 .commands
763 .into_iter()
764 .next()
765 .context("reconnect plan is empty")?;
766 let worker_root = targets::worker_root(&backend, session_id)?;
767 let harness_home = target_profile_home(&backend, session_id, profile);
768 let SessionExportLayout {
769 workspace_root,
770 primary_repository,
771 repositories,
772 ..
773 } = layout;
774 let target_path = |path: &str| match &backend {
775 targets::TargetLocator::AwsEc2 { .. } | targets::TargetLocator::SshBare { .. }
776 if !path.starts_with('/') =>
777 {
778 PathBuf::from(format!("~/{path}"))
779 }
780 _ => PathBuf::from(path),
781 };
782 let operation_id = new_command_id("checkpoint")?;
785 let remote_spec = format!("{worker_root}/{operation_id}-spec.json");
786 let remote_archive = format!("{worker_root}/{operation_id}.hel.zip");
787 let remote_stage = format!("{worker_root}/{operation_id}-stage");
788 let checkpointed_at = now();
789 let target_manifest = TargetManifest {
790 template_id: session.target_template_id.clone(),
791 target_kind: target_kind(&backend).into(),
792 details: Default::default(),
793 };
794 let bundle_manifest = BundleManifest {
795 id: session.bundle_id.clone(),
796 primary_repository,
797 };
798 let session_manifest = |native_session_id: &str| SessionManifest {
799 id: session.id.clone(),
800 title: session.title.clone(),
801 harness_kind: session.harness_kind,
802 profile_id: session.last_profile.clone(),
803 native_session_id: native_session_id.to_owned(),
804 created_at: session.created_at.clone(),
805 checkpointed_at: checkpointed_at.clone(),
806 hel_version: env!("CARGO_PKG_VERSION").into(),
807 relay_version: env!("CARGO_PKG_VERSION").into(),
808 adapter_version: "acp-v1".into(),
809 };
810 let releases_after_capture = exclusivity == LatchExclusivity::ReleaseAfterLatch;
811 if releases_after_capture
812 && let Some(native_session_id) = session.native_session_id.as_deref()
813 {
814 let prestage = CheckpointCaptureSpec {
815 protocol_version: CHECKPOINT_STAGING_PROTOCOL_VERSION,
816 session: session_manifest(native_session_id),
817 target: target_manifest.clone(),
818 bundle: bundle_manifest.clone(),
819 relay_root: target_path(&worker_root),
820 harness_home: target_path(&harness_home),
821 workspace_root: target_path(&workspace_root),
822 repositories: repositories.clone(),
823 allow_empty_native: false,
824 stage_path: target_path(&remote_stage),
825 refresh_existing: false,
826 };
827 let prestage_started = Instant::now();
828 let prestaged = {
829 let _recovery_copy = recovery_copy
830 .then(|| ProvisionStageGuard::new(executor, ProvisionStage::RecoveryCopy));
831 run_checkpoint_staging_command(
832 executor,
833 &backend,
834 session_id,
835 &prestage,
836 capture_stdin_command,
837 "prestage target checkpoint",
838 )
839 };
840 match prestaged {
841 Ok(output) => match serde_json::from_slice::<CapturedCheckpoint>(&output.stdout) {
842 Ok(captured) => tracing::info!(
843 session_id,
844 prestage_ms = prestage_started.elapsed().as_millis() as u64,
845 native_bytes = captured.native_bytes,
846 repository_bytes = captured.repository_bytes,
847 reused_native = captured.reused_native,
848 "checkpoint target state prestaged while ACP dispatch remained active"
849 ),
850 Err(error) => tracing::warn!(
851 session_id,
852 error = format!("{error:#}"),
853 "checkpoint prestage returned an invalid result; barrier capture will replace it"
854 ),
855 },
856 Err(error) => {
857 if executor.cancellation_requested() {
858 return Err(error.context("checkpoint prestage was cancelled"));
859 }
860 tracing::warn!(
861 session_id,
862 error = format!("{error:#}"),
863 "checkpoint prestage failed; barrier capture will collect a fresh generation"
864 );
865 }
866 }
867 }
868 let (mut relay, mut restarted_worker) = self
869 .open_checkpoint_relay(
870 session_id,
871 executor,
872 manager,
873 InstalledWorkerRestart {
874 backend: &backend,
875 worker_root: &worker_root,
876 reconnect: &reconnect,
877 launch: None,
878 messages: &RESTART_FOR_CHECKPOINT,
879 },
880 exclusivity == LatchExclusivity::HoldThroughClose
881 || session.harness_kind != HarnessKind::Kimi,
882 )
883 .await?;
884 let (barrier, barrier_command_id) = loop {
885 let checkpoint_only = relay
889 .connection_mut()
890 .sync()
891 .await?
892 .operational
893 .checkpoint_only;
894 if !checkpoint_only {
895 wait_for_native_session_in_stage(
896 relay.connection_mut(),
897 executor,
898 targets::ProvisionStage::Starting,
899 )
900 .await?;
901 }
902 if exclusivity == LatchExclusivity::HoldThroughClose {
903 let snapshot = relay.connection_mut().sync().await?;
904 if !checkpoint_only && snapshot.operational.capacity_retry.is_some() {
905 relay
908 .connection_mut()
909 .submit(
910 new_command_id("cancel-capacity-retry")?,
911 RelayCommand::CancelTurn,
912 )
913 .await?;
914 }
915 }
916 if exclusivity == LatchExclusivity::ReleaseAfterLatch {
917 let snapshot = relay.connection_mut().sync().await?;
918 if snapshot.operational.execution == RelayExecutionState::Running {
919 relay.release();
922 return Err(CheckpointDeferred::harness_busy().into());
923 }
924 if !snapshot
925 .operational
926 .safe_for_checkpoint(session.harness_kind)
927 {
928 relay.release();
933 return Err(CheckpointDeferred::background_snapshot(
934 &snapshot.operational,
935 session.harness_kind,
936 )
937 .into());
938 }
939 }
940 let barrier_command_id = new_command_id("checkpoint")?;
941 let timeout = if restarted_worker {
942 CHECKPOINT_BARRIER_TIMEOUT_AFTER_RESTART
943 } else {
944 CHECKPOINT_BARRIER_TIMEOUT
945 };
946 let result = {
947 let connection = relay.connection_mut();
948 connection
949 .submit(
950 barrier_command_id.clone(),
951 RelayCommand::BeginCheckpoint {
952 reason: Some("controller archive checkpoint".into()),
953 },
954 )
955 .await?;
956 wait_for_checkpoint_barrier(
957 connection,
958 session_id,
959 &barrier_command_id,
960 timeout,
961 BarrierBusyPolicy::of(exclusivity),
962 session.harness_kind,
963 )
964 .await
965 };
966 match result {
967 Ok(barrier) => break (barrier, barrier_command_id),
968 Err(error)
969 if !restarted_worker && checkpoint_barrier_needs_worker_restart(&error) =>
970 {
971 if exclusivity == LatchExclusivity::ReleaseAfterLatch
972 && matches!(session.harness_kind, HarnessKind::Kimi | HarnessKind::Codex)
973 {
974 let safe_to_restart =
975 relay.connection_mut().sync().await.is_ok_and(|snapshot| {
976 snapshot.operational.safe_to_replace(session.harness_kind)
977 });
978 if !safe_to_restart {
979 return Err(error.context(CheckpointDeferred::background_work()));
980 }
981 }
982 tracing::warn!(
983 session_id,
984 "checkpoint requires a worker restart; restarting and retrying: {error:#}"
985 );
986 let connection = self
987 .restart_worker_for_checkpoint(
988 session_id,
989 executor,
990 &backend,
991 &worker_root,
992 &reconnect,
993 )
994 .await?;
995 relay.replace_connection(connection);
996 restarted_worker = true;
997 }
998 Err(error) => return Err(error),
999 }
1000 };
1001 let barrier_ready_at = Instant::now();
1002 relay
1006 .connection_mut()
1007 .sync_project_memory()
1008 .await
1009 .context("synchronize project memory for checkpoint")?;
1010 let cursor = barrier
1011 .operational
1012 .checkpoint_ready
1013 .clone()
1014 .context("relay reported a checkpoint barrier without its ready cursor")?;
1015 let materialized = barrier.materialized;
1016 let expected_ordinal = materialized.applied_event_ordinal;
1017 let expected_digest = materialized.applied_event_digest.clone();
1018 ensure!(
1019 expected_ordinal == barrier.operational.latest_ordinal,
1020 "checkpoint projection frontier {expected_ordinal} does not match relay frontier {}",
1021 barrier.operational.latest_ordinal
1022 );
1023 ensure!(
1024 expected_digest == barrier.operational.latest_digest,
1025 "checkpoint projection digest does not match the relay frontier digest"
1026 );
1027 ensure_exact_checkpoint_cut(&cursor, expected_ordinal, &expected_digest)?;
1028 let canonical_session = canonical_session_from_materialized(&materialized)?;
1029 let native_session_id = barrier
1030 .operational
1031 .native_session_id
1032 .or_else(|| session.native_session_id.clone())
1033 .context("harness did not report its native session ID")?;
1034
1035 if exclusivity == LatchExclusivity::ReleaseAfterLatch {
1040 relay.end_latch();
1041 }
1042
1043 if export_policy == CheckpointExportPolicy::ReuseUnchangedArchive
1049 && session.managed_worktree.is_none()
1053 && let Some(artifact) = reusable_installed_checkpoint(
1054 session_id,
1055 session.checkpoint.as_ref(),
1056 &native_session_id,
1057 cursor.ordinal,
1058 &canonical_session,
1059 )
1060 {
1061 return Ok(LatchedCheckpoint {
1062 artifact,
1063 relay,
1064 barrier_command_id,
1065 cursor,
1066 completion: CheckpointCompletion::HeldBarrier,
1067 });
1068 }
1069
1070 let mut completion = CheckpointCompletion::HeldBarrier;
1075
1076 let exported: Result<CheckpointArtifact> = async {
1077 let spec = CheckpointExportSpec {
1078 protocol_version: CHECKPOINT_EXPORT_PROTOCOL_VERSION,
1079 session: session_manifest(&native_session_id),
1080 target: target_manifest,
1081 bundle: bundle_manifest,
1082 relay_root: target_path(&worker_root),
1083 harness_home: target_path(&harness_home),
1084 workspace_root: target_path(&workspace_root),
1085 repositories,
1086 canonical_session,
1087 output_path: target_path(&remote_archive),
1088 };
1089 let mut export_ms: Option<u64> = None;
1092 let exported = if releases_after_capture {
1093 let capture_spec = CheckpointCaptureSpec {
1094 protocol_version: CHECKPOINT_STAGING_PROTOCOL_VERSION,
1095 session: spec.session.clone(),
1096 target: spec.target.clone(),
1097 bundle: spec.bundle.clone(),
1098 relay_root: spec.relay_root.clone(),
1099 harness_home: spec.harness_home.clone(),
1100 workspace_root: spec.workspace_root.clone(),
1101 repositories: spec.repositories.clone(),
1102 allow_empty_native: !canonical_session_contains_prompt(&spec.canonical_session),
1103 stage_path: target_path(&remote_stage),
1104 refresh_existing: true,
1105 };
1106 let capture_started = Instant::now();
1107 let captured = {
1108 let _recovery_copy = recovery_copy.then(|| {
1109 ProvisionStageGuard::new(executor, ProvisionStage::RecoveryCopy)
1110 });
1111 run_checkpoint_staging_command(
1112 executor,
1113 &backend,
1114 session_id,
1115 &capture_spec,
1116 capture_stdin_command,
1117 "capture target checkpoint",
1118 )?
1119 };
1120 let captured: CapturedCheckpoint = serde_json::from_slice(&captured.stdout)
1121 .context("decode captured checkpoint result")?;
1122 tracing::info!(
1123 session_id,
1124 capture_ms = capture_started.elapsed().as_millis() as u64,
1125 barrier_held_ms = barrier_ready_at.elapsed().as_millis() as u64,
1126 native_bytes = captured.native_bytes,
1127 repository_bytes = captured.repository_bytes,
1128 reused_native = captured.reused_native,
1129 "checkpoint target state captured; releasing ACP dispatch"
1130 );
1131 if !session.harness_kind.captures_native_session() {
1132 tracing::info!(
1133 session_id,
1134 harness = %session.harness_kind.display_name(),
1135 "checkpoint holds repository state only; this harness keeps no per-session native files"
1136 );
1137 }
1138 completion = release_checkpoint_after_capture(
1139 &mut relay,
1140 session_id,
1141 &barrier_command_id,
1142 &cursor,
1143 session.harness_kind,
1144 )
1145 .await?;
1146 let pack_spec = CheckpointPackSpec {
1147 protocol_version: CHECKPOINT_STAGING_PROTOCOL_VERSION,
1148 relay_root: spec.relay_root.clone(),
1149 stage_path: target_path(&remote_stage),
1150 canonical_session: spec.canonical_session.clone(),
1151 output_path: spec.output_path.clone(),
1152 };
1153 let pack_started = Instant::now();
1154 let output = {
1155 let _recovery_copy = recovery_copy.then(|| {
1156 ProvisionStageGuard::new(executor, ProvisionStage::RecoveryCopy)
1157 });
1158 run_checkpoint_staging_command(
1159 executor,
1160 &backend,
1161 session_id,
1162 &pack_spec,
1163 pack_stdin_command,
1164 "pack target checkpoint",
1165 )?
1166 };
1167 tracing::info!(
1168 session_id,
1169 pack_ms = pack_started.elapsed().as_millis() as u64,
1170 "checkpoint archive packaged after ACP dispatch resumed"
1171 );
1172 output
1173 } else {
1174 let export_started = Instant::now();
1175 let output = {
1176 let _recovery_copy = recovery_copy.then(|| {
1177 ProvisionStageGuard::new(executor, ProvisionStage::RecoveryCopy)
1178 });
1179 export_target_checkpoint(
1180 executor,
1181 &backend,
1182 session_id,
1183 &spec,
1184 &remote_spec,
1185 )?
1186 };
1187 export_ms = Some(export_started.elapsed().as_millis() as u64);
1188 output
1189 };
1190 let target_checkpoint: mj_checkpoint::checkpoint::TargetCheckpoint =
1191 serde_json::from_slice(&exported.stdout)
1192 .context("decode target checkpoint result")?;
1193 if let Some(export_ms) = export_ms {
1194 let timings = target_checkpoint.timings.unwrap_or_default();
1197 tracing::info!(
1198 session_id,
1199 export_ms,
1200 timings_reported = target_checkpoint.timings.is_some(),
1201 native_ms = timings.native_ms,
1202 repositories_ms = timings.repositories_ms,
1203 archive_ms = timings.archive_ms,
1204 worker_total_ms = timings.total_ms,
1205 "checkpoint archive exported on the target"
1206 );
1207 }
1208 if target_checkpoint.event_frontier != expected_ordinal {
1209 bail!(
1210 "target checkpoint event frontier changed: expected {expected_ordinal}, found {}",
1211 target_checkpoint.event_frontier
1212 );
1213 }
1214 if target_checkpoint.event_frontier_digest != expected_digest {
1215 bail!("target checkpoint event frontier digest changed");
1216 }
1217
1218 let archive_id = new_command_id("archive")?;
1223 let destination = sessions_dir().join(format!(
1224 "{session_id}-{}-{archive_id}.hel.zip",
1225 target_checkpoint.event_frontier
1226 ));
1227 let transfer = CheckpointTransfer {
1228 locator: &backend,
1229 session_id,
1230 operation_id: &operation_id,
1231 remote_archive: &remote_archive,
1232 destination: &destination,
1233 expected_sha256: &target_checkpoint.sha256,
1234 expected_event_frontier: target_checkpoint.event_frontier,
1235 expected_event_frontier_digest: &target_checkpoint.event_frontier_digest,
1236 };
1237 let metadata = {
1238 let _verifying = ProvisionStageGuard::new(executor, ProvisionStage::Verifying);
1239 let transfer_started = Instant::now();
1240 let verified = transfer.execute(executor)?;
1241 tracing::info!(
1242 session_id,
1243 transfer_and_checksum_ms = transfer_started.elapsed().as_millis() as u64,
1244 "checkpoint archive transferred and checksum-verified"
1245 );
1246 let installed_archive = verified.archive_path().to_path_buf();
1247 let validate_transferred = || -> Result<()> {
1248 ensure!(
1249 verified.sha256() == target_checkpoint.sha256,
1250 "target and controller checkpoint checksums differ"
1251 );
1252 ensure!(
1253 verified.event_frontier_digest() == expected_digest,
1254 "verified checkpoint event frontier digest changed"
1255 );
1256 Ok(())
1257 };
1258 if let Err(error) = validate_transferred() {
1259 return Err(remove_uninstalled_checkpoint(&installed_archive, error));
1260 }
1261 if completion == CheckpointCompletion::HeldBarrier {
1265 let revalidated = relay.sync_snapshot().await.and_then(|snapshot| {
1266 if releases_after_capture {
1267 validate_automatic_checkpoint_barrier_snapshot(
1268 &snapshot,
1269 &barrier_command_id,
1270 &cursor,
1271 session.harness_kind,
1272 )
1273 } else {
1274 validate_checkpoint_barrier_snapshot(
1275 &snapshot,
1276 &barrier_command_id,
1277 &cursor,
1278 )
1279 }
1280 });
1281 if let Err(error) = revalidated {
1282 return Err(remove_uninstalled_checkpoint(
1283 &installed_archive,
1284 error.context(
1285 "checkpoint barrier changed while transferring its archive",
1286 ),
1287 ));
1288 }
1289 }
1290 if let Err(error) = transfer
1291 .cleanup_plan(&verified)
1292 .and_then(|plan| plan.execute(executor).map(|_| ()))
1293 {
1294 return Err(remove_uninstalled_checkpoint(
1295 &installed_archive,
1296 error.context("clean target checkpoint staging"),
1297 ));
1298 }
1299 CheckpointMetadata {
1300 archive_path: verified.archive_path().to_path_buf(),
1301 sha256: verified.sha256().to_string(),
1302 created_at: checkpointed_at.clone(),
1303 event_frontier: verified.event_frontier(),
1304 }
1305 };
1306 Ok(CheckpointArtifact {
1307 metadata,
1308 native_session_id,
1309 event_frontier_digest: expected_digest,
1310 })
1311 }
1312 .await;
1313
1314 let artifact = match exported {
1315 Ok(artifact) => artifact,
1316 Err(error) => {
1317 if completion == CheckpointCompletion::HeldBarrier
1323 && let Err(cancel_error) = relay.cancel_abandoned_barrier().await
1324 {
1325 tracing::warn!(
1326 session_id,
1327 "failed checkpoint could not cancel its relay barrier: {cancel_error:#}"
1328 );
1329 }
1330 return Err(error);
1331 }
1332 };
1333 Ok(LatchedCheckpoint {
1334 artifact,
1335 relay,
1336 barrier_command_id,
1337 cursor,
1338 completion,
1339 })
1340 }
1341
1342 pub(super) async fn prepare_move_source_checkpoint(
1343 &self,
1344 session_id: &str,
1345 executor: &(impl CommandExecutor + Sync),
1346 manager: &SessionManagerControl,
1347 operation: &mut mj_core::state::MoveOperation,
1348 ) -> Result<()> {
1349 let snapshot = super::move_session::refresh_move_source(manager, session_id).await?;
1350 if snapshot
1351 .as_ref()
1352 .is_some_and(|snapshot| snapshot.operational.checkpoint_only)
1353 {
1354 operation.source_checkpoint_only = true;
1355 crate::database::save_move_operation(operation)?;
1356 return Ok(());
1357 }
1358 if snapshot.as_ref().is_some_and(|snapshot| {
1359 matches!(
1360 snapshot.operational.execution,
1361 RelayExecutionState::Closing | RelayExecutionState::Closed
1362 )
1363 }) {
1364 return Ok(());
1365 }
1366 if !operation.source_checkpoint_only
1367 && snapshot
1368 .as_ref()
1369 .is_some_and(|snapshot| snapshot.operational.native_session_is_ready())
1370 {
1371 return Ok(());
1372 }
1373 ensure!(
1374 !executor.cancellation_requested() && !operation.cancellation_requested,
1375 "Move cancelled before source recovery"
1376 );
1377 operation.source_checkpoint_only = true;
1378 crate::database::save_move_operation(operation)?;
1379 executor.notify_notice("Recovering source data without starting its old harness");
1380 let (backend, worker_root) = self.worker_placement(session_id)?;
1381 let reconnect = targets::reconnect_plan(&backend, session_id)?
1382 .commands
1383 .into_iter()
1384 .next()
1385 .context("reconnect plan is empty")?;
1386 let launch = self.current_worker_launch_config(session_id, &backend)?;
1387 let connection = self
1388 .restart_worker_with_installed_binary(
1389 session_id,
1390 executor,
1391 InstalledWorkerRestart {
1392 backend: &backend,
1393 worker_root: &worker_root,
1394 reconnect: &reconnect,
1395 launch: Some(&launch),
1396 messages: &RESTART_FOR_CHECKPOINT,
1397 },
1398 )
1399 .await?;
1400 adopt_restarted_checkpoint_relay(session_id, Some(manager), connection)
1401 .await?
1402 .release();
1403 Ok(())
1404 }
1405
1406 async fn open_checkpoint_relay(
1410 &self,
1411 session_id: &str,
1412 executor: &(impl CommandExecutor + Sync),
1413 manager: Option<&SessionManagerControl>,
1414 target: InstalledWorkerRestart<'_>,
1415 restart_if_unreachable: bool,
1416 ) -> Result<(ControllerRelayLease, bool)> {
1417 let project_memory = match self.project_memory_sync_target(session_id) {
1418 Ok(target) => Some(target),
1419 Err(error) => {
1420 tracing::warn!(
1421 session_id,
1422 error = format!("{error:#}"),
1423 "project memory will not be synchronized during checkpoint reconnect"
1424 );
1425 None
1426 }
1427 };
1428 match connect_checkpoint_relay(
1429 session_id,
1430 manager,
1431 target.reconnect,
1432 project_memory.clone(),
1433 )
1434 .await
1435 {
1436 Ok(relay) => Ok((relay, false)),
1437 Err(error) if worker_connect_needs_restart(&error) && restart_if_unreachable => {
1438 tracing::warn!(
1439 session_id,
1440 "checkpoint could not reach the worker; restarting it: {error:#}"
1441 );
1442 let mut connection = self
1443 .restart_worker_for_checkpoint(
1444 session_id,
1445 executor,
1446 target.backend,
1447 target.worker_root,
1448 target.reconnect,
1449 )
1450 .await?;
1451 connection.set_project_memory_target(project_memory);
1452 let relay =
1453 adopt_restarted_checkpoint_relay(session_id, manager, connection).await?;
1454 Ok((relay, true))
1455 }
1456 Err(error) if worker_connect_needs_restart(&error) => {
1457 Err(error.context(CheckpointDeferred::background_work()))
1458 }
1459 Err(error) => Err(error).context("connect to the session worker for checkpoint"),
1460 }
1461 }
1462
1463 async fn restart_worker_for_checkpoint(
1467 &self,
1468 session_id: &str,
1469 executor: &(impl CommandExecutor + Sync),
1470 backend: &targets::TargetLocator,
1471 worker_root: &str,
1472 reconnect: &targets::CommandSpec,
1473 ) -> Result<StandaloneSession> {
1474 self.restart_worker_with_installed_binary(
1475 session_id,
1476 executor,
1477 InstalledWorkerRestart {
1478 backend,
1479 worker_root,
1480 reconnect,
1481 launch: None,
1482 messages: &RESTART_FOR_CHECKPOINT,
1483 },
1484 )
1485 .await
1486 }
1487}
1488
1489async fn connect_checkpoint_relay(
1490 session_id: &str,
1491 manager: Option<&SessionManagerControl>,
1492 reconnect: &targets::CommandSpec,
1493 project_memory: Option<crate::session_manager::ProjectMemorySyncTarget>,
1494) -> Result<ControllerRelayLease> {
1495 if let Some(manager) = manager {
1496 let handle = manager
1497 .wait_for_session(session_id, Duration::from_secs(5))
1498 .await?;
1499 let mut lease = handle.lease_connection().await?;
1500 lease
1501 .connection_mut()
1502 .set_project_memory_target(project_memory);
1503 Ok(ControllerRelayLease::Managed {
1504 handle,
1505 lease: Some(lease),
1506 })
1507 } else {
1508 let target = crate::session_manager::RelaySessionTarget {
1509 session_id: session_id.to_owned(),
1510 spec: reconnect.clone(),
1511 worker_recovery: None,
1512 project_memory,
1513 };
1514 Ok(ControllerRelayLease::Standalone(
1515 StandaloneSession::connect(&target).await?,
1516 ))
1517 }
1518}
1519
1520async fn adopt_restarted_checkpoint_relay(
1521 session_id: &str,
1522 manager: Option<&SessionManagerControl>,
1523 connection: StandaloneSession,
1524) -> Result<ControllerRelayLease> {
1525 let Some(manager) = manager else {
1526 return Ok(ControllerRelayLease::Standalone(connection));
1527 };
1528 let handle = manager
1529 .wait_for_session(session_id, Duration::from_secs(5))
1530 .await?;
1531 match handle.lease_connection().await {
1532 Ok(mut lease) => {
1533 lease.replace_connection(connection);
1534 Ok(ControllerRelayLease::Managed {
1535 handle,
1536 lease: Some(lease),
1537 })
1538 }
1539 Err(error) => {
1540 tracing::warn!(
1541 session_id,
1542 "session actor could not lease after worker restart; using the restarted proxy: {error:#}"
1543 );
1544 Ok(ControllerRelayLease::Standalone(connection))
1545 }
1546 }
1547}
1548
1549#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1551enum BarrierBusyPolicy {
1552 DeferWhileRunning,
1558 InterruptWhileRunning,
1562}
1563
1564impl BarrierBusyPolicy {
1565 fn of(exclusivity: LatchExclusivity) -> Self {
1566 match exclusivity {
1567 LatchExclusivity::ReleaseAfterLatch => Self::DeferWhileRunning,
1568 LatchExclusivity::HoldThroughClose => Self::InterruptWhileRunning,
1569 }
1570 }
1571}
1572
1573async fn wait_for_checkpoint_barrier(
1574 relay: &mut StandaloneSession,
1575 session_id: &str,
1576 command_id: &str,
1577 timeout: Duration,
1578 busy: BarrierBusyPolicy,
1579 harness: HarnessKind,
1580) -> Result<ManagedSessionSnapshot> {
1581 let deadline = tokio::time::Instant::now() + timeout;
1582 let mut cancel_submitted = false;
1583 let mut cancel_deadline = None;
1584 let mut cancel_started_at: Option<Instant> = None;
1585 loop {
1586 let snapshot = relay.sync().await?;
1587 if busy == BarrierBusyPolicy::DeferWhileRunning
1588 && !snapshot.operational.safe_for_checkpoint(harness)
1589 {
1590 return Err(
1595 CheckpointDeferred::background_snapshot(&snapshot.operational, harness).into(),
1596 );
1597 }
1598 if checkpoint_barrier_is_ready(&snapshot, command_id) {
1599 if let Some(started_at) = cancel_started_at {
1600 tracing::info!(
1601 session_id,
1602 barrier_command_id = command_id,
1603 cancellation_ms = started_at.elapsed().as_millis() as u64,
1604 "active turn cancellation settled before checkpoint barrier"
1605 );
1606 }
1607 return Ok(snapshot);
1608 }
1609 if busy == BarrierBusyPolicy::InterruptWhileRunning
1610 && snapshot.operational.execution == RelayExecutionState::Running
1611 && !cancel_submitted
1612 {
1613 let cancel_turn = RelayCommand::CancelTurn;
1614 if relay.protocol_version() < cancel_turn.minimum_protocol() {
1615 return Err(CheckpointBarrierUnreachable::cancel_turn_unavailable(
1616 command_id,
1617 relay.protocol_version(),
1618 )
1619 .into());
1620 }
1621 let cancel_command_id = new_command_id("checkpoint-cancel-turn")?;
1622 match relay.submit(cancel_command_id, cancel_turn).await {
1623 Ok(_) => {
1624 cancel_submitted = true;
1625 cancel_started_at = Some(Instant::now());
1626 cancel_deadline = Some(tokio::time::Instant::now() + CHECKPOINT_CANCEL_TIMEOUT);
1627 tracing::info!(
1628 session_id,
1629 barrier_command_id = command_id,
1630 "requested active turn cancellation before checkpoint barrier"
1631 );
1632 }
1633 Err(error) if checkpoint_cancel_turn_needs_worker_restart(&error) => {
1634 return Err(error.context(
1635 CheckpointBarrierUnreachable::cancel_turn_unavailable(
1636 command_id,
1637 relay.protocol_version(),
1638 ),
1639 ));
1640 }
1641 Err(error) if worker_connect_needs_restart(&error) => {
1642 return Err(error.context(
1643 CheckpointBarrierUnreachable::cancel_turn_unreachable(command_id),
1644 ));
1645 }
1646 Err(error) => {
1647 if let Ok(snapshot) = relay.sync().await
1651 && checkpoint_barrier_is_ready(&snapshot, command_id)
1652 {
1653 tracing::info!(
1654 session_id,
1655 barrier_command_id = command_id,
1656 "active turn settled while submitting checkpoint cancellation"
1657 );
1658 return Ok(snapshot);
1659 }
1660 return Err(error.context("cancel active ACP turn before checkpoint barrier"));
1661 }
1662 }
1663 continue;
1664 }
1665 let out_of_time = tokio::time::Instant::now() >= cancel_deadline.unwrap_or(deadline);
1666 if let Some(error) = checkpoint_barrier_wait_ended(
1667 &snapshot,
1668 command_id,
1669 busy,
1670 out_of_time,
1671 cancel_submitted,
1672 ) {
1673 return Err(error);
1674 }
1675 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
1676 }
1677}
1678
1679fn checkpoint_barrier_wait_ended(
1686 snapshot: &ManagedSessionSnapshot,
1687 command_id: &str,
1688 busy: BarrierBusyPolicy,
1689 out_of_time: bool,
1690 cancel_submitted: bool,
1691) -> Option<anyhow::Error> {
1692 if snapshot.operational.execution == RelayExecutionState::Closed {
1693 return Some(CheckpointBarrierUnreachable::runtime_stopped().into());
1694 }
1695 if snapshot.operational.execution == RelayExecutionState::Running {
1696 return Some(match busy {
1697 BarrierBusyPolicy::DeferWhileRunning => CheckpointDeferred::harness_busy().into(),
1698 BarrierBusyPolicy::InterruptWhileRunning if out_of_time && cancel_submitted => {
1699 CheckpointBarrierUnreachable::cancel_timed_out(command_id).into()
1700 }
1701 BarrierBusyPolicy::InterruptWhileRunning => return None,
1702 });
1703 }
1704 out_of_time.then(|| CheckpointBarrierUnreachable::not_admitted(command_id).into())
1705}
1706
1707#[derive(Debug)]
1714struct CheckpointBarrierUnreachable(String);
1715
1716impl CheckpointBarrierUnreachable {
1717 fn runtime_stopped() -> Self {
1718 Self("ACP runtime stopped before reaching the checkpoint barrier".to_owned())
1719 }
1720
1721 fn not_admitted(command_id: &str) -> Self {
1722 Self(format!(
1723 "ACP relay did not reach checkpoint barrier {command_id}"
1724 ))
1725 }
1726
1727 fn cancel_timed_out(command_id: &str) -> Self {
1728 Self(format!(
1729 "active ACP turn did not settle after cancellation before checkpoint barrier {command_id}"
1730 ))
1731 }
1732
1733 fn cancel_turn_unavailable(command_id: &str, protocol_version: u32) -> Self {
1734 Self(format!(
1735 "worker protocol {protocol_version} cannot cancel the active ACP turn before checkpoint barrier {command_id} (requires protocol {})",
1736 RelayCommand::CancelTurn.minimum_protocol(),
1737 ))
1738 }
1739
1740 fn cancel_turn_unreachable(command_id: &str) -> Self {
1741 Self(format!(
1742 "worker transport became unavailable while cancelling the active ACP turn before checkpoint barrier {command_id}"
1743 ))
1744 }
1745}
1746
1747impl std::fmt::Display for CheckpointBarrierUnreachable {
1748 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1749 formatter.write_str(&self.0)
1750 }
1751}
1752
1753impl std::error::Error for CheckpointBarrierUnreachable {}
1754
1755fn checkpoint_barrier_needs_worker_restart(error: &anyhow::Error) -> bool {
1756 error
1757 .downcast_ref::<CheckpointBarrierUnreachable>()
1758 .is_some()
1759}
1760
1761fn checkpoint_cancel_turn_needs_worker_restart(error: &anyhow::Error) -> bool {
1766 error.chain().any(|cause| {
1767 let Some(rejected) = cause.downcast_ref::<RelayRejected>() else {
1768 return false;
1769 };
1770 rejected.0.code == mj_core::relay::RelayErrorCode::IncompatibleProtocol
1771 })
1772}
1773
1774#[derive(Debug)]
1784pub struct CheckpointDeferred(String);
1785
1786impl CheckpointDeferred {
1787 pub fn harness_busy() -> Self {
1788 Self("the agent is working; try again when it is idle".to_owned())
1789 }
1790
1791 fn background_work() -> Self {
1792 Self("Kimi background-agent state could not be synchronized; checkpoint requires a synchronized empty task list".into())
1793 }
1794
1795 fn background_snapshot(
1796 state: &mj_core::relay::RelayOperationalState,
1797 harness: HarnessKind,
1798 ) -> Self {
1799 Self(
1800 state
1801 .checkpoint_background_blocker(harness)
1802 .unwrap_or("background state changed during checkpoint")
1803 .into(),
1804 )
1805 }
1806
1807 fn frontier_moved() -> Self {
1808 Self(
1809 "the session moved past the checkpoint-ready cursor before the barrier latched, so this checkpoint was deferred"
1810 .to_owned(),
1811 )
1812 }
1813
1814 fn harness_turn_during_capture() -> Self {
1815 Self(
1816 "the agent started a turn of its own while target state was captured, so this checkpoint was deferred"
1817 .to_owned(),
1818 )
1819 }
1820}
1821
1822impl std::fmt::Display for CheckpointDeferred {
1823 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1824 formatter.write_str(&self.0)
1825 }
1826}
1827
1828impl std::error::Error for CheckpointDeferred {}
1829
1830pub fn checkpoint_was_deferred(error: &anyhow::Error) -> bool {
1837 error.downcast_ref::<CheckpointDeferred>().is_some()
1838}
1839
1840pub struct IdleWorkspaceLease {
1844 lease: ManagedSessionLease,
1845 command_id: String,
1846 harness: HarnessKind,
1847}
1848
1849impl IdleWorkspaceLease {
1850 pub async fn acquire(handle: &ManagedSessionHandle, harness: HarnessKind) -> Result<Self> {
1851 tokio::time::timeout(Duration::from_secs(30), async {
1852 let mut lease = handle.lease_connection().await?;
1853 let snapshot = lease.connection_mut().sync().await?;
1854 ensure!(
1855 snapshot.operational.safe_to_replace(harness),
1856 "session must be live and idle with no queued or background work"
1857 );
1858 let command_id = new_command_id("workspace-write")?;
1859 lease
1860 .connection_mut()
1861 .submit(
1862 command_id.clone(),
1863 RelayCommand::BeginCheckpoint {
1864 reason: Some("API workspace file write".into()),
1865 },
1866 )
1867 .await?;
1868 loop {
1869 let snapshot = lease.connection_mut().sync().await?;
1870 if checkpoint_barrier_is_ready(&snapshot, &command_id) {
1871 let mut operation = Self {
1872 lease,
1873 command_id,
1874 harness,
1875 };
1876 operation.verify().await?;
1877 return Ok(operation);
1878 }
1879 ensure!(
1880 snapshot.operational.execution != RelayExecutionState::Running,
1881 "session started work before the file barrier was ready"
1882 );
1883 tokio::time::sleep(Duration::from_millis(25)).await;
1884 }
1885 })
1886 .await
1887 .context("session did not become available for a file write within 30 seconds")?
1888 }
1889
1890 pub async fn verify(&mut self) -> Result<()> {
1891 let mut snapshot = self.lease.connection_mut().sync().await?;
1892 ensure!(
1893 checkpoint_barrier_is_ready(&snapshot, &self.command_id),
1894 "file write lost its workspace barrier"
1895 );
1896 snapshot.operational.checkpoint_barrier = None;
1897 snapshot.operational.queued_prompts.clear();
1900 ensure!(
1901 snapshot.operational.safe_to_replace(self.harness),
1902 "session is no longer idle for the file write"
1903 );
1904 Ok(())
1905 }
1906
1907 pub async fn release(mut self) -> Result<()> {
1908 tokio::time::timeout(Duration::from_secs(30), async {
1909 self.lease
1910 .connection_mut()
1911 .submit(
1912 new_command_id("workspace-release")?,
1913 RelayCommand::ReleaseCheckpoint {
1914 barrier_command_id: self.command_id.clone(),
1915 },
1916 )
1917 .await?;
1918 loop {
1919 let snapshot = self.lease.connection_mut().sync().await?;
1920 if snapshot.operational.checkpoint_barrier.as_deref() != Some(&self.command_id) {
1921 return Ok::<_, anyhow::Error>(());
1922 }
1923 tokio::time::sleep(Duration::from_millis(25)).await;
1924 }
1925 })
1926 .await
1927 .context("release file write barrier timed out")??;
1928 self.lease.release();
1929 Ok(())
1930 }
1931}
1932
1933fn checkpoint_barrier_is_ready(snapshot: &ManagedSessionSnapshot, command_id: &str) -> bool {
1934 snapshot.operational.checkpoint_barrier.as_deref() == Some(command_id)
1935 && snapshot.operational.checkpoint_ready.is_some()
1936}
1937
1938fn ensure_exact_checkpoint_cut(
1946 cursor: &RelayCursor,
1947 expected_ordinal: u64,
1948 expected_digest: &str,
1949) -> Result<()> {
1950 if cursor.ordinal != expected_ordinal || cursor.digest != expected_digest {
1951 bail!(CheckpointDeferred::frontier_moved());
1952 }
1953 Ok(())
1954}
1955
1956fn validate_checkpoint_barrier_snapshot(
1971 snapshot: &ManagedSessionSnapshot,
1972 command_id: &str,
1973 expected: &RelayCursor,
1974) -> Result<()> {
1975 ensure!(
1976 snapshot.operational.checkpoint_barrier.as_deref() == Some(command_id),
1977 "checkpoint barrier {command_id} is no longer active"
1978 );
1979 ensure!(
1980 snapshot.operational.checkpoint_ready.as_ref() == Some(expected),
1981 "checkpoint barrier {command_id} has a different ready cursor"
1982 );
1983 if snapshot
1984 .operational
1985 .last_harness_turn_started_ordinal
1986 .is_some_and(|ordinal| ordinal > expected.ordinal)
1987 {
1988 bail!(CheckpointDeferred::harness_turn_during_capture());
1989 }
1990 Ok(())
1991}
1992
1993fn validate_automatic_checkpoint_barrier_snapshot(
1997 snapshot: &ManagedSessionSnapshot,
1998 command_id: &str,
1999 expected: &RelayCursor,
2000 harness: HarnessKind,
2001) -> Result<()> {
2002 validate_checkpoint_barrier_snapshot(snapshot, command_id, expected)?;
2003 ensure!(
2004 snapshot.operational.safe_for_checkpoint(harness),
2005 CheckpointDeferred::background_snapshot(&snapshot.operational, harness)
2006 );
2007 Ok(())
2008}
2009
2010fn remove_uninstalled_checkpoint(path: &Path, error: anyhow::Error) -> anyhow::Error {
2011 match std::fs::remove_file(path) {
2012 Ok(()) => error,
2013 Err(remove_error) if remove_error.kind() == std::io::ErrorKind::NotFound => error,
2014 Err(remove_error) => error.context(format!(
2015 "also failed to remove uninstalled checkpoint {}: {remove_error}",
2016 path.display()
2017 )),
2018 }
2019}
2020
2021pub(super) async fn wait_for_relay_closed(relay: &mut StandaloneSession) -> Result<()> {
2022 let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(30);
2023 loop {
2024 if relay.sync().await?.operational.execution == RelayExecutionState::Closed {
2025 return Ok(());
2026 }
2027 if tokio::time::Instant::now() >= deadline {
2028 bail!("ACP runtime did not close within 30 seconds");
2029 }
2030 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
2031 }
2032}
2033
2034async fn release_checkpoint_after_capture(
2046 relay: &mut ControllerRelayLease,
2047 session_id: &str,
2048 barrier_command_id: &str,
2049 cursor: &RelayCursor,
2050 harness: HarnessKind,
2051) -> Result<CheckpointCompletion> {
2052 relay
2053 .sync_snapshot()
2054 .await
2055 .and_then(|snapshot| {
2056 validate_automatic_checkpoint_barrier_snapshot(
2057 &snapshot,
2058 barrier_command_id,
2059 cursor,
2060 harness,
2061 )
2062 })
2063 .context("checkpoint barrier changed while capturing target state")?;
2064 match relay
2065 .submit(
2066 new_command_id("checkpoint-release")?,
2067 RelayCommand::ReleaseCheckpoint {
2068 barrier_command_id: barrier_command_id.to_owned(),
2069 },
2070 )
2071 .await
2072 {
2073 Ok(_) => Ok(CheckpointCompletion::ReleasedAfterCapture),
2074 Err(error) => {
2075 tracing::debug!(
2076 session_id,
2077 "relay kept the checkpoint barrier through the transfer: {error:#}"
2078 );
2079 Ok(CheckpointCompletion::HeldBarrier)
2080 }
2081 }
2082}
2083
2084fn run_checkpoint_staging_command<T: serde::Serialize>(
2085 executor: &impl CommandExecutor,
2086 locator: &targets::TargetLocator,
2087 session_id: &str,
2088 spec: &T,
2089 command: fn(&targets::TargetLocator, &str) -> Result<CommandSpec>,
2090 operation: &str,
2091) -> Result<CommandOutput> {
2092 let body = serde_json::to_vec(spec).with_context(|| format!("serialize {operation} spec"))?;
2093 let mut replaced_worker = false;
2094 loop {
2095 let command = command(locator, session_id)?;
2096 let output = executor.execute_with_stdin(&command, &mut body.as_slice())?;
2097 if output.status == 0 {
2098 return Ok(output);
2099 }
2100 let failure = String::from_utf8_lossy(&output.stderr).into_owned();
2101 if staging_protocol_unsupported(&failure)
2102 && replace_stale_export_worker(
2103 executor,
2104 locator,
2105 session_id,
2106 None,
2107 &failure,
2108 &mut replaced_worker,
2109 )?
2110 {
2111 continue;
2112 }
2113 bail!(
2114 "{operation} failed with status {}: {failure}",
2115 output.status
2116 );
2117 }
2118}
2119
2120fn export_target_checkpoint(
2125 executor: &impl CommandExecutor,
2126 locator: &targets::TargetLocator,
2127 session_id: &str,
2128 spec: &CheckpointExportSpec,
2129 remote_spec: &str,
2130) -> Result<CommandOutput> {
2131 export_target_checkpoint_with_worker(executor, locator, session_id, spec, remote_spec, None)
2132}
2133
2134fn export_target_checkpoint_with_worker(
2135 executor: &impl CommandExecutor,
2136 locator: &targets::TargetLocator,
2137 session_id: &str,
2138 spec: &CheckpointExportSpec,
2139 remote_spec: &str,
2140 worker_binary: Option<&Path>,
2141) -> Result<CommandOutput> {
2142 let body = serde_json::to_vec(spec).context("serialize checkpoint export spec")?;
2143 let mut replaced_worker = false;
2144 loop {
2145 let streamed = export_stdin_command(locator, session_id)?;
2146 let output = executor.execute_with_stdin(&streamed, &mut body.as_slice())?;
2147 if output.status == 0 {
2148 return Ok(output);
2149 }
2150 let failure = String::from_utf8_lossy(&output.stderr).into_owned();
2151 if export_spec_stdin_unsupported(&failure) {
2152 tracing::debug!(
2153 session_id,
2154 "target worker predates streamed checkpoint specs; uploading the spec file instead"
2155 );
2156 let output = export_uploaded_spec(executor, locator, session_id, spec, remote_spec)?;
2157 if output.status == 0 {
2158 return Ok(output);
2159 }
2160 let failure = String::from_utf8_lossy(&output.stderr).into_owned();
2161 if replace_stale_export_worker(
2162 executor,
2163 locator,
2164 session_id,
2165 worker_binary,
2166 &failure,
2167 &mut replaced_worker,
2168 )? {
2169 continue;
2170 }
2171 bail!(
2172 "export target checkpoint failed with status {}: {failure}",
2173 output.status
2174 );
2175 }
2176 if replace_stale_export_worker(
2177 executor,
2178 locator,
2179 session_id,
2180 worker_binary,
2181 &failure,
2182 &mut replaced_worker,
2183 )? {
2184 continue;
2185 }
2186 bail!(
2187 "{} failed with status {}: {failure}",
2188 streamed.purpose,
2189 output.status
2190 );
2191 }
2192}
2193
2194fn export_uploaded_spec(
2195 executor: &impl CommandExecutor,
2196 locator: &targets::TargetLocator,
2197 session_id: &str,
2198 spec: &CheckpointExportSpec,
2199 remote_spec: &str,
2200) -> Result<CommandOutput> {
2201 let staging = tempfile::tempdir().context("create checkpoint staging")?;
2202 let local_spec = staging.path().join("checkpoint-spec.json");
2203 spec.write(&local_spec)?;
2204 upload_checkpoint_spec(executor, locator, session_id, &local_spec, remote_spec)?;
2205 let output = executor.execute(&export_command(locator, session_id, remote_spec)?)?;
2206 if output.status == 0 {
2207 execute_checked(
2208 executor,
2209 targets::command_on_locator(
2210 locator,
2211 session_id,
2212 ["rm", "-f", "--", remote_spec].map(str::to_owned).to_vec(),
2213 "remove uploaded checkpoint specification",
2214 )?,
2215 )
2216 .context("clean successful checkpoint export specification")?;
2217 }
2218 Ok(output)
2219}
2220
2221fn replace_stale_export_worker(
2226 executor: &impl CommandExecutor,
2227 locator: &targets::TargetLocator,
2228 session_id: &str,
2229 worker_binary: Option<&Path>,
2230 failure: &str,
2231 replaced_worker: &mut bool,
2232) -> Result<bool> {
2233 if *replaced_worker || !staging_protocol_unsupported(failure) {
2234 return Ok(false);
2235 }
2236 tracing::debug!(
2237 session_id,
2238 "target worker does not support this checkpoint export protocol; replacing the installed Mjolnir binary and retrying"
2239 );
2240 let owned_binary;
2241 let binary = if let Some(path) = worker_binary {
2242 path
2243 } else {
2244 owned_binary = super::worker_binary::worker_binary_for(locator, executor)?;
2245 owned_binary.as_path()
2246 };
2247 super::worker_binary::replace_installed_worker_binary(executor, locator, session_id, binary)?;
2248 *replaced_worker = true;
2249 Ok(true)
2250}
2251
2252fn export_spec_stdin_unsupported(failure: &str) -> bool {
2260 failure.contains("read checkpoint export spec -")
2261 || failure.contains("unexpected argument")
2262 || failure.contains("invalid value")
2263}
2264
2265fn export_spec_schema_unsupported(failure: &str) -> bool {
2270 failure.contains("parse checkpoint")
2271 && (failure.contains("unknown field") || failure.contains("unknown variant"))
2272}
2273
2274fn export_protocol_unsupported(failure: &str) -> bool {
2275 export_spec_schema_unsupported(failure)
2276 || failure.contains("unsupported checkpoint export protocol version")
2277}
2278
2279fn staging_protocol_unsupported(failure: &str) -> bool {
2280 export_protocol_unsupported(failure)
2281 || failure.contains("unsupported checkpoint staging protocol version")
2282 || failure.contains("unrecognized subcommand")
2283 || failure.contains("unexpected argument")
2284}
2285
2286pub(super) fn upload_checkpoint_spec(
2287 executor: &impl CommandExecutor,
2288 locator: &targets::TargetLocator,
2289 session_id: &str,
2290 local: &Path,
2291 remote: &str,
2292) -> Result<()> {
2293 match locator {
2294 targets::TargetLocator::LocalBare { .. } => {
2295 std::fs::copy(local, remote)
2296 .with_context(|| format!("copy checkpoint specification to {remote}"))?;
2297 Ok(())
2298 }
2299 targets::TargetLocator::LocalPodman { container_id, .. } => execute_checked(
2300 executor,
2301 CommandSpec::new(
2302 "podman",
2303 [
2304 "cp".into(),
2305 local.to_string_lossy().into_owned(),
2306 format!("{container_id}:{remote}"),
2307 ],
2308 )
2309 .purpose("upload checkpoint specification"),
2310 )
2311 .map(|_| ()),
2312 targets::TargetLocator::LocalDocker { container_id } => execute_checked(
2313 executor,
2314 CommandSpec::new(
2315 "docker",
2316 [
2317 "cp".into(),
2318 local.to_string_lossy().into_owned(),
2319 format!("{container_id}:{remote}"),
2320 ],
2321 )
2322 .purpose("upload checkpoint specification"),
2323 )
2324 .map(|_| ()),
2325 targets::TargetLocator::AppleContainer { container_id } => execute_checked(
2326 executor,
2327 CommandSpec::new(
2328 "container",
2329 [
2330 "cp".into(),
2331 local.to_string_lossy().into_owned(),
2332 format!("{container_id}:{remote}"),
2333 ],
2334 )
2335 .purpose("upload checkpoint specification"),
2336 )
2337 .map(|_| ()),
2338 targets::TargetLocator::AwsEc2 { ssh, .. }
2339 | targets::TargetLocator::SshBare { ssh, .. } => execute_checked(
2340 executor,
2341 scp_command_spec(ssh, local, remote, false).purpose("upload checkpoint specification"),
2342 )
2343 .map(|_| ()),
2344 targets::TargetLocator::SshPodman {
2345 ssh, container_id, ..
2346 }
2347 | targets::TargetLocator::SshDocker { ssh, container_id } => {
2348 let engine = match locator {
2349 targets::TargetLocator::SshPodman { .. } => "podman",
2350 targets::TargetLocator::SshDocker { .. } => "docker",
2351 _ => unreachable!("matched remote container target"),
2352 };
2353 let staging = format!(".local/share/hel/uploads/{session_id}-checkpoint.json");
2354 execute_checked(
2355 executor,
2356 ssh_command_spec(ssh, ["mkdir", "-p", ".local/share/hel/uploads"])
2357 .purpose("create remote checkpoint staging"),
2358 )?;
2359 execute_checked(
2360 executor,
2361 scp_command_spec(ssh, local, &staging, false)
2362 .purpose("upload remote container checkpoint specification"),
2363 )?;
2364 execute_checked(
2365 executor,
2366 ssh_command_spec(
2367 ssh,
2368 [engine, "cp", &staging, &format!("{container_id}:{remote}")],
2369 )
2370 .purpose("install remote container checkpoint specification"),
2371 )?;
2372 execute_checked(
2373 executor,
2374 ssh_command_spec(ssh, ["rm", "-f", "--", &staging])
2375 .purpose("remove remote checkpoint staging"),
2376 )?;
2377 Ok(())
2378 }
2379 }?;
2380 Ok(())
2381}
2382
2383fn reusable_installed_checkpoint(
2391 session_id: &str,
2392 installed: Option<&CheckpointMetadata>,
2393 native_session_id: &str,
2394 latched_ordinal: u64,
2395 latched_session: &CanonicalSessionSnapshot,
2396) -> Option<CheckpointArtifact> {
2397 let installed = installed?;
2398 if installed.event_frontier > latched_ordinal {
2399 tracing::warn!(
2400 session_id,
2401 installed_frontier = installed.event_frontier,
2402 latched_ordinal,
2403 "installed checkpoint is ahead of the latched cursor; exporting a fresh archive"
2404 );
2405 return None;
2406 }
2407 let verified = match verify_archive_streaming(&installed.archive_path) {
2408 Ok(verified) => verified,
2409 Err(error) => {
2410 tracing::warn!(
2411 session_id,
2412 path = %installed.archive_path.display(),
2413 "installed checkpoint could not be verified for reuse: {error:#}"
2414 );
2415 return None;
2416 }
2417 };
2418 if verified.archive_sha256 != installed.sha256
2419 || verified.manifest.session.id != session_id
2420 || verified.canonical_session.event_frontier != installed.event_frontier
2421 {
2422 tracing::warn!(
2423 session_id,
2424 path = %installed.archive_path.display(),
2425 "installed checkpoint no longer matches its controller metadata; exporting a fresh archive"
2426 );
2427 return None;
2428 }
2429 if !verified.canonical_session.content_matches(latched_session) {
2430 tracing::info!(
2431 session_id,
2432 archive_frontier = verified.canonical_session.event_frontier,
2433 latched_ordinal,
2434 "session content changed since the installed checkpoint; exporting a fresh archive"
2435 );
2436 return None;
2437 }
2438 tracing::info!(
2439 session_id,
2440 archive_frontier = verified.canonical_session.event_frontier,
2441 latched_ordinal,
2442 "reusing the installed checkpoint archive; only relay bookkeeping moved"
2443 );
2444 Some(CheckpointArtifact {
2445 metadata: installed.clone(),
2446 native_session_id: native_session_id.to_owned(),
2447 event_frontier_digest: verified.canonical_session.event_frontier_digest,
2448 })
2449}
2450
2451pub(super) fn verify_installed_checkpoint_gate(
2452 session_id: &str,
2453 checkpoint: &CheckpointMetadata,
2454) -> Result<()> {
2455 let sha256 = checkpoint_sha256(&checkpoint.archive_path).with_context(|| {
2456 format!(
2457 "hash installed checkpoint {} before target cleanup",
2458 checkpoint.archive_path.display()
2459 )
2460 })?;
2461 ensure!(
2462 sha256 == checkpoint.sha256,
2463 "refusing target cleanup for session {session_id}: installed checkpoint SHA changed"
2464 );
2465 Ok(())
2466}
2467
2468fn verify_checkpoint_artifact(session_id: &str, artifact: &CheckpointArtifact) -> Result<()> {
2469 let sha256 = checkpoint_sha256(&artifact.metadata.archive_path).with_context(|| {
2470 format!(
2471 "hash completed checkpoint {}",
2472 artifact.metadata.archive_path.display()
2473 )
2474 })?;
2475 ensure!(
2476 sha256 == artifact.metadata.sha256,
2477 "completed checkpoint SHA changed before persistence for session {session_id}"
2478 );
2479 Ok(())
2480}
2481
2482pub(super) fn release_projection_behind_checkpoint(session_id: &str, current: &CheckpointMetadata) {
2490 match crate::database::compact_materialized_transcript_through(
2491 session_id,
2492 current.event_frontier,
2493 ) {
2494 Ok(retention) if retention.items == 0 => {}
2495 Ok(retention) => tracing::info!(
2496 session_id,
2497 items = retention.items,
2498 bytes = retention.bytes,
2499 remaining = retention.remaining,
2500 event_frontier = current.event_frontier,
2501 "released projection history the checkpoint covers"
2502 ),
2503 Err(error) => tracing::warn!(
2504 session_id,
2505 "checkpoint was saved, but the projection history it covers could not be released: {error:#}"
2506 ),
2507 }
2508}
2509
2510pub(super) fn prune_replaced_checkpoint(
2511 previous: Option<&CheckpointMetadata>,
2512 current: &CheckpointMetadata,
2513) {
2514 let Some(previous) = previous.filter(|old| old.archive_path != current.archive_path) else {
2515 return;
2516 };
2517 match crate::database::move_checkpoint_is_retained(&previous.archive_path) {
2518 Ok(true) => return,
2519 Ok(false) => {}
2520 Err(error) => {
2521 tracing::warn!(%error, "could not check move retention; keeping superseded checkpoint");
2522 return;
2523 }
2524 }
2525 if let Err(error) = std::fs::remove_file(&previous.archive_path)
2526 && error.kind() != std::io::ErrorKind::NotFound
2527 {
2528 tracing::warn!(
2529 path = %previous.archive_path.display(),
2530 "could not remove superseded recovery copy: {error}"
2531 );
2532 }
2533}
2534
2535#[cfg(test)]
2536mod tests {
2537 use std::cell::{Cell, RefCell};
2538 use std::collections::BTreeMap;
2539 use std::fs::OpenOptions;
2540 use std::path::{Path, PathBuf};
2541 #[cfg(unix)]
2542 use std::process::Command;
2543 #[cfg(unix)]
2544 use std::time::Duration;
2545
2546 #[cfg(unix)]
2547 use agent_client_protocol::schema::v1::{ContentBlock, TextContent};
2548 use anyhow::Result;
2549
2550 #[cfg(unix)]
2551 use crate::controller::now;
2552 use crate::controller::restore_session_after_persistence_failure;
2553 use crate::controller::test_support::{checkpoint_test_session, write_checkpoint_gate_archive};
2554 #[cfg(unix)]
2555 use crate::session_manager::{ManagedSessionHandle, new_command_id};
2556 use crate::worker_client::RelayTransportDead;
2557 use mj_checkpoint::archive::{
2558 BundleManifest, CanonicalTranscriptBody, CanonicalTranscriptItem, TargetManifest,
2559 };
2560 use mj_checkpoint::checkpoint::CheckpointExportSpec;
2561 #[cfg(unix)]
2562 use mj_core::config::{
2563 Config, HarnessProfile, ProjectBundle, ProjectRepository, TargetTemplate,
2564 };
2565 #[cfg(unix)]
2566 use mj_core::state::TargetLocator;
2567 use mj_core::state::{
2568 CheckpointMetadata, ManagedSessionSnapshot, MaterializedSession, SessionState, State,
2569 };
2570 use mj_transcript::projection::canonical_session_from_materialized;
2571
2572 #[cfg(unix)]
2573 use crate::targets::ProvisionStage;
2574 use crate::targets::{self, CommandExecutor, CommandOutput, CommandSpec};
2575 #[cfg(unix)]
2576 use mj_core::relay::RelayCommandOutcome;
2577 use mj_core::relay::{RelayCommand, RelayCursor, RelayExecutionState};
2578
2579 use super::*;
2580
2581 struct UnusedExecutor;
2585
2586 impl CommandExecutor for UnusedExecutor {
2587 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
2588 panic!("the export layout ran {command:?}");
2589 }
2590 }
2591
2592 #[test]
2593 fn the_export_layout_places_each_session_kind_in_its_workspace() {
2594 let session_id = "1123456789abcdef0123456789abcdef";
2595 let mut config = crate::controller::test_support::resume_compatibility_config();
2596 config.bundles.insert(
2597 "app-bundle".into(),
2598 mj_core::config::ProjectBundle {
2599 primary_repo: "app".into(),
2600 repositories: vec![mj_core::config::ProjectRepository {
2601 id: "app".into(),
2602 github: None,
2603 local: None,
2604 destination: PathBuf::from("app"),
2605 git_ref: None,
2606 }],
2607 },
2608 );
2609
2610 let mut session = checkpoint_test_session(session_id);
2613 session.bundle_id = "app-bundle".into();
2614 session.target = Some(mj_core::state::TargetLocator::LocalPodman {
2615 container_id: "hel-session".into(),
2616 workspace_storage: Default::default(),
2617 });
2618 let mut state = State::default();
2619 state.sessions.insert(session_id.into(), session.clone());
2620 let controller = Controller {
2621 config: config.clone(),
2622 state,
2623 };
2624
2625 let layout = controller
2626 .session_export_layout(session_id, &UnusedExecutor)
2627 .unwrap();
2628 assert_eq!(layout.workspace_root, "/workspace");
2629 assert_eq!(layout.primary_repository, "app");
2630 assert_eq!(
2631 layout
2632 .repositories
2633 .iter()
2634 .map(|repository| (
2635 repository.id.clone(),
2636 repository.relative_destination.clone()
2637 ))
2638 .collect::<Vec<_>>(),
2639 [("app".to_owned(), PathBuf::from("app"))]
2640 );
2641 assert!(matches!(
2642 layout.repositories[0].capture,
2643 CheckpointRepositoryCapture::RemoteWorkspace
2644 ));
2645 assert!(layout.managed_worktree.is_none());
2646
2647 let mut raw = session;
2650 raw.target_template_id = "local-bare".into();
2651 raw.project_directory = Some(PathBuf::from("/home/dev/project"));
2652 raw.target = Some(mj_core::state::TargetLocator::LocalBare {
2653 worker_root: PathBuf::from("/home/dev/.local/share/hel/workers/session"),
2654 });
2655 let mut state = State::default();
2656 state.sessions.insert(session_id.into(), raw);
2657 let controller = Controller { config, state };
2658
2659 let layout = controller
2660 .session_export_layout(session_id, &UnusedExecutor)
2661 .unwrap();
2662 assert_eq!(layout.workspace_root, "/home/dev");
2663 assert_eq!(layout.primary_repository, "project");
2664 assert_eq!(
2665 layout.repositories[0].relative_destination,
2666 PathBuf::from("project")
2667 );
2668 assert!(matches!(
2669 layout.repositories[0].capture,
2670 CheckpointRepositoryCapture::MetadataOnly
2671 ));
2672 }
2673
2674 fn managed_worktree_export_capture(
2678 clear_recorded_base: bool,
2679 ) -> (CheckpointRepositoryCapture, String, String) {
2680 let session_id = "2123456789abcdef0123456789abcdef";
2681 let repository = crate::controller::test_support::committed_repository();
2682 let mut session = crate::controller::test_support::managed_worktree_session(
2683 repository.path(),
2684 session_id,
2685 );
2686 let creation_commit =
2687 crate::controller::test_support::test_git(repository.path(), &["rev-parse", "HEAD"]);
2688 if clear_recorded_base {
2689 session.managed_worktree.as_mut().unwrap().base_commit = None;
2690 }
2691
2692 let worktree_root = session
2693 .managed_worktree
2694 .as_ref()
2695 .unwrap()
2696 .worktree_root
2697 .clone();
2698 std::fs::write(worktree_root.join("session.txt"), "work\n").unwrap();
2699 crate::controller::test_support::test_git(&worktree_root, &["add", "."]);
2700 crate::controller::test_support::test_git(
2701 &worktree_root,
2702 &["commit", "-m", "session work"],
2703 );
2704 let worktree_head =
2705 crate::controller::test_support::test_git(&worktree_root, &["rev-parse", "HEAD"]);
2706
2707 session.target = Some(mj_core::state::TargetLocator::LocalBare {
2708 worker_root: PathBuf::from("/home/dev/.local/share/hel/workers/session"),
2709 });
2710 let mut state = State::default();
2711 state.sessions.insert(session_id.into(), session);
2712 let controller = Controller {
2713 config: crate::controller::test_support::resume_compatibility_config(),
2714 state,
2715 };
2716
2717 let mut layout = controller
2718 .session_export_layout(session_id, &targets::ProcessExecutor)
2719 .unwrap();
2720 (
2721 layout.repositories.remove(0).capture,
2722 creation_commit,
2723 worktree_head,
2724 )
2725 }
2726
2727 #[test]
2728 fn a_managed_worktree_checkpoint_bundles_from_the_recorded_base() {
2729 let (capture, creation_commit, worktree_head) = managed_worktree_export_capture(false);
2730 let CheckpointRepositoryCapture::DeltaFrom { base_commit } = capture else {
2731 panic!("a managed worktree must be captured as a delta, got {capture:?}");
2732 };
2733 assert_eq!(base_commit, creation_commit);
2734 assert_ne!(base_commit, worktree_head);
2735 }
2736
2737 #[test]
2738 fn a_managed_worktree_without_a_recorded_base_uses_its_branch_creation_commit() {
2739 let (capture, creation_commit, worktree_head) = managed_worktree_export_capture(true);
2740 let CheckpointRepositoryCapture::DeltaFrom { base_commit } = capture else {
2741 panic!("a managed worktree must be captured as a delta, got {capture:?}");
2742 };
2743 assert_eq!(base_commit, creation_commit);
2744 assert_ne!(base_commit, worktree_head);
2745 }
2746
2747 #[test]
2748 fn startup_reconciliation_only_removes_unreferenced_controller_checkpoints() {
2749 let directory = tempfile::tempdir().unwrap();
2750 let session_id = "1123456789abcdef0123456789abcdef";
2751 let referenced_name =
2752 format!("{session_id}-7-archive-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.hel.zip");
2753 let orphan_name =
2754 format!("{session_id}-8-archive-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb.hel.zip");
2755 let imported_name = format!("{session_id}.hel.zip");
2756 for name in [
2757 &referenced_name,
2758 &orphan_name,
2759 &imported_name,
2760 "notes.hel.zip",
2761 ] {
2762 std::fs::write(directory.path().join(name), b"test").unwrap();
2763 }
2764 let mut state = State::default();
2765 let mut session = checkpoint_test_session(session_id);
2766 session.checkpoint = Some(CheckpointMetadata {
2767 archive_path: directory.path().join(&referenced_name),
2768 sha256: "c".repeat(64),
2769 created_at: "2026-08-12T00:00:00Z".into(),
2770 event_frontier: 7,
2771 });
2772 state.sessions.insert(session_id.into(), session);
2773
2774 assert_eq!(
2775 reconcile_managed_checkpoint_archives_in(directory.path(), &state).unwrap(),
2776 1
2777 );
2778 assert!(directory.path().join(referenced_name).exists());
2779 assert!(!directory.path().join(orphan_name).exists());
2780 assert!(directory.path().join(imported_name).exists());
2781 assert!(directory.path().join("notes.hel.zip").exists());
2782 }
2783 #[test]
2784 fn recovery_artifact_final_verification_checks_the_archive_digest() {
2785 let directory = tempfile::tempdir().unwrap();
2786 let session_id = "1123456789abcdef0123456789abcdef";
2787 let metadata = write_checkpoint_gate_archive(directory.path(), session_id, 7);
2788 let mut artifact = CheckpointArtifact {
2789 metadata,
2790 native_session_id: "native-session".into(),
2791 event_frontier_digest: "a".repeat(64),
2792 };
2793
2794 verify_checkpoint_artifact(session_id, &artifact).unwrap();
2795 artifact.metadata.sha256 = "b".repeat(64);
2796 assert!(
2797 verify_checkpoint_artifact(session_id, &artifact)
2798 .unwrap_err()
2799 .to_string()
2800 .contains("checkpoint SHA changed")
2801 );
2802 }
2803 fn checkpoint_barrier_snapshot(cursor: &RelayCursor) -> ManagedSessionSnapshot {
2806 let mut materialized = MaterializedSession::empty("session-1");
2807 materialized.applied_event_ordinal = cursor.ordinal;
2808 materialized.applied_event_digest = cursor.digest.clone();
2809 ManagedSessionSnapshot {
2810 subagent_requests: Vec::new(),
2811 subagent_results: Vec::new(),
2812 window: mj_core::state::ProjectionWindow::of(&materialized),
2813 materialized,
2814 latest_credential_sync_signal: None,
2815 worker_build: None,
2816 operational: mj_core::relay::RelayOperationalState {
2817 goal: serde_json::from_value(
2818 serde_json::json!({"known":true,"execution":{"version":1,"status":"idle"}}),
2819 )
2820 .unwrap(),
2821 capacity_retry: None,
2822 activity_turn_started_at_ms: None,
2823 checkpoint_only: false,
2824 acp_ready: None,
2825 store_id: None,
2826 idle_since_ms: None,
2827 session_id: "session-1".into(),
2828 execution: RelayExecutionState::Idle,
2829 latest_ordinal: cursor.ordinal,
2830 latest_digest: cursor.digest.clone(),
2831 acknowledged_through: cursor.ordinal,
2832 acknowledged_digest: cursor.digest.clone(),
2833 recovery_floor_ordinal: 0,
2834 recovery_floor_digest: mj_core::relay::RELAY_EVENT_GENESIS_DIGEST.into(),
2835 native_session_id: Some("native-session".into()),
2836 agent_capabilities: None,
2837 agent_info: None,
2838 steering_supported: None,
2839 config_options: Vec::new(),
2840 modes: None,
2841 available_commands: Vec::new(),
2842 config: BTreeMap::new(),
2843 active_prompt: None,
2844 queued_prompts: Vec::new(),
2845 active_user_shells: Vec::new(),
2846 active_agent_terminals: Vec::new(),
2847 checkpoint_barrier: Some("checkpoint-1".into()),
2848 checkpoint_ready: None,
2849 last_acp_activity_at_ms: None,
2850 current_step_started_at_ms: None,
2851 foreground_tool_started_at_ms: None,
2852 harness_turn: None,
2853 last_harness_turn_started_ordinal: None,
2854 background_commands: Vec::new(),
2855 background_work_known: None,
2856 },
2857 }
2858 }
2859 #[test]
2860 fn checkpoint_barrier_is_not_reached_until_its_ready_cursor_is_projected() {
2861 let cursor = RelayCursor {
2862 ordinal: 7,
2863 digest: "a".repeat(64),
2864 };
2865 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2866
2867 assert!(!checkpoint_barrier_is_ready(&snapshot, "checkpoint-1"));
2868 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2869 assert!(checkpoint_barrier_is_ready(&snapshot, "checkpoint-1"));
2870 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).unwrap();
2871 }
2872 #[test]
2873 fn checkpoint_revalidation_accepts_a_frontier_that_moved_past_the_ready_cursor() {
2874 let cursor = RelayCursor {
2875 ordinal: 7,
2876 digest: "a".repeat(64),
2877 };
2878 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2879 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2880
2881 snapshot.operational.latest_ordinal = cursor.ordinal + 2;
2885 snapshot.operational.latest_digest = "b".repeat(64);
2886 snapshot.materialized.applied_event_ordinal = cursor.ordinal + 2;
2887 snapshot.materialized.applied_event_digest = "b".repeat(64);
2888 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).unwrap();
2889
2890 snapshot.operational.checkpoint_ready = Some(RelayCursor {
2892 ordinal: cursor.ordinal + 1,
2893 digest: "c".repeat(64),
2894 });
2895 assert!(validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).is_err());
2896 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2897 snapshot.operational.checkpoint_barrier = None;
2898 assert!(validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).is_err());
2899 }
2900
2901 #[test]
2902 fn routine_kimi_checkpoint_defers_when_background_liveness_is_not_safe() {
2903 let cursor = RelayCursor {
2904 ordinal: 7,
2905 digest: "a".repeat(64),
2906 };
2907 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2908 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2909
2910 for (known, has_task, label) in [
2911 (Some(false), false, "tracker reported a failure"),
2912 (Some(true), true, "a native task is still active"),
2913 (None, false, "an older worker omitted the tracker field"),
2914 ] {
2915 snapshot.operational.background_work_known = known;
2916 snapshot.operational.background_commands = has_task
2917 .then(|| mj_core::relay::BackgroundCommand {
2918 id: "kimi:agent-1".into(),
2919 started_at_ms: 1,
2920 command: "background agent".into(),
2921 can_stop: false,
2922 })
2923 .into_iter()
2924 .collect();
2925 let error = validate_automatic_checkpoint_barrier_snapshot(
2926 &snapshot,
2927 "checkpoint-1",
2928 &cursor,
2929 HarnessKind::Kimi,
2930 )
2931 .expect_err(label);
2932 assert!(checkpoint_was_deferred(&error), "{label}: {error:#}");
2933 assert!(!checkpoint_barrier_needs_worker_restart(&error));
2934 }
2935
2936 snapshot.operational.background_work_known = None;
2940 snapshot.operational.background_commands = vec![mj_core::relay::BackgroundCommand {
2941 id: "legacy-task".into(),
2942 started_at_ms: 1,
2943 command: "legacy background work".into(),
2944 can_stop: false,
2945 }];
2946 validate_automatic_checkpoint_barrier_snapshot(
2947 &snapshot,
2948 "checkpoint-1",
2949 &cursor,
2950 HarnessKind::Codex,
2951 )
2952 .expect("non-Kimi checkpoint compatibility");
2953 }
2954 fn exported_checkpoint_json() -> Vec<u8> {
2956 serde_json::to_vec(&mj_checkpoint::checkpoint::TargetCheckpoint {
2957 path: PathBuf::from("/var/lib/hel/workers/session/checkpoint.hel.zip"),
2958 sha256: "c".repeat(64),
2959 event_frontier: 7,
2960 event_frontier_digest: "d".repeat(64),
2961 timings: None,
2962 })
2963 .unwrap()
2964 }
2965 fn export_spec_fixture() -> CheckpointExportSpec {
2966 CheckpointExportSpec {
2967 protocol_version: CHECKPOINT_EXPORT_PROTOCOL_VERSION,
2968 session: mj_checkpoint::archive::SessionManifest {
2969 id: LATCH_RELAY_SESSION.into(),
2970 title: "streamed spec".into(),
2971 harness_kind: mj_core::config::HarnessKind::Codex,
2972 profile_id: "codex".into(),
2973 native_session_id: "native-session".into(),
2974 created_at: "2026-08-12T00:00:00Z".into(),
2975 checkpointed_at: "2026-08-16T00:00:00Z".into(),
2976 hel_version: "test".into(),
2977 relay_version: "test".into(),
2978 adapter_version: "acp-v1".into(),
2979 },
2980 target: TargetManifest {
2981 template_id: "podman".into(),
2982 target_kind: "local-podman".into(),
2983 details: BTreeMap::new(),
2984 },
2985 bundle: BundleManifest {
2986 id: "project".into(),
2987 primary_repository: "app".into(),
2988 },
2989 relay_root: PathBuf::from("/var/lib/hel/workers/session"),
2990 harness_home: PathBuf::from("/var/lib/hel/profiles/codex"),
2991 workspace_root: PathBuf::from("/workspace"),
2992 repositories: Vec::new(),
2993 canonical_session: canonical_session_from_materialized(&MaterializedSession::empty(
2994 LATCH_RELAY_SESSION.to_owned(),
2995 ))
2996 .unwrap(),
2997 output_path: PathBuf::from("/var/lib/hel/workers/session/checkpoint.hel.zip"),
2998 }
2999 }
3000 struct ExportExecutor {
3003 streamed_status: i32,
3004 streamed_stderr: String,
3005 retry_stdin_after_failure: bool,
3006 stdin_calls: Cell<usize>,
3007 purposes: RefCell<Vec<String>>,
3008 streamed_spec: RefCell<Vec<u8>>,
3009 }
3010 impl ExportExecutor {
3011 fn new(streamed_status: i32, streamed_stderr: &str) -> Self {
3012 Self {
3013 streamed_status,
3014 streamed_stderr: streamed_stderr.to_owned(),
3015 retry_stdin_after_failure: false,
3016 stdin_calls: Cell::new(0),
3017 purposes: RefCell::new(Vec::new()),
3018 streamed_spec: RefCell::new(Vec::new()),
3019 }
3020 }
3021
3022 fn retry_stdin_after_failure(mut self) -> Self {
3023 self.retry_stdin_after_failure = true;
3024 self
3025 }
3026 }
3027 impl CommandExecutor for ExportExecutor {
3028 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
3029 self.purposes.borrow_mut().push(command.purpose.clone());
3030 Ok(CommandOutput {
3031 status: 0,
3032 stdout: exported_checkpoint_json(),
3033 stderr: Vec::new(),
3034 })
3035 }
3036
3037 fn execute_with_stdin(
3038 &self,
3039 command: &CommandSpec,
3040 input: &mut (dyn std::io::Read + Send),
3041 ) -> Result<CommandOutput> {
3042 self.purposes.borrow_mut().push(command.purpose.clone());
3043 let mut spec = Vec::new();
3044 input.read_to_end(&mut spec)?;
3045 *self.streamed_spec.borrow_mut() = spec;
3046 let attempt = self.stdin_calls.get();
3047 self.stdin_calls.set(attempt + 1);
3048 let failed =
3049 self.streamed_status != 0 && (attempt == 0 || !self.retry_stdin_after_failure);
3050 Ok(CommandOutput {
3051 status: if failed { self.streamed_status } else { 0 },
3052 stdout: if failed {
3053 Vec::new()
3054 } else {
3055 exported_checkpoint_json()
3056 },
3057 stderr: if failed {
3058 self.streamed_stderr.clone().into_bytes()
3059 } else {
3060 Vec::new()
3061 },
3062 })
3063 }
3064 }
3065 #[test]
3066 fn docker_checkpoint_fallback_upload_uses_docker_cp() {
3067 struct RecordingExecutor {
3068 commands: RefCell<Vec<CommandSpec>>,
3069 }
3070 impl CommandExecutor for RecordingExecutor {
3071 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
3072 self.commands.borrow_mut().push(command.clone());
3073 Ok(CommandOutput {
3074 status: 0,
3075 stdout: Vec::new(),
3076 stderr: Vec::new(),
3077 })
3078 }
3079 }
3080
3081 let executor = RecordingExecutor {
3082 commands: RefCell::new(Vec::new()),
3083 };
3084 let locator = targets::TargetLocator::LocalDocker {
3085 container_id: "hel-session-12345678".to_owned(),
3086 };
3087 upload_checkpoint_spec(
3088 &executor,
3089 &locator,
3090 LATCH_RELAY_SESSION,
3091 Path::new("checkpoint-spec.json"),
3092 "/var/lib/hel/workers/session/checkpoint-spec.json",
3093 )
3094 .unwrap();
3095
3096 let commands = executor.commands.borrow();
3097 assert_eq!(commands.len(), 1);
3098 assert_eq!(commands[0].program, "docker");
3099 assert_eq!(
3100 commands[0].args,
3101 [
3102 "cp",
3103 "checkpoint-spec.json",
3104 "hel-session-12345678:/var/lib/hel/workers/session/checkpoint-spec.json"
3105 ]
3106 );
3107 assert_eq!(commands[0].purpose, "upload checkpoint specification");
3108 }
3109 #[test]
3110 fn checkpoint_export_streams_its_spec_instead_of_uploading_it() {
3111 let locator = targets::TargetLocator::LocalPodman {
3112 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3113 workspace_storage: Default::default(),
3114 };
3115 let spec = export_spec_fixture();
3116 let executor = ExportExecutor::new(0, "");
3117
3118 let output = export_target_checkpoint(
3119 &executor,
3120 &locator,
3121 LATCH_RELAY_SESSION,
3122 &spec,
3123 "/var/lib/hel/workers/session/checkpoint-spec.json",
3124 )
3125 .unwrap();
3126
3127 assert_eq!(output.stdout, exported_checkpoint_json());
3128 assert_eq!(
3129 serde_json::from_slice::<CheckpointExportSpec>(&executor.streamed_spec.borrow())
3130 .unwrap(),
3131 spec
3132 );
3133 assert_eq!(
3134 executor.purposes.into_inner(),
3135 vec!["export target checkpoint".to_owned()]
3136 );
3137 }
3138 #[test]
3141 fn an_export_that_cannot_read_stdin_falls_back_to_uploading_the_spec() {
3142 let locator = targets::TargetLocator::LocalPodman {
3143 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3144 workspace_storage: Default::default(),
3145 };
3146 let executor = ExportExecutor::new(
3147 1,
3148 "Error: read checkpoint export spec -\n\nCaused by:\n \
3149 No such file or directory (os error 2)\n",
3150 );
3151
3152 let output = export_target_checkpoint(
3153 &executor,
3154 &locator,
3155 LATCH_RELAY_SESSION,
3156 &export_spec_fixture(),
3157 "/var/lib/hel/workers/session/checkpoint-spec.json",
3158 )
3159 .unwrap();
3160
3161 assert_eq!(output.stdout, exported_checkpoint_json());
3162 assert_eq!(
3163 executor.purposes.into_inner(),
3164 vec![
3165 "export target checkpoint".to_owned(),
3166 "upload checkpoint specification".to_owned(),
3167 "export target checkpoint".to_owned(),
3168 "remove uploaded checkpoint specification".to_owned(),
3169 ]
3170 );
3171 }
3172 #[test]
3173 fn a_failing_export_is_not_retried_as_an_old_worker() {
3174 let locator = targets::TargetLocator::LocalPodman {
3175 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3176 workspace_storage: Default::default(),
3177 };
3178 let executor = ExportExecutor::new(1, "Error: repository 'app' is missing\n");
3179
3180 let error = export_target_checkpoint(
3181 &executor,
3182 &locator,
3183 LATCH_RELAY_SESSION,
3184 &export_spec_fixture(),
3185 "/var/lib/hel/workers/session/checkpoint-spec.json",
3186 )
3187 .unwrap_err();
3188
3189 assert!(
3190 format!("{error:#}").contains("repository 'app' is missing"),
3191 "{error:#}"
3192 );
3193 assert_eq!(
3194 executor.purposes.into_inner(),
3195 vec!["export target checkpoint".to_owned()]
3196 );
3197 }
3198 #[test]
3201 fn a_legacy_export_worker_is_replaced_before_it_runs_obsolete_behavior() {
3202 let locator = targets::TargetLocator::LocalPodman {
3203 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3204 workspace_storage: Default::default(),
3205 };
3206 let spec = export_spec_fixture();
3207 let executor = ExportExecutor::new(
3208 1,
3209 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3210 unknown field `protocol_version`, expected `session` at line 1 column 20\n",
3211 )
3212 .retry_stdin_after_failure();
3213 let worker_binary = Path::new("/hel-test-worker");
3214
3215 let output = export_target_checkpoint_with_worker(
3216 &executor,
3217 &locator,
3218 LATCH_RELAY_SESSION,
3219 &spec,
3220 "/var/lib/hel/workers/session/checkpoint-spec.json",
3221 Some(worker_binary),
3222 )
3223 .unwrap();
3224
3225 assert_eq!(output.stdout, exported_checkpoint_json());
3226 assert_eq!(
3227 serde_json::from_slice::<CheckpointExportSpec>(&executor.streamed_spec.borrow())
3228 .unwrap(),
3229 spec
3230 );
3231 assert_eq!(
3232 executor.purposes.into_inner(),
3233 vec![
3234 "export target checkpoint".to_owned(),
3235 "stage replacement Mjolnir worker".to_owned(),
3236 "assign replacement worker to the worker user".to_owned(),
3237 "replace installed Mjolnir worker".to_owned(),
3238 "make replaced Mjolnir worker executable".to_owned(),
3239 "export target checkpoint".to_owned(),
3240 ]
3241 );
3242 }
3243 #[test]
3244 fn a_schema_mismatch_after_uploading_the_spec_still_replaces_the_worker_binary() {
3245 let locator = targets::TargetLocator::LocalPodman {
3246 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3247 workspace_storage: Default::default(),
3248 };
3249 struct FileThenRefreshExecutor {
3250 purposes: RefCell<Vec<String>>,
3251 file_export_calls: Cell<usize>,
3252 }
3253 impl CommandExecutor for FileThenRefreshExecutor {
3254 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
3255 self.purposes.borrow_mut().push(command.purpose.clone());
3256 if command.purpose == "export target checkpoint" {
3257 let attempt = self.file_export_calls.get();
3258 self.file_export_calls.set(attempt + 1);
3259 if attempt == 0 {
3260 return Ok(CommandOutput {
3261 status: 1,
3262 stdout: Vec::new(),
3263 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(),
3264 });
3265 }
3266 }
3267 Ok(CommandOutput {
3268 status: 0,
3269 stdout: exported_checkpoint_json(),
3270 stderr: Vec::new(),
3271 })
3272 }
3273
3274 fn execute_with_stdin(
3275 &self,
3276 command: &CommandSpec,
3277 input: &mut (dyn std::io::Read + Send),
3278 ) -> Result<CommandOutput> {
3279 self.purposes.borrow_mut().push(command.purpose.clone());
3280 let mut discarded = Vec::new();
3281 input.read_to_end(&mut discarded)?;
3282 let stdin_calls = self
3283 .purposes
3284 .borrow()
3285 .iter()
3286 .filter(|purpose| *purpose == "export target checkpoint")
3287 .count();
3288 if stdin_calls == 1 {
3289 return Ok(CommandOutput {
3290 status: 1,
3291 stdout: Vec::new(),
3292 stderr: b"Error: read checkpoint export spec -\n\nCaused by:\n No such file or directory (os error 2)\n".to_vec(),
3293 });
3294 }
3295 Ok(CommandOutput {
3296 status: 0,
3297 stdout: exported_checkpoint_json(),
3298 stderr: Vec::new(),
3299 })
3300 }
3301 }
3302
3303 let executor = FileThenRefreshExecutor {
3304 purposes: RefCell::new(Vec::new()),
3305 file_export_calls: Cell::new(0),
3306 };
3307 let output = export_target_checkpoint_with_worker(
3308 &executor,
3309 &locator,
3310 LATCH_RELAY_SESSION,
3311 &export_spec_fixture(),
3312 "/var/lib/hel/workers/session/checkpoint-spec.json",
3313 Some(Path::new("/hel-test-worker")),
3314 )
3315 .unwrap();
3316
3317 assert_eq!(output.stdout, exported_checkpoint_json());
3318 assert_eq!(
3319 executor.purposes.into_inner(),
3320 vec![
3321 "export target checkpoint".to_owned(),
3322 "upload checkpoint specification".to_owned(),
3323 "export target checkpoint".to_owned(),
3324 "stage replacement Mjolnir worker".to_owned(),
3325 "assign replacement worker to the worker user".to_owned(),
3326 "replace installed Mjolnir worker".to_owned(),
3327 "make replaced Mjolnir worker executable".to_owned(),
3328 "export target checkpoint".to_owned(),
3329 ]
3330 );
3331 }
3332 #[test]
3336 fn a_deferral_attached_as_context_under_more_context_is_still_a_deferral() {
3337 let deferred = anyhow::anyhow!("relay proxy disconnected during hello")
3338 .context(CheckpointDeferred::background_work())
3339 .context("connect to the session worker for checkpoint");
3340 assert!(checkpoint_was_deferred(&deferred), "{deferred:#}");
3341
3342 let plain = anyhow::anyhow!("relay proxy disconnected during hello")
3343 .context("connect to the session worker for checkpoint");
3344 assert!(!checkpoint_was_deferred(&plain), "{plain:#}");
3345 }
3346
3347 #[test]
3348 fn a_working_session_defers_but_close_waits_for_cancellation_before_recovery() {
3349 let cursor = RelayCursor {
3350 ordinal: 7,
3351 digest: "a".repeat(64),
3352 };
3353 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
3354 snapshot.operational.execution = RelayExecutionState::Running;
3355
3356 let deferred = checkpoint_barrier_wait_ended(
3357 &snapshot,
3358 "checkpoint-1",
3359 BarrierBusyPolicy::DeferWhileRunning,
3360 false,
3361 false,
3362 )
3363 .expect("a working session ends the wait at once");
3364 assert!(checkpoint_was_deferred(&deferred), "{deferred:#}");
3365 assert!(
3366 !checkpoint_barrier_needs_worker_restart(&deferred),
3367 "a deferred copy must never restart the worker: {deferred:#}"
3368 );
3369 assert_eq!(
3370 BarrierBusyPolicy::of(LatchExclusivity::HoldThroughClose),
3371 BarrierBusyPolicy::InterruptWhileRunning
3372 );
3373
3374 assert!(
3377 checkpoint_barrier_wait_ended(
3378 &snapshot,
3379 "checkpoint-1",
3380 BarrierBusyPolicy::InterruptWhileRunning,
3381 false,
3382 false,
3383 )
3384 .is_none()
3385 );
3386 let interrupted = checkpoint_barrier_wait_ended(
3387 &snapshot,
3388 "checkpoint-1",
3389 BarrierBusyPolicy::InterruptWhileRunning,
3390 true,
3391 true,
3392 )
3393 .expect("an unresponsive cancellation ends the wait at the deadline");
3394 assert!(
3395 checkpoint_barrier_needs_worker_restart(&interrupted),
3396 "{interrupted:#}"
3397 );
3398 assert!(!checkpoint_was_deferred(&interrupted), "{interrupted:#}");
3399
3400 snapshot.operational.execution = RelayExecutionState::Idle;
3403 let wedged = checkpoint_barrier_wait_ended(
3404 &snapshot,
3405 "checkpoint-1",
3406 BarrierBusyPolicy::DeferWhileRunning,
3407 true,
3408 false,
3409 )
3410 .expect("the deadline ends the wait");
3411 assert!(
3412 checkpoint_barrier_needs_worker_restart(&wedged),
3413 "{wedged:#}"
3414 );
3415 assert!(!checkpoint_was_deferred(&wedged), "{wedged:#}");
3416 }
3417
3418 #[test]
3421 fn a_frontier_that_moved_before_the_latch_defers_the_checkpoint() {
3422 let cursor = RelayCursor {
3423 ordinal: 220,
3424 digest: "a".repeat(64),
3425 };
3426 ensure_exact_checkpoint_cut(&cursor, cursor.ordinal, &cursor.digest)
3427 .expect("a projection latched at the ready cursor is an exact cut");
3428
3429 for (ordinal, digest) in [(223, "a".repeat(64)), (220, "b".repeat(64))] {
3430 let error = ensure_exact_checkpoint_cut(&cursor, ordinal, &digest)
3431 .expect_err("a projection past the ready cursor is not an exact cut");
3432 assert!(checkpoint_was_deferred(&error), "{error:#}");
3433 assert!(
3434 !checkpoint_barrier_needs_worker_restart(&error),
3435 "{error:#}"
3436 );
3437 }
3438 }
3439
3440 #[test]
3444 fn a_harness_turn_started_during_capture_abandons_the_archive() {
3445 let cursor = RelayCursor {
3446 ordinal: 220,
3447 digest: "a".repeat(64),
3448 };
3449 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
3450 snapshot.operational.checkpoint_ready = Some(cursor.clone());
3451
3452 snapshot.operational.last_harness_turn_started_ordinal = Some(cursor.ordinal);
3453 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor)
3454 .expect("a turn that started at or before the cursor is covered by the archive");
3455
3456 snapshot.operational.last_harness_turn_started_ordinal = Some(cursor.ordinal + 1);
3457 let error = validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor)
3458 .expect_err("a turn that started after the cursor invalidates the capture");
3459 assert!(checkpoint_was_deferred(&error), "{error:#}");
3460 }
3461
3462 #[test]
3463 fn a_stuck_checkpoint_barrier_is_retried_by_restarting_the_worker() {
3464 for failure in [
3467 CheckpointBarrierUnreachable::not_admitted(
3468 "checkpoint-976f6746887c5ccd93b9d8bbe120ef06",
3469 ),
3470 CheckpointBarrierUnreachable::runtime_stopped(),
3471 ] {
3472 let error = anyhow::Error::new(failure).context("latch a session checkpoint");
3473 assert!(checkpoint_barrier_needs_worker_restart(&error), "{error:#}");
3474 }
3475 assert!(!checkpoint_barrier_needs_worker_restart(&anyhow::anyhow!(
3476 "export target checkpoint failed with status 1"
3477 )));
3478 assert!(!checkpoint_barrier_needs_worker_restart(&anyhow::anyhow!(
3481 "ACP relay did not reach checkpoint barrier checkpoint-1"
3482 )));
3483 }
3484
3485 #[test]
3486 fn an_incompatible_cancel_turn_requests_worker_recovery() {
3487 let error = anyhow::Error::new(RelayRejected(mj_core::relay::RelayProtocolError {
3488 code: mj_core::relay::RelayErrorCode::IncompatibleProtocol,
3489 message: "request uses protocol 6".into(),
3490 retryable: false,
3491 detail: None,
3492 }))
3493 .context("cancel active ACP turn before checkpoint barrier");
3494 assert!(
3495 checkpoint_cancel_turn_needs_worker_restart(&error),
3496 "{error:#}"
3497 );
3498 assert!(checkpoint_barrier_needs_worker_restart(&error.context(
3499 CheckpointBarrierUnreachable::cancel_turn_unavailable("checkpoint-1", 6,)
3500 )));
3501 }
3502 #[test]
3503 fn a_dead_worker_hello_failure_is_retried_by_restarting_the_worker() {
3504 let dead = anyhow::Error::new(RelayTransportDead::new("the proxy is gone"))
3505 .context("connect to the session worker for checkpoint");
3506 assert!(worker_connect_needs_restart(&dead), "{dead:#}");
3507 assert!(!worker_connect_needs_restart(&anyhow::anyhow!(
3508 "unknown session"
3509 )));
3510 }
3511 #[cfg(unix)]
3512 #[tokio::test]
3513 async fn checkpoint_restart_stop_failure_names_mjolnir() {
3514 struct FailingStop;
3515
3516 impl CommandExecutor for FailingStop {
3517 fn execute(&self, _command: &CommandSpec) -> Result<CommandOutput> {
3518 Ok(CommandOutput {
3519 status: 1,
3520 stdout: Vec::new(),
3521 stderr: b"permission denied".to_vec(),
3522 })
3523 }
3524 }
3525
3526 let session_id = "0123456789abcdef0123456789abcdef";
3527 let worker_root = format!("/tmp/mjolnir-checkpoint-test/{session_id}");
3528 let backend = targets::TargetLocator::LocalBare {
3529 worker_root: worker_root.clone(),
3530 };
3531 let controller = Controller {
3532 config: Config::default(),
3533 state: State::default(),
3534 };
3535 let reconnect = CommandSpec::new("unused", std::iter::empty::<&str>());
3536
3537 let result = controller
3538 .restart_worker_for_checkpoint(
3539 session_id,
3540 &FailingStop,
3541 &backend,
3542 &worker_root,
3543 &reconnect,
3544 )
3545 .await;
3546 let error = match result {
3547 Ok(_) => panic!("a failed worker stop unexpectedly restarted the checkpoint worker"),
3548 Err(error) => error,
3549 };
3550 let detail = format!("{error:#}");
3551 assert!(
3552 detail.starts_with("stop wedged Mjolnir worker before retrying checkpoint"),
3553 "{detail}"
3554 );
3555 assert!(detail.contains("permission denied"), "{detail}");
3556 }
3557 #[test]
3558 fn export_spec_schema_mismatch_is_detected_from_the_parse_error() {
3559 assert!(export_spec_schema_unsupported(
3560 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3561 unknown field `terminal_refs`, expected `call` at line 1 column 7276552\n"
3562 ));
3563 assert!(export_spec_schema_unsupported(
3564 "Error: parse checkpoint export spec /spec.json\n\nCaused by:\n \
3565 unknown variant `terminal_output`, expected one of `user`, `agent`\n"
3566 ));
3567 assert!(!export_spec_schema_unsupported(
3568 "Error: repository 'app' is missing\n"
3569 ));
3570 assert!(!export_spec_schema_unsupported(
3571 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3572 missing field `relay_root`\n"
3573 ));
3574 assert!(export_protocol_unsupported(
3575 "Error: unsupported checkpoint export protocol version 3; worker supports 2\n"
3576 ));
3577 }
3578 const LATCH_RELAY_ROOT: &str = "MJ_TEST_LATCH_RELAY_ROOT";
3579 const LATCH_RELAY_STARTS: &str = "MJ_TEST_LATCH_RELAY_STARTS";
3580 const LATCH_RELAY_REJECT_RELEASE: &str = "MJ_TEST_LATCH_REJECT_RELEASE";
3581 #[cfg(unix)]
3582 const LATCH_RELAY_RUNNING: &str = "MJ_TEST_LATCH_RELAY_RUNNING";
3583 #[cfg(unix)]
3584 const LATCH_TEST_CHILD: &str = "MJ_TEST_LATCH_CHILD";
3585 #[cfg(unix)]
3586 const ABANDON_TEST_CHILD: &str = "MJ_TEST_ABANDON_LATCH_CHILD";
3587 #[cfg(unix)]
3588 const RELEASE_TEST_CHILD: &str = "MJ_TEST_RELEASE_LATCH_CHILD";
3589 #[cfg(unix)]
3590 const LEGACY_RELEASE_TEST_CHILD: &str = "MJ_TEST_LEGACY_RELEASE_LATCH_CHILD";
3591 #[cfg(unix)]
3592 const REUSE_TEST_CHILD: &str = "MJ_TEST_REUSE_LATCH_CHILD";
3593 const LATCH_CHECKPOINT_ONLY: &str = "MJ_TEST_LATCH_CHECKPOINT_ONLY";
3594 const LATCH_RELAY_STARTUP_DELAY_MS: &str = "MJ_TEST_LATCH_STARTUP_DELAY_MS";
3595 const LATCH_RELAY_SESSION: &str = "018f9dd2-a3b4-7c8d-9000-0123456789ab";
3596 #[cfg(unix)]
3598 #[derive(Clone, Copy, PartialEq, Eq)]
3599 enum ReleaseSupport {
3600 Supported,
3601 Rejected,
3604 }
3605 #[test]
3612 fn latch_relay_child_serves_stdio() {
3613 let Some(root) = std::env::var_os(LATCH_RELAY_ROOT) else {
3614 return;
3615 };
3616 println!();
3620 if let Some(starts) = std::env::var_os(LATCH_RELAY_STARTS) {
3623 use std::io::Write;
3624 let mut log = OpenOptions::new()
3625 .create(true)
3626 .append(true)
3627 .open(starts)
3628 .expect("open the relay start log");
3629 writeln!(log, "{}", std::process::id()).expect("record this relay start");
3630 }
3631 let checkpoint_only = std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some();
3632 let mut relay = if checkpoint_only {
3633 mj_worker::relay::DurableRelay::open_for_checkpoint(
3634 Path::new(&root),
3635 LATCH_RELAY_SESSION,
3636 "1.0.0",
3637 )
3638 } else {
3639 mj_worker::relay::DurableRelay::open(Path::new(&root), LATCH_RELAY_SESSION, "1.0.0")
3640 }
3641 .expect("open the test relay journal");
3642 if relay.operational_state().native_session_id.is_none() {
3643 relay
3644 .record_observation(mj_core::relay::RelayObservation::SessionOpened {
3645 native_session_id: "native-session".into(),
3646 resumed: true,
3647 })
3648 .unwrap();
3649 }
3650 if !relay.operational_state().goal.synchronized() && !checkpoint_only {
3651 relay.record_session_update(serde_json::from_value(serde_json::json!({
3652 "sessionUpdate":"session_info_update", "_meta":{"goal":null,"execution":{"version":1,"status":"idle"}}
3653 })).unwrap()).unwrap();
3654 }
3655 let ready_at = Instant::now()
3656 + Duration::from_millis(
3657 std::env::var(LATCH_RELAY_STARTUP_DELAY_MS)
3658 .ok()
3659 .map(|value| value.parse::<u64>().unwrap())
3660 .unwrap_or(0),
3661 );
3662 let reject_release = std::env::var_os(LATCH_RELAY_REJECT_RELEASE).is_some();
3663 #[cfg(unix)]
3664 let running = std::env::var_os(LATCH_RELAY_RUNNING).is_some();
3665 #[cfg(unix)]
3666 if running && relay.operational_state().active_prompt.is_none() {
3667 let response = relay.handle(mj_core::relay::RelayRequestEnvelope {
3668 request_id: "seed-running-request".into(),
3669 protocol_version: mj_core::relay::RELAY_PROTOCOL_VERSION,
3670 request: mj_core::relay::RelayRequest::Submit {
3671 command_id: "seed-running-prompt".into(),
3672 command: RelayCommand::Prompt {
3673 prompt: vec![ContentBlock::Text(TextContent::new("running"))],
3674 },
3675 },
3676 });
3677 assert!(matches!(
3678 response.body,
3679 mj_core::relay::RelayResponseBody::Ok {
3680 payload: mj_core::relay::RelayResponsePayload::Accepted { .. }
3681 }
3682 ));
3683 let claimed = relay
3684 .claim_pending_commands(true)
3685 .expect("seed the running prompt");
3686 assert_eq!(claimed.len(), 1);
3687 assert_eq!(claimed[0].command_id, "seed-running-prompt");
3688 }
3689 let mut reader = std::io::stdin().lock();
3690 let mut writer = std::io::stdout().lock();
3691 let mut configured = false;
3692 while let Some(request) =
3693 mj_core::relay::read_relay_frame(&mut reader).expect("read a relay request")
3694 {
3695 if !checkpoint_only && !configured && Instant::now() >= ready_at {
3696 relay
3697 .record_observation(mj_core::relay::RelayObservation::SessionConfigured {
3698 config_options: Vec::new(),
3699 })
3700 .unwrap();
3701 configured = true;
3702 }
3703 if matches!(
3704 &request.request,
3705 mj_core::relay::RelayRequest::Submit {
3706 command: RelayCommand::BeginCheckpoint { .. },
3707 ..
3708 }
3709 ) {
3710 assert!(
3711 checkpoint_only || relay.operational_state().native_session_is_ready(),
3712 "checkpoint submitted before current ACP startup finished"
3713 );
3714 }
3715 let response = if reject_release && requests_checkpoint_release(&request) {
3716 unparseable_request_response(&request)
3717 } else {
3718 relay.handle(request)
3719 };
3720 mj_core::relay::write_relay_frame(&mut writer, &response)
3721 .expect("answer a relay request");
3722 if checkpoint_only {
3723 relay.dispatch_checkpoint_only().unwrap();
3724 }
3725 for claimed in relay
3726 .claim_pending_commands(true)
3727 .expect("claim relay commands")
3728 {
3729 match claimed.command {
3730 RelayCommand::BeginCheckpoint { .. } => {
3731 relay
3732 .record_checkpoint_ready(&claimed.command_id)
3733 .expect("report the checkpoint barrier ready");
3734 }
3735 #[cfg(unix)]
3736 RelayCommand::CancelTurn => {
3737 let prompt_id = relay
3738 .operational_state()
3739 .active_prompt
3740 .as_ref()
3741 .map(|prompt| prompt.command_id.clone())
3742 .expect("a prompt to cancel");
3743 relay
3744 .record_command_completed(
3745 &claimed.command_id,
3746 RelayCommandOutcome::Cancelled,
3747 )
3748 .expect("complete the cancellation");
3749 relay
3750 .record_command_completed(
3751 &prompt_id,
3752 RelayCommandOutcome::Prompt {
3753 diagnostic: None,
3754 stop_reason: "cancelled".into(),
3755 usage: None,
3756 },
3757 )
3758 .expect("complete the cancelled prompt");
3759 }
3760 _ => {}
3761 }
3762 }
3763 }
3764 }
3765 fn requests_checkpoint_release(request: &mj_core::relay::RelayRequestEnvelope) -> bool {
3766 matches!(
3767 &request.request,
3768 mj_core::relay::RelayRequest::Submit {
3769 command: RelayCommand::ReleaseCheckpoint { .. },
3770 ..
3771 }
3772 )
3773 }
3774 fn unparseable_request_response(
3778 request: &mj_core::relay::RelayRequestEnvelope,
3779 ) -> mj_core::relay::RelayResponseEnvelope {
3780 mj_core::relay::RelayResponseEnvelope {
3781 request_id: request.request_id.clone(),
3782 protocol_version: request.protocol_version,
3783 body: mj_core::relay::RelayResponseBody::Error {
3784 error: mj_core::relay::RelayProtocolError {
3785 code: mj_core::relay::RelayErrorCode::InvalidRequest,
3786 message: "unknown variant `release_checkpoint`".into(),
3787 retryable: false,
3788 detail: None,
3789 },
3790 },
3791 }
3792 }
3793 #[cfg(unix)]
3796 fn latch_relay_target(
3797 relay_root: &Path,
3798 starts: Option<&Path>,
3799 release: ReleaseSupport,
3800 running: bool,
3801 ) -> crate::session_manager::RelaySessionTarget {
3802 let script = format!(
3805 "\"$0\" --exact {}::latch_relay_child_serves_stdio --nocapture | \
3806 grep --line-buffered '^{{'",
3807 module_path!()
3808 .strip_prefix("mj_controller::")
3809 .unwrap_or(module_path!())
3810 );
3811 let mut spec = CommandSpec::new(
3812 "sh",
3813 [
3814 "-c".to_owned(),
3815 script,
3816 std::env::current_exe()
3817 .unwrap()
3818 .to_string_lossy()
3819 .into_owned(),
3820 ],
3821 )
3822 .purpose("test latch relay");
3823 spec.env.insert(
3824 LATCH_RELAY_ROOT.to_owned(),
3825 relay_root.to_string_lossy().into_owned(),
3826 );
3827 if let Some(starts) = starts {
3828 spec.env.insert(
3829 LATCH_RELAY_STARTS.to_owned(),
3830 starts.to_string_lossy().into_owned(),
3831 );
3832 }
3833 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
3834 spec.env.insert(LATCH_CHECKPOINT_ONLY.into(), "1".into());
3835 }
3836 if release == ReleaseSupport::Rejected {
3837 spec.env
3838 .insert(LATCH_RELAY_REJECT_RELEASE.to_owned(), "1".to_owned());
3839 }
3840 if running {
3841 spec.env
3842 .insert(LATCH_RELAY_RUNNING.to_owned(), "1".to_owned());
3843 }
3844 crate::session_manager::RelaySessionTarget {
3845 session_id: LATCH_RELAY_SESSION.to_owned(),
3846 spec,
3847 worker_recovery: None,
3848 project_memory: None,
3849 }
3850 }
3851 #[cfg(unix)]
3854 async fn latch_a_live_checkpoint(
3855 relay_root: &Path,
3856 starts: Option<&Path>,
3857 release: ReleaseSupport,
3858 running: bool,
3859 ) -> (
3860 crate::session_manager::SessionManagerChannels,
3861 ManagedSessionHandle,
3862 ControllerRelayLease,
3863 String,
3864 RelayCursor,
3865 ) {
3866 crate::database::save_session(&checkpoint_test_session(LATCH_RELAY_SESSION)).unwrap();
3869 let channels = crate::session_manager::spawn_session_manager().unwrap();
3870 channels
3871 .targets
3872 .send(vec![latch_relay_target(
3873 relay_root, starts, release, running,
3874 )])
3875 .unwrap();
3876 let handle = channels
3877 .control
3878 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
3879 .await
3880 .unwrap();
3881
3882 let lease = handle.lease_connection().await.unwrap();
3883 let mut relay = ControllerRelayLease::Managed {
3884 handle: handle.clone(),
3885 lease: Some(lease),
3886 };
3887 let barrier_command_id = new_command_id("checkpoint").unwrap();
3888 let connection = relay.connection_mut();
3889 connection
3890 .submit(
3891 barrier_command_id.clone(),
3892 RelayCommand::BeginCheckpoint { reason: None },
3893 )
3894 .await
3895 .unwrap();
3896 let barrier = wait_for_checkpoint_barrier(
3897 connection,
3898 LATCH_RELAY_SESSION,
3899 &barrier_command_id,
3900 CHECKPOINT_BARRIER_TIMEOUT,
3901 BarrierBusyPolicy::InterruptWhileRunning,
3902 HarnessKind::Codex,
3903 )
3904 .await
3905 .unwrap();
3906 assert_eq!(
3907 barrier.materialized.applied_event_ordinal,
3908 barrier.operational.latest_ordinal
3909 );
3910 let cursor = barrier.operational.checkpoint_ready.clone().unwrap();
3911 (channels, handle, relay, barrier_command_id, cursor)
3912 }
3913
3914 #[cfg(unix)]
3919 #[tokio::test]
3920 async fn a_close_checkpoint_cancels_a_running_turn_without_restarting_the_worker() {
3921 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
3924 let directory = tempfile::tempdir().unwrap();
3925 let test_name = format!(
3926 "{}::a_close_checkpoint_cancels_a_running_turn_without_restarting_the_worker",
3927 module_path!()
3928 .strip_prefix("mj_controller::")
3929 .unwrap_or(module_path!())
3930 );
3931 let output = Command::new(std::env::current_exe().unwrap())
3932 .args(["--exact", &test_name, "--nocapture"])
3933 .env(LATCH_TEST_CHILD, "1")
3934 .env("MJ_DATA_DIR", directory.path())
3935 .output()
3936 .unwrap();
3937 assert!(
3938 output.status.success(),
3939 "isolated cancellation checkpoint test failed\nstdout:\n{}\nstderr:\n{}",
3940 String::from_utf8_lossy(&output.stdout),
3941 String::from_utf8_lossy(&output.stderr)
3942 );
3943 return;
3944 }
3945 let _writer = crate::database::install_isolated_test_writer();
3946 let relay_root = tempfile::tempdir().unwrap();
3947 let start_log_directory = tempfile::tempdir().unwrap();
3948 let start_log = start_log_directory.path().join("relay-starts");
3949 let (_channels, _handle, mut relay, _barrier_command_id, _cursor) =
3950 latch_a_live_checkpoint(
3951 relay_root.path(),
3952 Some(&start_log),
3953 ReleaseSupport::Supported,
3954 true,
3955 )
3956 .await;
3957 let snapshot = relay.sync_snapshot().await.unwrap();
3958 assert_eq!(
3959 snapshot.operational.execution,
3960 RelayExecutionState::Idle,
3961 "the close wait returned before the cancelled turn became idle"
3962 );
3963 assert!(
3964 snapshot.operational.active_prompt.is_none(),
3965 "the close wait returned before the cancelled prompt settled"
3966 );
3967 assert_eq!(
3968 relay_starts(&start_log),
3969 1,
3970 "responsive cancellation restarted worker"
3971 );
3972 }
3973 #[cfg(unix)]
3976 async fn wait_until_the_actor_serves_again(handle: &ManagedSessionHandle) {
3977 for attempt in 0.. {
3978 if handle.sync_now().await.is_ok() {
3979 return;
3980 }
3981 assert!(attempt < 200, "the actor never took its connection back");
3982 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
3983 }
3984 }
3985 #[cfg(unix)]
3989 #[tokio::test]
3990 async fn ending_the_checkpoint_latch_returns_the_connection_to_its_actor() {
3991 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
3994 let directory = tempfile::tempdir().unwrap();
3995 let test_name = format!(
3996 "{}::ending_the_checkpoint_latch_returns_the_connection_to_its_actor",
3997 module_path!()
3998 .strip_prefix("mj_controller::")
3999 .unwrap_or(module_path!())
4000 );
4001 let output = Command::new(std::env::current_exe().unwrap())
4002 .args(["--exact", &test_name, "--nocapture"])
4003 .env(LATCH_TEST_CHILD, "1")
4004 .env("MJ_DATA_DIR", directory.path())
4005 .output()
4006 .unwrap();
4007 assert!(
4008 output.status.success(),
4009 "isolated checkpoint latch test failed\nstdout:\n{}\nstderr:\n{}",
4010 String::from_utf8_lossy(&output.stdout),
4011 String::from_utf8_lossy(&output.stderr)
4012 );
4013 return;
4014 }
4015 let _writer = crate::database::install_isolated_test_writer();
4017
4018 std::thread::spawn(|| {
4021 std::thread::sleep(std::time::Duration::from_secs(120));
4022 eprintln!("the checkpoint latch never returned its connection");
4023 std::process::exit(101);
4024 });
4025
4026 let relay_root = tempfile::tempdir().unwrap();
4027 let (_channels, handle, mut relay, barrier_command_id, cursor) =
4028 latch_a_live_checkpoint(relay_root.path(), None, ReleaseSupport::Supported, false)
4029 .await;
4030
4031 assert!(
4034 handle.sync_now().await.is_err(),
4035 "a latched projection must not be advanced by its own actor"
4036 );
4037
4038 relay.end_latch();
4039 wait_until_the_actor_serves_again(&handle).await;
4040
4041 let latched = relay.sync_snapshot().await.unwrap();
4045 validate_checkpoint_barrier_snapshot(&latched, &barrier_command_id, &cursor).unwrap();
4046
4047 let prompt_ordinal = relay
4050 .submit(
4051 new_command_id("prompt").unwrap(),
4052 RelayCommand::Prompt {
4053 prompt: vec![ContentBlock::Text(TextContent::new("hello"))],
4054 },
4055 )
4056 .await
4057 .unwrap();
4058 assert!(prompt_ordinal > cursor.ordinal);
4059 let snapshot = relay.sync_snapshot().await.unwrap();
4060 assert!(snapshot.operational.latest_ordinal > cursor.ordinal);
4061 validate_checkpoint_barrier_snapshot(&snapshot, &barrier_command_id, &cursor).unwrap();
4062
4063 latched_checkpoint(
4064 relay,
4065 barrier_command_id,
4066 cursor,
4067 CheckpointCompletion::HeldBarrier,
4068 )
4069 .complete()
4070 .await
4071 .unwrap();
4072 handle.sync_now().await.unwrap();
4073 assert_eq!(
4074 handle
4075 .view()
4076 .snapshot
4077 .expect("the actor published the completed barrier")
4078 .operational
4079 .checkpoint_barrier,
4080 None
4081 );
4082 }
4083 #[cfg(unix)]
4087 #[tokio::test]
4088 async fn releasing_a_checkpoint_after_capture_defers_only_the_recovery_floor() {
4089 if std::env::var_os(RELEASE_TEST_CHILD).is_none() {
4092 let directory = tempfile::tempdir().unwrap();
4093 let test_name = format!(
4094 "{}::releasing_a_checkpoint_after_capture_defers_only_the_recovery_floor",
4095 module_path!()
4096 .strip_prefix("mj_controller::")
4097 .unwrap_or(module_path!())
4098 );
4099 let output = Command::new(std::env::current_exe().unwrap())
4100 .args(["--exact", &test_name, "--nocapture"])
4101 .env(RELEASE_TEST_CHILD, "1")
4102 .env("MJ_DATA_DIR", directory.path())
4103 .output()
4104 .unwrap();
4105 assert!(
4106 output.status.success(),
4107 "isolated checkpoint release test failed\nstdout:\n{}\nstderr:\n{}",
4108 String::from_utf8_lossy(&output.stdout),
4109 String::from_utf8_lossy(&output.stderr)
4110 );
4111 return;
4112 }
4113 let _writer = crate::database::install_isolated_test_writer();
4115
4116 std::thread::spawn(|| {
4119 std::thread::sleep(std::time::Duration::from_secs(120));
4120 eprintln!("the captured checkpoint never released its barrier");
4121 std::process::exit(101);
4122 });
4123
4124 let relay_root = tempfile::tempdir().unwrap();
4125 let (_channels, handle, mut relay, barrier_command_id, cursor) =
4126 latch_a_live_checkpoint(relay_root.path(), None, ReleaseSupport::Supported, false)
4127 .await;
4128 relay.end_latch();
4129 wait_until_the_actor_serves_again(&handle).await;
4130
4131 let completion = release_checkpoint_after_capture(
4134 &mut relay,
4135 LATCH_RELAY_SESSION,
4136 &barrier_command_id,
4137 &cursor,
4138 HarnessKind::Codex,
4139 )
4140 .await
4141 .unwrap();
4142 assert_eq!(completion, CheckpointCompletion::ReleasedAfterCapture);
4143 let released = relay.sync_snapshot().await.unwrap();
4144 assert_eq!(released.operational.checkpoint_barrier, None);
4145 assert_eq!(released.operational.checkpoint_ready, None);
4146 assert_eq!(
4147 released.operational.recovery_floor_ordinal, 0,
4148 "an exported archive that is not installed may not release journal history"
4149 );
4150
4151 relay
4154 .submit(
4155 new_command_id("prompt").unwrap(),
4156 RelayCommand::Prompt {
4157 prompt: vec![ContentBlock::Text(TextContent::new("during transfer"))],
4158 },
4159 )
4160 .await
4161 .unwrap();
4162 let mut dispatched = None;
4163 for attempt in 0.. {
4164 let snapshot = relay.sync_snapshot().await.unwrap();
4165 if let Some(active) = snapshot.operational.active_prompt {
4166 dispatched = Some(active);
4167 break;
4168 }
4169 assert!(attempt < 200, "a released barrier still froze ACP dispatch");
4170 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
4171 }
4172 assert!(dispatched.is_some());
4173
4174 latched_checkpoint(
4177 relay,
4178 barrier_command_id,
4179 cursor.clone(),
4180 CheckpointCompletion::ReleasedAfterCapture,
4181 )
4182 .complete()
4183 .await
4184 .unwrap();
4185 handle.sync_now().await.unwrap();
4186 let installed = handle
4187 .view()
4188 .snapshot
4189 .expect("the actor published the advanced recovery floor");
4190 assert_eq!(installed.operational.recovery_floor_ordinal, cursor.ordinal);
4191 assert_eq!(installed.operational.recovery_floor_digest, cursor.digest);
4192 }
4193 #[cfg(unix)]
4196 #[tokio::test]
4197 async fn a_worker_that_rejects_the_release_keeps_its_barrier_through_the_transfer() {
4198 if std::env::var_os(LEGACY_RELEASE_TEST_CHILD).is_none() {
4201 let directory = tempfile::tempdir().unwrap();
4202 let test_name = format!(
4203 "{}::a_worker_that_rejects_the_release_keeps_its_barrier_through_the_transfer",
4204 module_path!()
4205 .strip_prefix("mj_controller::")
4206 .unwrap_or(module_path!())
4207 );
4208 let output = Command::new(std::env::current_exe().unwrap())
4209 .args(["--exact", &test_name, "--nocapture"])
4210 .env(LEGACY_RELEASE_TEST_CHILD, "1")
4211 .env("MJ_DATA_DIR", directory.path())
4212 .output()
4213 .unwrap();
4214 assert!(
4215 output.status.success(),
4216 "isolated legacy checkpoint release test failed\nstdout:\n{}\nstderr:\n{}",
4217 String::from_utf8_lossy(&output.stdout),
4218 String::from_utf8_lossy(&output.stderr)
4219 );
4220 return;
4221 }
4222 let _writer = crate::database::install_isolated_test_writer();
4224
4225 std::thread::spawn(|| {
4228 std::thread::sleep(std::time::Duration::from_secs(120));
4229 eprintln!("the rejected release never finished its checkpoint");
4230 std::process::exit(101);
4231 });
4232
4233 let relay_root = tempfile::tempdir().unwrap();
4234 let start_log = tempfile::tempdir().unwrap();
4235 let start_log = start_log.path().join("relay-starts");
4236 let (_channels, handle, mut relay, barrier_command_id, cursor) = latch_a_live_checkpoint(
4237 relay_root.path(),
4238 Some(&start_log),
4239 ReleaseSupport::Rejected,
4240 false,
4241 )
4242 .await;
4243 relay.end_latch();
4244 wait_until_the_actor_serves_again(&handle).await;
4245
4246 let completion = release_checkpoint_after_capture(
4247 &mut relay,
4248 LATCH_RELAY_SESSION,
4249 &barrier_command_id,
4250 &cursor,
4251 HarnessKind::Codex,
4252 )
4253 .await
4254 .unwrap();
4255 assert_eq!(completion, CheckpointCompletion::HeldBarrier);
4256 assert_eq!(relay_starts(&start_log), 1);
4259
4260 let transferring = relay.sync_snapshot().await.unwrap();
4264 validate_checkpoint_barrier_snapshot(&transferring, &barrier_command_id, &cursor).unwrap();
4265 latched_checkpoint(relay, barrier_command_id, cursor.clone(), completion)
4266 .complete()
4267 .await
4268 .unwrap();
4269 handle.sync_now().await.unwrap();
4270 let completed = handle
4271 .view()
4272 .snapshot
4273 .expect("the actor published the completed barrier");
4274 assert_eq!(completed.operational.checkpoint_barrier, None);
4275 assert_eq!(completed.operational.recovery_floor_ordinal, cursor.ordinal);
4276 }
4277 #[cfg(unix)]
4282 #[tokio::test]
4283 async fn abandoning_a_latched_checkpoint_drops_the_connection_that_opened_its_barrier() {
4284 if std::env::var_os(ABANDON_TEST_CHILD).is_none() {
4287 let directory = tempfile::tempdir().unwrap();
4288 let test_name = format!(
4289 "{}::abandoning_a_latched_checkpoint_drops_the_connection_that_opened_its_barrier",
4290 module_path!()
4291 .strip_prefix("mj_controller::")
4292 .unwrap_or(module_path!())
4293 );
4294 let output = Command::new(std::env::current_exe().unwrap())
4295 .args(["--exact", &test_name, "--nocapture"])
4296 .env(ABANDON_TEST_CHILD, "1")
4297 .env("MJ_DATA_DIR", directory.path())
4298 .output()
4299 .unwrap();
4300 assert!(
4301 output.status.success(),
4302 "isolated abandoned checkpoint test failed\nstdout:\n{}\nstderr:\n{}",
4303 String::from_utf8_lossy(&output.stdout),
4304 String::from_utf8_lossy(&output.stderr)
4305 );
4306 return;
4307 }
4308 let _writer = crate::database::install_isolated_test_writer();
4310
4311 std::thread::spawn(|| {
4314 std::thread::sleep(std::time::Duration::from_secs(120));
4315 eprintln!("an abandoned checkpoint never released its relay connection");
4316 std::process::exit(101);
4317 });
4318
4319 let relay_root = tempfile::tempdir().unwrap();
4320 let start_log = tempfile::tempdir().unwrap();
4321 let start_log = start_log.path().join("relay-starts");
4322 let (_channels, handle, mut relay, barrier_command_id, cursor) = latch_a_live_checkpoint(
4323 relay_root.path(),
4324 Some(&start_log),
4325 ReleaseSupport::Supported,
4326 false,
4327 )
4328 .await;
4329 relay.end_latch();
4330 wait_until_the_actor_serves_again(&handle).await;
4331 assert_eq!(relay_starts(&start_log), 1);
4332
4333 latched_checkpoint(
4334 relay,
4335 barrier_command_id,
4336 cursor,
4337 CheckpointCompletion::HeldBarrier,
4338 )
4339 .abandon(LATCH_RELAY_SESSION)
4340 .await;
4341
4342 wait_until_the_actor_serves_again(&handle).await;
4347 assert_eq!(relay_starts(&start_log), 2);
4348 }
4349 #[cfg(unix)]
4354 #[test]
4355 fn a_move_checkpoint_can_verify_its_archive_without_source_harness_readiness() {
4356 let directory = tempfile::tempdir().unwrap();
4357 let name = format!(
4358 "{}::a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content",
4359 module_path!()
4360 .strip_prefix("mj_controller::")
4361 .unwrap_or(module_path!())
4362 );
4363 let output = Command::new(std::env::current_exe().unwrap())
4364 .args(["--exact", &name, "--nocapture"])
4365 .env(REUSE_TEST_CHILD, "1")
4366 .env(LATCH_CHECKPOINT_ONLY, "1")
4367 .env("MJ_DATA_DIR", directory.path())
4368 .output()
4369 .unwrap();
4370 assert!(
4371 output.status.success(),
4372 "checkpoint-only capture failed: {}\n{}",
4373 String::from_utf8_lossy(&output.stdout),
4374 String::from_utf8_lossy(&output.stderr)
4375 );
4376 }
4377
4378 #[cfg(unix)]
4379 #[tokio::test]
4380 async fn a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content() {
4381 if std::env::var_os(REUSE_TEST_CHILD).is_none() {
4384 let directory = tempfile::tempdir().unwrap();
4385 let test_name = format!(
4386 "{}::a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content",
4387 module_path!()
4388 .strip_prefix("mj_controller::")
4389 .unwrap_or(module_path!())
4390 );
4391 let output = Command::new(std::env::current_exe().unwrap())
4392 .args(["--exact", &test_name, "--nocapture"])
4393 .env(REUSE_TEST_CHILD, "1")
4394 .env(LATCH_RELAY_STARTUP_DELAY_MS, "31000")
4397 .env("MJ_DATA_DIR", directory.path())
4398 .output()
4399 .unwrap();
4400 assert!(
4401 output.status.success(),
4402 "isolated checkpoint reuse test failed\nstdout:\n{}\nstderr:\n{}",
4403 String::from_utf8_lossy(&output.stdout),
4404 String::from_utf8_lossy(&output.stderr)
4405 );
4406 return;
4407 }
4408 let _writer = crate::database::install_isolated_test_writer();
4410
4411 std::thread::spawn(|| {
4414 std::thread::sleep(std::time::Duration::from_secs(120));
4415 eprintln!("the reuse checkpoint never finished its latch");
4416 std::process::exit(101);
4417 });
4418
4419 #[derive(Default)]
4420 struct RecordingExecutor {
4421 purposes: std::sync::Mutex<Vec<String>>,
4422 active_stages: std::sync::Mutex<Vec<ProvisionStage>>,
4423 stage_events: std::sync::Mutex<Vec<(ProvisionStage, bool)>>,
4424 observed_stages: std::sync::Mutex<Vec<(String, Vec<ProvisionStage>)>>,
4425 }
4426
4427 impl RecordingExecutor {
4428 fn refused(&self, command: &CommandSpec) -> Result<CommandOutput> {
4429 self.purposes.lock().unwrap().push(command.purpose.clone());
4430 self.observed_stages.lock().unwrap().push((
4431 command.purpose.clone(),
4432 self.active_stages.lock().unwrap().clone(),
4433 ));
4434 Ok(CommandOutput {
4435 status: 1,
4436 stdout: Vec::new(),
4437 stderr: b"no target is provisioned for this test".to_vec(),
4438 })
4439 }
4440
4441 fn purposes(&self) -> Vec<String> {
4442 self.purposes.lock().unwrap().clone()
4443 }
4444
4445 fn observed_stages(&self) -> Vec<(String, Vec<ProvisionStage>)> {
4446 self.observed_stages.lock().unwrap().clone()
4447 }
4448
4449 fn stage_events(&self) -> Vec<(ProvisionStage, bool)> {
4450 self.stage_events.lock().unwrap().clone()
4451 }
4452 }
4453
4454 impl CommandExecutor for RecordingExecutor {
4455 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
4456 self.refused(command)
4457 }
4458
4459 fn execute_with_stdin(
4460 &self,
4461 command: &CommandSpec,
4462 _input: &mut (dyn std::io::Read + Send),
4463 ) -> Result<CommandOutput> {
4464 self.refused(command)
4465 }
4466
4467 fn stage_started(&self, stage: ProvisionStage) {
4468 self.active_stages.lock().unwrap().push(stage);
4469 self.stage_events.lock().unwrap().push((stage, true));
4470 }
4471
4472 fn stage_finished(&self, stage: ProvisionStage) {
4473 let mut active = self.active_stages.lock().unwrap();
4474 let position = active
4475 .iter()
4476 .position(|active_stage| *active_stage == stage)
4477 .expect("stage finished without a matching start");
4478 active.remove(position);
4479 self.stage_events.lock().unwrap().push((stage, false));
4480 }
4481 }
4482
4483 let data_directory = PathBuf::from(std::env::var_os("MJ_DATA_DIR").unwrap());
4484 let relay_root = data_directory.join("relay");
4485 let profile_home = data_directory.join("profile");
4486 let archive_directory = data_directory.join("archives");
4487 for directory in [&relay_root, &profile_home, &archive_directory] {
4488 std::fs::create_dir_all(directory).unwrap();
4489 }
4490 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
4491 let mut seed =
4492 mj_worker::relay::DurableRelay::open(&relay_root, LATCH_RELAY_SESSION, "1.0.0")
4493 .unwrap();
4494 seed.record_observation(mj_core::relay::RelayObservation::SessionOpened {
4495 native_session_id: "native-session".into(),
4496 resumed: true,
4497 })
4498 .unwrap();
4499 seed.record_observation(mj_core::relay::RelayObservation::SessionConfigured {
4500 config_options: Vec::new(),
4501 })
4502 .unwrap();
4503 }
4504 let checkpoint = write_checkpoint_gate_archive(&archive_directory, LATCH_RELAY_SESSION, 2);
4507
4508 let mut session = checkpoint_test_session(LATCH_RELAY_SESSION);
4509 session.target_template_id = "local".into();
4510 session.target = Some(TargetLocator::LocalBare {
4511 worker_root: data_directory.join("workers").join(LATCH_RELAY_SESSION),
4512 });
4513 session.checkpoint = Some(checkpoint.clone());
4514 crate::database::save_session(&session).unwrap();
4515
4516 let mut config = Config::default();
4517 config.profiles.insert(
4518 "codex".into(),
4519 HarnessProfile {
4520 enabled: true,
4521 kind: mj_core::config::HarnessKind::Codex,
4522 home: profile_home,
4523 environment: BTreeMap::new(),
4524 context_window_bytes: None,
4525 },
4526 );
4527 config
4528 .targets
4529 .insert("local".into(), TargetTemplate::LocalBare);
4530 config.bundles.insert(
4531 "project".into(),
4532 ProjectBundle {
4533 primary_repo: "project".into(),
4534 repositories: vec![ProjectRepository {
4535 id: "project".into(),
4536 github: Some("example/project".into()),
4537 local: None,
4538 destination: "project".into(),
4539 git_ref: None,
4540 }],
4541 },
4542 );
4543 let controller = Controller {
4544 config,
4545 state: State {
4546 sessions: BTreeMap::from([(LATCH_RELAY_SESSION.into(), session)]),
4547 ..State::default()
4548 },
4549 };
4550
4551 let channels = crate::session_manager::spawn_session_manager().unwrap();
4552 channels
4553 .targets
4554 .send(vec![latch_relay_target(
4555 &relay_root,
4556 None,
4557 ReleaseSupport::Supported,
4558 false,
4559 )])
4560 .unwrap();
4561 let handle = channels
4562 .control
4563 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
4564 .await
4565 .unwrap();
4566
4567 let executor = RecordingExecutor::default();
4568 let latched = controller
4569 .checkpoint_session_latched(
4570 LATCH_RELAY_SESSION,
4571 &executor,
4572 Some(&channels.control),
4573 LatchExclusivity::HoldThroughClose,
4574 CheckpointExportPolicy::ReuseUnchangedArchive,
4575 )
4576 .await
4577 .unwrap();
4578
4579 assert!(
4580 executor.purposes().is_empty(),
4581 "an unchanged session exported an archive anyway: {:?}",
4582 executor.purposes()
4583 );
4584 assert_eq!(latched.artifact.metadata, checkpoint);
4585 assert!(checkpoint.archive_path.exists());
4586
4587 assert!(latched.cursor.ordinal > checkpoint.event_frontier);
4590 let cursor = latched.cursor.clone();
4591 latched.complete().await.unwrap();
4592 wait_until_the_actor_serves_again(&handle).await;
4593
4594 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
4595 let snapshot = handle.view().snapshot.unwrap();
4596 assert!(snapshot.operational.checkpoint_only);
4597 assert!(!snapshot.operational.native_session_is_ready());
4598 assert_eq!(
4599 verify_archive_streaming(&checkpoint.archive_path)
4600 .unwrap()
4601 .manifest
4602 .session
4603 .native_session_id,
4604 "native-session"
4605 );
4606 channels.shutdown.shutdown().await.unwrap();
4607 return;
4608 }
4609
4610 handle
4614 .submit(
4615 new_command_id("busy-prompt").unwrap(),
4616 RelayCommand::Prompt {
4617 prompt: vec![ContentBlock::Text(TextContent::new("keep working"))],
4618 },
4619 )
4620 .await
4621 .unwrap();
4622 let mut connection = handle.lease_connection().await.unwrap();
4623 let before = connection.connection_mut().sync().await.unwrap();
4624 assert_eq!(before.operational.execution, RelayExecutionState::Running);
4625 connection.release();
4626 let deferred = controller
4627 .checkpoint_session_latched(
4628 LATCH_RELAY_SESSION,
4629 &executor,
4630 Some(&channels.control),
4631 LatchExclusivity::ReleaseAfterLatch,
4632 CheckpointExportPolicy::ReuseUnchangedArchive,
4633 )
4634 .await;
4635 assert!(
4636 matches!(deferred, Err(ref error) if error.downcast_ref::<CheckpointDeferred>().is_some())
4637 );
4638 wait_until_the_actor_serves_again(&handle).await;
4639 let mut connection = handle.lease_connection().await.unwrap();
4640 let after = connection.connection_mut().sync().await.unwrap();
4641 assert_eq!(after.operational.execution, RelayExecutionState::Running);
4642 assert!(after.operational.checkpoint_barrier.is_none());
4643 let journal =
4644 std::fs::read_to_string(relay_root.join("relay-journal/active.jsonl")).unwrap();
4645 for line in journal.lines() {
4646 let event: mj_core::relay::RelayEvent = serde_json::from_str(line).unwrap();
4647 if event.ordinal > before.operational.latest_ordinal {
4648 assert!(
4649 !matches!(
4650 event.observation,
4651 mj_core::relay::RelayObservation::CommandQueued {
4652 command: RelayCommand::BeginCheckpoint { .. },
4653 ..
4654 } | mj_core::relay::RelayObservation::CommandInterrupted {
4655 command: mj_core::relay::RelayCommandKind::BeginCheckpoint,
4656 ..
4657 }
4658 ),
4659 "busy deferral journaled checkpoint activity: {event:?}"
4660 );
4661 }
4662 }
4663 connection.release();
4664 handle
4665 .submit(
4666 new_command_id("finish-busy-prompt").unwrap(),
4667 RelayCommand::CancelTurn,
4668 )
4669 .await
4670 .unwrap();
4671 handle.sync_now().await.unwrap();
4672
4673 handle
4675 .submit(
4676 new_command_id("resume-notice").unwrap(),
4677 RelayCommand::RecordNotice {
4678 text: "the session changed".into(),
4679 },
4680 )
4681 .await
4682 .unwrap();
4683 for attempt in 0.. {
4684 handle.sync_now().await.unwrap();
4685 let materialized = handle.view().snapshot.map(|snapshot| snapshot.materialized);
4686 if materialized.is_some_and(|materialized| {
4687 materialized.applied_event_ordinal > cursor.ordinal
4688 && !materialized.transcript.is_empty()
4689 }) {
4690 break;
4691 }
4692 assert!(attempt < 200, "the notice never reached the projection");
4693 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
4694 }
4695
4696 let changed = controller
4697 .checkpoint_session_latched(
4698 LATCH_RELAY_SESSION,
4699 &executor,
4700 Some(&channels.control),
4701 LatchExclusivity::HoldThroughClose,
4702 CheckpointExportPolicy::ReuseUnchangedArchive,
4703 )
4704 .await;
4705 let Err(error) = changed else {
4706 panic!("a changed session reused its installed archive");
4707 };
4708
4709 assert!(
4710 executor
4711 .purposes()
4712 .contains(&"export target checkpoint".to_owned()),
4713 "a changed session skipped its export: {:?}",
4714 executor.purposes()
4715 );
4716 assert!(
4717 format!("{error:#}").contains("no target is provisioned for this test"),
4718 "{error:#}"
4719 );
4720 assert!(
4721 executor.observed_stages().iter().any(|(purpose, stages)| {
4722 purpose == "export target checkpoint"
4723 && stages.contains(&ProvisionStage::RecoveryCopy)
4724 }),
4725 "close checkpoint export did not run inside RecoveryCopy: {:?}",
4726 executor.observed_stages()
4727 );
4728 assert_eq!(
4729 executor
4730 .stage_events()
4731 .into_iter()
4732 .filter(|(stage, _)| *stage == ProvisionStage::RecoveryCopy)
4733 .collect::<Vec<_>>(),
4734 vec![
4735 (ProvisionStage::RecoveryCopy, true),
4736 (ProvisionStage::RecoveryCopy, false)
4737 ]
4738 );
4739 assert!(executor.active_stages.lock().unwrap().is_empty());
4740 assert!(checkpoint.archive_path.exists());
4741 }
4742 #[cfg(unix)]
4743 fn relay_starts(path: &Path) -> usize {
4744 std::fs::read_to_string(path)
4745 .unwrap_or_default()
4746 .lines()
4747 .count()
4748 }
4749 #[cfg(unix)]
4752 fn latched_checkpoint(
4753 relay: ControllerRelayLease,
4754 barrier_command_id: String,
4755 cursor: RelayCursor,
4756 completion: CheckpointCompletion,
4757 ) -> LatchedCheckpoint {
4758 LatchedCheckpoint {
4759 artifact: CheckpointArtifact {
4760 metadata: CheckpointMetadata {
4761 archive_path: PathBuf::from("checkpoint.hel.zip"),
4762 sha256: "a".repeat(64),
4763 created_at: now(),
4764 event_frontier: cursor.ordinal,
4765 },
4766 native_session_id: "native-session".into(),
4767 event_frontier_digest: cursor.digest.clone(),
4768 },
4769 relay,
4770 barrier_command_id,
4771 cursor,
4772 completion,
4773 }
4774 }
4775 #[test]
4776 fn checkpoint_persistence_rollback_restores_memory_and_reports_both_failures() {
4777 let session_id = "0123456789abcdef0123456789abcdef";
4778 let previous = checkpoint_test_session(session_id);
4779 let mut changed = previous.clone();
4780 changed.state = SessionState::Closing;
4781 changed.last_checkpoint_error = Some("partially installed checkpoint".into());
4782 let mut state = State::default();
4783 state.sessions.insert(session_id.into(), changed);
4784
4785 let error = restore_session_after_persistence_failure(
4786 &mut state,
4787 session_id,
4788 &previous,
4789 anyhow::anyhow!("verified checkpoint persistence failed"),
4790 |record| {
4791 assert_eq!(record, &previous);
4792 Err(anyhow::anyhow!("rollback database write failed"))
4793 },
4794 );
4795
4796 assert_eq!(state.sessions.get(session_id), Some(&previous));
4797 let detail = format!("{error:#}");
4798 assert!(detail.contains("verified checkpoint persistence failed"));
4799 assert!(detail.contains("rollback database write failed"));
4800 }
4801 #[test]
4802 fn installed_checkpoint_gate_reopens_and_checks_sha() {
4803 let directory = tempfile::tempdir().unwrap();
4804 let session_id = "0123456789abcdef0123456789abcdef";
4805 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4806 verify_installed_checkpoint_gate(session_id, &checkpoint).unwrap();
4807
4808 let mut wrong_sha = checkpoint.clone();
4809 wrong_sha.sha256 = "b".repeat(64);
4810 assert!(
4811 verify_installed_checkpoint_gate(session_id, &wrong_sha)
4812 .unwrap_err()
4813 .to_string()
4814 .contains("SHA changed")
4815 );
4816 std::fs::write(
4817 &checkpoint.archive_path,
4818 b"changed after first verification",
4819 )
4820 .unwrap();
4821 assert!(
4822 format!(
4823 "{:#}",
4824 verify_installed_checkpoint_gate(session_id, &checkpoint).unwrap_err()
4825 )
4826 .contains("installed checkpoint SHA changed")
4827 );
4828 }
4829 #[test]
4830 fn an_installed_archive_is_reused_when_only_relay_bookkeeping_moved() {
4831 let directory = tempfile::tempdir().unwrap();
4832 let session_id = "0123456789abcdef0123456789abcdef";
4833 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4834 let archived = verify_archive_streaming(&checkpoint.archive_path)
4835 .unwrap()
4836 .canonical_session;
4837
4838 let mut latched = archived.clone();
4841 latched.event_frontier += 6;
4842 latched.event_frontier_digest = "b".repeat(64);
4843 latched.session.last_activity_at_ms = Some(9_999);
4844
4845 let artifact = reusable_installed_checkpoint(
4846 session_id,
4847 Some(&checkpoint),
4848 "native-session",
4849 latched.event_frontier,
4850 &latched,
4851 )
4852 .expect("an unchanged session reuses its installed archive");
4853
4854 assert_eq!(artifact.metadata, checkpoint);
4855 assert_eq!(artifact.native_session_id, "native-session");
4856 assert_eq!(
4857 artifact.event_frontier_digest,
4858 archived.event_frontier_digest
4859 );
4860 verify_checkpoint_artifact(session_id, &artifact).unwrap();
4862 verify_installed_checkpoint_gate(session_id, &artifact.metadata).unwrap();
4863 }
4864 #[test]
4865 fn archive_reuse_falls_back_to_a_full_export_for_anything_but_bookkeeping() {
4866 let directory = tempfile::tempdir().unwrap();
4867 let session_id = "0123456789abcdef0123456789abcdef";
4868 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4869 let archived = verify_archive_streaming(&checkpoint.archive_path)
4870 .unwrap()
4871 .canonical_session;
4872 let mut latched = archived.clone();
4873 latched.event_frontier += 6;
4874 let reuse = |installed: Option<&CheckpointMetadata>,
4875 ordinal: u64,
4876 session: &CanonicalSessionSnapshot| {
4877 reusable_installed_checkpoint(session_id, installed, "native-session", ordinal, session)
4878 };
4879
4880 assert!(reuse(None, latched.event_frontier, &latched).is_none());
4881
4882 let mut with_new_content = latched.clone();
4883 with_new_content.transcript.push(CanonicalTranscriptItem {
4884 stable_id: "system:notice:notice-1".into(),
4885 position: latched.event_frontier,
4886 latest_content_event_ordinal: None,
4887 created_at_ms: 2_000,
4888 last_changed_at_ms: 2_000,
4889 body: CanonicalTranscriptBody::System {
4890 text: "resumed".into(),
4891 },
4892 });
4893 assert!(reuse(Some(&checkpoint), latched.event_frontier, &with_new_content).is_none());
4894
4895 assert!(reuse(Some(&checkpoint), checkpoint.event_frontier - 1, &latched).is_none());
4897
4898 let mut wrong_sha = checkpoint.clone();
4899 wrong_sha.sha256 = "b".repeat(64);
4900 assert!(reuse(Some(&wrong_sha), latched.event_frontier, &latched).is_none());
4901
4902 let another_session =
4903 write_checkpoint_gate_archive(directory.path(), "1123456789abcdef0123456789abcdef", 7);
4904 assert!(reuse(Some(&another_session), latched.event_frontier, &latched).is_none());
4905
4906 std::fs::write(&checkpoint.archive_path, b"not an archive any more").unwrap();
4907 assert!(reuse(Some(&checkpoint), latched.event_frontier, &latched).is_none());
4908 }
4909 #[cfg(unix)]
4910 #[tokio::test]
4911 async fn workspace_lease_blocks_prompts_and_releases_without_advancing_recovery() {
4912 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
4913 let directory = tempfile::tempdir().unwrap();
4914 let name = format!(
4915 "{}::workspace_lease_blocks_prompts_and_releases_without_advancing_recovery",
4916 module_path!()
4917 .strip_prefix("mj_controller::")
4918 .unwrap_or(module_path!())
4919 );
4920 let mut command = crate::targets::CommandSpec::new(
4921 std::env::current_exe().unwrap().to_string_lossy(),
4922 ["--exact", &name, "--nocapture"],
4923 );
4924 command.env.insert(LATCH_TEST_CHILD.into(), "1".into());
4925 command.env.insert(
4926 "MJ_DATA_DIR".into(),
4927 directory.path().to_string_lossy().into(),
4928 );
4929 let result =
4930 crate::targets::CancellableProcessExecutor::with_timeout(Duration::from_secs(60))
4931 .execute(&command)
4932 .unwrap();
4933 assert_eq!(
4934 result.status,
4935 0,
4936 "{}\n{}",
4937 String::from_utf8_lossy(&result.stdout),
4938 String::from_utf8_lossy(&result.stderr)
4939 );
4940 assert!(
4941 String::from_utf8_lossy(&result.stdout).contains("1 passed"),
4942 "child did not run its test"
4943 );
4944 return;
4945 }
4946 let _writer = crate::database::install_isolated_test_writer();
4947 let root = tempfile::tempdir().unwrap();
4948 let (_channels, handle, mut relay, barrier, _cursor) =
4949 latch_a_live_checkpoint(root.path(), None, ReleaseSupport::Supported, false).await;
4950 relay
4951 .connection_mut()
4952 .submit(
4953 new_command_id("release-initial").unwrap(),
4954 RelayCommand::ReleaseCheckpoint {
4955 barrier_command_id: barrier,
4956 },
4957 )
4958 .await
4959 .unwrap();
4960 relay.release();
4961 wait_until_the_actor_serves_again(&handle).await;
4962 let before = handle
4963 .view()
4964 .snapshot
4965 .unwrap()
4966 .operational
4967 .recovery_floor_ordinal;
4968 let mut workspace = IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
4969 .await
4970 .unwrap();
4971 workspace.verify().await.unwrap();
4972 drop(workspace);
4973 wait_until_the_actor_serves_again(&handle).await;
4974 assert!(
4975 handle
4976 .view()
4977 .snapshot
4978 .unwrap()
4979 .operational
4980 .checkpoint_barrier
4981 .is_none()
4982 );
4983 let mut workspace = IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
4984 .await
4985 .unwrap();
4986 let submitting = handle.clone();
4987 let mut prompt = tokio::spawn(async move {
4988 submitting
4989 .submit(
4990 new_command_id("after-write").unwrap(),
4991 RelayCommand::Prompt {
4992 prompt: vec![ContentBlock::Text(TextContent::new("go"))],
4993 },
4994 )
4995 .await
4996 });
4997 assert!(
4998 tokio::time::timeout(Duration::from_millis(50), &mut prompt)
4999 .await
5000 .is_err(),
5001 "prompt must wait for the workspace owner"
5002 );
5003 workspace.verify().await.unwrap();
5004 workspace.release().await.unwrap();
5005 tokio::time::timeout(Duration::from_secs(10), prompt)
5006 .await
5007 .unwrap()
5008 .unwrap()
5009 .unwrap();
5010 wait_until_the_actor_serves_again(&handle).await;
5011 let after = handle.view().snapshot.unwrap();
5012 assert!(after.operational.checkpoint_barrier.is_none());
5013 assert_eq!(after.operational.recovery_floor_ordinal, before);
5014 assert!(
5015 IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
5016 .await
5017 .is_err(),
5018 "a queued or running prompt must prevent file injection"
5019 );
5020 }
5021}