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