Skip to main content

mj_controller/controller/
backend.rs

1//! Backend target, locator, and capacity conversion for provisioned sessions.
2
3use std::collections::BTreeMap;
4use std::path::PathBuf;
5use std::process::{Command, Stdio};
6use std::time::{Duration, Instant};
7
8use anyhow::{Context, Result, bail, ensure};
9
10use mj_core::config::{AwsAddressSource, Config, ProjectBundle, TargetTemplate, data_dir};
11use mj_core::state::{
12    PodmanWorkspaceLocator, ResourceAllocationKind, SessionRecord, SessionResourceAllocation,
13    TargetLocator, allocation_cpus,
14};
15
16use crate::targets::{
17    self, AwsTemplate, CommandExecutor, CommandOutput, CommandSpec, ContainerTemplate,
18    ImageRefresh, ProjectBundleSpec, ProvisionStage, RepositorySpec, SshTarget,
19};
20
21use super::{Controller, execute_checked};
22
23impl Controller {
24    /// Inspect the actual execution checkout in a background worker.
25    pub fn session_working_context(
26        &self,
27        session_id: &str,
28        executor: &impl CommandExecutor,
29    ) -> Result<(PathBuf, String)> {
30        let session = self
31            .state
32            .sessions
33            .get(session_id)
34            .context("session is missing")?;
35        let locator = session
36            .target
37            .as_ref()
38            .context("target is still starting")?;
39        let backend = backend_locator(locator, session, &self.config)?;
40        let launch = self.current_worker_launch_config(session_id, &backend)?;
41        let output = executor.execute(&targets::command_on_locator(
42            &backend,
43            session_id,
44            vec![
45                "git".into(),
46                "-C".into(),
47                launch.cwd.to_string_lossy().into_owned(),
48                "rev-parse".into(),
49                "--abbrev-ref".into(),
50                "HEAD".into(),
51            ],
52            "read current session branch",
53        )?)?;
54        let branch = if output.status == 0 {
55            let branch = String::from_utf8(output.stdout).context("decode session branch")?;
56            if branch.trim() == "HEAD" {
57                "detached HEAD".to_owned()
58            } else {
59                branch.trim().to_owned()
60            }
61        } else {
62            let stderr = String::from_utf8_lossy(&output.stderr);
63            if stderr.contains("not a git repository") {
64                // A plain folder opened through `mj go` has no branch. Say so
65                // briefly rather than echoing multi-line git stderr, which wraps
66                // the banner.
67                "not a git checkout".to_owned()
68            } else {
69                format!(
70                    "unavailable: {}",
71                    stderr.lines().next().unwrap_or("").trim()
72                )
73            }
74        };
75        Ok((launch.cwd, branch))
76    }
77
78    /// The session checkout's branch, distance from upstream, and changed
79    /// files, read with the target's own `git`. Builds on
80    /// [`Self::session_working_context`], so it works wherever that does: a
81    /// local checkout, a container, or an SSH host.
82    pub fn session_git_status(
83        &self,
84        session_id: &str,
85        executor: &impl CommandExecutor,
86    ) -> Result<mj_core::local_git::SessionGitStatus> {
87        let (cwd, branch) = self.session_working_context(session_id, executor)?;
88        if branch.starts_with("not a git") || branch.starts_with("unavailable") {
89            return Ok(mj_core::local_git::parse_git_status(
90                cwd, &branch, None, "", "",
91            ));
92        }
93        let session = self
94            .state
95            .sessions
96            .get(session_id)
97            .context("session is missing")?;
98        let locator = session
99            .target
100            .as_ref()
101            .context("target is still starting")?;
102        let backend = backend_locator(locator, session, &self.config)?;
103        let cwd_text = cwd.to_string_lossy().into_owned();
104        let run = |args: &[&str], purpose: &str| -> Result<Option<String>> {
105            let mut command = vec!["git".to_owned(), "-C".to_owned(), cwd_text.clone()];
106            command.extend(args.iter().map(|arg| (*arg).to_owned()));
107            let output = executor.execute(&targets::command_on_locator(
108                &backend, session_id, command, purpose,
109            )?)?;
110            Ok((output.status == 0).then(|| String::from_utf8_lossy(&output.stdout).into_owned()))
111        };
112        // No upstream is an ordinary state, so a failing count is `None`
113        // rather than an error.
114        let ahead_behind = run(
115            &["rev-list", "--left-right", "--count", "@{upstream}...HEAD"],
116            "count commits against upstream",
117        )?;
118        // A repository with no commit yet has no HEAD to diff against.
119        let numstat = run(
120            &["--no-optional-locks", "diff", "--numstat", "HEAD"],
121            "count changed lines",
122        )?
123        .unwrap_or_default();
124        let porcelain = run(
125            &[
126                "--no-optional-locks",
127                "status",
128                "--porcelain",
129                "--untracked-files=normal",
130            ],
131            "list changed files",
132        )?
133        .unwrap_or_default();
134        Ok(mj_core::local_git::parse_git_status(
135            cwd,
136            &branch,
137            ahead_behind.as_deref(),
138            &numstat,
139            &porcelain,
140        ))
141    }
142
143    pub fn resolve_aws_resource_options(
144        &self,
145        target_id: &str,
146        executor: &impl CommandExecutor,
147    ) -> Result<Vec<SessionResourceAllocation>> {
148        let TargetTemplate::AwsEc2 {
149            aws_profile,
150            region,
151            launch_template,
152            launch_template_version,
153            ..
154        } = self
155            .config
156            .targets
157            .get(target_id)
158            .with_context(|| format!("unknown target template {target_id:?}"))?
159        else {
160            bail!("target {target_id:?} is not an AWS EC2 target");
161        };
162        let profile = aws_profile.as_deref().unwrap_or("default");
163        let launch_key = if launch_template.starts_with("lt-") {
164            "--launch-template-id"
165        } else {
166            "--launch-template-name"
167        };
168        let version = launch_template_version.as_deref().unwrap_or("$Default");
169        let describe_template = CommandSpec::new(
170            "aws",
171            [
172                "--profile",
173                profile,
174                "--region",
175                region,
176                "ec2",
177                "describe-launch-template-versions",
178                launch_key,
179                launch_template,
180                "--versions",
181                version,
182                "--output",
183                "json",
184            ],
185        )
186        .purpose("resolve EC2 launch template instance family");
187        let output = executor.execute(&describe_template)?;
188        if output.status != 0 {
189            bail!(
190                "{} failed with status {}: {}",
191                describe_template.purpose,
192                output.status,
193                String::from_utf8_lossy(&output.stderr).trim()
194            );
195        }
196        let response: serde_json::Value =
197            serde_json::from_slice(&output.stdout).context("parse EC2 launch template response")?;
198        let instance_type = response
199            .pointer("/LaunchTemplateVersions/0/LaunchTemplateData/InstanceType")
200            .and_then(serde_json::Value::as_str)
201            .context("launch template does not specify a concrete instance type")?;
202        let family = instance_type
203            .rsplit_once('.')
204            .map(|(family, _)| family)
205            .context("launch template instance type has no size suffix")?;
206        let filter = format!("Name=instance-type,Values={family}.*");
207        let describe_types = CommandSpec::new(
208            "aws",
209            [
210                "--profile",
211                profile,
212                "--region",
213                region,
214                "ec2",
215                "describe-instance-types",
216                "--filters",
217                &filter,
218                "--output",
219                "json",
220            ],
221        )
222        .purpose("discover EC2 instance sizes");
223        let output = executor.execute(&describe_types)?;
224        if output.status != 0 {
225            bail!(
226                "{} failed with status {}: {}",
227                describe_types.purpose,
228                output.status,
229                String::from_utf8_lossy(&output.stderr).trim()
230            );
231        }
232        let response: serde_json::Value =
233            serde_json::from_slice(&output.stdout).context("parse EC2 instance type response")?;
234        let mut options = response
235            .get("InstanceTypes")
236            .and_then(serde_json::Value::as_array)
237            .context("EC2 instance type response omitted InstanceTypes")?
238            .iter()
239            .filter_map(|entry| {
240                Some(SessionResourceAllocation::AwsEc2 {
241                    instance_type: entry.get("InstanceType")?.as_str()?.to_owned(),
242                    vcpus: entry.pointer("/VCpuInfo/DefaultVCpus")?.as_u64()?,
243                    memory_bytes: entry
244                        .pointer("/MemoryInfo/SizeInMiB")?
245                        .as_u64()?
246                        .checked_mul(1024 * 1024)?,
247                })
248            })
249            .collect::<Vec<_>>();
250        options.sort_by_key(allocation_cpus);
251        if !options.iter().any(|option| allocation_cpus(option) == 8) {
252            bail!("EC2 family {family:?} has no exact 8-vCPU baseline size");
253        }
254        Ok(options)
255    }
256
257    pub fn reconnect_command(&self, session_id: &str) -> Result<CommandSpec> {
258        let session = self
259            .state
260            .sessions
261            .get(session_id)
262            .with_context(|| format!("unknown session {session_id}"))?;
263        session.validate_configuration(&self.config)?;
264        let locator = session.target.as_ref().context("session has no target")?;
265        let backend = backend_locator(locator, session, &self.config)?;
266        targets::reconnect_plan(&backend, session_id)?
267            .commands
268            .into_iter()
269            .next()
270            .context("reconnect plan is empty")
271    }
272
273    /// The directories the live sessions on one target write to.
274    fn session_storage_paths(&self, target_id: &str) -> Vec<String> {
275        self.state
276            .sessions
277            .values()
278            .filter(|session| session.target_template_id == target_id && session.state.is_active())
279            .filter_map(|session| {
280                let checkout = self.state.checkout(&session.id).ok()?;
281                let paths = mj_core::targets::storage::session_storage_paths(
282                    session.target.as_ref()?,
283                    &session.id,
284                    checkout.project_directory(),
285                );
286                Some(paths.all().map(str::to_owned).collect::<Vec<_>>())
287            })
288            .flatten()
289            .collect()
290    }
291
292    pub fn deployment_capacity_targets(&self) -> Vec<targets::DeploymentCapacityTarget> {
293        use targets::{DeploymentCapacityKind, DeploymentCapacityTarget};
294
295        let mut local_ids = Vec::new();
296        // Per host: its target ids, and per SSH connection the paths to measure.
297        type HostProbes = (Vec<String>, Vec<(SshTarget, Vec<String>)>);
298        let mut ssh_hosts: BTreeMap<String, HostProbes> = BTreeMap::new();
299        let mut targets = Vec::new();
300        for (target_id, template) in &self.config.targets {
301            match template {
302                TargetTemplate::LocalBare
303                | TargetTemplate::LocalPodman { .. }
304                | TargetTemplate::LocalDocker { .. }
305                | TargetTemplate::AppleContainer { .. } => {
306                    local_ids.push(target_id.clone());
307                }
308                TargetTemplate::SshBare { ssh, .. }
309                | TargetTemplate::SshPodman { ssh, .. }
310                | TargetTemplate::SshDocker { ssh, .. } => {
311                    let entry = ssh_hosts.entry(ssh.host.clone()).or_default();
312                    entry.0.push(target_id.clone());
313                    let connection = SshTarget::from(ssh);
314                    let paths = match entry.1.iter_mut().find(|(known, _)| *known == connection) {
315                        Some((_, paths)) => paths,
316                        None => {
317                            entry.1.push((connection, Vec::new()));
318                            &mut entry.1.last_mut().expect("just pushed").1
319                        }
320                    };
321                    push_unique(paths, ssh_storage_paths(template));
322                    // Each live session's own directories: a project or
323                    // clone can sit on a filesystem of its own.
324                    push_unique(paths, self.session_storage_paths(target_id));
325                }
326                TargetTemplate::AwsEc2 { .. } => {
327                    let mut probes = Vec::new();
328                    let mut probe_error = None;
329                    for session in self.state.sessions.values().filter(|session| {
330                        session.target_template_id == *target_id
331                            && session.state.is_active()
332                            && session.target.is_some()
333                    }) {
334                        let result = backend_locator(
335                            session.target.as_ref().expect("filtered target"),
336                            session,
337                            &self.config,
338                        )
339                        .and_then(|locator| {
340                            targets::aws_allocated_capacity_command(&locator, &session.id)
341                        });
342                        match result {
343                            Ok(command) => probes.push(command),
344                            Err(error) => probe_error = Some(format!("{error:#}")),
345                        }
346                    }
347                    targets.push(DeploymentCapacityTarget {
348                        id: format!("aws:{target_id}"),
349                        host: target_id.clone(),
350                        target_ids: vec![target_id.clone()],
351                        kind: DeploymentCapacityKind::AwsFleet,
352                        local: false,
353                        probes,
354                        local_storage_paths: Vec::new(),
355                        probe_error,
356                    });
357                }
358            }
359        }
360        if !local_ids.is_empty() {
361            // CPU and memory come from sysinfo; storage is measured over the
362            // paths by `targets::measure_local_storage`.
363            let mut paths = Vec::new();
364            for target_id in &local_ids {
365                if let Some(template) = self.config.targets.get(target_id) {
366                    push_unique(&mut paths, local_storage_paths(template));
367                }
368                push_unique(&mut paths, self.session_storage_paths(target_id));
369            }
370            targets.push(DeploymentCapacityTarget {
371                id: "local".into(),
372                host: "local".into(),
373                target_ids: local_ids,
374                kind: DeploymentCapacityKind::Host,
375                local: true,
376                probes: Vec::new(),
377                local_storage_paths: paths,
378                probe_error: None,
379            });
380        }
381        targets.extend(
382            ssh_hosts
383                .into_iter()
384                .map(
385                    |(host, (target_ids, connections))| DeploymentCapacityTarget {
386                        id: format!("ssh:{host}"),
387                        host,
388                        target_ids,
389                        kind: DeploymentCapacityKind::Host,
390                        local: false,
391                        probes: connections
392                            .iter()
393                            .map(|(ssh, paths)| targets::ssh_host_capacity_command(ssh, paths))
394                            .collect(),
395                        local_storage_paths: Vec::new(),
396                        probe_error: None,
397                    },
398                ),
399        );
400        targets.sort_by(|left, right| left.id.cmp(&right.id));
401        targets
402    }
403
404    /// The full check behind the Targets pane's Test action.
405    pub fn test_target(&self, target_id: &str, executor: &impl CommandExecutor) -> Result<()> {
406        verify_target(
407            self.configured_target(target_id)?,
408            executor,
409            TargetCheck::BeforeLaunch,
410        )
411    }
412
413    /// The check a session wizard runs before it offers a target: the same
414    /// one a launch runs, worded for a launch that has not happened yet.
415    pub fn check_target_readiness(
416        &self,
417        target_id: &str,
418        executor: &impl CommandExecutor,
419    ) -> Result<()> {
420        preflight_target(
421            self.configured_target(target_id)?,
422            executor,
423            TargetCheck::BeforeLaunch,
424        )
425    }
426
427    fn configured_target(&self, target_id: &str) -> Result<&TargetTemplate> {
428        self.config
429            .targets
430            .get(target_id)
431            .with_context(|| format!("unknown target template {target_id:?}"))
432    }
433}
434
435/// When a target is checked, which decides how its refusal ends.
436#[derive(Debug, Clone, Copy, PartialEq, Eq)]
437pub(super) enum TargetCheck {
438    /// Before anything is launched: the session wizard's target row, the
439    /// Targets pane.
440    BeforeLaunch,
441    /// Inside a launch, resume, or move, whose failure the launch-failure
442    /// dialog reports with its Retry launch button.
443    Launch,
444}
445
446impl TargetCheck {
447    /// The clause that points at the failure dialog's Retry launch. Before a
448    /// launch there is no such dialog to point at (launch finding R6-3).
449    const fn then_retry(self) -> &'static str {
450        match self {
451            Self::BeforeLaunch => "",
452            Self::Launch => ", then Retry launch",
453        }
454    }
455}
456
457/// The directories Mjolnir writes to on an SSH host for one target, relative
458/// to the SSH user's home where they live there: worker roots, the upload
459/// staging and binary caches, and the workspace prefix or the container
460/// engine's storage. Each is measured at its nearest existing ancestor.
461pub(crate) fn ssh_storage_paths(template: &TargetTemplate) -> Vec<String> {
462    use mj_core::targets::storage::{
463        ContainerStorage, DEFAULT_BUILD_CACHE_DIRECTORY, REMOTE_CACHE_DIRECTORY,
464        REMOTE_PROFILES_DIRECTORY, REMOTE_WORKERS_DIRECTORY, TEMPORARY_DIRECTORY,
465    };
466    let mut paths = vec![
467        REMOTE_CACHE_DIRECTORY.to_owned(),
468        DEFAULT_BUILD_CACHE_DIRECTORY.to_owned(),
469        TEMPORARY_DIRECTORY.to_owned(),
470    ];
471    match template {
472        TargetTemplate::SshBare {
473            workspace_prefix, ..
474        } => {
475            paths.push(REMOTE_WORKERS_DIRECTORY.to_owned());
476            paths.push(REMOTE_PROFILES_DIRECTORY.to_owned());
477            paths.push(workspace_prefix.to_string_lossy().into_owned());
478        }
479        TargetTemplate::SshPodman { .. } => paths.push(ContainerStorage::Podman.path().to_owned()),
480        TargetTemplate::SshDocker { .. } => paths.push(ContainerStorage::Docker.path().to_owned()),
481        _ => {}
482    }
483    paths.extend(build_cache_directory(template));
484    paths
485}
486
487/// A target's configured mbx cache directory, when it names one.
488fn build_cache_directory(template: &TargetTemplate) -> Option<String> {
489    template
490        .container()?
491        .build_cache
492        .as_ref()?
493        .directory
494        .as_ref()
495        .map(|directory| directory.to_string_lossy().into_owned())
496}
497
498/// The paths the local probe measures for one local target.
499fn local_storage_paths(template: &TargetTemplate) -> Vec<String> {
500    use mj_core::targets::storage::{
501        ContainerStorage, DEFAULT_BUILD_CACHE_DIRECTORY, TEMPORARY_DIRECTORY, local_home_path,
502    };
503    let mut paths = vec![
504        data_dir().to_string_lossy().into_owned(),
505        TEMPORARY_DIRECTORY.to_owned(),
506        local_home_path(DEFAULT_BUILD_CACHE_DIRECTORY),
507    ];
508    match template {
509        TargetTemplate::LocalPodman { .. } => {
510            paths.push(local_home_path(ContainerStorage::Podman.path()));
511        }
512        TargetTemplate::LocalDocker { .. } => {
513            paths.push(ContainerStorage::Docker.path().to_owned());
514        }
515        _ => {}
516    }
517    paths.extend(build_cache_directory(template));
518    paths
519}
520
521fn push_unique(paths: &mut Vec<String>, more: impl IntoIterator<Item = String>) {
522    for path in more {
523        if !path.is_empty() && !paths.contains(&path) {
524            paths.push(path);
525        }
526    }
527}
528
529/// Check a target before a launch, resume, or move uses it.
530///
531/// An SSH host's container runtime is verified by `mj setup`, `mj doctor`,
532/// and the Targets pane's Test action; here it is assumed to still be as
533/// they left it, and only the host's reachability is checked. That probe
534/// joins an open shared connection, so it usually costs one SSH channel.
535pub(super) fn preflight_target(
536    template: &TargetTemplate,
537    executor: &impl CommandExecutor,
538    check: TargetCheck,
539) -> Result<()> {
540    match template {
541        TargetTemplate::SshBare { ssh, .. }
542        | TargetTemplate::SshPodman { ssh, .. }
543        | TargetTemplate::SshDocker { ssh, .. } => {
544            verify_ssh_connectivity(&SshTarget::from(ssh), executor)
545        }
546        _ => verify_target(template, executor, check),
547    }
548}
549
550/// Whether a local container engine can run sessions, as the launch
551/// preflight and the dashboard's Targets pane judge it.
552#[derive(Debug, Clone, Copy, PartialEq, Eq)]
553pub(crate) enum LocalEngineReadiness {
554    Ready,
555    /// The engine's command is not on this host.
556    NotInstalled,
557    /// The command is there, but the engine did not answer its check.
558    NotReady,
559}
560
561/// Run the launch preflight's engine check for a local container target
562/// kind (`local-podman`, `local-docker`, or `apple-container`).
563///
564/// Answers `None` for any other kind: their readiness is not a local engine's.
565pub(crate) fn local_engine_readiness(
566    kind: &str,
567    executor: &impl CommandExecutor,
568) -> Option<LocalEngineReadiness> {
569    let result = match kind {
570        "local-podman" => targets::verify_local_podman(executor).map(|_| ()),
571        "local-docker" => targets::verify_local_docker(executor).map(|_| ()),
572        // The same command `verify_target` runs, without its sentence, which
573        // flattens the cause this needs to tell a missing command apart.
574        "apple-container" => executor
575            .execute(
576                &CommandSpec::new("container", ["system", "status"])
577                    .purpose("preflight Apple container runtime")
578                    .stage(ProvisionStage::Provisioning),
579            )
580            .and_then(|output| {
581                ensure!(output.status == 0, "container system status failed");
582                Ok(())
583            }),
584        _ => return None,
585    };
586    Some(match result {
587        Ok(()) => LocalEngineReadiness::Ready,
588        Err(error) if is_missing_command(&error) => LocalEngineReadiness::NotInstalled,
589        Err(_) => LocalEngineReadiness::NotReady,
590    })
591}
592
593/// Whether a command failed because its program is not installed.
594fn is_missing_command(error: &anyhow::Error) -> bool {
595    error.chain().any(|cause| {
596        cause
597            .downcast_ref::<std::io::Error>()
598            .is_some_and(|io| io.kind() == std::io::ErrorKind::NotFound)
599    })
600}
601
602/// Check that a target's host and runtime can run sessions.
603pub(super) fn verify_target(
604    template: &TargetTemplate,
605    executor: &impl CommandExecutor,
606    check: TargetCheck,
607) -> Result<()> {
608    let then_retry = check.then_retry();
609    match template {
610        TargetTemplate::LocalPodman { .. } => targets::verify_local_podman(executor)
611            .map(|_| ())
612            .map_err(|error| {
613                anyhow::anyhow!(
614                    "local Podman is not ready. Fix the problem below{then_retry}: {error:#}"
615                )
616            }),
617        TargetTemplate::LocalDocker { .. } => targets::verify_local_docker(executor)
618            .map(|_| ())
619            .map_err(
620                |error| match error.downcast_ref::<targets::DockerUnavailable>() {
621                    // Leads with the sentence the launch options use, so a
622                    // missing Docker reads "not installed" in the wizard's
623                    // row too (launch finding R5-3).
624                    Some(problem) => anyhow::anyhow!(
625                        "{problem} {}",
626                        match check {
627                            TargetCheck::BeforeLaunch => problem.remedy(),
628                            TargetCheck::Launch => problem.launch_remedy(),
629                        }
630                    ),
631                    None => anyhow::anyhow!(
632                        "local Docker is not ready. Start Docker or fix the problem below{then_retry}: {error:#}"
633                    ),
634                },
635            ),
636        TargetTemplate::SshPodman { ssh, .. } => {
637            let ssh = SshTarget::from(ssh);
638            targets::verify_ssh_podman(&ssh, executor)
639                .map(|preflight| {
640                    for warning in preflight.warnings {
641                        executor.notify_notice(&warning.notice());
642                    }
643                })
644                .map_err(|error| {
645                    anyhow::anyhow!(
646                        "remote Podman is not ready on {}. Fix the problem below{then_retry}: {error:#}",
647                        ssh.destination
648                    )
649                })
650        }
651        TargetTemplate::SshDocker { ssh, .. } => {
652            let ssh = SshTarget::from(ssh);
653            targets::verify_ssh_docker(&ssh, executor)
654                .map(|_| ())
655                .map_err(|error| {
656                    anyhow::anyhow!(
657                        "remote Docker preflight failed for {}. Fix the problem below{then_retry}: {error:#}",
658                        ssh.destination
659                    )
660                })
661        }
662        TargetTemplate::AppleContainer { .. } => {
663            let command = CommandSpec::new("container", ["system", "status"])
664                .purpose("preflight Apple container runtime")
665                .stage(ProvisionStage::Provisioning);
666            let output = executor.execute(&command).map_err(|error| {
667                anyhow::anyhow!(
668                    "Apple container is not ready. Fix the problem below{then_retry}: {error}"
669                )
670            })?;
671            if output.status != 0 {
672                bail!(
673                    "Apple container is not ready. Start the runtime with `container system start`{then_retry}: container system status exited {}: {}",
674                    output.status,
675                    [
676                        String::from_utf8_lossy(&output.stdout).trim(),
677                        String::from_utf8_lossy(&output.stderr).trim(),
678                    ]
679                    .into_iter()
680                    .filter(|message| !message.is_empty())
681                    .collect::<Vec<_>>()
682                    .join("\n")
683                );
684            }
685            Ok(())
686        }
687        TargetTemplate::SshBare { ssh, .. } => {
688            verify_ssh_connectivity(&SshTarget::from(ssh), executor)
689        }
690        TargetTemplate::AwsEc2 {
691            aws_profile,
692            region,
693            launch_template,
694            launch_template_version,
695            ..
696        } => {
697            let mut identity_args = vec!["sts".into(), "get-caller-identity".into()];
698            if let Some(profile) = aws_profile {
699                identity_args.extend(["--profile".into(), profile.clone()]);
700            }
701            let identity = CommandSpec::new("aws", identity_args)
702                .purpose("verify AWS credentials")
703                .stage(ProvisionStage::Provisioning);
704            let output = executor.execute(&identity)?;
705            ensure!(
706                output.status == 0,
707                "AWS credential test failed with status {}: {}",
708                output.status,
709                String::from_utf8_lossy(&output.stderr).trim()
710            );
711
712            let mut launch_args = vec![
713                "ec2".into(),
714                "describe-launch-template-versions".into(),
715                "--region".into(),
716                region.clone(),
717                if launch_template.starts_with("lt-") { "--launch-template-id" } else { "--launch-template-name" }.into(),
718                launch_template.clone(),
719                "--versions".into(),
720                launch_template_version
721                    .clone()
722                    .unwrap_or_else(|| "$Default".into()),
723            ];
724            if let Some(profile) = aws_profile {
725                launch_args.extend(["--profile".into(), profile.clone()]);
726            }
727            let launch = CommandSpec::new("aws", launch_args)
728                .purpose("verify AWS launch template")
729                .stage(ProvisionStage::Provisioning);
730            let output = executor.execute(&launch)?;
731            ensure!(
732                output.status == 0,
733                "AWS launch-template test failed with status {}: {}",
734                output.status,
735                String::from_utf8_lossy(&output.stderr).trim()
736            );
737            Ok(())
738        }
739        TargetTemplate::LocalBare => Ok(()),
740    }
741}
742
743fn verify_ssh_connectivity(ssh: &SshTarget, executor: &impl CommandExecutor) -> Result<()> {
744    let output = executor.execute(&targets::ssh_connectivity_probe(ssh))?;
745    ensure!(
746        output.status == 0,
747        "SSH connectivity test failed for {} with status {}: {}",
748        ssh.destination,
749        output.status,
750        String::from_utf8_lossy(&output.stderr).trim()
751    );
752    Ok(())
753}
754
755#[cfg(test)]
756pub(super) fn backend_bundle(
757    bundle: &ProjectBundle,
758    executor: &impl CommandExecutor,
759) -> Result<ProjectBundleSpec> {
760    backend_bundle_with_sources(bundle, None, executor)
761}
762
763pub(super) fn backend_session_bundle(
764    session: &SessionRecord,
765    config: &mj_core::config::Config,
766    executor: &impl CommandExecutor,
767) -> Result<ProjectBundleSpec> {
768    backend_bundle_with_sources(
769        session
770            .project_bundle(config)
771            .context("session bundle is missing")?,
772        session
773            .project
774            .as_ref()
775            .map(|project| &project.network_sources),
776        executor,
777    )
778}
779
780fn backend_bundle_with_sources(
781    bundle: &ProjectBundle,
782    sources: Option<&std::collections::BTreeMap<String, mj_core::remote_git::NetworkGitSource>>,
783    executor: &impl CommandExecutor,
784) -> Result<ProjectBundleSpec> {
785    let primary = bundle.primary().context("bundle primary is missing")?;
786    Ok(ProjectBundleSpec {
787        primary: primary.destination.to_string_lossy().into_owned(),
788        repositories: bundle
789            .repositories
790            .iter()
791            .map(|repository| {
792                let source = match sources.and_then(|sources| sources.get(&repository.id)) {
793                    Some(source) => source.clone(),
794                    None => mj_core::remote_git::resolve_repository(repository, executor)
795                        .with_context(|| format!("repository {:?}", repository.id))?,
796                };
797                Ok(RepositorySpec {
798                    url: Some(source.fetch_url),
799                    push_urls: source.push_urls,
800                    destination: repository.destination.to_string_lossy().into_owned(),
801                    git_ref: None,
802                    reference: None,
803                })
804            })
805            .collect::<Result<Vec<_>>>()?,
806    })
807}
808
809/// Per-session container size overrides. They win over both the target
810/// template's values and any recorded resource allocation, and they are read
811/// only while a container is being created.
812#[derive(Debug, Clone, Copy, Default)]
813pub(super) struct ContainerOverrides<'a> {
814    pub cpus: Option<&'a str>,
815    pub memory: Option<&'a str>,
816}
817
818impl<'a> ContainerOverrides<'a> {
819    pub(super) fn for_session(session: &'a SessionRecord) -> Self {
820        Self {
821            cpus: session.container_cpus.as_deref(),
822            memory: session.container_memory.as_deref(),
823        }
824    }
825}
826
827pub(super) fn backend_target(
828    template: &TargetTemplate,
829    allocation: Option<&SessionResourceAllocation>,
830    overrides: ContainerOverrides<'_>,
831) -> Result<targets::TargetTemplate> {
832    Ok(match template {
833        TargetTemplate::LocalBare => targets::TargetTemplate::LocalBare,
834        TargetTemplate::LocalPodman { container } => {
835            let mut backend = backend_container(container, allocation, overrides);
836            backend.workspace_storage = (&container.workspace_storage).into();
837            targets::TargetTemplate::LocalPodman(backend)
838        }
839        TargetTemplate::LocalDocker { container } => targets::TargetTemplate::LocalDocker(
840            backend_container(container, allocation, overrides),
841        ),
842        TargetTemplate::AppleContainer { container } => targets::TargetTemplate::AppleContainer(
843            backend_container(container, allocation, overrides),
844        ),
845        TargetTemplate::AwsEc2 {
846            aws_profile,
847            region,
848            launch_template,
849            launch_template_version,
850            ssh_user,
851            identity_file,
852            ssh_args,
853            ..
854        } => targets::TargetTemplate::AwsEc2(AwsTemplate {
855            profile: aws_profile.clone().unwrap_or_else(|| "default".into()),
856            region: region.clone(),
857            launch_template: launch_template.clone(),
858            launch_template_version: launch_template_version.clone(),
859            instance_type: match allocation {
860                Some(SessionResourceAllocation::AwsEc2 { instance_type, .. }) => {
861                    Some(instance_type.clone())
862                }
863                _ => None,
864            },
865            // The address is filled after describe-instances.
866            ssh: SshTarget {
867                destination: format!("{ssh_user}@pending.invalid"),
868                ssh_args: targets::ssh_args_with_identity(ssh_args, identity_file.as_deref()),
869            },
870        }),
871        TargetTemplate::SshBare {
872            ssh,
873            workspace_prefix,
874            ..
875        } => targets::TargetTemplate::SshBare {
876            ssh: SshTarget::from(ssh),
877            workspace_prefix: workspace_prefix.to_string_lossy().into_owned(),
878        },
879        TargetTemplate::SshPodman { ssh, container, .. } => {
880            let mut backend = backend_container(container, allocation, overrides);
881            backend.workspace_storage = (&container.workspace_storage).into();
882            targets::TargetTemplate::SshPodman {
883                ssh: SshTarget::from(ssh),
884                container: backend,
885            }
886        }
887        TargetTemplate::SshDocker { ssh, container, .. } => targets::TargetTemplate::SshDocker {
888            ssh: SshTarget::from(ssh),
889            container: backend_container(container, allocation, overrides),
890        },
891    })
892}
893
894/// Every container image the daemon downloads in the background, once per
895/// (host, image, platform).
896///
897/// Every container target is covered, including Apple's `container` engine.
898/// A `never` policy is the one opt-out. The rest differ only in when they
899/// download: `always` and `newer` pull on every refresh, while the others
900/// pull only when the host has no copy of the image.
901///
902/// Several targets often share one image on one host, and that needs one
903/// download. When two such targets disagree about when to pull, the merged
904/// entry takes the more eager of the two.
905pub fn image_refresh_plan(config: &Config) -> Vec<ImageRefresh> {
906    let mut plan: Vec<ImageRefresh> = Vec::new();
907    for target in config.targets.values() {
908        let Some((host, container)) = target.image_host() else {
909            continue;
910        };
911        let Some(refresh) = targets::image_refresh(
912            host,
913            &container.image,
914            container.platform.as_deref(),
915            container.pull_policy,
916        ) else {
917            continue;
918        };
919        // Commands are decided by the host, the image, and the platform alone,
920        // so those three identify the duplicates worth collapsing.
921        if let Some(existing) = plan.iter_mut().find(|entry| {
922            entry.host == refresh.host
923                && entry.image == refresh.image
924                && entry.platform == refresh.platform
925        }) {
926            existing.when = existing.when.max(refresh.when);
927            continue;
928        }
929        plan.push(refresh);
930    }
931    plan
932}
933
934pub(crate) fn controller_github_token() -> Option<String> {
935    for name in ["GH_TOKEN", "GITHUB_TOKEN"] {
936        if let Ok(token) = std::env::var(name)
937            && let Some(token) = usable_github_token(&token)
938        {
939            return Some(token.to_owned());
940        }
941    }
942    let output = match Command::new("gh")
943        .args(["auth", "token", "--hostname", "github.com"])
944        .stdin(Stdio::null())
945        .stderr(Stdio::null())
946        .output()
947    {
948        Ok(output) => output,
949        Err(error) => {
950            tracing::debug!(%error, "could not query the GitHub CLI for a token");
951            return None;
952        }
953    };
954    if !output.status.success() {
955        tracing::debug!(status = ?output.status, "GitHub CLI did not return an authenticated token");
956        return None;
957    }
958    let token = match std::str::from_utf8(&output.stdout) {
959        Ok(token) => token,
960        Err(error) => {
961            tracing::debug!(%error, "GitHub CLI returned a non-UTF-8 token");
962            return None;
963        }
964    };
965    let Some(token) = usable_github_token(token) else {
966        tracing::debug!("GitHub CLI returned an empty or invalid token");
967        return None;
968    };
969    Some(token.to_owned())
970}
971
972fn usable_github_token(token: &str) -> Option<&str> {
973    let token = token.trim();
974    (!token.is_empty() && !token.chars().any(char::is_whitespace)).then_some(token)
975}
976
977pub(super) fn configure_github_token_environment(target: &mut targets::TargetTemplate) -> bool {
978    let container = match target {
979        targets::TargetTemplate::LocalPodman(container)
980        | targets::TargetTemplate::LocalDocker(container)
981        | targets::TargetTemplate::AppleContainer(container)
982        | targets::TargetTemplate::SshPodman { container, .. }
983        | targets::TargetTemplate::SshDocker { container, .. } => container,
984        targets::TargetTemplate::LocalBare
985        | targets::TargetTemplate::AwsEc2(_)
986        | targets::TargetTemplate::SshBare { .. } => return false,
987    };
988    container
989        .extra_run_args
990        .extend(["--env".to_owned(), "GH_TOKEN".to_owned()]);
991    true
992}
993
994pub(super) fn use_github_https_urls(bundle: &mut targets::ProjectBundleSpec) {
995    for repository in &mut bundle.repositories {
996        for source in repository
997            .url
998            .iter_mut()
999            .chain(repository.push_urls.iter_mut())
1000        {
1001            if let Some(github) = crate::setup::github_repository_from_origin(source) {
1002                *source = format!(
1003                    "https://github.com/{}/{}.git",
1004                    github.owner, github.repository
1005                );
1006            }
1007        }
1008    }
1009}
1010
1011fn backend_container(
1012    container: &mj_core::config::ContainerTemplate,
1013    allocation: Option<&SessionResourceAllocation>,
1014    overrides: ContainerOverrides<'_>,
1015) -> ContainerTemplate {
1016    let mut extra_run_args = Vec::new();
1017    if let Some(platform) = &container.platform {
1018        extra_run_args.push(format!("--platform={platform}"));
1019    }
1020    let (cpus, memory) = match allocation {
1021        Some(SessionResourceAllocation::Container { cpus, memory_bytes }) => {
1022            (Some(cpus.to_string()), Some(memory_bytes.to_string()))
1023        }
1024        _ => (container.cpus.clone(), container.memory.clone()),
1025    };
1026    // The session's own overrides are the last word on size.
1027    let cpus = overrides.cpus.map(str::to_owned).or(cpus);
1028    let memory = overrides.memory.map(str::to_owned).or(memory);
1029    if let Some(cpus) = cpus {
1030        extra_run_args.push(format!("--cpus={cpus}"));
1031    }
1032    if let Some(memory) = memory {
1033        extra_run_args.push(format!("--memory={memory}"));
1034    }
1035    for (key, value) in &container.environment {
1036        extra_run_args.extend(["--env".to_string(), format!("{key}={value}")]);
1037    }
1038    ContainerTemplate {
1039        image: container.image.clone(),
1040        pull_policy: container.pull_policy,
1041        extra_run_args,
1042        workspace_storage: targets::PodmanWorkspaceStorage::ContainerLayer,
1043        build_cache: container.build_cache.clone(),
1044    }
1045}
1046
1047pub(crate) fn validate_resource_allocation(
1048    target_kind: impl Into<ResourceAllocationKind>,
1049    allocation: Option<&SessionResourceAllocation>,
1050) -> Result<()> {
1051    if let Some(allocation) = allocation {
1052        allocation.validate()?;
1053    }
1054    match (target_kind.into(), allocation) {
1055        (_, None)
1056        | (ResourceAllocationKind::Container, Some(SessionResourceAllocation::Container { .. }))
1057        | (ResourceAllocationKind::AwsEc2, Some(SessionResourceAllocation::AwsEc2 { .. })) => {
1058            Ok(())
1059        }
1060        (ResourceAllocationKind::Fixed, Some(_)) => {
1061            bail!(mj_core::state::BARE_TARGET_FIXED_RESOURCES)
1062        }
1063        _ => bail!("resource allocation does not match the selected target kind"),
1064    }
1065}
1066
1067/// How long a freshly launched EC2 instance may take to accept SSH.
1068const AWS_SSH_READY_TIMEOUT: Duration = Duration::from_secs(300);
1069
1070const AWS_SSH_READY_RETRY_DELAY: Duration = Duration::from_secs(3);
1071
1072/// Poll a remote host until it accepts SSH, or until the deadline passes.
1073///
1074/// `now` and `sleep` are injected so tests can drive the deadline without
1075/// waiting in real time.
1076fn wait_for_ssh_ready(
1077    executor: &impl CommandExecutor,
1078    probe: &CommandSpec,
1079    timeout: Duration,
1080    mut now: impl FnMut() -> Instant,
1081    mut sleep: impl FnMut(Duration),
1082) -> Result<()> {
1083    let started = now();
1084    loop {
1085        if executor.cancellation_requested() {
1086            bail!("cancelled while waiting for SSH on the new instance");
1087        }
1088        let failure = match executor.execute(probe) {
1089            Ok(output) if output.status == 0 => return Ok(()),
1090            Ok(output) => String::from_utf8_lossy(&output.stderr).trim().to_string(),
1091            Err(error) => error.to_string(),
1092        };
1093        if now().duration_since(started) >= timeout {
1094            bail!(
1095                "{} timed out after {}s: {}",
1096                probe.purpose,
1097                timeout.as_secs(),
1098                if failure.is_empty() {
1099                    "the SSH probe reported no error output"
1100                } else {
1101                    failure.as_str()
1102                }
1103            );
1104        }
1105        sleep(AWS_SSH_READY_RETRY_DELAY);
1106    }
1107}
1108
1109pub(super) fn locator_after_provision_named(
1110    canonical: &TargetTemplate,
1111    backend: &targets::TargetTemplate,
1112    session_id: &str,
1113    first_output: Option<&CommandOutput>,
1114    executor: &(impl CommandExecutor + Sync),
1115    name: &str,
1116) -> Result<TargetLocator> {
1117    let generated = name.to_owned();
1118
1119    Ok(match canonical {
1120        TargetTemplate::LocalBare => TargetLocator::LocalBare {
1121            worker_root: data_dir().join("workers").join(session_id),
1122        },
1123        TargetTemplate::LocalPodman { .. } => {
1124            let targets::TargetTemplate::LocalPodman(container) = backend else {
1125                bail!("session locator/template mismatch")
1126            };
1127            TargetLocator::LocalPodman {
1128                borrowed_from: None,
1129                container_id: generated,
1130                workspace_storage: PodmanWorkspaceLocator::from(
1131                    targets::podman_workspace_locator_named(container, name)?,
1132                ),
1133            }
1134        }
1135        TargetTemplate::LocalDocker { .. } => TargetLocator::LocalDocker {
1136            borrowed_from: None,
1137            container_id: generated,
1138        },
1139        TargetTemplate::AppleContainer { .. } => TargetLocator::AppleContainer {
1140            borrowed_from: None,
1141            container_id: generated,
1142        },
1143        TargetTemplate::SshBare { ssh, .. } => TargetLocator::SshBare {
1144            host: ssh.host.clone(),
1145            workspace: PathBuf::from(targets::workspace_for(backend, session_id)?),
1146            worker_id: None,
1147        },
1148        TargetTemplate::SshPodman { ssh, .. } => {
1149            let targets::TargetTemplate::SshPodman { container, .. } = backend else {
1150                bail!("session locator/template mismatch")
1151            };
1152            TargetLocator::SshPodman {
1153                borrowed_from: None,
1154                host: ssh.host.clone(),
1155                container_id: generated,
1156                workspace_storage: PodmanWorkspaceLocator::from(
1157                    targets::podman_workspace_locator_named(container, name)?,
1158                ),
1159            }
1160        }
1161        TargetTemplate::SshDocker { ssh, .. } => TargetLocator::SshDocker {
1162            borrowed_from: None,
1163            host: ssh.host.clone(),
1164            container_id: generated,
1165        },
1166        TargetTemplate::AwsEc2 { .. } => {
1167            let output = first_output.context("AWS launch produced no output")?;
1168            let json: serde_json::Value = serde_json::from_slice(&output.stdout)
1169                .context("parse aws ec2 run-instances response")?;
1170            let instance_id = json
1171                .pointer("/Instances/0/InstanceId")
1172                .and_then(serde_json::Value::as_str)
1173                .context("AWS response omitted instance ID")?
1174                .to_string();
1175            return ec2_locator_after_launch(canonical, instance_id, executor);
1176        }
1177    })
1178}
1179
1180/// Resolve a created instance. Callers persist its ID before these long waits.
1181pub(super) fn ec2_locator_after_launch(
1182    canonical: &TargetTemplate,
1183    instance_id: String,
1184    executor: &(impl CommandExecutor + Sync),
1185) -> Result<TargetLocator> {
1186    let TargetTemplate::AwsEc2 {
1187        aws_profile,
1188        region,
1189        ssh_user,
1190        address_source,
1191        identity_file,
1192        ssh_args,
1193        ..
1194    } = canonical
1195    else {
1196        bail!("EC2 locator requires an EC2 target");
1197    };
1198    let profile = aws_profile.clone().unwrap_or_else(|| "default".into());
1199    execute_checked(
1200        executor,
1201        CommandSpec::new(
1202            "aws",
1203            [
1204                "--profile".into(),
1205                profile.clone(),
1206                "--region".into(),
1207                region.clone(),
1208                "ec2".into(),
1209                "wait".into(),
1210                "instance-running".into(),
1211                "--instance-ids".into(),
1212                instance_id.clone(),
1213            ],
1214        )
1215        .purpose("wait for EC2 session instance to run")
1216        .stage(ProvisionStage::Booting),
1217    )?;
1218    let field = match address_source {
1219        AwsAddressSource::PublicDns => "PublicDnsName",
1220        AwsAddressSource::PublicIp => "PublicIpAddress",
1221        AwsAddressSource::PrivateDns => "PrivateDnsName",
1222        AwsAddressSource::PrivateIp => "PrivateIpAddress",
1223    };
1224    let address = execute_checked(
1225        executor,
1226        CommandSpec::new(
1227            "aws",
1228            [
1229                "--profile".into(),
1230                profile.clone(),
1231                "--region".into(),
1232                region.clone(),
1233                "ec2".into(),
1234                "describe-instances".into(),
1235                "--instance-ids".into(),
1236                instance_id.clone(),
1237                "--query".into(),
1238                format!("Reservations[0].Instances[0].{field}"),
1239                "--output".into(),
1240                "text".into(),
1241            ],
1242        )
1243        .purpose("resolve EC2 session address")
1244        .stage(ProvisionStage::Booting),
1245    )?;
1246    let address = String::from_utf8(address.stdout)
1247        .context("AWS address was not UTF-8")?
1248        .trim()
1249        .to_string();
1250    if address.is_empty() || address == "None" {
1251        bail!("AWS instance {instance_id} has no configured address");
1252    }
1253    let ssh = SshTarget {
1254        destination: format!("{ssh_user}@{address}"),
1255        ssh_args: targets::ssh_args_with_identity(ssh_args, identity_file.as_deref()),
1256    };
1257    wait_for_ssh_ready(
1258        executor,
1259        &crate::targets::ssh_command(&ssh, ["true"])
1260            .purpose("wait for EC2 SSH availability")
1261            .stage(ProvisionStage::Booting),
1262        AWS_SSH_READY_TIMEOUT,
1263        Instant::now,
1264        std::thread::sleep,
1265    )?;
1266    Ok(TargetLocator::AwsEc2 {
1267        instance_id,
1268        address: Some(address),
1269    })
1270}
1271
1272/// The execution-plan locator for a session's stored target.
1273///
1274/// The mapping itself lives in `mj_core::targets`; this only pairs the stored
1275/// locator with the target template the session was created against.
1276pub(crate) fn backend_locator(
1277    locator: &TargetLocator,
1278    session: &SessionRecord,
1279    config: &Config,
1280) -> Result<targets::TargetLocator> {
1281    let runtime = if session.target_runtime.is_some() || targets::locator_needs_connection(locator)
1282    {
1283        Some(session.target_runtime_settings(config)?)
1284    } else {
1285        None
1286    };
1287    Ok(targets::TargetLocator::try_from(targets::RecordedTarget {
1288        locator,
1289        runtime: runtime.as_deref(),
1290        session_id: &session.id,
1291    })?)
1292}
1293
1294#[cfg(test)]
1295mod tests;