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