Skip to main content

mj_controller/
hel_controller.rs

1//! Controller-side lifecycle transitions and canonical-to-backend conversion.
2
3mod backend;
4mod checkpoint;
5mod git_cache;
6mod lifecycle;
7pub mod move_session;
8mod provisioning;
9mod readiness;
10mod recovery_scan;
11mod resume;
12mod reviewer;
13#[cfg(test)]
14mod test_support;
15mod worker_binary;
16mod worker_restart;
17mod worktree;
18
19use std::collections::BTreeMap;
20use std::fs::{self, File, OpenOptions};
21use std::path::{Path, PathBuf};
22
23use anyhow::{Context, Result, bail, ensure};
24use chrono::Utc;
25
26use hel::hel_config::{
27    HelConfig, ProjectBundle, ProjectRepository, SshConnection, TargetTemplate, atomic_write,
28    container_size_host, data_dir, is_bare_project_target, mount_history_host,
29};
30
31use crate::hel_import::{configured_bundle_for_local, configured_bundle_for_origin};
32use crate::hel_setup::github_repository_from_origin;
33
34const CONFIG_RENAME_JOURNAL: &str = "config-rename.json";
35
36#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
37#[serde(rename_all = "snake_case")]
38enum ConfigRenameKind {
39    Profile,
40    Target,
41}
42
43#[derive(Debug, serde::Serialize, serde::Deserialize)]
44#[serde(deny_unknown_fields)]
45struct ConfigRenameJournal {
46    kind: ConfigRenameKind,
47    old_id: String,
48    new_id: String,
49}
50use hel::hel_local_git::dirty_local_repositories;
51use hel::hel_state::{
52    HelState, HostContainerSize, SessionRecord, SessionResourceAllocation, SessionState,
53    new_session_id, normalize_session_title,
54};
55use hel::hel_targets::{
56    self, AdditionalMount, CommandExecutor, CommandOutput, CommandSpec, SshTarget,
57};
58
59pub(crate) use backend::controller_github_token;
60pub use backend::image_refresh_plan;
61use backend::validate_resource_allocation;
62use provisioning::apply_failed_new_session_rollback;
63pub(crate) use worker_binary::refresh_remote_worker_binary_if_stale;
64pub(crate) use worktree::path_exists_on_managed_target;
65
66pub use checkpoint::{
67    CheckpointArtifact, CheckpointDeferred, checkpoint_was_deferred,
68    reconcile_managed_checkpoint_archives,
69};
70pub use recovery_scan::{RecoveryCandidate, RecoveryScan};
71pub use resume::{
72    ResumeRepositorySourceMismatch, ResumeRepositorySourcePreflight, ResumeRepositorySourceReceipt,
73};
74pub use worker_binary::{WorkerBinaryAvailability, worker_binary_prerequisite_for_arch};
75pub use worker_restart::WorkerUpgradeOutcome;
76pub use worktree::{ResumePlan, local_project_repository, resume_compatibility};
77
78pub struct Controller {
79    pub config: HelConfig,
80    pub state: HelState,
81}
82
83/// Machine-wide advisory lock for one controller data store. This prevents a
84/// dashboard, server, or CLI lifecycle command from concurrently acting as a
85/// second controller against the same SQLite state and relay sessions.
86#[derive(Debug)]
87pub struct ControllerStoreGuard {
88    file: File,
89}
90
91impl ControllerStoreGuard {
92    pub fn acquire() -> Result<Self> {
93        let directory = data_dir();
94        Self::acquire_at(&directory)
95    }
96
97    fn acquire_at(directory: &Path) -> Result<Self> {
98        Self::try_acquire_at(directory)?.with_context(|| {
99            format!(
100                "another Mjolnir controller is already using {}; stop it before starting this command",
101                directory.display()
102            )
103        })
104    }
105
106    /// Probe exclusivity without treating an owner that is still exiting as an error.
107    pub fn try_acquire() -> Result<Option<Self>> {
108        Self::try_acquire_at(&data_dir())
109    }
110
111    fn try_acquire_at(directory: &Path) -> Result<Option<Self>> {
112        std::fs::create_dir_all(directory)
113            .with_context(|| format!("create controller data directory {}", directory.display()))?;
114        let path = directory.join("controller.lock");
115        let mut options = OpenOptions::new();
116        options.create(true).read(true).write(true);
117        #[cfg(unix)]
118        {
119            use std::os::unix::fs::OpenOptionsExt;
120            options.mode(0o600);
121        }
122        let file = options
123            .open(&path)
124            .with_context(|| format!("open controller lock {}", path.display()))?;
125        match file.try_lock() {
126            Ok(()) => {}
127            Err(std::fs::TryLockError::WouldBlock) => return Ok(None),
128            Err(std::fs::TryLockError::Error(error)) => {
129                return Err(error)
130                    .with_context(|| format!("lock controller store {}", directory.display()));
131            }
132        }
133        Ok(Some(Self { file }))
134    }
135
136    /// Start the sole production SQLite writer after controller exclusivity
137    /// has been established by this guard.
138    pub fn start_database_writer(&self) -> Result<hel::hel_database::DatabaseWriterOwner> {
139        hel::hel_database::start_database_writer()
140    }
141}
142
143impl Drop for ControllerStoreGuard {
144    fn drop(&mut self) {
145        // Make release explicit. `File` also unlocks on close, but an explicit
146        // unlock keeps same-process handoff deterministic across platforms.
147        let _ = self.file.unlock();
148    }
149}
150
151/// The durable result of creating a quick bundle. The returned config is the
152/// same fresh config that was written, allowing a serving projection to publish
153/// the new bundle before acknowledging the request that created it.
154#[derive(Debug)]
155pub struct QuickBundleCreation {
156    pub config: HelConfig,
157    pub bundle_id: String,
158}
159
160/// Failure stages exposed to a viewer request without exposing the underlying
161/// filesystem/configuration error. The detailed error remains available to
162/// the caller for logs and terminal notices.
163#[derive(Debug)]
164pub enum QuickBundleFailure {
165    InvalidSource(anyhow::Error),
166    Persistence(anyhow::Error),
167}
168
169impl std::fmt::Display for QuickBundleFailure {
170    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
171        match self {
172            Self::InvalidSource(error) => write!(formatter, "invalid repository source: {error}"),
173            Self::Persistence(error) => write!(formatter, "persist quick bundle: {error}"),
174        }
175    }
176}
177
178impl std::error::Error for QuickBundleFailure {
179    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
180        match self {
181            Self::InvalidSource(error) | Self::Persistence(error) => Some(error.root_cause()),
182        }
183    }
184}
185
186/// Create a quick bundle from a local repository or GitHub source and persist
187/// it as one serialized fresh-config transaction. Identical sources reuse the
188/// existing configured bundle, matching the terminal's behavior. The returned
189/// config is the same fresh config that was written, allowing a serving
190/// projection to publish the new bundle before acknowledging the request.
191pub fn create_quick_bundle(
192    source: &str,
193) -> std::result::Result<QuickBundleCreation, QuickBundleFailure> {
194    let (config, bundle_id) = HelConfig::update(|config| {
195        create_quick_bundle_in_config(config, source)
196            .map_err(|error| anyhow::Error::new(QuickBundleFailure::InvalidSource(error)))
197    })
198    .map_err(|error| {
199        error
200            .downcast::<QuickBundleFailure>()
201            .unwrap_or_else(QuickBundleFailure::Persistence)
202    })?;
203    Ok(QuickBundleCreation { config, bundle_id })
204}
205
206/// Add a quick bundle to an already-loaded config. The helper still performs
207/// the local repository canonicalization/GitHub-source parsing, but callers
208/// that persist a config should use [`create_quick_bundle`] so concurrent saves
209/// cannot clobber one another.
210pub fn create_quick_bundle_in_config(config: &mut HelConfig, source: &str) -> Result<String> {
211    let source = source.trim();
212    if source.is_empty() {
213        bail!("repository source cannot be empty");
214    }
215    let candidate = Path::new(source);
216    let (name, github, local) = if candidate.exists() {
217        let root = hel::hel_local_git::canonical_repository(candidate)?;
218        if let Some(existing) = configured_bundle_for_local(config, &root) {
219            return Ok(existing);
220        }
221        let name = root
222            .file_name()
223            .and_then(|name| name.to_str())
224            .context("local repository has no usable directory name")?
225            .to_owned();
226        (name, None, Some(root))
227    } else {
228        if candidate.is_absolute() || source.starts_with('.') || source.starts_with('~') {
229            bail!("local repository path {source:?} does not exist");
230        }
231        let repository = github_repository_from_origin(source).context(format!(
232            "{source:?} is not a GitHub owner/repository or URL"
233        ))?;
234        if let Some(existing) = configured_bundle_for_origin(config, &repository) {
235            return Ok(existing);
236        }
237        let name = repository.repository.clone();
238        let github = format!("{}/{}", repository.owner, repository.repository);
239        (name, Some(github), None)
240    };
241    let repository_id = quick_config_id(&name);
242    let mut bundle_id = repository_id.clone();
243    for suffix in 2_u32.. {
244        if !config.bundles.contains_key(&bundle_id) {
245            break;
246        }
247        bundle_id = format!("{repository_id}-{suffix}");
248    }
249    config.bundles.insert(
250        bundle_id.clone(),
251        ProjectBundle {
252            primary_repo: repository_id.clone(),
253            repositories: vec![ProjectRepository {
254                id: repository_id.clone(),
255                github,
256                local,
257                destination: PathBuf::from(repository_id),
258                git_ref: None,
259            }],
260        },
261    );
262    config.validate()?;
263    Ok(bundle_id)
264}
265
266fn quick_config_id(value: &str) -> String {
267    let id = value
268        .chars()
269        .filter(|character| {
270            character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.')
271        })
272        .take(64)
273        .collect::<String>();
274    if id.is_empty() || matches!(id.as_str(), "." | "..") {
275        "repository".into()
276    } else {
277        id
278    }
279}
280
281pub struct SessionLaunchOptions {
282    pub initial_prompt: Option<String>,
283    pub workspace_id: String,
284    pub additional_mounts: Vec<AdditionalMount>,
285    pub allow_dirty_local: bool,
286    pub resource_allocation: Option<SessionResourceAllocation>,
287    pub project_directory: Option<PathBuf>,
288    pub session_title_override: Option<String>,
289}
290
291pub struct SessionResumeOptions {
292    pub additional_mounts: Option<Vec<AdditionalMount>>,
293    pub resource_allocation: Option<SessionResourceAllocation>,
294    pub discard_queue: bool,
295}
296
297fn selected_host_container_size(
298    template: &TargetTemplate,
299    allocation: Option<&SessionResourceAllocation>,
300) -> Option<(String, HostContainerSize)> {
301    let host = container_size_host(template)?;
302    let SessionResourceAllocation::Container { cpus, memory_bytes } = allocation? else {
303        return None;
304    };
305    Some((
306        host.to_owned(),
307        HostContainerSize {
308            cpus: *cpus,
309            memory_bytes: *memory_bytes,
310        },
311    ))
312}
313
314impl Controller {
315    pub fn load() -> Result<Self> {
316        let config = HelConfig::load()?;
317        let state = HelState::load()?;
318        state.validate_against_config(&config)?;
319        Ok(Self { config, state })
320    }
321
322    pub fn reload(&mut self) -> Result<()> {
323        *self = Self::load()?;
324        Ok(())
325    }
326
327    fn persist_session_state(&self, session_id: &str) -> Result<()> {
328        match self.state.sessions.get(session_id) {
329            Some(session) => hel::hel_database::save_lifecycle_session(session),
330            None => hel::hel_database::delete_session(session_id),
331        }
332    }
333
334    fn persist_session_transition_or_restore(
335        &mut self,
336        session_id: &str,
337        previous: &SessionRecord,
338        context: &'static str,
339    ) -> Result<()> {
340        persist_session_record_transition_or_restore(
341            &mut self.state,
342            session_id,
343            previous,
344            context,
345            &hel::hel_database::save_lifecycle_session,
346        )
347    }
348
349    fn restore_prior_session_after_persistence_failure(
350        &mut self,
351        session_id: &str,
352        previous: &SessionRecord,
353        primary: anyhow::Error,
354    ) -> anyhow::Error {
355        restore_session_after_persistence_failure(
356            &mut self.state,
357            session_id,
358            previous,
359            primary,
360            hel::hel_database::save_lifecycle_session,
361        )
362    }
363
364    /// Complete a mount source at the same host that will run the container.
365    pub fn complete_mount_source(
366        &self,
367        target_id: &str,
368        prefix: &str,
369        executor: &impl CommandExecutor,
370    ) -> Result<Vec<String>> {
371        let target = self
372            .config
373            .targets
374            .get(target_id)
375            .with_context(|| format!("unknown target template {target_id:?}"))?;
376        match target {
377            TargetTemplate::LocalPodman { .. }
378            | TargetTemplate::LocalDocker { .. }
379            | TargetTemplate::AppleContainer { .. }
380            | TargetTemplate::AwsEc2 { .. } => Ok(hel_targets::local_directory_completions(prefix)),
381            TargetTemplate::SshPodman { ssh, .. } | TargetTemplate::SshDocker { ssh, .. } => {
382                hel_targets::ssh_directory_completions(&backend_ssh(ssh), prefix, executor)
383            }
384            TargetTemplate::LocalBare | TargetTemplate::SshBare { .. } => {
385                bail!("resource path completion is unsupported for bare targets")
386            }
387        }
388    }
389
390    /// Verify a mount source on the host where Mjolnir will consume it, and report
391    /// the filesystem reason it must be attached read-only, if there is one.
392    ///
393    /// The probe runs in the same round trip as the existence check so the
394    /// editor learns both answers without a second wait. A probe that cannot
395    /// answer reports no reason: provisioning decides that authoritatively.
396    pub fn validate_mount_source(
397        &self,
398        target_id: &str,
399        source: &Path,
400        executor: &impl CommandExecutor,
401    ) -> Result<Option<String>> {
402        let target = self
403            .config
404            .targets
405            .get(target_id)
406            .with_context(|| format!("unknown target template {target_id:?}"))?;
407        let exists = match target {
408            TargetTemplate::LocalPodman { .. }
409            | TargetTemplate::LocalDocker { .. }
410            | TargetTemplate::AppleContainer { .. }
411            | TargetTemplate::AwsEc2 { .. } => std::fs::metadata(source)
412                .map(|metadata| metadata.is_dir())
413                .or_else(|error| {
414                    if error.kind() == std::io::ErrorKind::NotFound {
415                        Ok(false)
416                    } else {
417                        Err(error)
418                    }
419                })
420                .with_context(|| format!("inspect resource source {}", source.display()))?,
421            TargetTemplate::SshPodman { ssh, .. } | TargetTemplate::SshDocker { ssh, .. } => {
422                hel_targets::ssh_directory_exists(&backend_ssh(ssh), source, executor)?
423            }
424            TargetTemplate::LocalBare | TargetTemplate::SshBare { .. } => {
425                bail!("resource attachments are unsupported for bare targets")
426            }
427        };
428        ensure!(
429            exists,
430            "source path {} does not exist or is not a directory",
431            source.display()
432        );
433        Ok(self.forced_read_only_reason(target, source, executor))
434    }
435
436    /// The `filesystem (reason)` label for a source the runtime cannot overlay.
437    fn forced_read_only_reason(
438        &self,
439        target: &TargetTemplate,
440        source: &Path,
441        executor: &impl CommandExecutor,
442    ) -> Option<String> {
443        let ssh = match target {
444            TargetTemplate::LocalPodman { .. } | TargetTemplate::LocalDocker { .. } => None,
445            TargetTemplate::SshPodman { ssh, .. } | TargetTemplate::SshDocker { ssh, .. } => {
446                Some(backend_ssh(ssh))
447            }
448            // Apple Container already mounts read-only, and EC2 copies instead
449            // of mounting, so neither has an overlay to lose.
450            _ => return None,
451        };
452        let filesystem = hel_targets::probe_filesystem_types(
453            ssh.as_ref(),
454            std::slice::from_ref(&source.to_path_buf()),
455            executor,
456        )
457        .map_err(|error| {
458            tracing::debug!(
459                source = %source.display(),
460                error = format!("{error:#}"),
461                "could not probe the filesystem under a mount source"
462            );
463        })
464        .ok()?
465        .pop()?;
466        let reason = hel_targets::overlay_unsupported_filesystem(&filesystem)?;
467        Some(format!("{filesystem} ({reason})"))
468    }
469
470    fn fail_new_session_with_cleanup(
471        &mut self,
472        session_id: &str,
473        error: anyhow::Error,
474        executor: &impl CommandExecutor,
475    ) -> Result<anyhow::Error> {
476        let original = format!("{error:#}");
477        let cleanup_error = self
478            .cleanup_new_session_worktree_after_failure(session_id, executor)
479            .err()
480            .map(|cleanup_error| format!("{cleanup_error:#}"));
481        if let Some(cleanup_error) = &cleanup_error {
482            tracing::warn!(
483                session_id,
484                error = %cleanup_error,
485                "new-session worktree rollback reported a cleanup failure"
486            );
487        }
488        let failure = apply_failed_new_session_rollback(
489            &mut self.state,
490            session_id,
491            &original,
492            cleanup_error,
493        );
494        self.persist_session_state(session_id)?;
495        Ok(failure)
496    }
497
498    pub fn register_session_with_resources(
499        &mut self,
500        profile_id: &str,
501        bundle_id: &str,
502        target_id: &str,
503        title: impl Into<String>,
504        options: SessionLaunchOptions,
505    ) -> Result<String> {
506        let SessionLaunchOptions {
507            initial_prompt,
508            workspace_id,
509            additional_mounts,
510            allow_dirty_local,
511            resource_allocation,
512            project_directory,
513            session_title_override,
514        } = options;
515        let session_title_override = match session_title_override {
516            Some(title) => {
517                Some(normalize_session_title(&title).context("session name cannot be empty")?)
518            }
519            None => None,
520        };
521        let profile = self
522            .config
523            .profiles
524            .get(profile_id)
525            .with_context(|| format!("unknown profile {profile_id:?}"))?;
526        let template = self
527            .config
528            .targets
529            .get(target_id)
530            .with_context(|| format!("unknown target template {target_id:?}"))?;
531        if project_directory.is_some() != is_bare_project_target(template) {
532            bail!("raw project directories require a bare target, and bare targets require one");
533        }
534        if let Some(path) = &project_directory
535            && (!path.is_absolute()
536                || path
537                    .components()
538                    .any(|part| part == std::path::Component::ParentDir))
539        {
540            bail!("bare project directory must be an absolute safe path");
541        }
542        let bundle = project_directory
543            .is_none()
544            .then(|| self.config.bundles.get(bundle_id))
545            .flatten();
546        if project_directory.is_none() && bundle.is_none() {
547            bail!("unknown bundle {bundle_id:?}");
548        }
549        if matches!(
550            profile.kind,
551            hel::hel_config::HarnessKind::Deepseek | hel::hel_config::HarnessKind::Muse
552        ) && (!additional_mounts.is_empty()
553            || bundle.is_some_and(|bundle| bundle.repositories.len() > 1))
554        {
555            bail!(
556                "{} ACP supports one workspace root; use a single-repository bundle without attached directories",
557                profile.kind.display_name()
558            );
559        }
560        let dirty = bundle
561            .map(dirty_local_repositories)
562            .transpose()?
563            .unwrap_or_default();
564        if !allow_dirty_local && !dirty.is_empty() {
565            let repositories = dirty
566                .iter()
567                .map(|repository| format!("{} ({})", repository.path.display(), repository.summary))
568                .collect::<Vec<_>>()
569                .join(", ");
570            bail!(
571                "local repositories have uncommitted changes: {repositories}; explicit confirmation is required"
572            );
573        }
574        validate_resource_allocation(template, resource_allocation.as_ref())?;
575        let selected_container_size =
576            selected_host_container_size(template, resource_allocation.as_ref());
577        if !additional_mounts.is_empty() && mount_history_host(template).is_none() {
578            bail!("attached resources are unsupported for this target");
579        }
580        hel_targets::validate_additional_mounts(&additional_mounts)?;
581        let id = new_session_id()?;
582        let now = now();
583        let record = SessionRecord {
584            archived: false,
585            container_cpus: None,
586            container_memory: None,
587            id: id.clone(),
588            workspace_id,
589            title: title.into(),
590            harness_kind: profile.kind,
591            last_profile: profile_id.to_string(),
592            bundle_id: bundle_id.to_string(),
593            project_directory,
594            managed_worktree: None,
595            target_template_id: target_id.to_string(),
596            resource_allocation,
597            additional_mounts: additional_mounts.clone(),
598            state: SessionState::Provisioning,
599            target: None,
600            native_session_id: None,
601            acp_session_title: None,
602            session_title_override,
603            created_at: now.clone(),
604            updated_at: now,
605            viewed_through_event_ordinal: 0,
606            draft_input: initial_prompt.unwrap_or_default(),
607            last_error: None,
608            last_checkpoint_error: None,
609            checkpoint: None,
610        };
611        // Creation authors the whole record, so it writes the whole row. The
612        // record reaches memory only once it is durable: a session this process
613        // alone knows about is one the database can never resume or clean up.
614        if let Some((host, size)) = selected_container_size.as_ref() {
615            hel::hel_database::save_session_with_container_size(&record, host, *size)?;
616        } else {
617            hel::hel_database::save_session(&record)?;
618        }
619        self.state.sessions.insert(id.clone(), record);
620        if let Some((host, size)) = selected_container_size {
621            self.state.remember_container_size(&host, size);
622        }
623        if let Some(host) = mount_history_host(template) {
624            // Mount history only seeds the attach dialog's suggestions. The
625            // session row is already committed, so a failed suggestion write is
626            // reported rather than turned into a failed registration.
627            match hel::hel_database::remember_mount_sources(host, &additional_mounts) {
628                Ok(()) => self.state.remember_mount_sources(host, &additional_mounts),
629                Err(error) => tracing::warn!(
630                    session_id = id,
631                    error = format!("{error:#}"),
632                    "could not remember the attached resource directories for later suggestions"
633                ),
634            }
635        }
636        Ok(id)
637    }
638
639    pub fn rename_session(&mut self, session_id: &str, title: &str) -> Result<String> {
640        let title = normalize_session_title(title).context("session name cannot be empty")?;
641        ensure!(
642            self.state.sessions.contains_key(session_id),
643            "unknown session {session_id}"
644        );
645        let updated_at = now();
646        hel::hel_database::set_session_title_override(session_id, &title, &updated_at)?;
647        let record = self
648            .state
649            .sessions
650            .get_mut(session_id)
651            .expect("session was checked before updating its title");
652        record.session_title_override = Some(title.clone());
653        record.updated_at = updated_at;
654        Ok(title)
655    }
656
657    pub fn rename_profile_id(&mut self, old_id: &str, new_id: &str) -> Result<()> {
658        hel::hel_config::validate_id("profile", new_id)?;
659        if old_id == new_id {
660            ensure!(
661                self.config.profiles.contains_key(old_id),
662                "unknown profile {old_id:?}"
663            );
664            return Ok(());
665        }
666        let journal = ConfigRenameJournal {
667            kind: ConfigRenameKind::Profile,
668            old_id: old_id.to_owned(),
669            new_id: new_id.to_owned(),
670        };
671        write_config_rename_journal(&journal)?;
672        let (config, ()) = match HelConfig::update(|config| {
673            ensure!(
674                config.profiles.contains_key(old_id),
675                "unknown profile {old_id:?}"
676            );
677            ensure!(
678                !config.profiles.contains_key(new_id),
679                "profile {new_id:?} already exists"
680            );
681            let profile = config
682                .profiles
683                .remove(old_id)
684                .expect("profile was checked in the transaction");
685            config.profiles.insert(new_id.to_owned(), profile);
686            if config.startup.profile.as_deref() == Some(old_id) {
687                config.startup.profile = Some(new_id.to_owned());
688            }
689            Ok(())
690        }) {
691            Ok(result) => result,
692            Err(error) => {
693                remove_config_rename_journal()
694                    .context("remove profile rename journal after config save failed")?;
695                return Err(error).context("save renamed profile configuration");
696            }
697        };
698        self.config = config;
699        hel::hel_test_hooks::reach_test_hook("config_replacement_before_reference_migration")?;
700        if let Err(error) = hel::hel_database::rename_profile_references(old_id, new_id) {
701            let restore = HelConfig::update(|config| {
702                let profile = config
703                    .profiles
704                    .remove(new_id)
705                    .with_context(|| format!("renamed profile {new_id:?} is missing"))?;
706                ensure!(
707                    !config.profiles.contains_key(old_id),
708                    "cannot restore profile rename: both {old_id:?} and {new_id:?} exist"
709                );
710                config.profiles.insert(old_id.to_owned(), profile);
711                if config.startup.profile.as_deref() == Some(new_id) {
712                    config.startup.profile = Some(old_id.to_owned());
713                }
714                Ok(())
715            });
716            let restored = match restore {
717                Ok((config, ())) => config,
718                Err(restore_error) => {
719                    return Err(error).context(format!(
720                        "rename profile references; additionally failed to restore config: {restore_error:#}"
721                    ));
722                }
723            };
724            self.config = restored;
725            if let Err(restore_error) = remove_config_rename_journal() {
726                return Err(error).context(format!(
727                    "rename profile references; additionally failed to remove rename journal: {restore_error:#}"
728                ));
729            }
730            return Err(error).context("rename profile references");
731        }
732        for session in self.state.sessions.values_mut() {
733            if session.last_profile == old_id {
734                session.last_profile = new_id.to_owned();
735            }
736        }
737        remove_config_rename_journal()?;
738        Ok(())
739    }
740
741    pub fn rename_target_id(&mut self, old_id: &str, new_id: &str) -> Result<()> {
742        hel::hel_config::validate_id("target template", new_id)?;
743        if old_id == new_id {
744            ensure!(
745                self.config.targets.contains_key(old_id),
746                "unknown target {old_id:?}"
747            );
748            return Ok(());
749        }
750        let journal = ConfigRenameJournal {
751            kind: ConfigRenameKind::Target,
752            old_id: old_id.to_owned(),
753            new_id: new_id.to_owned(),
754        };
755        write_config_rename_journal(&journal)?;
756        let (config, ()) = match HelConfig::update(|config| {
757            ensure!(
758                config.targets.contains_key(old_id),
759                "unknown target {old_id:?}"
760            );
761            ensure!(
762                !config.targets.contains_key(new_id),
763                "target {new_id:?} already exists"
764            );
765            let target = config
766                .targets
767                .remove(old_id)
768                .expect("target was checked in the transaction");
769            config.targets.insert(new_id.to_owned(), target);
770            if config.startup.target.as_deref() == Some(old_id) {
771                config.startup.target = Some(new_id.to_owned());
772            }
773            Ok(())
774        }) {
775            Ok(result) => result,
776            Err(error) => {
777                remove_config_rename_journal()
778                    .context("remove target rename journal after config save failed")?;
779                return Err(error).context("save renamed target configuration");
780            }
781        };
782        self.config = config;
783        hel::hel_test_hooks::reach_test_hook("config_replacement_before_reference_migration")?;
784        if let Err(error) = hel::hel_database::rename_target_references(old_id, new_id) {
785            let restore = HelConfig::update(|config| {
786                let target = config
787                    .targets
788                    .remove(new_id)
789                    .with_context(|| format!("renamed target {new_id:?} is missing"))?;
790                ensure!(
791                    !config.targets.contains_key(old_id),
792                    "cannot restore target rename: both {old_id:?} and {new_id:?} exist"
793                );
794                config.targets.insert(old_id.to_owned(), target);
795                if config.startup.target.as_deref() == Some(new_id) {
796                    config.startup.target = Some(old_id.to_owned());
797                }
798                Ok(())
799            });
800            let restored = match restore {
801                Ok((config, ())) => config,
802                Err(restore_error) => {
803                    return Err(error).context(format!(
804                        "rename target references; additionally failed to restore config: {restore_error:#}"
805                    ));
806                }
807            };
808            self.config = restored;
809            if let Err(restore_error) = remove_config_rename_journal() {
810                return Err(error).context(format!(
811                    "rename target references; additionally failed to remove rename journal: {restore_error:#}"
812                ));
813            }
814            return Err(error).context("rename target references");
815        }
816        for session in self.state.sessions.values_mut() {
817            if session.target_template_id == old_id {
818                session.target_template_id = new_id.to_owned();
819            }
820        }
821        remove_config_rename_journal()?;
822        Ok(())
823    }
824
825    /// Finish a profile/target id rename interrupted between the atomic config
826    /// replacement and SQLite transaction. Each step is idempotent, so a
827    /// second crash leaves the same intent available for the next startup.
828    pub fn recover_config_id_rename() -> Result<bool> {
829        let path = config_rename_journal_path();
830        let body = match fs::read(&path) {
831            Ok(body) => body,
832            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
833            Err(error) => return Err(error).context(format!("read {}", path.display())),
834        };
835        let journal: ConfigRenameJournal =
836            serde_json::from_slice(&body).with_context(|| format!("parse {}", path.display()))?;
837        match journal.kind {
838            ConfigRenameKind::Profile => {
839                HelConfig::update(|config| {
840                    finish_config_map_rename(
841                        &mut config.profiles,
842                        &journal.old_id,
843                        &journal.new_id,
844                        "profile",
845                    )?;
846                    Ok(())
847                })?;
848                hel::hel_database::rename_profile_references(&journal.old_id, &journal.new_id)?;
849            }
850            ConfigRenameKind::Target => {
851                HelConfig::update(|config| {
852                    finish_config_map_rename(
853                        &mut config.targets,
854                        &journal.old_id,
855                        &journal.new_id,
856                        "target",
857                    )?;
858                    Ok(())
859                })?;
860                hel::hel_database::rename_target_references(&journal.old_id, &journal.new_id)?;
861            }
862        }
863        remove_config_rename_journal()?;
864        Ok(true)
865    }
866
867    /// Record the per-session container size overrides and attached
868    /// directories. Nothing is applied to a running container: the values are
869    /// read the next time the session's container is created.
870    pub fn update_session_container_settings(
871        &mut self,
872        session_id: &str,
873        cpus: Option<String>,
874        memory: Option<String>,
875        additional_mounts: Vec<hel_targets::AdditionalMount>,
876        mount_history: Vec<std::path::PathBuf>,
877    ) -> Result<()> {
878        ensure!(
879            self.state.sessions.contains_key(session_id),
880            "unknown session {session_id}"
881        );
882        let cpus = cpus.filter(|value| !value.trim().is_empty());
883        let memory = memory.filter(|value| !value.trim().is_empty());
884        let updated_at = now();
885        hel::hel_database::set_session_container_settings(
886            session_id,
887            cpus.as_deref(),
888            memory.as_deref(),
889            &additional_mounts,
890            &updated_at,
891        )?;
892        if let Some(host) = self
893            .config
894            .targets
895            .get(
896                &self.state.sessions[session_id]
897                    .target_template_id
898                    .to_owned(),
899            )
900            .and_then(hel::hel_config::mount_history_host)
901        {
902            let host = host.to_owned();
903            // The dialog owns the suggestion list, so forgetting a directory
904            // there has to survive the mounts being remembered right after.
905            hel::hel_database::replace_mount_history(&host, &mount_history)?;
906            hel::hel_database::remember_mount_sources(&host, &additional_mounts)?;
907            self.state.mount_history.insert(host.clone(), mount_history);
908            self.state.remember_mount_sources(&host, &additional_mounts);
909        }
910        let record = self
911            .state
912            .sessions
913            .get_mut(session_id)
914            .expect("session was checked before updating its container settings");
915        record.container_cpus = cpus;
916        record.container_memory = memory;
917        record.additional_mounts = additional_mounts;
918        record.updated_at = updated_at;
919        Ok(())
920    }
921}
922
923fn config_rename_journal_path() -> PathBuf {
924    data_dir().join(CONFIG_RENAME_JOURNAL)
925}
926
927fn write_config_rename_journal(journal: &ConfigRenameJournal) -> Result<()> {
928    let path = config_rename_journal_path();
929    let body = serde_json::to_vec(journal).context("serialize config rename journal")?;
930    atomic_write(&path, &body).with_context(|| format!("write {}", path.display()))
931}
932
933fn remove_config_rename_journal() -> Result<()> {
934    let path = config_rename_journal_path();
935    match fs::remove_file(&path) {
936        Ok(()) => Ok(()),
937        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
938        Err(error) => Err(error).with_context(|| format!("remove {}", path.display())),
939    }
940}
941
942fn finish_config_map_rename<T>(
943    entries: &mut BTreeMap<String, T>,
944    old_id: &str,
945    new_id: &str,
946    kind: &str,
947) -> Result<()> {
948    if let Some(entry) = entries.remove(old_id) {
949        ensure!(
950            !entries.contains_key(new_id),
951            "cannot recover {kind} rename: both {old_id:?} and {new_id:?} exist"
952        );
953        entries.insert(new_id.to_owned(), entry);
954    } else {
955        ensure!(
956            entries.contains_key(new_id),
957            "cannot recover {kind} rename: neither {old_id:?} nor {new_id:?} exists"
958        );
959    }
960    Ok(())
961}
962
963fn target_kind(locator: &hel_targets::TargetLocator) -> &'static str {
964    match locator {
965        hel_targets::TargetLocator::LocalBare { .. } => "local-bare",
966        hel_targets::TargetLocator::LocalPodman { .. } => "local-podman",
967        hel_targets::TargetLocator::LocalDocker { .. } => "local-docker",
968        hel_targets::TargetLocator::AppleContainer { .. } => "apple-container",
969        hel_targets::TargetLocator::AwsEc2 { .. } => "aws-ec2",
970        hel_targets::TargetLocator::SshBare { .. } => "ssh-bare",
971        hel_targets::TargetLocator::SshPodman { .. } => "ssh-podman",
972        hel_targets::TargetLocator::SshDocker { .. } => "ssh-docker",
973    }
974}
975
976fn target_profile_home(
977    locator: &hel_targets::TargetLocator,
978    session_id: &str,
979    profile: &hel::hel_config::HarnessProfile,
980) -> String {
981    let home = match locator {
982        hel_targets::TargetLocator::LocalBare { .. } => profile.home.to_string_lossy().into_owned(),
983        hel_targets::TargetLocator::LocalPodman { .. }
984        | hel_targets::TargetLocator::LocalDocker { .. }
985        | hel_targets::TargetLocator::AppleContainer { .. }
986        | hel_targets::TargetLocator::SshPodman { .. }
987        | hel_targets::TargetLocator::SshDocker { .. } => {
988            format!("/var/lib/hel/profiles/{session_id}")
989        }
990        hel_targets::TargetLocator::AwsEc2 { .. } | hel_targets::TargetLocator::SshBare { .. } => {
991            format!(".local/share/hel/profiles/{session_id}")
992        }
993    };
994    if profile.kind == hel::hel_config::HarnessKind::Muse {
995        let root = if matches!(locator, hel_targets::TargetLocator::LocalBare { .. }) {
996            hel::hel_config::data_dir()
997                .join("profiles")
998                .join(session_id)
999        } else {
1000            PathBuf::from(home)
1001        };
1002        root.join("muse").to_string_lossy().into_owned()
1003    } else {
1004        home
1005    }
1006}
1007
1008pub(crate) fn backend_ssh(ssh: &SshConnection) -> SshTarget {
1009    let destination = match &ssh.user {
1010        Some(user) => format!("{user}@{}", ssh.host),
1011        None => ssh.host.clone(),
1012    };
1013    SshTarget {
1014        destination,
1015        ssh_args: ssh_args_with_identity(&ssh.extra_args, ssh.identity_file.as_deref()),
1016    }
1017}
1018
1019fn ssh_command_spec(
1020    ssh: &SshTarget,
1021    args: impl IntoIterator<Item = impl AsRef<str>>,
1022) -> CommandSpec {
1023    let remote = args
1024        .into_iter()
1025        .map(|arg| arg.as_ref().to_string())
1026        .collect::<Vec<_>>();
1027    let mut command_args = ssh.ssh_args.clone();
1028    command_args.push(ssh.destination.clone());
1029    command_args.push(hel_targets::join_remote_command(&remote));
1030    CommandSpec::new("ssh", command_args)
1031}
1032
1033fn scp_command_spec(ssh: &SshTarget, source: &Path, remote: &str, recursive: bool) -> CommandSpec {
1034    let mut args = ssh.ssh_args.clone();
1035    if recursive {
1036        args.push("-r".into());
1037    }
1038    args.push(source.to_string_lossy().into_owned());
1039    args.push(format!("{}:{remote}", ssh.destination));
1040    CommandSpec::new("scp", args)
1041}
1042
1043fn ssh_args_with_identity(args: &[String], identity: Option<&Path>) -> Vec<String> {
1044    // Mjolnir drives ssh non-interactively from a TUI; a host-key or password
1045    // prompt would steal the terminal and wedge provisioning. BatchMode fails
1046    // fast instead of prompting, and accept-new trusts a first-seen host key
1047    // (fresh EC2 instances are always first-seen) while still rejecting
1048    // changed keys. User-supplied ssh_args come last so they can override.
1049    let mut result = vec![
1050        "-o".into(),
1051        "BatchMode=yes".into(),
1052        "-o".into(),
1053        "StrictHostKeyChecking=accept-new".into(),
1054        "-o".into(),
1055        "ConnectTimeout=15".into(),
1056    ];
1057    result.extend(args.iter().cloned());
1058    if let Some(identity) = identity {
1059        result.push("-i".into());
1060        result.push(identity.to_string_lossy().into_owned());
1061    }
1062    result
1063}
1064
1065fn execute_checked(executor: &impl CommandExecutor, command: CommandSpec) -> Result<CommandOutput> {
1066    let output = executor.execute(&command)?;
1067    if output.status != 0 {
1068        let detail = command_error_detail(&output.stderr);
1069        if detail.is_empty() {
1070            bail!("{} failed with status {}", command.purpose, output.status);
1071        }
1072        bail!("{detail}");
1073    }
1074    Ok(output)
1075}
1076
1077fn command_error_detail(stderr: &[u8]) -> String {
1078    let reported = String::from_utf8_lossy(stderr);
1079    let reported = reported.trim();
1080    let detail = reported
1081        .rsplit_once("\nCaused by:\n")
1082        .map_or(reported, |(_, causes)| causes);
1083    let detail = detail.strip_prefix("Error: ").unwrap_or(detail);
1084    detail
1085        .lines()
1086        .map(|line| line.strip_prefix("    ").unwrap_or(line))
1087        .collect::<Vec<_>>()
1088        .join("\n")
1089        .trim()
1090        .to_owned()
1091}
1092
1093fn now() -> String {
1094    Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true)
1095}
1096
1097fn restore_session_after_persistence_failure(
1098    state: &mut HelState,
1099    session_id: &str,
1100    previous: &SessionRecord,
1101    primary: anyhow::Error,
1102    persist: impl FnOnce(&SessionRecord) -> Result<()>,
1103) -> anyhow::Error {
1104    state
1105        .sessions
1106        .insert(session_id.to_owned(), previous.clone());
1107    let restored = state
1108        .sessions
1109        .get(session_id)
1110        .expect("restored session record disappeared");
1111    match persist(restored) {
1112        Ok(()) => primary,
1113        Err(error) => primary.context(format!(
1114            "restored prior session state in memory, but failed to persist the rollback: {error:#}"
1115        )),
1116    }
1117}
1118
1119fn persist_session_record_transition_or_restore(
1120    state: &mut HelState,
1121    session_id: &str,
1122    previous: &SessionRecord,
1123    context: &'static str,
1124    persist: &impl Fn(&SessionRecord) -> Result<()>,
1125) -> Result<()> {
1126    let result = persist(
1127        state
1128            .sessions
1129            .get(session_id)
1130            .expect("checkpoint session disappeared before persistence"),
1131    );
1132    match result {
1133        Ok(()) => Ok(()),
1134        Err(error) => Err(restore_session_after_persistence_failure(
1135            state,
1136            session_id,
1137            previous,
1138            error.context(context),
1139            persist,
1140        )),
1141    }
1142}
1143
1144#[cfg(test)]
1145mod tests {
1146    use std::collections::BTreeMap;
1147    use std::path::Path;
1148
1149    use hel::hel_config::{
1150        ContainerTemplate as ConfigContainer, HarnessKind, HarnessProfile, HelConfig,
1151        ProjectBundle, ProjectRepository, TargetTemplate,
1152    };
1153    use hel::hel_state::HelState;
1154    use hel::hel_targets::ProcessExecutor;
1155
1156    use super::*;
1157
1158    /// One profile, one bundle with nothing checked out locally, and one
1159    /// container target, which is all `register_session_with_resources` reads.
1160    fn registration_config() -> HelConfig {
1161        let mut config = HelConfig::default();
1162        config.profiles.insert(
1163            "codex".into(),
1164            HarnessProfile {
1165                kind: HarnessKind::Codex,
1166                home: PathBuf::from("/home/dev/.codex"),
1167                environment: BTreeMap::new(),
1168                context_window_bytes: None,
1169            },
1170        );
1171        config.bundles.insert(
1172            "project".into(),
1173            ProjectBundle {
1174                primary_repo: "project".into(),
1175                repositories: vec![ProjectRepository {
1176                    id: "project".into(),
1177                    github: Some("owner/project".into()),
1178                    local: None,
1179                    destination: PathBuf::from("project"),
1180                    git_ref: None,
1181                }],
1182            },
1183        );
1184        config.targets.insert(
1185            "podman".into(),
1186            TargetTemplate::LocalPodman {
1187                container: ConfigContainer {
1188                    image: "example.invalid/hel-test:latest".into(),
1189                    pull_policy: Default::default(),
1190                    platform: None,
1191                    cpus: None,
1192                    memory: None,
1193                    environment: BTreeMap::new(),
1194                    workspace_storage: Default::default(),
1195                },
1196            },
1197        );
1198        config
1199    }
1200
1201    fn launch_options(additional_mounts: Vec<AdditionalMount>) -> SessionLaunchOptions {
1202        SessionLaunchOptions {
1203            initial_prompt: None,
1204            workspace_id: hel::hel_workspace::DEFAULT_WORKSPACE_ID.to_owned(),
1205            additional_mounts,
1206            allow_dirty_local: false,
1207            resource_allocation: None,
1208            project_directory: None,
1209            session_title_override: None,
1210        }
1211    }
1212
1213    #[test]
1214    fn deepseek_registration_rejects_more_than_one_workspace_root_before_persisting() {
1215        let mut config = registration_config();
1216        config.profiles.get_mut("codex").unwrap().kind = HarnessKind::Deepseek;
1217        let second = config.bundles["project"].repositories[0].clone();
1218        config
1219            .bundles
1220            .get_mut("project")
1221            .unwrap()
1222            .repositories
1223            .push(hel::hel_config::ProjectRepository {
1224                id: "second".into(),
1225                destination: "second".into(),
1226                ..second
1227            });
1228        let mut controller = Controller {
1229            config,
1230            state: HelState::default(),
1231        };
1232
1233        let error = controller
1234            .register_session_with_resources(
1235                "codex",
1236                "project",
1237                "podman",
1238                "unsupported",
1239                launch_options(Vec::new()),
1240            )
1241            .unwrap_err();
1242
1243        assert!(error.to_string().contains("one workspace root"));
1244        assert!(controller.state.sessions.is_empty());
1245    }
1246
1247    /// MJ_DATA_DIR is process-global, so every test that reaches the
1248    /// controller database runs in an exact child with its own data directory.
1249    fn run_registration_child(marker: &str, test: &str, data_directory: &Path) {
1250        let output = std::process::Command::new(std::env::current_exe().unwrap())
1251            .args([
1252                "--exact",
1253                &format!("hel_controller::tests::{test}"),
1254                "--nocapture",
1255            ])
1256            .env(marker, "1")
1257            .env("MJ_DATA_DIR", data_directory)
1258            .env("MJ_CONFIG_DIR", data_directory)
1259            .output()
1260            .unwrap();
1261        assert!(
1262            output.status.success(),
1263            "isolated {test} failed\nstdout:\n{}\nstderr:\n{}",
1264            String::from_utf8_lossy(&output.stdout),
1265            String::from_utf8_lossy(&output.stderr)
1266        );
1267    }
1268
1269    #[test]
1270    fn registration_saves_the_initial_task_before_provisioning() {
1271        const MARKER: &str = "MJ_TEST_INITIAL_TASK_CHILD";
1272        if std::env::var_os(MARKER).is_none() {
1273            let directory = tempfile::tempdir().unwrap();
1274            run_registration_child(
1275                MARKER,
1276                "registration_saves_the_initial_task_before_provisioning",
1277                directory.path(),
1278            );
1279            return;
1280        }
1281        let _writer = hel::hel_database::install_isolated_test_writer();
1282        let mut controller = Controller {
1283            config: registration_config(),
1284            state: HelState::default(),
1285        };
1286        let prompt = format!(
1287            "Initial task\n{}\n\tPreserve indentation and λ",
1288            "x".repeat(70_000)
1289        );
1290        let mut options = launch_options(Vec::new());
1291        options.initial_prompt = Some(prompt.clone());
1292        let id = controller
1293            .register_session_with_resources("codex", "project", "podman", "fresh task", options)
1294            .unwrap();
1295        let saved = hel::hel_database::load_state().unwrap();
1296        assert_eq!(saved.sessions[&id].draft_input, prompt);
1297        assert_eq!(saved.sessions[&id].state, SessionState::Provisioning);
1298        hel::hel_database::set_session_draft_input(&id, "a newer draft").unwrap();
1299        hel::hel_database::clear_session_draft_input_if_matches(&id, &prompt).unwrap();
1300        assert_eq!(
1301            hel::hel_database::load_state().unwrap().sessions[&id].draft_input,
1302            "a newer draft"
1303        );
1304        hel::hel_database::clear_session_draft_input_if_matches(&id, "a newer draft").unwrap();
1305        assert!(
1306            hel::hel_database::load_state().unwrap().sessions[&id]
1307                .draft_input
1308                .is_empty()
1309        );
1310    }
1311
1312    const UNPERSISTABLE_SESSION_CHILD: &str = "MJ_TEST_UNPERSISTABLE_SESSION_CHILD";
1313
1314    const CONFIG_ID_RENAME_CHILD: &str = "MJ_TEST_CONFIG_ID_RENAME_CHILD";
1315
1316    #[test]
1317    fn configuration_id_rename_rewrites_durable_session_references() {
1318        if std::env::var_os(CONFIG_ID_RENAME_CHILD).is_none() {
1319            let directory = tempfile::tempdir().unwrap();
1320            run_registration_child(
1321                CONFIG_ID_RENAME_CHILD,
1322                "configuration_id_rename_rewrites_durable_session_references",
1323                directory.path(),
1324            );
1325            return;
1326        }
1327        // Alone in this child process, so it installs the one writer.
1328        let _writer = hel::hel_database::install_isolated_test_writer();
1329
1330        let mut controller = Controller {
1331            config: registration_config(),
1332            state: HelState::default(),
1333        };
1334        controller.config.startup.profile = Some("codex".into());
1335        controller.config.startup.target = Some("podman".into());
1336        controller.config.save().unwrap();
1337        let session_id = controller
1338            .register_session_with_resources(
1339                "codex",
1340                "project",
1341                "podman",
1342                "rename references",
1343                launch_options(Vec::new()),
1344            )
1345            .unwrap();
1346
1347        controller
1348            .rename_profile_id("codex", "codex-renamed")
1349            .unwrap();
1350        controller
1351            .rename_target_id("podman", "podman-renamed")
1352            .unwrap();
1353
1354        let loaded = Controller::load().unwrap();
1355        let session = &loaded.state.sessions[&session_id];
1356        assert_eq!(session.last_profile, "codex-renamed");
1357        assert_eq!(session.target_template_id, "podman-renamed");
1358        assert!(loaded.config.profiles.contains_key("codex-renamed"));
1359        assert!(loaded.config.targets.contains_key("podman-renamed"));
1360        assert_eq!(
1361            loaded.config.startup.profile.as_deref(),
1362            Some("codex-renamed")
1363        );
1364        assert_eq!(
1365            loaded.config.startup.target.as_deref(),
1366            Some("podman-renamed")
1367        );
1368        assert!(!config_rename_journal_path().exists());
1369    }
1370
1371    #[test]
1372    fn a_session_the_database_rejects_is_never_left_in_memory() {
1373        if std::env::var_os(UNPERSISTABLE_SESSION_CHILD).is_none() {
1374            let directory = tempfile::tempdir().unwrap();
1375            run_registration_child(
1376                UNPERSISTABLE_SESSION_CHILD,
1377                "a_session_the_database_rejects_is_never_left_in_memory",
1378                directory.path(),
1379            );
1380            return;
1381        }
1382        // Alone in this child process, so it installs the one writer.
1383        let _writer = hel::hel_database::install_isolated_test_writer();
1384
1385        let mut controller = Controller {
1386            config: registration_config(),
1387            state: HelState::default(),
1388        };
1389        // The store has to be healthy enough to open before it can reject a
1390        // write: this test is about a write the database refuses, not about a
1391        // store that cannot be opened at all, which now fails earlier and
1392        // louder when the writer is installed. The first registration builds
1393        // the schema the second one then loses.
1394        controller
1395            .register_session_with_resources(
1396                "codex",
1397                "project",
1398                "podman",
1399                "first",
1400                launch_options(Vec::new()),
1401            )
1402            .expect("a healthy store registers a session");
1403        rusqlite::Connection::open(hel::hel_database::database_path())
1404            .unwrap()
1405            .execute_batch("DROP TABLE sessions")
1406            .unwrap();
1407
1408        let error = controller
1409            .register_session_with_resources(
1410                "codex",
1411                "project",
1412                "podman",
1413                "unpersistable",
1414                launch_options(Vec::new()),
1415            )
1416            .expect_err("a store that rejects the write cannot register a session");
1417        assert!(
1418            format!("{error:#}").contains("sessions"),
1419            "unexpected error: {error:#}"
1420        );
1421        assert_eq!(
1422            controller.state.sessions.len(),
1423            1,
1424            "a session the database never accepted stayed in controller memory"
1425        );
1426        assert!(
1427            controller
1428                .state
1429                .sessions
1430                .values()
1431                .all(|session| session.title != "unpersistable"),
1432            "the rejected session is the one that stayed"
1433        );
1434    }
1435
1436    const MOUNT_HISTORY_FAILURE_CHILD: &str = "MJ_TEST_MOUNT_HISTORY_FAILURE_CHILD";
1437
1438    const CONTAINER_SIZE_HISTORY_CHILD: &str = "MJ_TEST_CONTAINER_SIZE_HISTORY_CHILD";
1439
1440    #[test]
1441    fn registration_remembers_launch_size_but_session_overrides_do_not_replace_it() {
1442        if std::env::var_os(CONTAINER_SIZE_HISTORY_CHILD).is_none() {
1443            let directory = tempfile::tempdir().unwrap();
1444            run_registration_child(
1445                CONTAINER_SIZE_HISTORY_CHILD,
1446                "registration_remembers_launch_size_but_session_overrides_do_not_replace_it",
1447                directory.path(),
1448            );
1449            return;
1450        }
1451        // Alone in this child process, so it installs the one writer.
1452        let _writer = hel::hel_database::install_isolated_test_writer();
1453
1454        let mut controller = Controller {
1455            config: registration_config(),
1456            state: HelState::default(),
1457        };
1458        let mut options = launch_options(Vec::new());
1459        options.resource_allocation = Some(SessionResourceAllocation::Container {
1460            cpus: 12,
1461            memory_bytes: 48 * 1024 * 1024 * 1024,
1462        });
1463        let id = controller
1464            .register_session_with_resources("codex", "project", "podman", "sized", options)
1465            .unwrap();
1466        let expected = HostContainerSize {
1467            cpus: 12,
1468            memory_bytes: 48 * 1024 * 1024 * 1024,
1469        };
1470        assert_eq!(controller.state.container_sizes["local"], expected);
1471        assert_eq!(HelState::load().unwrap().container_sizes["local"], expected);
1472
1473        controller
1474            .update_session_container_settings(
1475                &id,
1476                Some("2".into()),
1477                Some("4g".into()),
1478                Vec::new(),
1479                Vec::new(),
1480            )
1481            .unwrap();
1482        assert_eq!(HelState::load().unwrap().container_sizes["local"], expected);
1483    }
1484
1485    #[test]
1486    fn a_failed_mount_history_write_does_not_fail_the_registered_session() {
1487        if std::env::var_os(MOUNT_HISTORY_FAILURE_CHILD).is_none() {
1488            let directory = tempfile::tempdir().unwrap();
1489            run_registration_child(
1490                MOUNT_HISTORY_FAILURE_CHILD,
1491                "a_failed_mount_history_write_does_not_fail_the_registered_session",
1492                directory.path(),
1493            );
1494            return;
1495        }
1496        // Alone in this child process, so it installs the one writer.
1497        let _writer = hel::hel_database::install_isolated_test_writer();
1498
1499        let mut controller = Controller {
1500            config: registration_config(),
1501            state: HelState::default(),
1502        };
1503        // The first registration builds the schema this test then breaks.
1504        controller
1505            .register_session_with_resources(
1506                "codex",
1507                "project",
1508                "podman",
1509                "first",
1510                launch_options(Vec::new()),
1511            )
1512            .expect("a healthy store registers a session");
1513        let database = hel::hel_database::database_path();
1514        rusqlite::Connection::open(&database)
1515            .unwrap()
1516            .execute_batch("DROP TABLE mount_history")
1517            .unwrap();
1518
1519        let id = controller
1520            .register_session_with_resources(
1521                "codex",
1522                "project",
1523                "podman",
1524                "attached",
1525                launch_options(vec![AdditionalMount {
1526                    source: PathBuf::from("/host/models"),
1527                    destination: PathBuf::from("/mnt/models"),
1528                    read_only: false,
1529                }]),
1530            )
1531            .expect("a suggestion list that cannot be written must not fail a registration");
1532
1533        let stored: i64 = rusqlite::Connection::open(&database)
1534            .unwrap()
1535            .query_row(
1536                "SELECT count(*) FROM sessions WHERE session_id = ?1",
1537                [&id],
1538                |row| row.get(0),
1539            )
1540            .unwrap();
1541        assert_eq!(stored, 1, "the registered session was not committed");
1542        assert!(
1543            controller.state.mount_history.is_empty(),
1544            "controller memory remembered mount sources the database never stored"
1545        );
1546    }
1547
1548    #[test]
1549    fn command_errors_report_the_root_cause_without_worker_wrappers() {
1550        let stderr = b"Error: restore target checkpoint failed with status 1: Error: restore repository \"bifrost\"\n\nCaused by:\n    checkpoint base b41dc78 is absent from configured source\n    repository may have moved\n";
1551
1552        assert_eq!(
1553            command_error_detail(stderr),
1554            "checkpoint base b41dc78 is absent from configured source\nrepository may have moved"
1555        );
1556    }
1557
1558    #[test]
1559    fn controller_store_lock_excludes_a_second_process_owner() {
1560        let directory = tempfile::tempdir().unwrap();
1561        let first = ControllerStoreGuard::acquire_at(directory.path()).unwrap();
1562        run_controller_lock_probe(directory.path(), true);
1563        drop(first);
1564        run_controller_lock_probe(directory.path(), false);
1565    }
1566    fn run_controller_lock_probe(directory: &Path, expect_locked: bool) {
1567        let output = std::process::Command::new(std::env::current_exe().unwrap())
1568            .args([
1569                "--exact",
1570                "hel_controller::tests::controller_store_lock_subprocess_probe",
1571                "--nocapture",
1572            ])
1573            .env("MJ_CONTROLLER_LOCK_PROBE", directory)
1574            .env(
1575                "MJ_CONTROLLER_LOCK_EXPECTED",
1576                if expect_locked { "locked" } else { "available" },
1577            )
1578            .output()
1579            .unwrap();
1580        assert!(
1581            output.status.success(),
1582            "controller lock subprocess failed:\nstdout:\n{}\nstderr:\n{}",
1583            String::from_utf8_lossy(&output.stdout),
1584            String::from_utf8_lossy(&output.stderr)
1585        );
1586    }
1587    #[test]
1588    fn controller_store_lock_subprocess_probe() {
1589        let Some(directory) = std::env::var_os("MJ_CONTROLLER_LOCK_PROBE") else {
1590            return;
1591        };
1592        let expected = std::env::var("MJ_CONTROLLER_LOCK_EXPECTED").unwrap();
1593        let acquired = ControllerStoreGuard::acquire_at(Path::new(&directory));
1594        match expected.as_str() {
1595            "locked" => {
1596                let error = acquired.expect_err("a second process acquired the controller store");
1597                assert!(error.to_string().contains("another Mjolnir controller"));
1598            }
1599            "available" => {
1600                acquired.expect("released controller store stayed locked");
1601            }
1602            value => panic!("unexpected lock probe expectation {value:?}"),
1603        }
1604    }
1605    #[test]
1606    fn local_mount_source_must_be_an_existing_directory() {
1607        let directory = tempfile::tempdir().unwrap();
1608        let file = directory.path().join("file");
1609        std::fs::write(&file, "not a directory").unwrap();
1610        let mut config = HelConfig::default();
1611        config.targets.insert(
1612            "local".into(),
1613            TargetTemplate::LocalPodman {
1614                container: ConfigContainer {
1615                    image: "ubuntu:24.04".into(),
1616                    pull_policy: Default::default(),
1617                    platform: None,
1618                    cpus: None,
1619                    memory: None,
1620                    environment: BTreeMap::new(),
1621                    workspace_storage: Default::default(),
1622                },
1623            },
1624        );
1625        let controller = Controller {
1626            config,
1627            state: HelState::default(),
1628        };
1629
1630        assert!(
1631            controller
1632                .validate_mount_source("local", directory.path(), &ProcessExecutor)
1633                .is_ok()
1634        );
1635        for invalid in [file, directory.path().join("missing")] {
1636            let error = controller
1637                .validate_mount_source("local", &invalid, &ProcessExecutor)
1638                .unwrap_err();
1639            assert!(
1640                error
1641                    .to_string()
1642                    .contains("does not exist or is not a directory")
1643            );
1644        }
1645    }
1646}