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