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