Skip to main content

mj_controller/controller/
checkpoint.rs

1//! Checkpoint export, latching, verification, and archive bookkeeping.
2
3use 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_stdin_command, pack_stdin_command,
12};
13use crate::session_manager::{
14    ManagedSessionHandle, ManagedSessionLease, SessionManagerControl, StandaloneSession,
15    new_command_id, worker_connect_needs_restart,
16};
17use crate::worker_client::RelayRejected;
18use mj_checkpoint::archive::{
19    BundleManifest, CanonicalSessionSnapshot, SessionManifest, TargetManifest,
20    verify_archive_streaming,
21};
22use mj_checkpoint::checkpoint::{
23    CHECKPOINT_EXPORT_PROTOCOL_VERSION, CHECKPOINT_STAGING_PROTOCOL_VERSION, CapturedCheckpoint,
24    CheckpointCaptureSpec, CheckpointExportSpec, CheckpointPackSpec, CheckpointRepositoryCapture,
25    CheckpointRepositorySpec, canonical_session_contains_prompt, checkpoint_sha256,
26};
27use mj_core::config::{HarnessKind, sessions_dir};
28use mj_core::state::{
29    CheckpointMetadata, ManagedSessionSnapshot, SessionRecord, SessionState, State,
30};
31use mj_transcript::projection::canonical_session_from_materialized;
32
33use crate::targets::{
34    self, CommandExecutor, CommandOutput, CommandSpec, ProcessExecutor, ProvisionStage,
35    ProvisionStageGuard,
36};
37use mj_core::relay::{RelayCommand, RelayCursor, RelayExecutionState};
38
39use super::backend::backend_locator;
40use super::readiness::wait_for_native_session_in_stage;
41use super::worker_restart::{InstalledWorkerRestart, RESTART_FOR_CHECKPOINT};
42use super::{
43    Controller, execute_checked, now, persist_session_record_transition_or_restore,
44    target_profile_home,
45};
46
47/// Where one session's work lives on its target.
48///
49/// Produced by [`Controller::session_export_layout`] and used both to build a
50/// checkpoint export specification and to run the worker's export commands
51/// against the right repository.
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct SessionExportLayout {
54    /// The provisioned target the session's commands run on.
55    pub backend: targets::TargetLocator,
56    /// The directory on the target that every repository is relative to.
57    pub workspace_root: String,
58    /// The id, within `repositories`, of the repository a caller means when it
59    /// names no repository.
60    pub primary_repository: String,
61    pub repositories: Vec<CheckpointRepositorySpec>,
62    /// Set when the session works in a Hel-owned worktree of the user's own
63    /// checkout, which is what records the branch and base commit an export
64    /// compares against.
65    pub managed_worktree: Option<mj_core::state::ManagedWorktree>,
66}
67
68/// How long an idle relay may fail to admit a barrier before its worker is
69/// treated as wedged. Busy recovery checkpoints defer immediately. A close
70/// sends a non-steering turn cancellation and gives the worker this same
71/// bounded interval to settle before recovery restarts it.
72const CHECKPOINT_BARRIER_TIMEOUT: Duration = Duration::from_secs(30);
73/// A close gets a fresh cancellation grace period once the worker accepts the
74/// request. This keeps an expensive status sync from consuming the whole
75/// cancellation budget before the worker has had a chance to settle.
76const CHECKPOINT_CANCEL_TIMEOUT: Duration = Duration::from_secs(30);
77/// After a wedged ACP forces a worker restart, wait as long as native-session
78/// startup: session/load of a long kimi transcript can outlast 30s.
79const CHECKPOINT_BARRIER_TIMEOUT_AFTER_RESTART: Duration = Duration::from_secs(300);
80
81mod archives;
82pub use archives::*;
83mod lease;
84pub use lease::*;
85mod barrier;
86mod capture;
87mod latched;
88mod layout;
89mod persist;
90mod relay;
91pub use barrier::*;
92mod workspace_lease;
93pub use workspace_lease::*;
94mod validate;
95use validate::*;
96mod staging;
97pub(super) use staging::*;
98
99#[cfg(test)]
100mod tests;