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 native_continuity_lost: false,
2837 agent_capabilities: None,
2838 agent_info: None,
2839 steering_supported: None,
2840 config_options: Vec::new(),
2841 modes: None,
2842 available_commands: Vec::new(),
2843 config: BTreeMap::new(),
2844 active_prompt: None,
2845 queued_prompts: Vec::new(),
2846 active_user_shells: Vec::new(),
2847 active_agent_terminals: Vec::new(),
2848 checkpoint_barrier: Some("checkpoint-1".into()),
2849 checkpoint_ready: None,
2850 last_acp_activity_at_ms: None,
2851 current_step_started_at_ms: None,
2852 foreground_tool_started_at_ms: None,
2853 harness_turn: None,
2854 last_harness_turn_started_ordinal: None,
2855 background_commands: Vec::new(),
2856 background_work_known: None,
2857 },
2858 }
2859 }
2860 #[test]
2861 fn checkpoint_barrier_is_not_reached_until_its_ready_cursor_is_projected() {
2862 let cursor = RelayCursor {
2863 ordinal: 7,
2864 digest: "a".repeat(64),
2865 };
2866 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2867
2868 assert!(!checkpoint_barrier_is_ready(&snapshot, "checkpoint-1"));
2869 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2870 assert!(checkpoint_barrier_is_ready(&snapshot, "checkpoint-1"));
2871 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).unwrap();
2872 }
2873 #[test]
2874 fn checkpoint_revalidation_accepts_a_frontier_that_moved_past_the_ready_cursor() {
2875 let cursor = RelayCursor {
2876 ordinal: 7,
2877 digest: "a".repeat(64),
2878 };
2879 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2880 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2881
2882 snapshot.operational.latest_ordinal = cursor.ordinal + 2;
2886 snapshot.operational.latest_digest = "b".repeat(64);
2887 snapshot.materialized.applied_event_ordinal = cursor.ordinal + 2;
2888 snapshot.materialized.applied_event_digest = "b".repeat(64);
2889 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).unwrap();
2890
2891 snapshot.operational.checkpoint_ready = Some(RelayCursor {
2893 ordinal: cursor.ordinal + 1,
2894 digest: "c".repeat(64),
2895 });
2896 assert!(validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).is_err());
2897 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2898 snapshot.operational.checkpoint_barrier = None;
2899 assert!(validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor).is_err());
2900 }
2901
2902 #[test]
2903 fn routine_kimi_checkpoint_defers_when_background_liveness_is_not_safe() {
2904 let cursor = RelayCursor {
2905 ordinal: 7,
2906 digest: "a".repeat(64),
2907 };
2908 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
2909 snapshot.operational.checkpoint_ready = Some(cursor.clone());
2910
2911 for (known, has_task, label) in [
2912 (Some(false), false, "tracker reported a failure"),
2913 (Some(true), true, "a native task is still active"),
2914 (None, false, "an older worker omitted the tracker field"),
2915 ] {
2916 snapshot.operational.background_work_known = known;
2917 snapshot.operational.background_commands = has_task
2918 .then(|| mj_core::relay::BackgroundCommand {
2919 id: "kimi:agent-1".into(),
2920 started_at_ms: 1,
2921 command: "background agent".into(),
2922 can_stop: false,
2923 })
2924 .into_iter()
2925 .collect();
2926 let error = validate_automatic_checkpoint_barrier_snapshot(
2927 &snapshot,
2928 "checkpoint-1",
2929 &cursor,
2930 HarnessKind::Kimi,
2931 )
2932 .expect_err(label);
2933 assert!(checkpoint_was_deferred(&error), "{label}: {error:#}");
2934 assert!(!checkpoint_barrier_needs_worker_restart(&error));
2935 }
2936
2937 snapshot.operational.background_work_known = None;
2941 snapshot.operational.background_commands = vec![mj_core::relay::BackgroundCommand {
2942 id: "legacy-task".into(),
2943 started_at_ms: 1,
2944 command: "legacy background work".into(),
2945 can_stop: false,
2946 }];
2947 validate_automatic_checkpoint_barrier_snapshot(
2948 &snapshot,
2949 "checkpoint-1",
2950 &cursor,
2951 HarnessKind::Codex,
2952 )
2953 .expect("non-Kimi checkpoint compatibility");
2954 }
2955 fn exported_checkpoint_json() -> Vec<u8> {
2957 serde_json::to_vec(&mj_checkpoint::checkpoint::TargetCheckpoint {
2958 path: PathBuf::from("/var/lib/hel/workers/session/checkpoint.hel.zip"),
2959 sha256: "c".repeat(64),
2960 event_frontier: 7,
2961 event_frontier_digest: "d".repeat(64),
2962 timings: None,
2963 })
2964 .unwrap()
2965 }
2966 fn export_spec_fixture() -> CheckpointExportSpec {
2967 CheckpointExportSpec {
2968 protocol_version: CHECKPOINT_EXPORT_PROTOCOL_VERSION,
2969 session: mj_checkpoint::archive::SessionManifest {
2970 id: LATCH_RELAY_SESSION.into(),
2971 title: "streamed spec".into(),
2972 harness_kind: mj_core::config::HarnessKind::Codex,
2973 profile_id: "codex".into(),
2974 native_session_id: "native-session".into(),
2975 created_at: "2026-08-12T00:00:00Z".into(),
2976 checkpointed_at: "2026-08-16T00:00:00Z".into(),
2977 hel_version: "test".into(),
2978 relay_version: "test".into(),
2979 adapter_version: "acp-v1".into(),
2980 },
2981 target: TargetManifest {
2982 template_id: "podman".into(),
2983 target_kind: "local-podman".into(),
2984 details: BTreeMap::new(),
2985 },
2986 bundle: BundleManifest {
2987 id: "project".into(),
2988 primary_repository: "app".into(),
2989 },
2990 relay_root: PathBuf::from("/var/lib/hel/workers/session"),
2991 harness_home: PathBuf::from("/var/lib/hel/profiles/codex"),
2992 workspace_root: PathBuf::from("/workspace"),
2993 repositories: Vec::new(),
2994 canonical_session: canonical_session_from_materialized(&MaterializedSession::empty(
2995 LATCH_RELAY_SESSION.to_owned(),
2996 ))
2997 .unwrap(),
2998 output_path: PathBuf::from("/var/lib/hel/workers/session/checkpoint.hel.zip"),
2999 }
3000 }
3001 struct ExportExecutor {
3004 streamed_status: i32,
3005 streamed_stderr: String,
3006 retry_stdin_after_failure: bool,
3007 stdin_calls: Cell<usize>,
3008 purposes: RefCell<Vec<String>>,
3009 streamed_spec: RefCell<Vec<u8>>,
3010 }
3011 impl ExportExecutor {
3012 fn new(streamed_status: i32, streamed_stderr: &str) -> Self {
3013 Self {
3014 streamed_status,
3015 streamed_stderr: streamed_stderr.to_owned(),
3016 retry_stdin_after_failure: false,
3017 stdin_calls: Cell::new(0),
3018 purposes: RefCell::new(Vec::new()),
3019 streamed_spec: RefCell::new(Vec::new()),
3020 }
3021 }
3022
3023 fn retry_stdin_after_failure(mut self) -> Self {
3024 self.retry_stdin_after_failure = true;
3025 self
3026 }
3027 }
3028 impl CommandExecutor for ExportExecutor {
3029 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
3030 self.purposes.borrow_mut().push(command.purpose.clone());
3031 Ok(CommandOutput {
3032 status: 0,
3033 stdout: exported_checkpoint_json(),
3034 stderr: Vec::new(),
3035 })
3036 }
3037
3038 fn execute_with_stdin(
3039 &self,
3040 command: &CommandSpec,
3041 input: &mut (dyn std::io::Read + Send),
3042 ) -> Result<CommandOutput> {
3043 self.purposes.borrow_mut().push(command.purpose.clone());
3044 let mut spec = Vec::new();
3045 input.read_to_end(&mut spec)?;
3046 *self.streamed_spec.borrow_mut() = spec;
3047 let attempt = self.stdin_calls.get();
3048 self.stdin_calls.set(attempt + 1);
3049 let failed =
3050 self.streamed_status != 0 && (attempt == 0 || !self.retry_stdin_after_failure);
3051 Ok(CommandOutput {
3052 status: if failed { self.streamed_status } else { 0 },
3053 stdout: if failed {
3054 Vec::new()
3055 } else {
3056 exported_checkpoint_json()
3057 },
3058 stderr: if failed {
3059 self.streamed_stderr.clone().into_bytes()
3060 } else {
3061 Vec::new()
3062 },
3063 })
3064 }
3065 }
3066 #[test]
3067 fn docker_checkpoint_fallback_upload_uses_docker_cp() {
3068 struct RecordingExecutor {
3069 commands: RefCell<Vec<CommandSpec>>,
3070 }
3071 impl CommandExecutor for RecordingExecutor {
3072 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
3073 self.commands.borrow_mut().push(command.clone());
3074 Ok(CommandOutput {
3075 status: 0,
3076 stdout: Vec::new(),
3077 stderr: Vec::new(),
3078 })
3079 }
3080 }
3081
3082 let executor = RecordingExecutor {
3083 commands: RefCell::new(Vec::new()),
3084 };
3085 let locator = targets::TargetLocator::LocalDocker {
3086 container_id: "hel-session-12345678".to_owned(),
3087 };
3088 upload_checkpoint_spec(
3089 &executor,
3090 &locator,
3091 LATCH_RELAY_SESSION,
3092 Path::new("checkpoint-spec.json"),
3093 "/var/lib/hel/workers/session/checkpoint-spec.json",
3094 )
3095 .unwrap();
3096
3097 let commands = executor.commands.borrow();
3098 assert_eq!(commands.len(), 1);
3099 assert_eq!(commands[0].program, "docker");
3100 assert_eq!(
3101 commands[0].args,
3102 [
3103 "cp",
3104 "checkpoint-spec.json",
3105 "hel-session-12345678:/var/lib/hel/workers/session/checkpoint-spec.json"
3106 ]
3107 );
3108 assert_eq!(commands[0].purpose, "upload checkpoint specification");
3109 }
3110 #[test]
3111 fn checkpoint_export_streams_its_spec_instead_of_uploading_it() {
3112 let locator = targets::TargetLocator::LocalPodman {
3113 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3114 workspace_storage: Default::default(),
3115 };
3116 let spec = export_spec_fixture();
3117 let executor = ExportExecutor::new(0, "");
3118
3119 let output = export_target_checkpoint(
3120 &executor,
3121 &locator,
3122 LATCH_RELAY_SESSION,
3123 &spec,
3124 "/var/lib/hel/workers/session/checkpoint-spec.json",
3125 )
3126 .unwrap();
3127
3128 assert_eq!(output.stdout, exported_checkpoint_json());
3129 assert_eq!(
3130 serde_json::from_slice::<CheckpointExportSpec>(&executor.streamed_spec.borrow())
3131 .unwrap(),
3132 spec
3133 );
3134 assert_eq!(
3135 executor.purposes.into_inner(),
3136 vec!["export target checkpoint".to_owned()]
3137 );
3138 }
3139 #[test]
3142 fn an_export_that_cannot_read_stdin_falls_back_to_uploading_the_spec() {
3143 let locator = targets::TargetLocator::LocalPodman {
3144 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3145 workspace_storage: Default::default(),
3146 };
3147 let executor = ExportExecutor::new(
3148 1,
3149 "Error: read checkpoint export spec -\n\nCaused by:\n \
3150 No such file or directory (os error 2)\n",
3151 );
3152
3153 let output = export_target_checkpoint(
3154 &executor,
3155 &locator,
3156 LATCH_RELAY_SESSION,
3157 &export_spec_fixture(),
3158 "/var/lib/hel/workers/session/checkpoint-spec.json",
3159 )
3160 .unwrap();
3161
3162 assert_eq!(output.stdout, exported_checkpoint_json());
3163 assert_eq!(
3164 executor.purposes.into_inner(),
3165 vec![
3166 "export target checkpoint".to_owned(),
3167 "upload checkpoint specification".to_owned(),
3168 "export target checkpoint".to_owned(),
3169 "remove uploaded checkpoint specification".to_owned(),
3170 ]
3171 );
3172 }
3173 #[test]
3174 fn a_failing_export_is_not_retried_as_an_old_worker() {
3175 let locator = targets::TargetLocator::LocalPodman {
3176 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3177 workspace_storage: Default::default(),
3178 };
3179 let executor = ExportExecutor::new(1, "Error: repository 'app' is missing\n");
3180
3181 let error = export_target_checkpoint(
3182 &executor,
3183 &locator,
3184 LATCH_RELAY_SESSION,
3185 &export_spec_fixture(),
3186 "/var/lib/hel/workers/session/checkpoint-spec.json",
3187 )
3188 .unwrap_err();
3189
3190 assert!(
3191 format!("{error:#}").contains("repository 'app' is missing"),
3192 "{error:#}"
3193 );
3194 assert_eq!(
3195 executor.purposes.into_inner(),
3196 vec!["export target checkpoint".to_owned()]
3197 );
3198 }
3199 #[test]
3202 fn a_legacy_export_worker_is_replaced_before_it_runs_obsolete_behavior() {
3203 let locator = targets::TargetLocator::LocalPodman {
3204 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3205 workspace_storage: Default::default(),
3206 };
3207 let spec = export_spec_fixture();
3208 let executor = ExportExecutor::new(
3209 1,
3210 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3211 unknown field `protocol_version`, expected `session` at line 1 column 20\n",
3212 )
3213 .retry_stdin_after_failure();
3214 let worker_binary = Path::new("/hel-test-worker");
3215
3216 let output = export_target_checkpoint_with_worker(
3217 &executor,
3218 &locator,
3219 LATCH_RELAY_SESSION,
3220 &spec,
3221 "/var/lib/hel/workers/session/checkpoint-spec.json",
3222 Some(worker_binary),
3223 )
3224 .unwrap();
3225
3226 assert_eq!(output.stdout, exported_checkpoint_json());
3227 assert_eq!(
3228 serde_json::from_slice::<CheckpointExportSpec>(&executor.streamed_spec.borrow())
3229 .unwrap(),
3230 spec
3231 );
3232 assert_eq!(
3233 executor.purposes.into_inner(),
3234 vec![
3235 "export target checkpoint".to_owned(),
3236 "stage replacement Mjolnir worker".to_owned(),
3237 "assign replacement worker to the worker user".to_owned(),
3238 "replace installed Mjolnir worker".to_owned(),
3239 "make replaced Mjolnir worker executable".to_owned(),
3240 "export target checkpoint".to_owned(),
3241 ]
3242 );
3243 }
3244 #[test]
3245 fn a_schema_mismatch_after_uploading_the_spec_still_replaces_the_worker_binary() {
3246 let locator = targets::TargetLocator::LocalPodman {
3247 container_id: targets::resource_name(LATCH_RELAY_SESSION).unwrap(),
3248 workspace_storage: Default::default(),
3249 };
3250 struct FileThenRefreshExecutor {
3251 purposes: RefCell<Vec<String>>,
3252 file_export_calls: Cell<usize>,
3253 }
3254 impl CommandExecutor for FileThenRefreshExecutor {
3255 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
3256 self.purposes.borrow_mut().push(command.purpose.clone());
3257 if command.purpose == "export target checkpoint" {
3258 let attempt = self.file_export_calls.get();
3259 self.file_export_calls.set(attempt + 1);
3260 if attempt == 0 {
3261 return Ok(CommandOutput {
3262 status: 1,
3263 stdout: Vec::new(),
3264 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(),
3265 });
3266 }
3267 }
3268 Ok(CommandOutput {
3269 status: 0,
3270 stdout: exported_checkpoint_json(),
3271 stderr: Vec::new(),
3272 })
3273 }
3274
3275 fn execute_with_stdin(
3276 &self,
3277 command: &CommandSpec,
3278 input: &mut (dyn std::io::Read + Send),
3279 ) -> Result<CommandOutput> {
3280 self.purposes.borrow_mut().push(command.purpose.clone());
3281 let mut discarded = Vec::new();
3282 input.read_to_end(&mut discarded)?;
3283 let stdin_calls = self
3284 .purposes
3285 .borrow()
3286 .iter()
3287 .filter(|purpose| *purpose == "export target checkpoint")
3288 .count();
3289 if stdin_calls == 1 {
3290 return Ok(CommandOutput {
3291 status: 1,
3292 stdout: Vec::new(),
3293 stderr: b"Error: read checkpoint export spec -\n\nCaused by:\n No such file or directory (os error 2)\n".to_vec(),
3294 });
3295 }
3296 Ok(CommandOutput {
3297 status: 0,
3298 stdout: exported_checkpoint_json(),
3299 stderr: Vec::new(),
3300 })
3301 }
3302 }
3303
3304 let executor = FileThenRefreshExecutor {
3305 purposes: RefCell::new(Vec::new()),
3306 file_export_calls: Cell::new(0),
3307 };
3308 let output = export_target_checkpoint_with_worker(
3309 &executor,
3310 &locator,
3311 LATCH_RELAY_SESSION,
3312 &export_spec_fixture(),
3313 "/var/lib/hel/workers/session/checkpoint-spec.json",
3314 Some(Path::new("/hel-test-worker")),
3315 )
3316 .unwrap();
3317
3318 assert_eq!(output.stdout, exported_checkpoint_json());
3319 assert_eq!(
3320 executor.purposes.into_inner(),
3321 vec![
3322 "export target checkpoint".to_owned(),
3323 "upload checkpoint specification".to_owned(),
3324 "export target checkpoint".to_owned(),
3325 "stage replacement Mjolnir worker".to_owned(),
3326 "assign replacement worker to the worker user".to_owned(),
3327 "replace installed Mjolnir worker".to_owned(),
3328 "make replaced Mjolnir worker executable".to_owned(),
3329 "export target checkpoint".to_owned(),
3330 ]
3331 );
3332 }
3333 #[test]
3337 fn a_deferral_attached_as_context_under_more_context_is_still_a_deferral() {
3338 let deferred = anyhow::anyhow!("relay proxy disconnected during hello")
3339 .context(CheckpointDeferred::background_work())
3340 .context("connect to the session worker for checkpoint");
3341 assert!(checkpoint_was_deferred(&deferred), "{deferred:#}");
3342
3343 let plain = anyhow::anyhow!("relay proxy disconnected during hello")
3344 .context("connect to the session worker for checkpoint");
3345 assert!(!checkpoint_was_deferred(&plain), "{plain:#}");
3346 }
3347
3348 #[test]
3349 fn a_working_session_defers_but_close_waits_for_cancellation_before_recovery() {
3350 let cursor = RelayCursor {
3351 ordinal: 7,
3352 digest: "a".repeat(64),
3353 };
3354 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
3355 snapshot.operational.execution = RelayExecutionState::Running;
3356
3357 let deferred = checkpoint_barrier_wait_ended(
3358 &snapshot,
3359 "checkpoint-1",
3360 BarrierBusyPolicy::DeferWhileRunning,
3361 false,
3362 false,
3363 )
3364 .expect("a working session ends the wait at once");
3365 assert!(checkpoint_was_deferred(&deferred), "{deferred:#}");
3366 assert!(
3367 !checkpoint_barrier_needs_worker_restart(&deferred),
3368 "a deferred copy must never restart the worker: {deferred:#}"
3369 );
3370 assert_eq!(
3371 BarrierBusyPolicy::of(LatchExclusivity::HoldThroughClose),
3372 BarrierBusyPolicy::InterruptWhileRunning
3373 );
3374
3375 assert!(
3378 checkpoint_barrier_wait_ended(
3379 &snapshot,
3380 "checkpoint-1",
3381 BarrierBusyPolicy::InterruptWhileRunning,
3382 false,
3383 false,
3384 )
3385 .is_none()
3386 );
3387 let interrupted = checkpoint_barrier_wait_ended(
3388 &snapshot,
3389 "checkpoint-1",
3390 BarrierBusyPolicy::InterruptWhileRunning,
3391 true,
3392 true,
3393 )
3394 .expect("an unresponsive cancellation ends the wait at the deadline");
3395 assert!(
3396 checkpoint_barrier_needs_worker_restart(&interrupted),
3397 "{interrupted:#}"
3398 );
3399 assert!(!checkpoint_was_deferred(&interrupted), "{interrupted:#}");
3400
3401 snapshot.operational.execution = RelayExecutionState::Idle;
3404 let wedged = checkpoint_barrier_wait_ended(
3405 &snapshot,
3406 "checkpoint-1",
3407 BarrierBusyPolicy::DeferWhileRunning,
3408 true,
3409 false,
3410 )
3411 .expect("the deadline ends the wait");
3412 assert!(
3413 checkpoint_barrier_needs_worker_restart(&wedged),
3414 "{wedged:#}"
3415 );
3416 assert!(!checkpoint_was_deferred(&wedged), "{wedged:#}");
3417 }
3418
3419 #[test]
3422 fn a_frontier_that_moved_before_the_latch_defers_the_checkpoint() {
3423 let cursor = RelayCursor {
3424 ordinal: 220,
3425 digest: "a".repeat(64),
3426 };
3427 ensure_exact_checkpoint_cut(&cursor, cursor.ordinal, &cursor.digest)
3428 .expect("a projection latched at the ready cursor is an exact cut");
3429
3430 for (ordinal, digest) in [(223, "a".repeat(64)), (220, "b".repeat(64))] {
3431 let error = ensure_exact_checkpoint_cut(&cursor, ordinal, &digest)
3432 .expect_err("a projection past the ready cursor is not an exact cut");
3433 assert!(checkpoint_was_deferred(&error), "{error:#}");
3434 assert!(
3435 !checkpoint_barrier_needs_worker_restart(&error),
3436 "{error:#}"
3437 );
3438 }
3439 }
3440
3441 #[test]
3445 fn a_harness_turn_started_during_capture_abandons_the_archive() {
3446 let cursor = RelayCursor {
3447 ordinal: 220,
3448 digest: "a".repeat(64),
3449 };
3450 let mut snapshot = checkpoint_barrier_snapshot(&cursor);
3451 snapshot.operational.checkpoint_ready = Some(cursor.clone());
3452
3453 snapshot.operational.last_harness_turn_started_ordinal = Some(cursor.ordinal);
3454 validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor)
3455 .expect("a turn that started at or before the cursor is covered by the archive");
3456
3457 snapshot.operational.last_harness_turn_started_ordinal = Some(cursor.ordinal + 1);
3458 let error = validate_checkpoint_barrier_snapshot(&snapshot, "checkpoint-1", &cursor)
3459 .expect_err("a turn that started after the cursor invalidates the capture");
3460 assert!(checkpoint_was_deferred(&error), "{error:#}");
3461 }
3462
3463 #[test]
3464 fn a_stuck_checkpoint_barrier_is_retried_by_restarting_the_worker() {
3465 for failure in [
3468 CheckpointBarrierUnreachable::not_admitted(
3469 "checkpoint-976f6746887c5ccd93b9d8bbe120ef06",
3470 ),
3471 CheckpointBarrierUnreachable::runtime_stopped(),
3472 ] {
3473 let error = anyhow::Error::new(failure).context("latch a session checkpoint");
3474 assert!(checkpoint_barrier_needs_worker_restart(&error), "{error:#}");
3475 }
3476 assert!(!checkpoint_barrier_needs_worker_restart(&anyhow::anyhow!(
3477 "export target checkpoint failed with status 1"
3478 )));
3479 assert!(!checkpoint_barrier_needs_worker_restart(&anyhow::anyhow!(
3482 "ACP relay did not reach checkpoint barrier checkpoint-1"
3483 )));
3484 }
3485
3486 #[test]
3487 fn an_incompatible_cancel_turn_requests_worker_recovery() {
3488 let error = anyhow::Error::new(RelayRejected(mj_core::relay::RelayProtocolError {
3489 code: mj_core::relay::RelayErrorCode::IncompatibleProtocol,
3490 message: "request uses protocol 6".into(),
3491 retryable: false,
3492 detail: None,
3493 }))
3494 .context("cancel active ACP turn before checkpoint barrier");
3495 assert!(
3496 checkpoint_cancel_turn_needs_worker_restart(&error),
3497 "{error:#}"
3498 );
3499 assert!(checkpoint_barrier_needs_worker_restart(&error.context(
3500 CheckpointBarrierUnreachable::cancel_turn_unavailable("checkpoint-1", 6,)
3501 )));
3502 }
3503 #[test]
3504 fn a_dead_worker_hello_failure_is_retried_by_restarting_the_worker() {
3505 let dead = anyhow::Error::new(RelayTransportDead::new("the proxy is gone"))
3506 .context("connect to the session worker for checkpoint");
3507 assert!(worker_connect_needs_restart(&dead), "{dead:#}");
3508 assert!(!worker_connect_needs_restart(&anyhow::anyhow!(
3509 "unknown session"
3510 )));
3511 }
3512 #[cfg(unix)]
3513 #[tokio::test]
3514 async fn checkpoint_restart_stop_failure_names_mjolnir() {
3515 struct FailingStop;
3516
3517 impl CommandExecutor for FailingStop {
3518 fn execute(&self, _command: &CommandSpec) -> Result<CommandOutput> {
3519 Ok(CommandOutput {
3520 status: 1,
3521 stdout: Vec::new(),
3522 stderr: b"permission denied".to_vec(),
3523 })
3524 }
3525 }
3526
3527 let session_id = "0123456789abcdef0123456789abcdef";
3528 let worker_root = format!("/tmp/mjolnir-checkpoint-test/{session_id}");
3529 let backend = targets::TargetLocator::LocalBare {
3530 worker_root: worker_root.clone(),
3531 };
3532 let controller = Controller {
3533 config: Config::default(),
3534 state: State::default(),
3535 };
3536 let reconnect = CommandSpec::new("unused", std::iter::empty::<&str>());
3537
3538 let result = controller
3539 .restart_worker_for_checkpoint(
3540 session_id,
3541 &FailingStop,
3542 &backend,
3543 &worker_root,
3544 &reconnect,
3545 )
3546 .await;
3547 let error = match result {
3548 Ok(_) => panic!("a failed worker stop unexpectedly restarted the checkpoint worker"),
3549 Err(error) => error,
3550 };
3551 let detail = format!("{error:#}");
3552 assert!(
3553 detail.starts_with("stop wedged Mjolnir worker before retrying checkpoint"),
3554 "{detail}"
3555 );
3556 assert!(detail.contains("permission denied"), "{detail}");
3557 }
3558 #[test]
3559 fn export_spec_schema_mismatch_is_detected_from_the_parse_error() {
3560 assert!(export_spec_schema_unsupported(
3561 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3562 unknown field `terminal_refs`, expected `call` at line 1 column 7276552\n"
3563 ));
3564 assert!(export_spec_schema_unsupported(
3565 "Error: parse checkpoint export spec /spec.json\n\nCaused by:\n \
3566 unknown variant `terminal_output`, expected one of `user`, `agent`\n"
3567 ));
3568 assert!(!export_spec_schema_unsupported(
3569 "Error: repository 'app' is missing\n"
3570 ));
3571 assert!(!export_spec_schema_unsupported(
3572 "Error: parse checkpoint export spec from standard input\n\nCaused by:\n \
3573 missing field `relay_root`\n"
3574 ));
3575 assert!(export_protocol_unsupported(
3576 "Error: unsupported checkpoint export protocol version 3; worker supports 2\n"
3577 ));
3578 }
3579 const LATCH_RELAY_ROOT: &str = "MJ_TEST_LATCH_RELAY_ROOT";
3580 const LATCH_RELAY_STARTS: &str = "MJ_TEST_LATCH_RELAY_STARTS";
3581 const LATCH_RELAY_REJECT_RELEASE: &str = "MJ_TEST_LATCH_REJECT_RELEASE";
3582 #[cfg(unix)]
3583 const LATCH_RELAY_RUNNING: &str = "MJ_TEST_LATCH_RELAY_RUNNING";
3584 #[cfg(unix)]
3585 const LATCH_TEST_CHILD: &str = "MJ_TEST_LATCH_CHILD";
3586 #[cfg(unix)]
3587 const ABANDON_TEST_CHILD: &str = "MJ_TEST_ABANDON_LATCH_CHILD";
3588 #[cfg(unix)]
3589 const RELEASE_TEST_CHILD: &str = "MJ_TEST_RELEASE_LATCH_CHILD";
3590 #[cfg(unix)]
3591 const LEGACY_RELEASE_TEST_CHILD: &str = "MJ_TEST_LEGACY_RELEASE_LATCH_CHILD";
3592 #[cfg(unix)]
3593 const REUSE_TEST_CHILD: &str = "MJ_TEST_REUSE_LATCH_CHILD";
3594 const LATCH_CHECKPOINT_ONLY: &str = "MJ_TEST_LATCH_CHECKPOINT_ONLY";
3595 const LATCH_RELAY_STARTUP_DELAY_MS: &str = "MJ_TEST_LATCH_STARTUP_DELAY_MS";
3596 const LATCH_RELAY_SESSION: &str = "018f9dd2-a3b4-7c8d-9000-0123456789ab";
3597 #[cfg(unix)]
3599 #[derive(Clone, Copy, PartialEq, Eq)]
3600 enum ReleaseSupport {
3601 Supported,
3602 Rejected,
3605 }
3606 #[test]
3613 fn latch_relay_child_serves_stdio() {
3614 let Some(root) = std::env::var_os(LATCH_RELAY_ROOT) else {
3615 return;
3616 };
3617 println!();
3621 if let Some(starts) = std::env::var_os(LATCH_RELAY_STARTS) {
3624 use std::io::Write;
3625 let mut log = OpenOptions::new()
3626 .create(true)
3627 .append(true)
3628 .open(starts)
3629 .expect("open the relay start log");
3630 writeln!(log, "{}", std::process::id()).expect("record this relay start");
3631 }
3632 let checkpoint_only = std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some();
3633 let mut relay = if checkpoint_only {
3634 mj_worker::relay::DurableRelay::open_for_checkpoint(
3635 Path::new(&root),
3636 LATCH_RELAY_SESSION,
3637 "1.0.0",
3638 )
3639 } else {
3640 mj_worker::relay::DurableRelay::open(Path::new(&root), LATCH_RELAY_SESSION, "1.0.0")
3641 }
3642 .expect("open the test relay journal");
3643 if relay.operational_state().native_session_id.is_none() {
3644 relay
3645 .record_observation(mj_core::relay::RelayObservation::SessionOpened {
3646 native_session_id: "native-session".into(),
3647 native_continuity_lost: false,
3648 resumed: true,
3649 })
3650 .unwrap();
3651 }
3652 if !relay.operational_state().goal.synchronized() && !checkpoint_only {
3653 relay.record_session_update(serde_json::from_value(serde_json::json!({
3654 "sessionUpdate":"session_info_update", "_meta":{"goal":null,"execution":{"version":1,"status":"idle"}}
3655 })).unwrap()).unwrap();
3656 }
3657 let ready_at = Instant::now()
3658 + Duration::from_millis(
3659 std::env::var(LATCH_RELAY_STARTUP_DELAY_MS)
3660 .ok()
3661 .map(|value| value.parse::<u64>().unwrap())
3662 .unwrap_or(0),
3663 );
3664 let reject_release = std::env::var_os(LATCH_RELAY_REJECT_RELEASE).is_some();
3665 #[cfg(unix)]
3666 let running = std::env::var_os(LATCH_RELAY_RUNNING).is_some();
3667 #[cfg(unix)]
3668 if running && relay.operational_state().active_prompt.is_none() {
3669 let response = relay.handle(mj_core::relay::RelayRequestEnvelope {
3670 request_id: "seed-running-request".into(),
3671 protocol_version: mj_core::relay::RELAY_PROTOCOL_VERSION,
3672 request: mj_core::relay::RelayRequest::Submit {
3673 command_id: "seed-running-prompt".into(),
3674 command: RelayCommand::Prompt {
3675 prompt: vec![ContentBlock::Text(TextContent::new("running"))],
3676 },
3677 },
3678 });
3679 assert!(matches!(
3680 response.body,
3681 mj_core::relay::RelayResponseBody::Ok {
3682 payload: mj_core::relay::RelayResponsePayload::Accepted { .. }
3683 }
3684 ));
3685 let claimed = relay
3686 .claim_pending_commands(true)
3687 .expect("seed the running prompt");
3688 assert_eq!(claimed.len(), 1);
3689 assert_eq!(claimed[0].command_id, "seed-running-prompt");
3690 }
3691 let mut reader = std::io::stdin().lock();
3692 let mut writer = std::io::stdout().lock();
3693 let mut configured = false;
3694 while let Some(request) =
3695 mj_core::relay::read_relay_frame(&mut reader).expect("read a relay request")
3696 {
3697 if !checkpoint_only && !configured && Instant::now() >= ready_at {
3698 relay
3699 .record_observation(mj_core::relay::RelayObservation::SessionConfigured {
3700 config_options: Vec::new(),
3701 })
3702 .unwrap();
3703 configured = true;
3704 }
3705 if matches!(
3706 &request.request,
3707 mj_core::relay::RelayRequest::Submit {
3708 command: RelayCommand::BeginCheckpoint { .. },
3709 ..
3710 }
3711 ) {
3712 assert!(
3713 checkpoint_only || relay.operational_state().native_session_is_ready(),
3714 "checkpoint submitted before current ACP startup finished"
3715 );
3716 }
3717 let response = if reject_release && requests_checkpoint_release(&request) {
3718 unparseable_request_response(&request)
3719 } else {
3720 relay.handle(request)
3721 };
3722 mj_core::relay::write_relay_frame(&mut writer, &response)
3723 .expect("answer a relay request");
3724 if checkpoint_only {
3725 relay.dispatch_checkpoint_only().unwrap();
3726 }
3727 for claimed in relay
3728 .claim_pending_commands(true)
3729 .expect("claim relay commands")
3730 {
3731 match claimed.command {
3732 RelayCommand::BeginCheckpoint { .. } => {
3733 relay
3734 .record_checkpoint_ready(&claimed.command_id)
3735 .expect("report the checkpoint barrier ready");
3736 }
3737 #[cfg(unix)]
3738 RelayCommand::CancelTurn => {
3739 let prompt_id = relay
3740 .operational_state()
3741 .active_prompt
3742 .as_ref()
3743 .map(|prompt| prompt.command_id.clone())
3744 .expect("a prompt to cancel");
3745 relay
3746 .record_command_completed(
3747 &claimed.command_id,
3748 RelayCommandOutcome::Cancelled,
3749 )
3750 .expect("complete the cancellation");
3751 relay
3752 .record_command_completed(
3753 &prompt_id,
3754 RelayCommandOutcome::Prompt {
3755 diagnostic: None,
3756 stop_reason: "cancelled".into(),
3757 usage: None,
3758 },
3759 )
3760 .expect("complete the cancelled prompt");
3761 }
3762 _ => {}
3763 }
3764 }
3765 }
3766 }
3767 fn requests_checkpoint_release(request: &mj_core::relay::RelayRequestEnvelope) -> bool {
3768 matches!(
3769 &request.request,
3770 mj_core::relay::RelayRequest::Submit {
3771 command: RelayCommand::ReleaseCheckpoint { .. },
3772 ..
3773 }
3774 )
3775 }
3776 fn unparseable_request_response(
3780 request: &mj_core::relay::RelayRequestEnvelope,
3781 ) -> mj_core::relay::RelayResponseEnvelope {
3782 mj_core::relay::RelayResponseEnvelope {
3783 request_id: request.request_id.clone(),
3784 protocol_version: request.protocol_version,
3785 body: mj_core::relay::RelayResponseBody::Error {
3786 error: mj_core::relay::RelayProtocolError {
3787 code: mj_core::relay::RelayErrorCode::InvalidRequest,
3788 message: "unknown variant `release_checkpoint`".into(),
3789 retryable: false,
3790 detail: None,
3791 },
3792 },
3793 }
3794 }
3795 #[cfg(unix)]
3798 fn latch_relay_target(
3799 relay_root: &Path,
3800 starts: Option<&Path>,
3801 release: ReleaseSupport,
3802 running: bool,
3803 ) -> crate::session_manager::RelaySessionTarget {
3804 let script = format!(
3807 "\"$0\" --exact {}::latch_relay_child_serves_stdio --nocapture | \
3808 grep --line-buffered '^{{'",
3809 module_path!()
3810 .strip_prefix("mj_controller::")
3811 .unwrap_or(module_path!())
3812 );
3813 let mut spec = CommandSpec::new(
3814 "sh",
3815 [
3816 "-c".to_owned(),
3817 script,
3818 std::env::current_exe()
3819 .unwrap()
3820 .to_string_lossy()
3821 .into_owned(),
3822 ],
3823 )
3824 .purpose("test latch relay");
3825 spec.env.insert(
3826 LATCH_RELAY_ROOT.to_owned(),
3827 relay_root.to_string_lossy().into_owned(),
3828 );
3829 if let Some(starts) = starts {
3830 spec.env.insert(
3831 LATCH_RELAY_STARTS.to_owned(),
3832 starts.to_string_lossy().into_owned(),
3833 );
3834 }
3835 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
3836 spec.env.insert(LATCH_CHECKPOINT_ONLY.into(), "1".into());
3837 }
3838 if release == ReleaseSupport::Rejected {
3839 spec.env
3840 .insert(LATCH_RELAY_REJECT_RELEASE.to_owned(), "1".to_owned());
3841 }
3842 if running {
3843 spec.env
3844 .insert(LATCH_RELAY_RUNNING.to_owned(), "1".to_owned());
3845 }
3846 crate::session_manager::RelaySessionTarget {
3847 session_id: LATCH_RELAY_SESSION.to_owned(),
3848 spec,
3849 worker_recovery: None,
3850 project_memory: None,
3851 }
3852 }
3853 #[cfg(unix)]
3856 async fn latch_a_live_checkpoint(
3857 relay_root: &Path,
3858 starts: Option<&Path>,
3859 release: ReleaseSupport,
3860 running: bool,
3861 ) -> (
3862 crate::session_manager::SessionManagerChannels,
3863 ManagedSessionHandle,
3864 ControllerRelayLease,
3865 String,
3866 RelayCursor,
3867 ) {
3868 crate::database::save_session(&checkpoint_test_session(LATCH_RELAY_SESSION)).unwrap();
3871 let channels = crate::session_manager::spawn_session_manager().unwrap();
3872 channels
3873 .targets
3874 .send(vec![latch_relay_target(
3875 relay_root, starts, release, running,
3876 )])
3877 .unwrap();
3878 let handle = channels
3879 .control
3880 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
3881 .await
3882 .unwrap();
3883
3884 let lease = handle.lease_connection().await.unwrap();
3885 let mut relay = ControllerRelayLease::Managed {
3886 handle: handle.clone(),
3887 lease: Some(lease),
3888 };
3889 let barrier_command_id = new_command_id("checkpoint").unwrap();
3890 let connection = relay.connection_mut();
3891 connection
3892 .submit(
3893 barrier_command_id.clone(),
3894 RelayCommand::BeginCheckpoint { reason: None },
3895 )
3896 .await
3897 .unwrap();
3898 let barrier = wait_for_checkpoint_barrier(
3899 connection,
3900 LATCH_RELAY_SESSION,
3901 &barrier_command_id,
3902 CHECKPOINT_BARRIER_TIMEOUT,
3903 BarrierBusyPolicy::InterruptWhileRunning,
3904 HarnessKind::Codex,
3905 )
3906 .await
3907 .unwrap();
3908 assert_eq!(
3909 barrier.materialized.applied_event_ordinal,
3910 barrier.operational.latest_ordinal
3911 );
3912 let cursor = barrier.operational.checkpoint_ready.clone().unwrap();
3913 (channels, handle, relay, barrier_command_id, cursor)
3914 }
3915
3916 #[cfg(unix)]
3921 #[tokio::test]
3922 async fn a_close_checkpoint_cancels_a_running_turn_without_restarting_the_worker() {
3923 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
3926 let directory = tempfile::tempdir().unwrap();
3927 let test_name = format!(
3928 "{}::a_close_checkpoint_cancels_a_running_turn_without_restarting_the_worker",
3929 module_path!()
3930 .strip_prefix("mj_controller::")
3931 .unwrap_or(module_path!())
3932 );
3933 let output = Command::new(std::env::current_exe().unwrap())
3934 .args(["--exact", &test_name, "--nocapture"])
3935 .env(LATCH_TEST_CHILD, "1")
3936 .env("MJ_DATA_DIR", directory.path())
3937 .output()
3938 .unwrap();
3939 assert!(
3940 output.status.success(),
3941 "isolated cancellation checkpoint test failed\nstdout:\n{}\nstderr:\n{}",
3942 String::from_utf8_lossy(&output.stdout),
3943 String::from_utf8_lossy(&output.stderr)
3944 );
3945 return;
3946 }
3947 let _writer = crate::database::install_isolated_test_writer();
3948 let relay_root = tempfile::tempdir().unwrap();
3949 let start_log_directory = tempfile::tempdir().unwrap();
3950 let start_log = start_log_directory.path().join("relay-starts");
3951 let (_channels, _handle, mut relay, _barrier_command_id, _cursor) =
3952 latch_a_live_checkpoint(
3953 relay_root.path(),
3954 Some(&start_log),
3955 ReleaseSupport::Supported,
3956 true,
3957 )
3958 .await;
3959 let snapshot = relay.sync_snapshot().await.unwrap();
3960 assert_eq!(
3961 snapshot.operational.execution,
3962 RelayExecutionState::Idle,
3963 "the close wait returned before the cancelled turn became idle"
3964 );
3965 assert!(
3966 snapshot.operational.active_prompt.is_none(),
3967 "the close wait returned before the cancelled prompt settled"
3968 );
3969 assert_eq!(
3970 relay_starts(&start_log),
3971 1,
3972 "responsive cancellation restarted worker"
3973 );
3974 }
3975 #[cfg(unix)]
3978 async fn wait_until_the_actor_serves_again(handle: &ManagedSessionHandle) {
3979 for attempt in 0.. {
3980 if handle.sync_now().await.is_ok() {
3981 return;
3982 }
3983 assert!(attempt < 200, "the actor never took its connection back");
3984 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
3985 }
3986 }
3987 #[cfg(unix)]
3991 #[tokio::test]
3992 async fn ending_the_checkpoint_latch_returns_the_connection_to_its_actor() {
3993 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
3996 let directory = tempfile::tempdir().unwrap();
3997 let test_name = format!(
3998 "{}::ending_the_checkpoint_latch_returns_the_connection_to_its_actor",
3999 module_path!()
4000 .strip_prefix("mj_controller::")
4001 .unwrap_or(module_path!())
4002 );
4003 let output = Command::new(std::env::current_exe().unwrap())
4004 .args(["--exact", &test_name, "--nocapture"])
4005 .env(LATCH_TEST_CHILD, "1")
4006 .env("MJ_DATA_DIR", directory.path())
4007 .output()
4008 .unwrap();
4009 assert!(
4010 output.status.success(),
4011 "isolated checkpoint latch test failed\nstdout:\n{}\nstderr:\n{}",
4012 String::from_utf8_lossy(&output.stdout),
4013 String::from_utf8_lossy(&output.stderr)
4014 );
4015 return;
4016 }
4017 let _writer = crate::database::install_isolated_test_writer();
4019
4020 std::thread::spawn(|| {
4023 std::thread::sleep(std::time::Duration::from_secs(120));
4024 eprintln!("the checkpoint latch never returned its connection");
4025 std::process::exit(101);
4026 });
4027
4028 let relay_root = tempfile::tempdir().unwrap();
4029 let (_channels, handle, mut relay, barrier_command_id, cursor) =
4030 latch_a_live_checkpoint(relay_root.path(), None, ReleaseSupport::Supported, false)
4031 .await;
4032
4033 assert!(
4036 handle.sync_now().await.is_err(),
4037 "a latched projection must not be advanced by its own actor"
4038 );
4039
4040 relay.end_latch();
4041 wait_until_the_actor_serves_again(&handle).await;
4042
4043 let latched = relay.sync_snapshot().await.unwrap();
4047 validate_checkpoint_barrier_snapshot(&latched, &barrier_command_id, &cursor).unwrap();
4048
4049 let prompt_ordinal = relay
4052 .submit(
4053 new_command_id("prompt").unwrap(),
4054 RelayCommand::Prompt {
4055 prompt: vec![ContentBlock::Text(TextContent::new("hello"))],
4056 },
4057 )
4058 .await
4059 .unwrap();
4060 assert!(prompt_ordinal > cursor.ordinal);
4061 let snapshot = relay.sync_snapshot().await.unwrap();
4062 assert!(snapshot.operational.latest_ordinal > cursor.ordinal);
4063 validate_checkpoint_barrier_snapshot(&snapshot, &barrier_command_id, &cursor).unwrap();
4064
4065 latched_checkpoint(
4066 relay,
4067 barrier_command_id,
4068 cursor,
4069 CheckpointCompletion::HeldBarrier,
4070 )
4071 .complete()
4072 .await
4073 .unwrap();
4074 handle.sync_now().await.unwrap();
4075 assert_eq!(
4076 handle
4077 .view()
4078 .snapshot
4079 .expect("the actor published the completed barrier")
4080 .operational
4081 .checkpoint_barrier,
4082 None
4083 );
4084 }
4085 #[cfg(unix)]
4089 #[tokio::test]
4090 async fn releasing_a_checkpoint_after_capture_defers_only_the_recovery_floor() {
4091 if std::env::var_os(RELEASE_TEST_CHILD).is_none() {
4094 let directory = tempfile::tempdir().unwrap();
4095 let test_name = format!(
4096 "{}::releasing_a_checkpoint_after_capture_defers_only_the_recovery_floor",
4097 module_path!()
4098 .strip_prefix("mj_controller::")
4099 .unwrap_or(module_path!())
4100 );
4101 let output = Command::new(std::env::current_exe().unwrap())
4102 .args(["--exact", &test_name, "--nocapture"])
4103 .env(RELEASE_TEST_CHILD, "1")
4104 .env("MJ_DATA_DIR", directory.path())
4105 .output()
4106 .unwrap();
4107 assert!(
4108 output.status.success(),
4109 "isolated checkpoint release test failed\nstdout:\n{}\nstderr:\n{}",
4110 String::from_utf8_lossy(&output.stdout),
4111 String::from_utf8_lossy(&output.stderr)
4112 );
4113 return;
4114 }
4115 let _writer = crate::database::install_isolated_test_writer();
4117
4118 std::thread::spawn(|| {
4121 std::thread::sleep(std::time::Duration::from_secs(120));
4122 eprintln!("the captured checkpoint never released its barrier");
4123 std::process::exit(101);
4124 });
4125
4126 let relay_root = tempfile::tempdir().unwrap();
4127 let (_channels, handle, mut relay, barrier_command_id, cursor) =
4128 latch_a_live_checkpoint(relay_root.path(), None, ReleaseSupport::Supported, false)
4129 .await;
4130 relay.end_latch();
4131 wait_until_the_actor_serves_again(&handle).await;
4132
4133 let completion = release_checkpoint_after_capture(
4136 &mut relay,
4137 LATCH_RELAY_SESSION,
4138 &barrier_command_id,
4139 &cursor,
4140 HarnessKind::Codex,
4141 )
4142 .await
4143 .unwrap();
4144 assert_eq!(completion, CheckpointCompletion::ReleasedAfterCapture);
4145 let released = relay.sync_snapshot().await.unwrap();
4146 assert_eq!(released.operational.checkpoint_barrier, None);
4147 assert_eq!(released.operational.checkpoint_ready, None);
4148 assert_eq!(
4149 released.operational.recovery_floor_ordinal, 0,
4150 "an exported archive that is not installed may not release journal history"
4151 );
4152
4153 relay
4156 .submit(
4157 new_command_id("prompt").unwrap(),
4158 RelayCommand::Prompt {
4159 prompt: vec![ContentBlock::Text(TextContent::new("during transfer"))],
4160 },
4161 )
4162 .await
4163 .unwrap();
4164 let mut dispatched = None;
4165 for attempt in 0.. {
4166 let snapshot = relay.sync_snapshot().await.unwrap();
4167 if let Some(active) = snapshot.operational.active_prompt {
4168 dispatched = Some(active);
4169 break;
4170 }
4171 assert!(attempt < 200, "a released barrier still froze ACP dispatch");
4172 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
4173 }
4174 assert!(dispatched.is_some());
4175
4176 latched_checkpoint(
4179 relay,
4180 barrier_command_id,
4181 cursor.clone(),
4182 CheckpointCompletion::ReleasedAfterCapture,
4183 )
4184 .complete()
4185 .await
4186 .unwrap();
4187 handle.sync_now().await.unwrap();
4188 let installed = handle
4189 .view()
4190 .snapshot
4191 .expect("the actor published the advanced recovery floor");
4192 assert_eq!(installed.operational.recovery_floor_ordinal, cursor.ordinal);
4193 assert_eq!(installed.operational.recovery_floor_digest, cursor.digest);
4194 }
4195 #[cfg(unix)]
4198 #[tokio::test]
4199 async fn a_worker_that_rejects_the_release_keeps_its_barrier_through_the_transfer() {
4200 if std::env::var_os(LEGACY_RELEASE_TEST_CHILD).is_none() {
4203 let directory = tempfile::tempdir().unwrap();
4204 let test_name = format!(
4205 "{}::a_worker_that_rejects_the_release_keeps_its_barrier_through_the_transfer",
4206 module_path!()
4207 .strip_prefix("mj_controller::")
4208 .unwrap_or(module_path!())
4209 );
4210 let output = Command::new(std::env::current_exe().unwrap())
4211 .args(["--exact", &test_name, "--nocapture"])
4212 .env(LEGACY_RELEASE_TEST_CHILD, "1")
4213 .env("MJ_DATA_DIR", directory.path())
4214 .output()
4215 .unwrap();
4216 assert!(
4217 output.status.success(),
4218 "isolated legacy checkpoint release test failed\nstdout:\n{}\nstderr:\n{}",
4219 String::from_utf8_lossy(&output.stdout),
4220 String::from_utf8_lossy(&output.stderr)
4221 );
4222 return;
4223 }
4224 let _writer = crate::database::install_isolated_test_writer();
4226
4227 std::thread::spawn(|| {
4230 std::thread::sleep(std::time::Duration::from_secs(120));
4231 eprintln!("the rejected release never finished its checkpoint");
4232 std::process::exit(101);
4233 });
4234
4235 let relay_root = tempfile::tempdir().unwrap();
4236 let start_log = tempfile::tempdir().unwrap();
4237 let start_log = start_log.path().join("relay-starts");
4238 let (_channels, handle, mut relay, barrier_command_id, cursor) = latch_a_live_checkpoint(
4239 relay_root.path(),
4240 Some(&start_log),
4241 ReleaseSupport::Rejected,
4242 false,
4243 )
4244 .await;
4245 relay.end_latch();
4246 wait_until_the_actor_serves_again(&handle).await;
4247
4248 let completion = release_checkpoint_after_capture(
4249 &mut relay,
4250 LATCH_RELAY_SESSION,
4251 &barrier_command_id,
4252 &cursor,
4253 HarnessKind::Codex,
4254 )
4255 .await
4256 .unwrap();
4257 assert_eq!(completion, CheckpointCompletion::HeldBarrier);
4258 assert_eq!(relay_starts(&start_log), 1);
4261
4262 let transferring = relay.sync_snapshot().await.unwrap();
4266 validate_checkpoint_barrier_snapshot(&transferring, &barrier_command_id, &cursor).unwrap();
4267 latched_checkpoint(relay, barrier_command_id, cursor.clone(), completion)
4268 .complete()
4269 .await
4270 .unwrap();
4271 handle.sync_now().await.unwrap();
4272 let completed = handle
4273 .view()
4274 .snapshot
4275 .expect("the actor published the completed barrier");
4276 assert_eq!(completed.operational.checkpoint_barrier, None);
4277 assert_eq!(completed.operational.recovery_floor_ordinal, cursor.ordinal);
4278 }
4279 #[cfg(unix)]
4284 #[tokio::test]
4285 async fn abandoning_a_latched_checkpoint_drops_the_connection_that_opened_its_barrier() {
4286 if std::env::var_os(ABANDON_TEST_CHILD).is_none() {
4289 let directory = tempfile::tempdir().unwrap();
4290 let test_name = format!(
4291 "{}::abandoning_a_latched_checkpoint_drops_the_connection_that_opened_its_barrier",
4292 module_path!()
4293 .strip_prefix("mj_controller::")
4294 .unwrap_or(module_path!())
4295 );
4296 let output = Command::new(std::env::current_exe().unwrap())
4297 .args(["--exact", &test_name, "--nocapture"])
4298 .env(ABANDON_TEST_CHILD, "1")
4299 .env("MJ_DATA_DIR", directory.path())
4300 .output()
4301 .unwrap();
4302 assert!(
4303 output.status.success(),
4304 "isolated abandoned checkpoint test failed\nstdout:\n{}\nstderr:\n{}",
4305 String::from_utf8_lossy(&output.stdout),
4306 String::from_utf8_lossy(&output.stderr)
4307 );
4308 return;
4309 }
4310 let _writer = crate::database::install_isolated_test_writer();
4312
4313 std::thread::spawn(|| {
4316 std::thread::sleep(std::time::Duration::from_secs(120));
4317 eprintln!("an abandoned checkpoint never released its relay connection");
4318 std::process::exit(101);
4319 });
4320
4321 let relay_root = tempfile::tempdir().unwrap();
4322 let start_log = tempfile::tempdir().unwrap();
4323 let start_log = start_log.path().join("relay-starts");
4324 let (_channels, handle, mut relay, barrier_command_id, cursor) = latch_a_live_checkpoint(
4325 relay_root.path(),
4326 Some(&start_log),
4327 ReleaseSupport::Supported,
4328 false,
4329 )
4330 .await;
4331 relay.end_latch();
4332 wait_until_the_actor_serves_again(&handle).await;
4333 assert_eq!(relay_starts(&start_log), 1);
4334
4335 latched_checkpoint(
4336 relay,
4337 barrier_command_id,
4338 cursor,
4339 CheckpointCompletion::HeldBarrier,
4340 )
4341 .abandon(LATCH_RELAY_SESSION)
4342 .await;
4343
4344 wait_until_the_actor_serves_again(&handle).await;
4349 assert_eq!(relay_starts(&start_log), 2);
4350 }
4351 #[cfg(unix)]
4356 #[test]
4357 fn a_move_checkpoint_can_verify_its_archive_without_source_harness_readiness() {
4358 let directory = tempfile::tempdir().unwrap();
4359 let name = format!(
4360 "{}::a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content",
4361 module_path!()
4362 .strip_prefix("mj_controller::")
4363 .unwrap_or(module_path!())
4364 );
4365 let output = Command::new(std::env::current_exe().unwrap())
4366 .args(["--exact", &name, "--nocapture"])
4367 .env(REUSE_TEST_CHILD, "1")
4368 .env(LATCH_CHECKPOINT_ONLY, "1")
4369 .env("MJ_DATA_DIR", directory.path())
4370 .output()
4371 .unwrap();
4372 assert!(
4373 output.status.success(),
4374 "checkpoint-only capture failed: {}\n{}",
4375 String::from_utf8_lossy(&output.stdout),
4376 String::from_utf8_lossy(&output.stderr)
4377 );
4378 }
4379
4380 #[cfg(unix)]
4381 #[tokio::test]
4382 async fn a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content() {
4383 if std::env::var_os(REUSE_TEST_CHILD).is_none() {
4386 let directory = tempfile::tempdir().unwrap();
4387 let test_name = format!(
4388 "{}::a_close_latch_reuses_an_unchanged_archive_and_exports_after_new_content",
4389 module_path!()
4390 .strip_prefix("mj_controller::")
4391 .unwrap_or(module_path!())
4392 );
4393 let output = Command::new(std::env::current_exe().unwrap())
4394 .args(["--exact", &test_name, "--nocapture"])
4395 .env(REUSE_TEST_CHILD, "1")
4396 .env(LATCH_RELAY_STARTUP_DELAY_MS, "31000")
4399 .env("MJ_DATA_DIR", directory.path())
4400 .output()
4401 .unwrap();
4402 assert!(
4403 output.status.success(),
4404 "isolated checkpoint reuse test failed\nstdout:\n{}\nstderr:\n{}",
4405 String::from_utf8_lossy(&output.stdout),
4406 String::from_utf8_lossy(&output.stderr)
4407 );
4408 return;
4409 }
4410 let _writer = crate::database::install_isolated_test_writer();
4412
4413 std::thread::spawn(|| {
4416 std::thread::sleep(std::time::Duration::from_secs(120));
4417 eprintln!("the reuse checkpoint never finished its latch");
4418 std::process::exit(101);
4419 });
4420
4421 #[derive(Default)]
4422 struct RecordingExecutor {
4423 purposes: std::sync::Mutex<Vec<String>>,
4424 active_stages: std::sync::Mutex<Vec<ProvisionStage>>,
4425 stage_events: std::sync::Mutex<Vec<(ProvisionStage, bool)>>,
4426 observed_stages: std::sync::Mutex<Vec<(String, Vec<ProvisionStage>)>>,
4427 }
4428
4429 impl RecordingExecutor {
4430 fn refused(&self, command: &CommandSpec) -> Result<CommandOutput> {
4431 self.purposes.lock().unwrap().push(command.purpose.clone());
4432 self.observed_stages.lock().unwrap().push((
4433 command.purpose.clone(),
4434 self.active_stages.lock().unwrap().clone(),
4435 ));
4436 Ok(CommandOutput {
4437 status: 1,
4438 stdout: Vec::new(),
4439 stderr: b"no target is provisioned for this test".to_vec(),
4440 })
4441 }
4442
4443 fn purposes(&self) -> Vec<String> {
4444 self.purposes.lock().unwrap().clone()
4445 }
4446
4447 fn observed_stages(&self) -> Vec<(String, Vec<ProvisionStage>)> {
4448 self.observed_stages.lock().unwrap().clone()
4449 }
4450
4451 fn stage_events(&self) -> Vec<(ProvisionStage, bool)> {
4452 self.stage_events.lock().unwrap().clone()
4453 }
4454 }
4455
4456 impl CommandExecutor for RecordingExecutor {
4457 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
4458 self.refused(command)
4459 }
4460
4461 fn execute_with_stdin(
4462 &self,
4463 command: &CommandSpec,
4464 _input: &mut (dyn std::io::Read + Send),
4465 ) -> Result<CommandOutput> {
4466 self.refused(command)
4467 }
4468
4469 fn stage_started(&self, stage: ProvisionStage) {
4470 self.active_stages.lock().unwrap().push(stage);
4471 self.stage_events.lock().unwrap().push((stage, true));
4472 }
4473
4474 fn stage_finished(&self, stage: ProvisionStage) {
4475 let mut active = self.active_stages.lock().unwrap();
4476 let position = active
4477 .iter()
4478 .position(|active_stage| *active_stage == stage)
4479 .expect("stage finished without a matching start");
4480 active.remove(position);
4481 self.stage_events.lock().unwrap().push((stage, false));
4482 }
4483 }
4484
4485 let data_directory = PathBuf::from(std::env::var_os("MJ_DATA_DIR").unwrap());
4486 let relay_root = data_directory.join("relay");
4487 let profile_home = data_directory.join("profile");
4488 let archive_directory = data_directory.join("archives");
4489 for directory in [&relay_root, &profile_home, &archive_directory] {
4490 std::fs::create_dir_all(directory).unwrap();
4491 }
4492 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
4493 let mut seed =
4494 mj_worker::relay::DurableRelay::open(&relay_root, LATCH_RELAY_SESSION, "1.0.0")
4495 .unwrap();
4496 seed.record_observation(mj_core::relay::RelayObservation::SessionOpened {
4497 native_session_id: "native-session".into(),
4498 native_continuity_lost: false,
4499 resumed: true,
4500 })
4501 .unwrap();
4502 seed.record_observation(mj_core::relay::RelayObservation::SessionConfigured {
4503 config_options: Vec::new(),
4504 })
4505 .unwrap();
4506 }
4507 let checkpoint = write_checkpoint_gate_archive(&archive_directory, LATCH_RELAY_SESSION, 2);
4510
4511 let mut session = checkpoint_test_session(LATCH_RELAY_SESSION);
4512 session.target_template_id = "local".into();
4513 session.target = Some(TargetLocator::LocalBare {
4514 worker_root: data_directory.join("workers").join(LATCH_RELAY_SESSION),
4515 });
4516 session.checkpoint = Some(checkpoint.clone());
4517 crate::database::save_session(&session).unwrap();
4518
4519 let mut config = Config::default();
4520 config.profiles.insert(
4521 "codex".into(),
4522 HarnessProfile {
4523 enabled: true,
4524 kind: mj_core::config::HarnessKind::Codex,
4525 home: profile_home,
4526 environment: BTreeMap::new(),
4527 context_window_bytes: None,
4528 },
4529 );
4530 config
4531 .targets
4532 .insert("local".into(), TargetTemplate::LocalBare);
4533 config.bundles.insert(
4534 "project".into(),
4535 ProjectBundle {
4536 primary_repo: "project".into(),
4537 repositories: vec![ProjectRepository {
4538 id: "project".into(),
4539 github: Some("example/project".into()),
4540 local: None,
4541 destination: "project".into(),
4542 git_ref: None,
4543 }],
4544 },
4545 );
4546 let controller = Controller {
4547 config,
4548 state: State {
4549 sessions: BTreeMap::from([(LATCH_RELAY_SESSION.into(), session)]),
4550 ..State::default()
4551 },
4552 };
4553
4554 let channels = crate::session_manager::spawn_session_manager().unwrap();
4555 channels
4556 .targets
4557 .send(vec![latch_relay_target(
4558 &relay_root,
4559 None,
4560 ReleaseSupport::Supported,
4561 false,
4562 )])
4563 .unwrap();
4564 let handle = channels
4565 .control
4566 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
4567 .await
4568 .unwrap();
4569
4570 let executor = RecordingExecutor::default();
4571 let latched = controller
4572 .checkpoint_session_latched(
4573 LATCH_RELAY_SESSION,
4574 &executor,
4575 Some(&channels.control),
4576 LatchExclusivity::HoldThroughClose,
4577 CheckpointExportPolicy::ReuseUnchangedArchive,
4578 )
4579 .await
4580 .unwrap();
4581
4582 assert!(
4583 executor.purposes().is_empty(),
4584 "an unchanged session exported an archive anyway: {:?}",
4585 executor.purposes()
4586 );
4587 assert_eq!(latched.artifact.metadata, checkpoint);
4588 assert!(checkpoint.archive_path.exists());
4589
4590 assert!(latched.cursor.ordinal > checkpoint.event_frontier);
4593 let cursor = latched.cursor.clone();
4594 latched.complete().await.unwrap();
4595 wait_until_the_actor_serves_again(&handle).await;
4596
4597 if std::env::var_os(LATCH_CHECKPOINT_ONLY).is_some() {
4598 let snapshot = handle.view().snapshot.unwrap();
4599 assert!(snapshot.operational.checkpoint_only);
4600 assert!(!snapshot.operational.native_session_is_ready());
4601 assert_eq!(
4602 verify_archive_streaming(&checkpoint.archive_path)
4603 .unwrap()
4604 .manifest
4605 .session
4606 .native_session_id,
4607 "native-session"
4608 );
4609 channels.shutdown.shutdown().await.unwrap();
4610 return;
4611 }
4612
4613 handle
4617 .submit(
4618 new_command_id("busy-prompt").unwrap(),
4619 RelayCommand::Prompt {
4620 prompt: vec![ContentBlock::Text(TextContent::new("keep working"))],
4621 },
4622 )
4623 .await
4624 .unwrap();
4625 let mut connection = handle.lease_connection().await.unwrap();
4626 let before = connection.connection_mut().sync().await.unwrap();
4627 assert_eq!(before.operational.execution, RelayExecutionState::Running);
4628 connection.release();
4629 let deferred = controller
4630 .checkpoint_session_latched(
4631 LATCH_RELAY_SESSION,
4632 &executor,
4633 Some(&channels.control),
4634 LatchExclusivity::ReleaseAfterLatch,
4635 CheckpointExportPolicy::ReuseUnchangedArchive,
4636 )
4637 .await;
4638 assert!(
4639 matches!(deferred, Err(ref error) if error.downcast_ref::<CheckpointDeferred>().is_some())
4640 );
4641 wait_until_the_actor_serves_again(&handle).await;
4642 let mut connection = handle.lease_connection().await.unwrap();
4643 let after = connection.connection_mut().sync().await.unwrap();
4644 assert_eq!(after.operational.execution, RelayExecutionState::Running);
4645 assert!(after.operational.checkpoint_barrier.is_none());
4646 let journal =
4647 std::fs::read_to_string(relay_root.join("relay-journal/active.jsonl")).unwrap();
4648 for line in journal.lines() {
4649 let event: mj_core::relay::RelayEvent = serde_json::from_str(line).unwrap();
4650 if event.ordinal > before.operational.latest_ordinal {
4651 assert!(
4652 !matches!(
4653 event.observation,
4654 mj_core::relay::RelayObservation::CommandQueued {
4655 command: RelayCommand::BeginCheckpoint { .. },
4656 ..
4657 } | mj_core::relay::RelayObservation::CommandInterrupted {
4658 command: mj_core::relay::RelayCommandKind::BeginCheckpoint,
4659 ..
4660 }
4661 ),
4662 "busy deferral journaled checkpoint activity: {event:?}"
4663 );
4664 }
4665 }
4666 connection.release();
4667 handle
4668 .submit(
4669 new_command_id("finish-busy-prompt").unwrap(),
4670 RelayCommand::CancelTurn,
4671 )
4672 .await
4673 .unwrap();
4674 handle.sync_now().await.unwrap();
4675
4676 handle
4678 .submit(
4679 new_command_id("resume-notice").unwrap(),
4680 RelayCommand::RecordNotice {
4681 text: "the session changed".into(),
4682 },
4683 )
4684 .await
4685 .unwrap();
4686 for attempt in 0.. {
4687 handle.sync_now().await.unwrap();
4688 let materialized = handle.view().snapshot.map(|snapshot| snapshot.materialized);
4689 if materialized.is_some_and(|materialized| {
4690 materialized.applied_event_ordinal > cursor.ordinal
4691 && !materialized.transcript.is_empty()
4692 }) {
4693 break;
4694 }
4695 assert!(attempt < 200, "the notice never reached the projection");
4696 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
4697 }
4698
4699 let changed = controller
4700 .checkpoint_session_latched(
4701 LATCH_RELAY_SESSION,
4702 &executor,
4703 Some(&channels.control),
4704 LatchExclusivity::HoldThroughClose,
4705 CheckpointExportPolicy::ReuseUnchangedArchive,
4706 )
4707 .await;
4708 let Err(error) = changed else {
4709 panic!("a changed session reused its installed archive");
4710 };
4711
4712 assert!(
4713 executor
4714 .purposes()
4715 .contains(&"export target checkpoint".to_owned()),
4716 "a changed session skipped its export: {:?}",
4717 executor.purposes()
4718 );
4719 assert!(
4720 format!("{error:#}").contains("no target is provisioned for this test"),
4721 "{error:#}"
4722 );
4723 assert!(
4724 executor.observed_stages().iter().any(|(purpose, stages)| {
4725 purpose == "export target checkpoint"
4726 && stages.contains(&ProvisionStage::RecoveryCopy)
4727 }),
4728 "close checkpoint export did not run inside RecoveryCopy: {:?}",
4729 executor.observed_stages()
4730 );
4731 assert_eq!(
4732 executor
4733 .stage_events()
4734 .into_iter()
4735 .filter(|(stage, _)| *stage == ProvisionStage::RecoveryCopy)
4736 .collect::<Vec<_>>(),
4737 vec![
4738 (ProvisionStage::RecoveryCopy, true),
4739 (ProvisionStage::RecoveryCopy, false)
4740 ]
4741 );
4742 assert!(executor.active_stages.lock().unwrap().is_empty());
4743 assert!(checkpoint.archive_path.exists());
4744 }
4745 #[cfg(unix)]
4746 fn relay_starts(path: &Path) -> usize {
4747 std::fs::read_to_string(path)
4748 .unwrap_or_default()
4749 .lines()
4750 .count()
4751 }
4752 #[cfg(unix)]
4755 fn latched_checkpoint(
4756 relay: ControllerRelayLease,
4757 barrier_command_id: String,
4758 cursor: RelayCursor,
4759 completion: CheckpointCompletion,
4760 ) -> LatchedCheckpoint {
4761 LatchedCheckpoint {
4762 artifact: CheckpointArtifact {
4763 metadata: CheckpointMetadata {
4764 archive_path: PathBuf::from("checkpoint.hel.zip"),
4765 sha256: "a".repeat(64),
4766 created_at: now(),
4767 event_frontier: cursor.ordinal,
4768 },
4769 native_session_id: "native-session".into(),
4770 event_frontier_digest: cursor.digest.clone(),
4771 },
4772 relay,
4773 barrier_command_id,
4774 cursor,
4775 completion,
4776 }
4777 }
4778 #[test]
4779 fn checkpoint_persistence_rollback_restores_memory_and_reports_both_failures() {
4780 let session_id = "0123456789abcdef0123456789abcdef";
4781 let previous = checkpoint_test_session(session_id);
4782 let mut changed = previous.clone();
4783 changed.state = SessionState::Closing;
4784 changed.last_checkpoint_error = Some("partially installed checkpoint".into());
4785 let mut state = State::default();
4786 state.sessions.insert(session_id.into(), changed);
4787
4788 let error = restore_session_after_persistence_failure(
4789 &mut state,
4790 session_id,
4791 &previous,
4792 anyhow::anyhow!("verified checkpoint persistence failed"),
4793 |record| {
4794 assert_eq!(record, &previous);
4795 Err(anyhow::anyhow!("rollback database write failed"))
4796 },
4797 );
4798
4799 assert_eq!(state.sessions.get(session_id), Some(&previous));
4800 let detail = format!("{error:#}");
4801 assert!(detail.contains("verified checkpoint persistence failed"));
4802 assert!(detail.contains("rollback database write failed"));
4803 }
4804 #[test]
4805 fn installed_checkpoint_gate_reopens_and_checks_sha() {
4806 let directory = tempfile::tempdir().unwrap();
4807 let session_id = "0123456789abcdef0123456789abcdef";
4808 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4809 verify_installed_checkpoint_gate(session_id, &checkpoint).unwrap();
4810
4811 let mut wrong_sha = checkpoint.clone();
4812 wrong_sha.sha256 = "b".repeat(64);
4813 assert!(
4814 verify_installed_checkpoint_gate(session_id, &wrong_sha)
4815 .unwrap_err()
4816 .to_string()
4817 .contains("SHA changed")
4818 );
4819 std::fs::write(
4820 &checkpoint.archive_path,
4821 b"changed after first verification",
4822 )
4823 .unwrap();
4824 assert!(
4825 format!(
4826 "{:#}",
4827 verify_installed_checkpoint_gate(session_id, &checkpoint).unwrap_err()
4828 )
4829 .contains("installed checkpoint SHA changed")
4830 );
4831 }
4832 #[test]
4833 fn an_installed_archive_is_reused_when_only_relay_bookkeeping_moved() {
4834 let directory = tempfile::tempdir().unwrap();
4835 let session_id = "0123456789abcdef0123456789abcdef";
4836 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4837 let archived = verify_archive_streaming(&checkpoint.archive_path)
4838 .unwrap()
4839 .canonical_session;
4840
4841 let mut latched = archived.clone();
4844 latched.event_frontier += 6;
4845 latched.event_frontier_digest = "b".repeat(64);
4846 latched.session.last_activity_at_ms = Some(9_999);
4847
4848 let artifact = reusable_installed_checkpoint(
4849 session_id,
4850 Some(&checkpoint),
4851 "native-session",
4852 latched.event_frontier,
4853 &latched,
4854 )
4855 .expect("an unchanged session reuses its installed archive");
4856
4857 assert_eq!(artifact.metadata, checkpoint);
4858 assert_eq!(artifact.native_session_id, "native-session");
4859 assert_eq!(
4860 artifact.event_frontier_digest,
4861 archived.event_frontier_digest
4862 );
4863 verify_checkpoint_artifact(session_id, &artifact).unwrap();
4865 verify_installed_checkpoint_gate(session_id, &artifact.metadata).unwrap();
4866 }
4867 #[test]
4868 fn archive_reuse_falls_back_to_a_full_export_for_anything_but_bookkeeping() {
4869 let directory = tempfile::tempdir().unwrap();
4870 let session_id = "0123456789abcdef0123456789abcdef";
4871 let checkpoint = write_checkpoint_gate_archive(directory.path(), session_id, 7);
4872 let archived = verify_archive_streaming(&checkpoint.archive_path)
4873 .unwrap()
4874 .canonical_session;
4875 let mut latched = archived.clone();
4876 latched.event_frontier += 6;
4877 let reuse = |installed: Option<&CheckpointMetadata>,
4878 ordinal: u64,
4879 session: &CanonicalSessionSnapshot| {
4880 reusable_installed_checkpoint(session_id, installed, "native-session", ordinal, session)
4881 };
4882
4883 assert!(reuse(None, latched.event_frontier, &latched).is_none());
4884
4885 let mut with_new_content = latched.clone();
4886 with_new_content.transcript.push(CanonicalTranscriptItem {
4887 stable_id: "system:notice:notice-1".into(),
4888 position: latched.event_frontier,
4889 latest_content_event_ordinal: None,
4890 created_at_ms: 2_000,
4891 last_changed_at_ms: 2_000,
4892 body: CanonicalTranscriptBody::System {
4893 text: "resumed".into(),
4894 },
4895 });
4896 assert!(reuse(Some(&checkpoint), latched.event_frontier, &with_new_content).is_none());
4897
4898 assert!(reuse(Some(&checkpoint), checkpoint.event_frontier - 1, &latched).is_none());
4900
4901 let mut wrong_sha = checkpoint.clone();
4902 wrong_sha.sha256 = "b".repeat(64);
4903 assert!(reuse(Some(&wrong_sha), latched.event_frontier, &latched).is_none());
4904
4905 let another_session =
4906 write_checkpoint_gate_archive(directory.path(), "1123456789abcdef0123456789abcdef", 7);
4907 assert!(reuse(Some(&another_session), latched.event_frontier, &latched).is_none());
4908
4909 std::fs::write(&checkpoint.archive_path, b"not an archive any more").unwrap();
4910 assert!(reuse(Some(&checkpoint), latched.event_frontier, &latched).is_none());
4911 }
4912 #[cfg(unix)]
4913 #[tokio::test]
4914 async fn workspace_lease_blocks_prompts_and_releases_without_advancing_recovery() {
4915 if std::env::var_os(LATCH_TEST_CHILD).is_none() {
4916 let directory = tempfile::tempdir().unwrap();
4917 let name = format!(
4918 "{}::workspace_lease_blocks_prompts_and_releases_without_advancing_recovery",
4919 module_path!()
4920 .strip_prefix("mj_controller::")
4921 .unwrap_or(module_path!())
4922 );
4923 let mut command = crate::targets::CommandSpec::new(
4924 std::env::current_exe().unwrap().to_string_lossy(),
4925 ["--exact", &name, "--nocapture"],
4926 );
4927 command.env.insert(LATCH_TEST_CHILD.into(), "1".into());
4928 command.env.insert(
4929 "MJ_DATA_DIR".into(),
4930 directory.path().to_string_lossy().into(),
4931 );
4932 let result =
4933 crate::targets::CancellableProcessExecutor::with_timeout(Duration::from_secs(60))
4934 .execute(&command)
4935 .unwrap();
4936 assert_eq!(
4937 result.status,
4938 0,
4939 "{}\n{}",
4940 String::from_utf8_lossy(&result.stdout),
4941 String::from_utf8_lossy(&result.stderr)
4942 );
4943 assert!(
4944 String::from_utf8_lossy(&result.stdout).contains("1 passed"),
4945 "child did not run its test"
4946 );
4947 return;
4948 }
4949 let _writer = crate::database::install_isolated_test_writer();
4950 let root = tempfile::tempdir().unwrap();
4951 let (_channels, handle, mut relay, barrier, _cursor) =
4952 latch_a_live_checkpoint(root.path(), None, ReleaseSupport::Supported, false).await;
4953 relay
4954 .connection_mut()
4955 .submit(
4956 new_command_id("release-initial").unwrap(),
4957 RelayCommand::ReleaseCheckpoint {
4958 barrier_command_id: barrier,
4959 },
4960 )
4961 .await
4962 .unwrap();
4963 relay.release();
4964 wait_until_the_actor_serves_again(&handle).await;
4965 let before = handle
4966 .view()
4967 .snapshot
4968 .unwrap()
4969 .operational
4970 .recovery_floor_ordinal;
4971 let mut workspace = IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
4972 .await
4973 .unwrap();
4974 workspace.verify().await.unwrap();
4975 drop(workspace);
4976 wait_until_the_actor_serves_again(&handle).await;
4977 assert!(
4978 handle
4979 .view()
4980 .snapshot
4981 .unwrap()
4982 .operational
4983 .checkpoint_barrier
4984 .is_none()
4985 );
4986 let mut workspace = IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
4987 .await
4988 .unwrap();
4989 let submitting = handle.clone();
4990 let mut prompt = tokio::spawn(async move {
4991 submitting
4992 .submit(
4993 new_command_id("after-write").unwrap(),
4994 RelayCommand::Prompt {
4995 prompt: vec![ContentBlock::Text(TextContent::new("go"))],
4996 },
4997 )
4998 .await
4999 });
5000 assert!(
5001 tokio::time::timeout(Duration::from_millis(50), &mut prompt)
5002 .await
5003 .is_err(),
5004 "prompt must wait for the workspace owner"
5005 );
5006 workspace.verify().await.unwrap();
5007 workspace.release().await.unwrap();
5008 tokio::time::timeout(Duration::from_secs(10), prompt)
5009 .await
5010 .unwrap()
5011 .unwrap()
5012 .unwrap();
5013 wait_until_the_actor_serves_again(&handle).await;
5014 let after = handle.view().snapshot.unwrap();
5015 assert!(after.operational.checkpoint_barrier.is_none());
5016 assert_eq!(after.operational.recovery_floor_ordinal, before);
5017 assert!(
5018 IdleWorkspaceLease::acquire(&handle, HarnessKind::Codex)
5019 .await
5020 .is_err(),
5021 "a queued or running prompt must prevent file injection"
5022 );
5023 }
5024}