Skip to main content

mj_controller/controller/
provisioning.rs

1//! Session provisioning, rollback, and worker-side Git bootstrap.
2
3use std::collections::BTreeMap;
4use std::path::{Path, PathBuf};
5use std::process::{Command, Stdio};
6use std::time::{Duration, Instant};
7
8use anyhow::{Context, Result, bail, ensure};
9
10use mj_core::config::{TargetTemplate, atomic_write, data_dir};
11use mj_core::state::{SessionState, State, TargetLocator};
12
13use crate::targets::{
14    self, CancellableProcessExecutor, CommandExecutor, CommandOutput, CommandSpec, ProvisionStage,
15    ProvisionStageGuard,
16};
17
18use super::backend::{
19    ContainerOverrides, backend_bundle, backend_locator, backend_target,
20    configure_github_token_environment, controller_github_token, locator_after_provision,
21    preflight_target, use_github_https_urls,
22};
23use super::git_cache;
24use super::readiness::{connect_started_worker, wait_for_native_session_in_stage};
25use super::worker_binary::{bridge_readiness_stage, start_worker, worker_probe_diagnosis};
26use super::{Controller, execute_checked, now};
27
28const INHERITED_GIT_SETTINGS: &[&str] = &[
29    "diff.algorithm",
30    "fetch.prune",
31    "fetch.prunetags",
32    "init.defaultbranch",
33    "merge.conflictstyle",
34    "pull.ff",
35    "pull.rebase",
36    "push.autosetupremote",
37    "push.default",
38    "rebase.autostash",
39    "rerere.autoupdate",
40    "rerere.enabled",
41    "user.email",
42    "user.name",
43];
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub(super) enum ProvisioningFailureDisposition {
47    /// A freshly registered session has no durable history to retain.
48    Discard,
49    /// Resume owns rollback to the archived record and checkpoint lineage.
50    Preserve,
51}
52
53impl Controller {
54    pub async fn provision_session_controlled_with_commit(
55        &mut self,
56        session_id: &str,
57        executor: &(impl CommandExecutor + Sync),
58        grant_commit: impl FnOnce() -> Result<()>,
59    ) -> Result<()> {
60        let github_token = controller_github_token();
61        let repositories = self
62            .provision_session_target_with_failure_disposition(
63                session_id,
64                executor,
65                github_token.as_deref(),
66                ProvisioningFailureDisposition::Discard,
67            )
68            .await?;
69        let setup = execute_concurrent_lanes(
70            || execute_repository_setup(&repositories, executor),
71            || self.install_worker_payload(session_id, executor),
72        );
73        let result = match setup {
74            Ok(((), (backend, worker_root))) => {
75                self.connect_and_start_worker(session_id, executor, &backend, &worker_root, true)
76                    .await
77            }
78            Err(error) => Err(error),
79        };
80        match result {
81            Ok(native_session_id) => {
82                if let Err(error) = grant_commit() {
83                    return Err(self.rollback_failed_new_session(session_id, error, executor)?);
84                }
85                self.mark_worker_connected(session_id, native_session_id)
86            }
87            Err(error) => Err(self.rollback_failed_new_session(session_id, error, executor)?),
88        }
89    }
90
91    /// Start a child worker inside an already-provisioned parent target.
92    /// Repository, target, and mount setup belong exclusively to the parent.
93    pub async fn provision_subagent_session_controlled(
94        &mut self,
95        session_id: &str,
96        executor: &(impl CommandExecutor + Sync),
97    ) -> Result<()> {
98        // Placement failures must reach the same failure arm as startup
99        // failures; otherwise the child record stays `Provisioning` forever.
100        let placement = self.worker_placement(session_id);
101        let (result, placement) = match placement {
102            Ok((backend, worker_root)) => {
103                let syncing = &StagedExecutor::new(executor, ProvisionStage::Syncing);
104                let prepared =
105                    self.prepare_worker_files(session_id, &backend, &worker_root, syncing);
106                let result = match prepared {
107                    Ok(()) => {
108                        self.connect_and_start_worker(
109                            session_id,
110                            executor,
111                            &backend,
112                            &worker_root,
113                            false,
114                        )
115                        .await
116                    }
117                    Err(error) => Err(error),
118                };
119                (result, Some((backend, worker_root)))
120            }
121            Err(error) => (Err(error), None),
122        };
123        match result {
124            Ok(native_session_id) => self.mark_worker_connected(session_id, native_session_id),
125            Err(error) => {
126                // Without placement there is no worker to stop.
127                if let Some((backend, worker_root)) = placement
128                    && let Err(stop_error) =
129                        super::worker_binary::stop_worker(executor, &backend, &worker_root)
130                {
131                    tracing::warn!(
132                        session_id,
133                        error = format!("{stop_error:#}"),
134                        "failed sub-agent worker could not be stopped cleanly"
135                    );
136                }
137                tracing::warn!(
138                    session_id,
139                    error = format!("{error:#}"),
140                    "sub-agent startup failed"
141                );
142                let record = self
143                    .state
144                    .sessions
145                    .get_mut(session_id)
146                    .context("failed sub-agent session disappeared")?;
147                record.state = SessionState::Error;
148                record.updated_at = super::now();
149                record.last_error = Some(format!("sub-agent startup failed: {error:#}"));
150                crate::database::save_lifecycle_session(record)?;
151                Err(error)
152            }
153        }
154    }
155
156    fn rollback_failed_new_session(
157        &mut self,
158        session_id: &str,
159        error: anyhow::Error,
160        executor: &impl CommandExecutor,
161    ) -> Result<anyhow::Error> {
162        let session = self
163            .state
164            .sessions
165            .get(session_id)
166            .with_context(|| format!("unknown session {session_id}"))?
167            .clone();
168        let target_cleanup = match session.target.as_ref() {
169            Some(locator) => (|| -> Result<()> {
170                let backend = backend_locator(locator, &session, &self.config)?;
171                targets::close_plan(&backend, session_id)?
172                    // Rollback must remain possible after the foreground
173                    // operation's cancellation token has been set.
174                    .execute(&CancellableProcessExecutor::with_timeout(
175                        Duration::from_secs(15),
176                    ))
177                    .map(|_| ())
178            })(),
179            None => Ok(()),
180        };
181        let worktree_cleanup =
182            self.cleanup_new_session_worktree_after_failure(session_id, executor);
183        let cleanup_error = [target_cleanup, worktree_cleanup]
184            .into_iter()
185            .filter_map(Result::err)
186            .map(|error| format!("{error:#}"))
187            .collect::<Vec<_>>()
188            .join("; ");
189        if !cleanup_error.is_empty() {
190            tracing::warn!(
191                session_id,
192                error = %cleanup_error,
193                "new-session rollback cleanup reported failures"
194            );
195        }
196        let original = note_new_session_launch_failure(session_id, &error);
197        let failure = apply_failed_new_session_rollback(
198            &mut self.state,
199            session_id,
200            &original,
201            (!cleanup_error.is_empty()).then_some(cleanup_error),
202        );
203        self.persist_session_state(session_id)?;
204        Ok(failure)
205    }
206
207    pub(super) async fn provision_session_with_failure_disposition(
208        &mut self,
209        session_id: &str,
210        executor: &(impl CommandExecutor + Sync),
211        github_token: Option<&str>,
212        failure_disposition: ProvisioningFailureDisposition,
213    ) -> Result<()> {
214        let repositories = self
215            .provision_session_target_with_failure_disposition(
216                session_id,
217                executor,
218                github_token,
219                failure_disposition,
220            )
221            .await?;
222        match execute_repository_setup(&repositories, executor) {
223            Ok(()) => Ok(()),
224            Err(error) if failure_disposition == ProvisioningFailureDisposition::Discard => {
225                Err(self.rollback_failed_new_session(session_id, error, executor)?)
226            }
227            Err(error) => Err(error),
228        }
229    }
230
231    async fn provision_session_target_with_failure_disposition(
232        &mut self,
233        session_id: &str,
234        executor: &(impl CommandExecutor + Sync),
235        github_token: Option<&str>,
236        failure_disposition: ProvisioningFailureDisposition,
237    ) -> Result<targets::CommandPlan> {
238        let session = self
239            .state
240            .sessions
241            .get(session_id)
242            .with_context(|| format!("unknown session {session_id}"))?
243            .clone();
244        if session.state != SessionState::Provisioning {
245            bail!("session {session_id} is not provisioning");
246        }
247        let preparation = (|| {
248            let template = self
249                .config
250                .targets
251                .get(&session.target_template_id)
252                .context("target template disappeared during provisioning")?;
253            let profile = self
254                .config
255                .profiles
256                .get(&session.last_profile)
257                .context("harness profile disappeared during provisioning")?;
258            super::worker_binary::preflight_harness(template, profile, executor)?;
259            self.prepare_managed_raw_worktree(session_id, executor)
260        })();
261        let created_worktree = match preparation {
262            Ok(created) => created,
263            Err(error) if failure_disposition == ProvisioningFailureDisposition::Discard => {
264                return Err(self.fail_new_session_with_cleanup(session_id, error, executor)?);
265            }
266            Err(error) => return Err(error),
267        };
268        let session = self
269            .state
270            .sessions
271            .get(session_id)
272            .expect("session retained after managed worktree preparation")
273            .clone();
274        // Keep planning, preflight, creation, and locator discovery in one
275        // result so the caller's failure disposition applies to every error.
276        let result = (|| {
277            let template = self
278                .config
279                .targets
280                .get(&session.target_template_id)
281                .context("target template disappeared during provisioning")?;
282            if matches!(template, TargetTemplate::AwsEc2 { .. }) {
283                for resource in &session.additional_mounts {
284                    ensure!(
285                        resource.source.is_dir(),
286                        "attached resource source is not a directory: {}",
287                        resource.source.display()
288                    );
289                }
290            }
291            let mut target = backend_target(
292                template,
293                session.resource_allocation.as_ref(),
294                ContainerOverrides::for_session(&session),
295            )?;
296            let mut runtime_mounts = if matches!(target, targets::TargetTemplate::AwsEc2(_)) {
297                Vec::new()
298            } else {
299                session.additional_mounts.clone()
300            };
301            // The mounts this container runs with, not the ones the session
302            // stores: a forced downgrade belongs to the host the container
303            // lands on, so it is decided here every time and never written
304            // over the user's choice.
305            for notice in enforce_overlay_capable_mounts(&target, &mut runtime_mounts, executor) {
306                executor.notify_notice(&notice);
307            }
308            // The image's user is a property of the host's copy of the image,
309            // so it is read here, once per image per daemon, and handed to the
310            // plan rather than stored on the session.
311            let image_user = podman_image_user(&target, executor);
312            let mut bundle = if session.project_directory.is_some() {
313                None
314            } else if failure_disposition == ProvisioningFailureDisposition::Preserve {
315                Some(super::network_git::checkpoint_bundle(&session)?)
316            } else {
317                Some(backend_bundle(
318                    self.config
319                        .bundles
320                        .get(&session.bundle_id)
321                        .context("session bundle is missing")?,
322                    executor,
323                )?)
324            };
325            let container_github_token =
326                github_token.filter(|_| configure_github_token_environment(&mut target));
327            if container_github_token.is_some()
328                && let Some(bundle) = bundle.as_mut()
329            {
330                use_github_https_urls(bundle);
331            }
332            preflight_target(template, executor)?;
333            let prepared_cache = bundle.as_mut().and_then(|bundle| {
334                git_cache::prepare(
335                    &target,
336                    session_id,
337                    bundle,
338                    &mut runtime_mounts,
339                    container_github_token,
340                    executor,
341                )
342            });
343            // Mounts are fixed when the container is created, so the build
344            // cache is decided here, before the provisioning plan is built.
345            let build_cache = super::mbx::prepare(
346                &target,
347                &self.config.build_cache,
348                &session,
349                bundle.as_ref(),
350                prepared_cache.as_ref(),
351                &mut runtime_mounts,
352                executor,
353            );
354            let provision = if let Some(project_directory) = &session.project_directory {
355                targets::provision_bare_project_plan(
356                    &target,
357                    session_id,
358                    &project_directory.to_string_lossy(),
359                )
360            } else {
361                bundle
362                    .as_ref()
363                    .context("project bundle disappeared during provisioning")
364                    .and_then(|bundle| {
365                        targets::provision_plan(
366                            &target,
367                            session_id,
368                            bundle,
369                            &runtime_mounts,
370                            image_user,
371                            session.container_workspace.as_deref(),
372                        )
373                    })
374            };
375            let mut provision = match provision {
376                Ok(provision) => provision,
377                Err(error) => {
378                    if let Some(cache) = &prepared_cache {
379                        let _ = cache.cleanup(executor);
380                    }
381                    return Err(error);
382                }
383            };
384            if let Some(token) = container_github_token
385                && let Err(error) =
386                    provision.provide_target_environment_secret(&target, "GH_TOKEN", token)
387            {
388                if let Some(cache) = &prepared_cache {
389                    let _ = cache.cleanup(executor);
390                }
391                return Err(error);
392            }
393
394            let started = Instant::now();
395            let result =
396                provision_target_creation(&provision, &target, session_id, executor, |outputs| {
397                    locator_after_provision(
398                        template,
399                        &target,
400                        session_id,
401                        outputs.first(),
402                        executor,
403                    )
404                })
405                .map(|(locator, remainder)| (locator, remainder, bundle, build_cache));
406            if result.is_err()
407                && let Some(cache) = &prepared_cache
408            {
409                if let Some(locator) = provisioned_locator(&target, session_id, None) {
410                    let _ = targets::close_plan(&locator, session_id)
411                        .and_then(|plan| plan.execute(executor).map(|_| ()));
412                } else {
413                    let _ = cache.cleanup(executor);
414                }
415            }
416            tracing::debug!(
417                session_id,
418                elapsed_ms = started.elapsed().as_millis(),
419                "provisioning plan execution completed"
420            );
421            result
422        })();
423        let result = match result {
424            Err(error)
425                if created_worktree
426                    && failure_disposition == ProvisioningFailureDisposition::Discard =>
427            {
428                return Err(self.fail_new_session_with_cleanup(session_id, error, executor)?);
429            }
430            Err(error) if failure_disposition == ProvisioningFailureDisposition::Preserve => {
431                Err(error)
432            }
433            Err(error) => {
434                // This arm (no managed worktree to unwind) is the one a
435                // provisioning failure such as a dropped target connection
436                // hits; record the diagnostic and the session-id log here too.
437                let detail = note_new_session_launch_failure(session_id, &error);
438                {
439                    let record = self.state.sessions.get_mut(session_id).unwrap();
440                    record.state = SessionState::Error;
441                    record.target = None;
442                    record.updated_at = super::now();
443                    record.last_error = Some(format!("session provisioning failed: {detail}"));
444                }
445                return match self.persist_session_state(session_id) {
446                    Ok(()) => Err(error),
447                    Err(persistence_error) => Err(error.context(format!(
448                        "persist removal of failed provisioning session {session_id}: {persistence_error:#}"
449                    ))),
450                };
451            }
452            Ok((locator, remainder, bundle, build_cache)) => {
453                apply_new_session_provisioning_result(&mut self.state, session_id, Ok(locator))?;
454                self.state
455                    .sessions
456                    .get_mut(session_id)
457                    .expect("session retained after provisioning")
458                    .build_cache = build_cache;
459                let session = &self.state.sessions[session_id];
460                let backend = backend_locator(
461                    session
462                        .target
463                        .as_ref()
464                        .context("provisioned target disappeared")?,
465                    session,
466                    &self.config,
467                )?;
468                if matches!(backend, targets::TargetLocator::AwsEc2 { .. }) {
469                    targets::provision_on_locator_plan(
470                        &backend,
471                        session_id,
472                        bundle
473                            .as_ref()
474                            .context("AWS provisioning requires a project bundle")?,
475                    )
476                } else {
477                    Ok(remainder)
478                }
479            }
480        };
481        let result = match result {
482            Err(error) if failure_disposition == ProvisioningFailureDisposition::Discard => {
483                return Err(self.rollback_failed_new_session(session_id, error, executor)?);
484            }
485            result => result,
486        };
487        if result.is_ok()
488            && let Some(session) = self.state.sessions.get(session_id)
489            && let Some(directory) = session
490                .managed_worktree
491                .as_ref()
492                .map(|worktree| worktree.source_project_directory.clone())
493                .or_else(|| session.project_directory.clone())
494            && let Some(template) = self.config.targets.get(&session.target_template_id)
495        {
496            let host = match template {
497                TargetTemplate::LocalBare => Some("local"),
498                TargetTemplate::SshBare { ssh, .. } => Some(ssh.host.as_str()),
499                _ => None,
500            };
501            if let Some(host) = host {
502                self.state.remember_project_directory(host, &directory);
503                crate::database::remember_project_directory(host, &directory)?;
504            }
505        }
506        self.persist_session_state(session_id)?;
507        result
508    }
509
510    pub fn mark_worker_connected(
511        &mut self,
512        session_id: &str,
513        native_session_id: Option<String>,
514    ) -> Result<()> {
515        let session = self
516            .state
517            .sessions
518            .get(session_id)
519            .with_context(|| format!("unknown session {session_id}"))?;
520        if session.target.is_none() {
521            bail!("session {session_id} has no provisioned target");
522        }
523        let updated_at = now();
524        crate::database::mark_session_worker_connected(
525            session_id,
526            native_session_id.as_deref(),
527            &updated_at,
528        )?;
529        let session = self
530            .state
531            .sessions
532            .get_mut(session_id)
533            .expect("session disappeared after its worker connection was saved");
534        session.state = SessionState::Running;
535        if native_session_id.is_some() {
536            session.native_session_id = native_session_id;
537        }
538        session.updated_at = updated_at;
539        session.last_error = None;
540        Ok(())
541    }
542
543    fn install_worker_payload(
544        &self,
545        session_id: &str,
546        executor: &impl CommandExecutor,
547    ) -> Result<(targets::TargetLocator, String)> {
548        // Worker/profile installation is independent of repository cloning.
549        let syncing = &StagedExecutor::new(executor, ProvisionStage::Syncing);
550        let (backend, worker_root) = self.worker_placement(session_id)?;
551        self.prepare_worker_files(session_id, &backend, &worker_root, syncing)?;
552        install_attached_resources(&self.state, session_id, &backend, &worker_root, syncing)?;
553        Ok((backend, worker_root))
554    }
555
556    async fn connect_and_start_worker(
557        &self,
558        session_id: &str,
559        executor: &impl CommandExecutor,
560        backend: &targets::TargetLocator,
561        worker_root: &str,
562        initialize_workspace: bool,
563    ) -> Result<Option<String>> {
564        let syncing = &StagedExecutor::new(executor, ProvisionStage::Syncing);
565        if initialize_workspace {
566            install_inherited_git_settings(executor, backend, session_id)?;
567            self.initialize_network_workspaces(session_id, backend, syncing)?;
568        }
569        let session = self
570            .state
571            .sessions
572            .get(session_id)
573            .with_context(|| format!("unknown session {session_id}"))?;
574        let profile = self
575            .config
576            .profiles
577            .get(&session.last_profile)
578            .with_context(|| format!("unknown profile {}", session.last_profile))?;
579        let readiness_stage = bridge_readiness_stage(profile);
580        let reconnect = &targets::reconnect_plan(backend, session_id)?.commands[0];
581        let readiness = async {
582            let mut relay = {
583                let _starting = ProvisionStageGuard::new(executor, ProvisionStage::Starting);
584                start_worker(executor, backend, worker_root)?;
585                connect_started_worker(reconnect, session_id, executor, backend, worker_root)
586                    .await?
587            };
588            let native_session_id =
589                wait_for_native_session_in_stage(&mut relay, executor, readiness_stage).await?;
590            Ok(Some(native_session_id))
591        }
592        .await;
593        match readiness {
594            Ok(native_session_id) => Ok(native_session_id),
595            Err(error) => Err(worker_probe_diagnosis(
596                executor,
597                backend,
598                worker_root,
599                error,
600            )),
601        }
602    }
603}
604
605const MAX_LAUNCH_DIAGNOSTIC_BYTES: usize = 64 * 1024;
606
607const RETAINED_LAUNCH_DIAGNOSTICS: usize = 20;
608
609/// Record a failed new-session launch consistently across every failure arm:
610/// log the failure with the session id (so `logs/mj-*.log` names the session)
611/// and save the local diagnostic file. Returns the underlying error chain
612/// annotated with the diagnostic path, which the caller stores in `last_error`
613/// so the reason travels to `mj sessions`, `mj wait`, and `mj events`.
614pub(super) fn note_new_session_launch_failure(session_id: &str, error: &anyhow::Error) -> String {
615    note_new_session_launch_failure_in(&data_dir().join("diagnostics"), session_id, error)
616}
617
618fn note_new_session_launch_failure_in(
619    directory: &Path,
620    session_id: &str,
621    error: &anyhow::Error,
622) -> String {
623    let original = format!("{error:#}");
624    tracing::warn!(session_id, error = %original, "session launch failed");
625    match persist_launch_failure_to(directory, session_id, &original) {
626        Ok(path) => format!("{original}; full diagnostic saved to {}", path.display()),
627        Err(save_error) => {
628            format!("{original}; saving the local diagnostic failed: {save_error:#}")
629        }
630    }
631}
632
633fn persist_launch_failure_to(directory: &Path, session_id: &str, detail: &str) -> Result<PathBuf> {
634    mj_core::config::validate_id("session", session_id)?;
635    std::fs::create_dir_all(directory).with_context(|| {
636        format!(
637            "create launch diagnostics directory {}",
638            directory.display()
639        )
640    })?;
641    #[cfg(unix)]
642    {
643        use std::os::unix::fs::PermissionsExt;
644        std::fs::set_permissions(directory, std::fs::Permissions::from_mode(0o700))?;
645    }
646    let path = directory.join(format!("{session_id}-launch-error.txt"));
647    let detail = bounded_launch_diagnostic(detail);
648    let body = format!(
649        "Hel session launch failure\nsession: {session_id}\nat: {}\n\n{detail}\n",
650        now()
651    );
652    atomic_write(&path, body.as_bytes())?;
653    prune_launch_diagnostics(directory)?;
654    Ok(path)
655}
656
657fn bounded_launch_diagnostic(detail: &str) -> String {
658    if detail.len() <= MAX_LAUNCH_DIAGNOSTIC_BYTES {
659        return detail.to_owned();
660    }
661    let mut head_end = MAX_LAUNCH_DIAGNOSTIC_BYTES / 4;
662    while !detail.is_char_boundary(head_end) {
663        head_end -= 1;
664    }
665    let tail_bytes = MAX_LAUNCH_DIAGNOSTIC_BYTES - head_end;
666    let mut tail_start = detail.len() - tail_bytes;
667    while !detail.is_char_boundary(tail_start) {
668        tail_start += 1;
669    }
670    format!(
671        "{}\n\n[... launch diagnostic truncated ...]\n\n{}",
672        &detail[..head_end],
673        &detail[tail_start..]
674    )
675}
676
677fn prune_launch_diagnostics(directory: &Path) -> Result<()> {
678    let mut diagnostics = Vec::new();
679    for entry in std::fs::read_dir(directory)? {
680        let entry = entry?;
681        if !entry
682            .file_name()
683            .to_str()
684            .is_some_and(|name| name.ends_with("-launch-error.txt"))
685        {
686            continue;
687        }
688        diagnostics.push((entry.metadata()?.modified()?, entry.path()));
689    }
690    diagnostics.sort_by_key(|entry| std::cmp::Reverse(entry.0));
691    for (_, path) in diagnostics.into_iter().skip(RETAINED_LAUNCH_DIAGNOSTICS) {
692        std::fs::remove_file(&path)
693            .with_context(|| format!("prune old launch diagnostic {}", path.display()))?;
694    }
695    Ok(())
696}
697
698fn apply_new_session_provisioning_result(
699    state: &mut State,
700    session_id: &str,
701    result: Result<TargetLocator>,
702) -> Result<()> {
703    match result {
704        Ok(locator) => {
705            let record = state.sessions.get_mut(session_id).unwrap();
706            record.target = Some(locator);
707            // Provisioning has completed, but Running is reserved for a
708            // successful worker handshake.
709            record.state = SessionState::Disconnected;
710            record.updated_at = now();
711            record.last_error = None;
712            Ok(())
713        }
714        Err(error) => {
715            let record = state.sessions.get_mut(session_id).unwrap();
716            record.state = SessionState::Error;
717            record.target = None;
718            record.updated_at = now();
719            record.last_error = Some(format!("session provisioning failed: {error:#}"));
720            Err(error)
721        }
722    }
723}
724
725pub(super) fn apply_failed_new_session_rollback(
726    state: &mut State,
727    session_id: &str,
728    original_error: &str,
729    cleanup_error: Option<String>,
730) -> anyhow::Error {
731    match cleanup_error {
732        None => {
733            let record = state.sessions.get_mut(session_id).unwrap();
734            record.state = SessionState::Error;
735            record.target = None;
736            record.updated_at = now();
737            record.last_error = Some(format!("worker bootstrap failed: {original_error}"));
738            anyhow::anyhow!("{original_error}; partial target removed and failed session retained")
739        }
740        Some(cleanup_error) => {
741            let failure = format!(
742                "{original_error}; cleanup of the failed session target failed: {cleanup_error}"
743            );
744            let record = state.sessions.get_mut(session_id).unwrap();
745            record.state = SessionState::Error;
746            record.updated_at = now();
747            record.last_error = Some(format!("worker bootstrap failed: {failure}"));
748            anyhow::anyhow!(failure)
749        }
750    }
751}
752
753pub(super) fn install_attached_resources(
754    state: &State,
755    session_id: &str,
756    backend: &targets::TargetLocator,
757    worker_root: &str,
758    executor: &impl CommandExecutor,
759) -> Result<()> {
760    let targets::TargetLocator::AwsEc2 { .. } = backend else {
761        return Ok(());
762    };
763    let session = state
764        .sessions
765        .get(session_id)
766        .with_context(|| format!("unknown session {session_id}"))?;
767    if session.additional_mounts.is_empty() {
768        return Ok(());
769    }
770    for resource in &session.additional_mounts {
771        let install = targets::command_on_locator(
772            backend,
773            session_id,
774            vec![
775                format!("{worker_root}/hel"),
776                "worker".into(),
777                "install-resource".into(),
778                "--destination".into(),
779                resource.destination.to_string_lossy().into_owned(),
780            ],
781            "stream attached resource",
782        )?;
783        mj_checkpoint::resources::stream_resource(&resource.source, |stream| {
784            execute_checked_with_stdin(executor, &install, stream).map(|_| ())
785        })
786        .with_context(|| format!("stream attached resource {}", resource.source.display()))?;
787    }
788    Ok(())
789}
790
791/// Run two independent target setup lanes at the same time and wait for both.
792/// The first lane's failure wins deterministically when both fail, and neither
793/// lane is abandoned while it may still own a transfer or subprocess.
794pub(super) fn execute_concurrent_lanes<A: Send, B: Send>(
795    first: impl FnOnce() -> Result<A> + Send,
796    second: impl FnOnce() -> Result<B> + Send,
797) -> Result<(A, B)> {
798    std::thread::scope(|scope| {
799        let second = scope.spawn(second);
800        let first = first();
801        let second = second.join().unwrap_or_else(|panic| {
802            Err(anyhow::anyhow!(
803                "concurrent target lane panicked: {}",
804                targets::command_thread_panic_message(panic.as_ref())
805            ))
806        });
807        match (first, second) {
808            (Err(error), _) => Err(error),
809            (Ok(_), Err(error)) => Err(error),
810            (Ok(first), Ok(second)) => Ok((first, second)),
811        }
812    })
813}
814
815fn execute_repository_setup(
816    plan: &targets::CommandPlan,
817    executor: &(impl CommandExecutor + Sync),
818) -> Result<()> {
819    if plan.commands.is_empty() {
820        return Ok(());
821    }
822    let _cloning = ProvisionStageGuard::new(executor, ProvisionStage::Cloning);
823    plan.execute_concurrent(executor).map(|_| ())
824}
825
826/// Run a provisioning plan and discover the locator it produced, tearing the
827/// target down again if anything after its creation fails.
828///
829/// Creation is the boundary that matters. A step that fails before the target
830/// exists has left nothing behind; every failure after it — a later plan step
831/// or locator discovery — owns a target no session record will point at.
832#[cfg(test)]
833fn provision_target(
834    plan: &targets::CommandPlan,
835    target: &targets::TargetTemplate,
836    session_id: &str,
837    executor: &(impl CommandExecutor + Sync),
838    discover: impl FnOnce(&[CommandOutput]) -> Result<TargetLocator>,
839) -> Result<TargetLocator> {
840    let Some((creation, remainder)) = plan.split_at_target_creation() else {
841        // Nothing this plan runs can leave a target behind.
842        return discover(&plan.execute_concurrent(executor)?);
843    };
844    let mut outputs = creation.execute_concurrent(executor)?;
845    let result = match remainder.execute_concurrent(executor) {
846        Ok(rest) => {
847            outputs.extend(rest);
848            discover(&outputs)
849        }
850        Err(error) => Err(error),
851    };
852    result.map_err(|error| {
853        match cleanup_failed_provision(target, session_id, outputs.first(), executor) {
854            Some(note) => error.context(note),
855            None => error,
856        }
857    })
858}
859
860/// Bring the target into existence and return the commands that populate its
861/// repositories. The caller may overlap that remainder with worker/profile
862/// installation once it has persisted the discovered locator.
863fn provision_target_creation(
864    plan: &targets::CommandPlan,
865    target: &targets::TargetTemplate,
866    session_id: &str,
867    executor: &(impl CommandExecutor + Sync),
868    discover: impl FnOnce(&[CommandOutput]) -> Result<TargetLocator>,
869) -> Result<(TargetLocator, targets::CommandPlan)> {
870    let Some((creation, remainder)) = plan.split_at_target_creation() else {
871        // Nothing this plan runs can leave a target behind, so its commands
872        // must still finish before the locator is usable.
873        let outputs = plan.execute_concurrent(executor)?;
874        return discover(&outputs).map(|locator| {
875            (
876                locator,
877                targets::CommandPlan {
878                    description: plan.description.clone(),
879                    commands: Vec::new(),
880                },
881            )
882        });
883    };
884    let outputs = creation.execute_concurrent(executor)?;
885    discover(&outputs)
886        .map(|locator| (locator, remainder))
887        .map_err(|error| {
888            match cleanup_failed_provision(target, session_id, outputs.first(), executor) {
889                Some(note) => error.context(note),
890                None => error,
891            }
892        })
893}
894
895/// Best-effort teardown of a target whose creation succeeded but whose
896/// provisioning failed before a locator was recorded. Returns a note
897/// describing what happened for inclusion in the session error.
898///
899/// The teardown is the session's own close plan, so a failed launch and an
900/// ordinary close can never disagree about what removing a target means.
901fn cleanup_failed_provision(
902    target: &targets::TargetTemplate,
903    session_id: &str,
904    create_output: Option<&CommandOutput>,
905    executor: &impl CommandExecutor,
906) -> Option<String> {
907    let locator = provisioned_locator(target, session_id, create_output)?;
908    let leak = format!(
909        "the resource may still exist; find it via its dev.mj.session={session_id} label/tag"
910    );
911    let plan = match targets::close_plan(&locator, session_id) {
912        Ok(plan) => plan,
913        Err(error) => {
914            tracing::warn!(
915                session_id,
916                error = format!("{error:#}"),
917                "could not build provisioning cleanup plan"
918            );
919            return Some(format!("cleanup FAILED: {error:#}; {leak}"));
920        }
921    };
922    let purpose = plan
923        .commands
924        .iter()
925        .map(|command| command.purpose.clone())
926        .collect::<Vec<_>>()
927        .join("; ");
928    let Err(error) = plan.execute(executor) else {
929        return Some(format!("cleanup succeeded: {purpose}"));
930    };
931    match targets::cleanup_target_is_confirmed_absent(&locator, session_id, executor) {
932        Ok(true) => Some(format!("cleanup succeeded: {purpose}")),
933        Ok(false) => {
934            tracing::warn!(
935                session_id,
936                error = format!("{error:#}"),
937                "provisioning cleanup failed and the target may still exist"
938            );
939            Some(format!("cleanup FAILED ({purpose}): {error:#}; {leak}"))
940        }
941        Err(confirm_error) => {
942            tracing::warn!(
943                session_id,
944                error = format!("{confirm_error:#}"),
945                "could not confirm whether the failed provisioning target was removed"
946            );
947            Some(format!(
948                "cleanup FAILED ({purpose}): {error:#}; checking whether it was removed also failed: {confirm_error:#}; {leak}"
949            ))
950        }
951    }
952}
953
954/// The locator a provisioning plan's creating command brought into existence.
955///
956/// Every target but AWS is named before its plan runs; an EC2 instance
957/// reports its own ID in the launch response.
958fn provisioned_locator(
959    target: &targets::TargetTemplate,
960    session_id: &str,
961    create_output: Option<&CommandOutput>,
962) -> Option<targets::TargetLocator> {
963    let container_id = || targets::resource_name(session_id).ok();
964    Some(match target {
965        // A bare project directory belongs to the user: provisioning creates
966        // nothing that a failure could leak.
967        targets::TargetTemplate::LocalBare => return None,
968        targets::TargetTemplate::LocalPodman(container) => targets::TargetLocator::LocalPodman {
969            borrowed_from: None,
970            container_id: container_id()?,
971            workspace_storage: targets::podman_workspace_locator(container, session_id).ok()?,
972        },
973        targets::TargetTemplate::LocalDocker(_) => targets::TargetLocator::LocalDocker {
974            borrowed_from: None,
975            container_id: container_id()?,
976        },
977        targets::TargetTemplate::AppleContainer(_) => targets::TargetLocator::AppleContainer {
978            borrowed_from: None,
979            container_id: container_id()?,
980        },
981        targets::TargetTemplate::SshPodman { ssh, container } => {
982            targets::TargetLocator::SshPodman {
983                borrowed_from: None,
984                ssh: ssh.clone(),
985                container_id: container_id()?,
986                workspace_storage: targets::podman_workspace_locator(container, session_id).ok()?,
987            }
988        }
989        targets::TargetTemplate::SshDocker { ssh, .. } => targets::TargetLocator::SshDocker {
990            borrowed_from: None,
991            ssh: ssh.clone(),
992            container_id: container_id()?,
993        },
994        targets::TargetTemplate::SshBare { ssh, .. } => targets::TargetLocator::SshBare {
995            ssh: ssh.clone(),
996            workspace: targets::workspace_for(target, session_id).ok()?,
997            worker_id: None,
998        },
999        targets::TargetTemplate::AwsEc2(aws) => targets::TargetLocator::AwsEc2 {
1000            profile: aws.profile.clone(),
1001            region: aws.region.clone(),
1002            instance_id: serde_json::from_slice::<serde_json::Value>(&create_output?.stdout)
1003                .ok()?
1004                .pointer("/Instances/0/InstanceId")?
1005                .as_str()?
1006                .to_owned(),
1007            ssh: aws.ssh.clone(),
1008            workspace: targets::workspace_for(target, session_id).ok()?,
1009        },
1010    })
1011}
1012
1013/// Attach read-only whatever the selected container overlay cannot hold, and
1014/// say so.
1015///
1016/// The filesystem is probed on the host that runs the container, because that
1017/// is where the overlay would be built. A probe that cannot answer leaves the
1018/// overlay alone: a failed probe is no evidence of an unsupported filesystem,
1019/// and refusing to provision over one would cost the user their session.
1020///
1021/// Apple's `container` engine already mounts every extra directory read-only,
1022/// and EC2 copies the directory instead of mounting it, so neither is probed.
1023pub(super) fn enforce_overlay_capable_mounts(
1024    target: &targets::TargetTemplate,
1025    mounts: &mut [targets::AdditionalMount],
1026    executor: &impl CommandExecutor,
1027) -> Vec<String> {
1028    let ssh = match target {
1029        targets::TargetTemplate::LocalPodman(_) | targets::TargetTemplate::LocalDocker(_) => None,
1030        targets::TargetTemplate::SshPodman { ssh, .. }
1031        | targets::TargetTemplate::SshDocker { ssh, .. } => Some(ssh),
1032        _ => return Vec::new(),
1033    };
1034    let overlaid = mounts
1035        .iter()
1036        .filter(|mount| mount.access == targets::MountAccess::Cow)
1037        .map(|mount| mount.source.clone())
1038        .collect::<Vec<_>>();
1039    if overlaid.is_empty() {
1040        return Vec::new();
1041    }
1042    let filesystems = match targets::probe_filesystem_types(ssh, &overlaid, executor) {
1043        Ok(filesystems) => filesystems,
1044        Err(error) => {
1045            tracing::warn!(
1046                error = format!("{error:#}"),
1047                "could not probe attached-directory filesystems; preserving overlay mounts"
1048            );
1049            return vec![format!(
1050                "Could not read the filesystem under the attached directories, so they keep the \
1051                 copy-on-write overlay: {error:#}"
1052            )];
1053        }
1054    };
1055    let mut notices = Vec::new();
1056    for (mount, filesystem) in mounts
1057        .iter_mut()
1058        .filter(|mount| mount.access == targets::MountAccess::Cow)
1059        .zip(filesystems)
1060    {
1061        let Some(reason) = targets::overlay_unsupported_filesystem(&filesystem) else {
1062            continue;
1063        };
1064        mount.access = mount.access.without_overlay();
1065        notices.push(format!(
1066            "Mounted {} read-only: the overlay is unreliable on {filesystem} ({reason}).",
1067            mount.source.display()
1068        ));
1069    }
1070    notices
1071}
1072
1073/// The image users already probed, keyed by container host and image
1074/// reference. An image's
1075/// configured user does not change under a fixed reference, and reading it
1076/// costs a container start, so each daemon asks a host once.
1077static IMAGE_USERS: std::sync::LazyLock<std::sync::Mutex<BTreeMap<String, targets::ImageUser>>> =
1078    std::sync::LazyLock::new(std::sync::Mutex::default);
1079
1080/// The uid and gid a Podman session container maps onto the host user.
1081///
1082/// Only Podman is asked: Docker and Apple's `container` engine are left with
1083/// their own defaults. A probe that cannot answer is not a launch failure —
1084/// the container falls back to plain `--userns=keep-id`, which maps the
1085/// image's default user, and the user is told what happened.
1086pub(super) fn podman_image_user(
1087    target: &targets::TargetTemplate,
1088    executor: &impl CommandExecutor,
1089) -> Option<targets::ImageUser> {
1090    let (ssh, container) = match target {
1091        targets::TargetTemplate::LocalPodman(container) => (None, container),
1092        targets::TargetTemplate::SshPodman { ssh, container } => (Some(ssh), container),
1093        _ => return None,
1094    };
1095    let image = container.image.as_str();
1096    let key = format!(
1097        "{}|{image}",
1098        ssh.map_or("local", |ssh| ssh.destination.as_str())
1099    );
1100    if let Some(cached) = IMAGE_USERS.lock().expect("image user cache").get(&key) {
1101        return Some(*cached);
1102    }
1103    match targets::probe_image_user(ssh, container, executor) {
1104        Ok(user) => {
1105            IMAGE_USERS
1106                .lock()
1107                .expect("image user cache")
1108                .insert(key, user);
1109            Some(user)
1110        }
1111        Err(error) => {
1112            tracing::warn!(
1113                image,
1114                error = format!("{error:#}"),
1115                "could not read the container image user; keeping Podman's default user mapping"
1116            );
1117            executor.notify_notice(&format!(
1118                "Could not read the user of image {image}, so the container runs with Podman's \
1119                 default user mapping and may not be able to write to an attached directory: \
1120                 {error:#}"
1121            ));
1122            None
1123        }
1124    }
1125}
1126
1127/// Reports every command an installer issues as one launch stage, so progress
1128/// stays accurate without threading the stage through each `CommandSpec`.
1129/// A command that already names a stage keeps it.
1130pub(super) struct StagedExecutor<'a, E: CommandExecutor> {
1131    inner: &'a E,
1132    stage: ProvisionStage,
1133    _guard: ProvisionStageGuard<'a, E>,
1134}
1135
1136impl<'a, E: CommandExecutor> StagedExecutor<'a, E> {
1137    pub(crate) fn new(inner: &'a E, stage: ProvisionStage) -> Self {
1138        Self {
1139            inner,
1140            stage,
1141            _guard: ProvisionStageGuard::new(inner, stage),
1142        }
1143    }
1144
1145    fn staged(&self, command: &CommandSpec) -> CommandSpec {
1146        if command.stage.is_some() {
1147            return command.clone();
1148        }
1149        command.clone().stage(self.stage)
1150    }
1151}
1152
1153impl<E: CommandExecutor> CommandExecutor for StagedExecutor<'_, E> {
1154    fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
1155        self.inner.execute(&self.staged(command))
1156    }
1157
1158    fn cancellation_requested(&self) -> bool {
1159        self.inner.cancellation_requested()
1160    }
1161
1162    fn stage_started(&self, stage: ProvisionStage) {
1163        self.inner.stage_started(stage);
1164    }
1165
1166    fn stage_finished(&self, stage: ProvisionStage) {
1167        self.inner.stage_finished(stage);
1168    }
1169
1170    fn notify_notice(&self, notice: &str) {
1171        self.inner.notify_notice(notice);
1172    }
1173
1174    fn execute_with_stdin(
1175        &self,
1176        command: &CommandSpec,
1177        input: &mut (dyn std::io::Read + Send),
1178    ) -> Result<CommandOutput> {
1179        self.inner.execute_with_stdin(&self.staged(command), input)
1180    }
1181}
1182
1183fn execute_checked_with_stdin(
1184    executor: &impl CommandExecutor,
1185    command: &CommandSpec,
1186    input: &mut (dyn std::io::Read + Send),
1187) -> Result<CommandOutput> {
1188    let output = executor.execute_with_stdin(command, input)?;
1189    if output.status != 0 {
1190        bail!(
1191            "{} failed with status {}: {}",
1192            command.purpose,
1193            output.status,
1194            String::from_utf8_lossy(&output.stderr)
1195        );
1196    }
1197    Ok(output)
1198}
1199
1200pub(super) fn install_inherited_git_settings(
1201    executor: &impl CommandExecutor,
1202    locator: &targets::TargetLocator,
1203    session_id: &str,
1204) -> Result<()> {
1205    let settings = if inherits_controller_git_settings(locator) {
1206        controller_git_settings()?
1207    } else {
1208        BTreeMap::new()
1209    };
1210    for command in inherited_git_setting_commands(locator, session_id, settings)? {
1211        execute_checked(executor, command)?;
1212    }
1213    Ok(())
1214}
1215
1216fn inherits_controller_git_settings(locator: &targets::TargetLocator) -> bool {
1217    !matches!(
1218        locator,
1219        targets::TargetLocator::LocalBare { .. } | targets::TargetLocator::SshBare { .. }
1220    )
1221}
1222
1223fn inherited_git_setting_commands(
1224    locator: &targets::TargetLocator,
1225    session_id: &str,
1226    settings: BTreeMap<String, String>,
1227) -> Result<Vec<CommandSpec>> {
1228    if matches!(locator, targets::TargetLocator::SshBare { .. }) {
1229        return Ok(Vec::new());
1230    }
1231    settings
1232        .into_iter()
1233        .map(|(key, value)| {
1234            targets::command_on_locator(
1235                locator,
1236                session_id,
1237                vec![
1238                    "git".into(),
1239                    "config".into(),
1240                    "--global".into(),
1241                    "--replace-all".into(),
1242                    "--".into(),
1243                    key.clone(),
1244                    value,
1245                ],
1246                format!("inherit Git setting {key}"),
1247            )
1248        })
1249        .collect()
1250}
1251
1252fn controller_git_settings() -> Result<BTreeMap<String, String>> {
1253    let output = match Command::new("git")
1254        .args(["config", "--global", "--includes", "--null", "--list"])
1255        .stdin(Stdio::null())
1256        .output()
1257    {
1258        Ok(output) => output,
1259        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(BTreeMap::new()),
1260        Err(error) => return Err(error).context("read controller Git configuration"),
1261    };
1262    if !output.status.success() {
1263        bail!(
1264            "read controller Git configuration failed with status {}: {}",
1265            output.status,
1266            String::from_utf8_lossy(&output.stderr).trim()
1267        );
1268    }
1269    parse_inherited_git_settings(&output.stdout)
1270}
1271
1272fn parse_inherited_git_settings(output: &[u8]) -> Result<BTreeMap<String, String>> {
1273    let mut settings = BTreeMap::new();
1274    for entry in output
1275        .split(|byte| *byte == 0)
1276        .filter(|entry| !entry.is_empty())
1277    {
1278        let entry = std::str::from_utf8(entry).context("decode controller Git configuration")?;
1279        let (key, value) = entry
1280            .split_once('\n')
1281            .with_context(|| format!("controller Git returned malformed entry {entry:?}"))?;
1282        let key = key.to_ascii_lowercase();
1283        if INHERITED_GIT_SETTINGS.contains(&key.as_str()) {
1284            settings.insert(key, value.to_owned());
1285        }
1286    }
1287    Ok(settings)
1288}
1289
1290#[cfg(test)]
1291mod tests;