use std::path::{Path, PathBuf};
use anyhow::{Context, Result, bail};
use crate::session_manager::StandaloneSession;
use mj_core::config::{AwsAddressSource, SshConnection, TargetTemplate};
use mj_core::state::{
PodmanWorkspaceLocator, SessionRecord, SessionState, TargetLocator, normalize_session_title,
};
use crate::targets::{self, CommandExecutor, CommandOutput, CommandSpec, SshTarget};
use mj_core::worker_launch::WorkerOwnership;
use super::backend::{ContainerOverrides, backend_locator, backend_target};
use super::readiness::wait_for_native_session;
use super::{Controller, now};
pub use mj_core::state::{RecoveryCandidate, RecoveryScan};
impl Controller {
pub fn scan_orphan_workers(
&self,
executor: &impl CommandExecutor,
all_instances: bool,
) -> RecoveryScan {
let mut scan = RecoveryScan {
instance_id: mj_core::config::instance_identity(),
..RecoveryScan::default()
};
for (target_id, template) in &self.config.targets {
match scan_target_workers(target_id, template, executor) {
Ok(candidates) => {
scan.candidates
.extend(candidates.into_iter().filter(|candidate| {
!self.state.sessions.contains_key(&candidate.session_id)
}))
}
Err(error) => scan.warnings.push(format!("target {target_id}: {error:#}")),
}
}
scan.candidates.sort_by(|left, right| {
(&left.session_id, &left.target_template_id)
.cmp(&(&right.session_id, &right.target_template_id))
});
scan.candidates.dedup_by(|left, right| {
left.session_id == right.session_id
&& left.target_template_id == right.target_template_id
});
if !all_instances {
restrict_to_instance(&mut scan);
}
scan
}
pub async fn adopt_orphan_worker(
&mut self,
session_id: &str,
target_id: &str,
profile_override: Option<&str>,
bundle_override: Option<&str>,
all_instances: bool,
executor: &impl CommandExecutor,
) -> Result<()> {
let (record, newly_adopted) = match self.state.sessions.get(session_id).cloned() {
Some(existing) if adoption_unfinished(&existing, target_id) => {
for (flag, requested, adopted) in [
("profile", profile_override, existing.last_profile.as_str()),
("bundle", bundle_override, existing.bundle_id.as_str()),
] {
if let Some(requested) = requested
&& requested != adopted
{
bail!(
"session {session_id} was already adopted with {flag} {adopted:?}; retry without --{flag}"
);
}
}
(existing, false)
}
Some(_) => bail!("session {session_id} is already tracked"),
None => {
let scan = self.scan_orphan_workers(executor, true);
let candidate = scan
.candidates
.into_iter()
.find(|candidate| {
candidate.session_id == session_id
&& candidate.target_template_id == target_id
})
.with_context(|| {
format!("no managed orphan {session_id} was found on target {target_id:?}")
})?;
require_instance_access(&candidate, &scan.instance_id, all_instances)?;
let profile_id = profile_override
.map(str::to_owned)
.or_else(|| {
candidate
.ownership
.as_ref()
.map(|marker| marker.profile_id.clone())
})
.context("orphan has no ownership marker; pass --profile")?;
let bundle_id = bundle_override
.map(str::to_owned)
.or_else(|| {
candidate
.ownership
.as_ref()
.map(|marker| marker.bundle_id.clone())
})
.context("orphan has no ownership marker; pass --bundle")?;
let profile = self
.config
.profiles
.get(&profile_id)
.with_context(|| format!("unknown profile {profile_id:?}"))?;
self.config
.bundles
.get(&bundle_id)
.with_context(|| format!("unknown bundle {bundle_id:?}"))?;
let workspace_id = resolve_recovery_workspace_id(
candidate
.ownership
.as_ref()
.map(|ownership| ownership.workspace_id.as_str())
.unwrap_or(mj_core::workspace::DEFAULT_WORKSPACE_ID),
)?;
let container_workspace = self
.config
.targets
.get(target_id)
.and_then(|template| {
recovery_backend_locator(template, &candidate.locator, session_id).ok()
})
.and_then(|backend| {
adopted_container_workspace(&backend, session_id, executor)
});
let mut record = adopted_session_record(
session_id,
target_id,
profile_id,
profile.kind,
bundle_id,
workspace_id,
candidate.locator,
);
record.container_workspace = container_workspace;
(record, true)
}
};
let locator = record
.target
.as_ref()
.context("adopted session has no target locator")?;
let backend = backend_locator(locator, &record, &self.config)?;
let spec = targets::reconnect_plan(&backend, session_id)?
.commands
.into_iter()
.next()
.context("reconnect plan is empty")?;
if newly_adopted {
crate::database::save_session(&record)?;
self.state.sessions.insert(session_id.to_owned(), record);
}
match self.complete_adoption(session_id, &spec, executor).await {
Ok(()) => Ok(()),
Err(error) => Err(self.record_adoption_failure(session_id, error)),
}
}
async fn complete_adoption(
&mut self,
session_id: &str,
spec: &CommandSpec,
executor: &impl CommandExecutor,
) -> Result<()> {
let mut relay = StandaloneSession::connect_command(spec, session_id)
.await
.context("orphan relay did not complete the v1 handshake")?;
let native_session_id = wait_for_native_session(&mut relay, executor).await?;
self.mark_worker_connected(session_id, Some(native_session_id))?;
if let Some(title) = relay
.snapshot()
.materialized
.session_title
.as_deref()
.and_then(normalize_session_title)
{
crate::database::set_session_acp_title(session_id, Some(&title))?;
self.state
.sessions
.get_mut(session_id)
.expect("adopted session disappeared while saving its ACP title")
.acp_session_title = Some(title);
}
Ok(())
}
fn record_adoption_failure(&mut self, session_id: &str, error: anyhow::Error) -> anyhow::Error {
let Some(record) = self.state.sessions.get_mut(session_id) else {
return error;
};
record.updated_at = now();
record.last_error = Some(format!("orphan adoption failed: {error:#}"));
match self.persist_session_state(session_id) {
Ok(()) => error,
Err(persist_error) => error.context(format!(
"recorded the adoption failure in memory, but failed to persist it: {persist_error:#}"
)),
}
}
pub fn destroy_orphan_worker(
&self,
session_id: &str,
target_id: &str,
confirmation: &str,
all_instances: bool,
executor: &impl CommandExecutor,
) -> Result<()> {
if confirmation != session_id {
bail!("refusing destructive recovery: --confirm must exactly match the session ID");
}
let scan = self.scan_orphan_workers(executor, true);
let candidate = scan
.candidates
.into_iter()
.find(|candidate| {
candidate.session_id == session_id && candidate.target_template_id == target_id
})
.with_context(|| {
format!("no managed orphan {session_id} was found on target {target_id:?}")
})?;
require_instance_access(&candidate, &scan.instance_id, all_instances)?;
let template = self.config.targets.get(target_id).unwrap();
let backend = recovery_backend_locator(template, &candidate.locator, session_id)?;
targets::close_plan(&backend, session_id)?
.execute(executor)
.map(|_| ())
}
}
fn restrict_to_instance(scan: &mut RecoveryScan) {
let before = scan.candidates.len();
scan.candidates
.retain(|candidate| candidate.instance_id.as_deref() == Some(scan.instance_id.as_str()));
scan.hidden_other_instances = before - scan.candidates.len();
}
fn require_instance_access(
candidate: &RecoveryCandidate,
scan_instance: &str,
all_instances: bool,
) -> Result<()> {
if all_instances {
return Ok(());
}
match candidate.instance_id.as_deref() {
Some(instance) if instance == scan_instance => Ok(()),
Some(other) => bail!(
"worker {} belongs to instance {other:?}, not this instance {scan_instance:?}; pass --all-instances to act on it",
candidate.session_id
),
None => bail!(
"worker {} has no instance stamp (created by an older build); pass --all-instances to act on it",
candidate.session_id
),
}
}
fn resolve_recovery_workspace_id(marked_workspace_id: &str) -> Result<String> {
if marked_workspace_id == mj_core::workspace::DEFAULT_WORKSPACE_ID {
return Ok(marked_workspace_id.to_owned());
}
if crate::database::list_workspaces()?
.into_iter()
.any(|workspace| workspace.id == marked_workspace_id)
{
return Ok(marked_workspace_id.to_owned());
}
Ok(crate::database::create_or_get_workspace("Recovered")?.id)
}
fn adopted_container_workspace(
backend: &targets::TargetLocator,
session_id: &str,
executor: &impl CommandExecutor,
) -> Option<PathBuf> {
if !matches!(
backend,
targets::TargetLocator::LocalPodman { .. }
| targets::TargetLocator::LocalDocker { .. }
| targets::TargetLocator::AppleContainer { .. }
| targets::TargetLocator::SshPodman { .. }
| targets::TargetLocator::SshDocker { .. }
) {
return None;
}
let workspace = targets::new_container_workspace(session_id).ok()?;
let command = targets::command_on_locator(
backend,
session_id,
vec![
"test".to_owned(),
"-d".to_owned(),
workspace.to_string_lossy().into_owned(),
],
"probe the adopted session workspace",
)
.ok()?;
match executor.execute(&command) {
Ok(output) if output.status == 0 => Some(workspace),
Ok(_) => None,
Err(error) => {
tracing::debug!(
session_id,
%error,
"could not probe the adopted session workspace; assuming the shared one"
);
None
}
}
}
fn adopted_session_record(
session_id: &str,
target_id: &str,
profile_id: String,
harness_kind: mj_core::config::HarnessKind,
bundle_id: String,
workspace_id: String,
locator: TargetLocator,
) -> SessionRecord {
let now = now();
SessionRecord {
build_cache: None,
mjolnir_subagents: None,
container_workspace: None,
create_managed_worktree: None,
workspace_id,
archived: false,
container_cpus: None,
container_memory: None,
id: session_id.to_owned(),
title: format!("Recovered {}", &session_id[..session_id.len().min(8)]),
harness_kind,
last_profile: profile_id,
bundle_id,
project_directory: None,
managed_worktree: None,
target_template_id: target_id.to_owned(),
resource_allocation: None,
additional_mounts: Vec::new(),
state: SessionState::Disconnected,
target: Some(locator),
native_session_id: None,
acp_session_title: None,
session_title_override: None,
created_at: now.clone(),
updated_at: now,
viewed_through_event_ordinal: 0,
draft_input: String::new(),
last_error: None,
last_checkpoint_error: None,
checkpoint: None,
}
}
fn adoption_unfinished(record: &SessionRecord, target_id: &str) -> bool {
record.state == SessionState::Disconnected
&& record.native_session_id.is_none()
&& record.target_template_id == target_id
&& record.target.is_some()
}
fn scan_target_workers(
target_id: &str,
template: &TargetTemplate,
executor: &impl CommandExecutor,
) -> Result<Vec<RecoveryCandidate>> {
let mut candidates = match template {
TargetTemplate::LocalBare => Vec::new(),
TargetTemplate::LocalPodman { .. } => scan_container_engine(
target_id,
template,
"podman",
vec![
"ps".into(),
"--all".into(),
"--filter".into(),
format!("label={}=true", targets::MANAGED_LABEL),
"--format".into(),
"json".into(),
],
executor,
)?,
TargetTemplate::LocalDocker { .. } => scan_container_engine(
target_id,
template,
"docker",
vec![
"ps".into(),
"--all".into(),
"--filter".into(),
format!("label={}=true", targets::MANAGED_LABEL),
"--format".into(),
"json".into(),
],
executor,
)?,
TargetTemplate::AppleContainer { .. } => scan_container_engine(
target_id,
template,
"container",
vec![
"list".into(),
"--all".into(),
"--format".into(),
"json".into(),
],
executor,
)?,
TargetTemplate::SshPodman { ssh, .. } => {
let remote = targets::join_remote_command(&[
"podman".into(),
"ps".into(),
"--all".into(),
"--filter".into(),
format!("label={}=true", targets::MANAGED_LABEL),
"--format".into(),
"json".into(),
]);
let output = execute_scan(
executor,
ssh_spec(ssh, [remote]),
"scan remote Podman workers",
)?;
candidates_from_container_json(target_id, template, &output.stdout)?
}
TargetTemplate::SshDocker { ssh, .. } => {
let remote = targets::join_remote_command(&[
"docker".into(),
"ps".into(),
"--all".into(),
"--filter".into(),
format!("label={}=true", targets::MANAGED_LABEL),
"--format".into(),
"json".into(),
]);
let output = execute_scan(
executor,
ssh_spec(ssh, [remote]),
"scan remote Docker workers",
)?;
candidates_from_container_json(target_id, template, &output.stdout)?
}
TargetTemplate::AwsEc2 {
aws_profile,
region,
address_source,
..
} => {
let profile = aws_profile.clone().unwrap_or_else(|| "default".into());
let output = execute_scan(
executor,
CommandSpec::new(
"aws",
[
"--profile".into(),
profile,
"--region".into(),
region.clone(),
"ec2".into(),
"describe-instances".into(),
"--filters".into(),
format!("Name=tag:{},Values=true", targets::MANAGED_TAG),
"Name=instance-state-name,Values=pending,running,stopping,stopped".into(),
"--output".into(),
"json".into(),
],
)
.purpose("scan managed EC2 workers"),
"scan managed EC2 workers",
)?;
candidates_from_aws_json(target_id, address_source.clone(), &output.stdout)?
}
TargetTemplate::SshBare { ssh, .. } => {
let output = execute_scan(
executor,
ssh_spec(
ssh,
[targets::join_remote_command(&[
"find".into(),
".local/share/hel/workers".into(),
"-mindepth".into(),
"2".into(),
"-maxdepth".into(),
"2".into(),
"-name".into(),
"ownership.json".into(),
"-print".into(),
])],
),
"scan bare SSH worker markers",
)?;
output
.stdout
.split(|byte| *byte == b'\n')
.filter_map(|line| {
let path = match std::str::from_utf8(line) {
Ok(path) => path.trim(),
Err(error) => {
tracing::debug!(%error, "recovery scan skipped a non-UTF-8 worker marker path");
return None;
}
};
let Some(session_id) = Path::new(path)
.parent()
.and_then(|parent| parent.file_name())
.and_then(|name| name.to_str())
else {
tracing::debug!(path, "recovery scan skipped a malformed worker marker path");
return None;
};
if let Err(error) = targets::resource_name(session_id) {
tracing::debug!(session_id, %error, "recovery scan skipped an invalid session id");
return None;
}
let backend = match backend_target(template, None, ContainerOverrides::default()) {
Ok(backend) => backend,
Err(error) => {
tracing::debug!(session_id, %error, "recovery scan could not construct the target backend");
return None;
}
};
let workspace = match targets::workspace_for(&backend, session_id) {
Ok(workspace) => workspace,
Err(error) => {
tracing::debug!(session_id, %error, "recovery scan could not derive the target workspace");
return None;
}
};
Some(RecoveryCandidate {
session_id: session_id.to_owned(),
target_template_id: target_id.to_owned(),
locator: TargetLocator::SshBare {
host: ssh.host.clone(),
workspace: PathBuf::from(workspace),
worker_id: None,
},
ownership: None,
instance_id: None,
})
})
.collect()
}
};
for candidate in &mut candidates {
candidate.ownership = read_recovery_ownership(template, candidate, executor);
if candidate.instance_id.is_none() {
candidate.instance_id = candidate
.ownership
.as_ref()
.and_then(|marker| marker.instance_id.clone());
}
}
Ok(candidates)
}
fn scan_container_engine(
target_id: &str,
template: &TargetTemplate,
engine: &str,
args: Vec<String>,
executor: &impl CommandExecutor,
) -> Result<Vec<RecoveryCandidate>> {
let output = execute_scan(
executor,
CommandSpec::new(engine, args).purpose("scan managed container workers"),
"scan managed container workers",
)?;
candidates_from_container_json(target_id, template, &output.stdout)
}
fn recovery_workspace_storage(
template: &TargetTemplate,
session_id: &str,
) -> Result<PodmanWorkspaceLocator> {
let backend = backend_target(template, None, ContainerOverrides::default())?;
let container = match &backend {
targets::TargetTemplate::LocalPodman(container) => container,
targets::TargetTemplate::SshPodman { container, .. } => container,
_ => bail!("target template is not a Podman target"),
};
Ok(PodmanWorkspaceLocator::from(
targets::podman_workspace_locator(container, session_id)?,
))
}
fn candidates_from_container_json(
target_id: &str,
template: &TargetTemplate,
stdout: &[u8],
) -> Result<Vec<RecoveryCandidate>> {
let sessions = managed_sessions_from_container_json(stdout)?;
Ok(sessions
.into_iter()
.filter_map(|(session_id, instance_id)| {
let generated = match targets::resource_name(&session_id) {
Ok(generated) => generated,
Err(error) => {
tracing::debug!(%session_id, %error, "recovery scan skipped an invalid managed session id");
return None;
}
};
let workspace_storage = match template {
TargetTemplate::LocalPodman { .. } | TargetTemplate::SshPodman { .. } => {
match recovery_workspace_storage(template, &session_id) {
Ok(storage) => storage,
Err(error) => {
tracing::debug!(%session_id, %error, "recovery scan could not derive the Podman workspace storage");
return None;
}
}
}
_ => Default::default(),
};
let locator = match template {
TargetTemplate::LocalPodman { .. } => TargetLocator::LocalPodman {
borrowed_from: None,
container_id: generated,
workspace_storage,
},
TargetTemplate::LocalDocker { .. } => TargetLocator::LocalDocker {
borrowed_from: None,
container_id: generated,
},
TargetTemplate::AppleContainer { .. } => TargetLocator::AppleContainer {
borrowed_from: None,
container_id: generated,
},
TargetTemplate::SshPodman { ssh, .. } => TargetLocator::SshPodman {
borrowed_from: None,
host: ssh.host.clone(),
container_id: generated,
workspace_storage,
},
TargetTemplate::SshDocker { ssh, .. } => TargetLocator::SshDocker {
borrowed_from: None,
host: ssh.host.clone(),
container_id: generated,
},
_ => return None,
};
Some(RecoveryCandidate {
session_id,
target_template_id: target_id.to_owned(),
locator,
ownership: None,
instance_id,
})
})
.collect())
}
pub(super) fn managed_sessions_from_container_json(
stdout: &[u8],
) -> Result<Vec<(String, Option<String>)>> {
let values = serde_json::Deserializer::from_slice(stdout)
.into_iter::<serde_json::Value>()
.collect::<std::result::Result<Vec<_>, _>>()
.context("parse container list JSON")?;
let mut sessions = Vec::new();
for value in &values {
collect_managed_sessions(value, &mut sessions);
}
sessions.sort();
sessions.dedup();
Ok(sessions)
}
pub(super) fn collect_managed_sessions(
value: &serde_json::Value,
sessions: &mut Vec<(String, Option<String>)>,
) {
match value {
serde_json::Value::Array(values) => {
for value in values {
collect_managed_sessions(value, sessions);
}
}
serde_json::Value::Object(object) => {
for label_key in ["Labels", "labels"] {
if let Some(labels) = object.get(label_key) {
let managed = label_value(labels, targets::MANAGED_LABEL)
.is_some_and(|value| value == "true");
if managed && let Some(session) = label_value(labels, targets::SESSION_LABEL) {
sessions.push((session, label_value(labels, targets::INSTANCE_LABEL)));
}
}
}
for value in object.values() {
collect_managed_sessions(value, sessions);
}
}
_ => {}
}
}
fn label_value(labels: &serde_json::Value, key: &str) -> Option<String> {
match labels {
serde_json::Value::Object(object) => object.get(key)?.as_str().map(str::to_owned),
serde_json::Value::String(text) => text
.split(',')
.find_map(|label| {
label
.trim()
.split_once('=')
.filter(|(name, _)| *name == key)
})
.map(|(_, value)| value.to_owned()),
_ => None,
}
}
fn candidates_from_aws_json(
target_id: &str,
address_source: AwsAddressSource,
stdout: &[u8],
) -> Result<Vec<RecoveryCandidate>> {
let value: serde_json::Value =
serde_json::from_slice(stdout).context("parse AWS instance JSON")?;
let mut result = Vec::new();
let reservations = value
.get("Reservations")
.and_then(serde_json::Value::as_array)
.cloned()
.unwrap_or_default();
for instance in reservations.iter().flat_map(|reservation| {
reservation
.get("Instances")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
}) {
let tags = instance
.get("Tags")
.and_then(serde_json::Value::as_array)
.cloned()
.unwrap_or_default();
let tag = |key: &str| {
tags.iter()
.find(|tag| tag.get("Key").and_then(serde_json::Value::as_str) == Some(key))
.and_then(|tag| tag.get("Value"))
.and_then(serde_json::Value::as_str)
};
if tag(targets::MANAGED_TAG) != Some("true") {
continue;
}
let Some(session_id) = tag(targets::SESSION_TAG).map(str::to_owned) else {
continue;
};
targets::resource_name(&session_id)?;
let instance_id = instance
.get("InstanceId")
.and_then(serde_json::Value::as_str)
.context("managed EC2 instance omitted InstanceId")?
.to_owned();
let field = match address_source {
AwsAddressSource::PublicDns => "PublicDnsName",
AwsAddressSource::PublicIp => "PublicIpAddress",
AwsAddressSource::PrivateDns => "PrivateDnsName",
AwsAddressSource::PrivateIp => "PrivateIpAddress",
};
let address = instance
.get(field)
.and_then(serde_json::Value::as_str)
.filter(|value| !value.is_empty())
.map(str::to_owned);
let created_by = tag(targets::INSTANCE_TAG).map(str::to_owned);
result.push(RecoveryCandidate {
session_id,
target_template_id: target_id.to_owned(),
locator: TargetLocator::AwsEc2 {
instance_id,
address,
},
ownership: None,
instance_id: created_by,
});
}
Ok(result)
}
fn execute_scan(
executor: &impl CommandExecutor,
command: CommandSpec,
operation: &str,
) -> Result<CommandOutput> {
let output = executor.execute(&command)?;
if output.status != 0 {
bail!(
"{operation} failed with status {}: {}",
output.status,
String::from_utf8_lossy(&output.stderr).trim()
);
}
Ok(output)
}
fn ssh_spec(ssh: &SshConnection, remote: impl IntoIterator<Item = String>) -> CommandSpec {
let backend = SshTarget::from(ssh);
let mut args = backend.ssh_args;
mj_core::targets::push_connection_sharing_args(&mut args);
args.push(backend.destination.clone());
args.extend(remote);
CommandSpec::new("ssh", args).ssh_destination(backend.destination)
}
fn read_recovery_ownership(
template: &TargetTemplate,
candidate: &RecoveryCandidate,
executor: &impl CommandExecutor,
) -> Option<WorkerOwnership> {
let backend =
match recovery_backend_locator(template, &candidate.locator, &candidate.session_id) {
Ok(backend) => backend,
Err(error) => {
tracing::debug!(
session_id = %candidate.session_id,
%error,
"could not construct a recovery ownership probe"
);
return None;
}
};
let root = match targets::worker_root(&backend, &candidate.session_id) {
Ok(root) => root,
Err(error) => {
tracing::debug!(
session_id = %candidate.session_id,
%error,
"could not derive a recovery worker root"
);
return None;
}
};
let command = match targets::command_on_locator(
&backend,
&candidate.session_id,
vec!["cat".into(), format!("{root}/ownership.json")],
"read worker ownership marker",
) {
Ok(command) => command,
Err(error) => {
tracing::debug!(
session_id = %candidate.session_id,
%error,
"could not construct a recovery ownership command"
);
return None;
}
};
let output = match executor.execute(&command) {
Ok(output) => output,
Err(error) => {
tracing::debug!(
session_id = %candidate.session_id,
%error,
"could not read a recovery worker ownership marker"
);
return None;
}
};
if output.status != 0 {
tracing::debug!(
session_id = %candidate.session_id,
status = output.status,
"recovery worker ownership probe returned a failure"
);
return None;
}
let marker: WorkerOwnership = match serde_json::from_slice(&output.stdout) {
Ok(marker) => marker,
Err(error) => {
tracing::debug!(
session_id = %candidate.session_id,
%error,
"recovery worker ownership marker was not valid JSON"
);
return None;
}
};
if !(1..=WorkerOwnership::VERSION).contains(&marker.version)
|| marker.session_id != candidate.session_id
|| marker.target_template_id != candidate.target_template_id
{
tracing::debug!(
session_id = %candidate.session_id,
marker_session_id = %marker.session_id,
marker_target_template_id = %marker.target_template_id,
"recovery worker ownership marker did not match the candidate"
);
return None;
}
Some(marker)
}
fn recovery_backend_locator(
template: &TargetTemplate,
locator: &TargetLocator,
session_id: &str,
) -> Result<targets::TargetLocator> {
Ok(match (template, locator) {
(TargetTemplate::LocalBare, TargetLocator::LocalBare { worker_root }) => {
targets::TargetLocator::LocalBare {
worker_root: worker_root.to_string_lossy().into_owned(),
}
}
(
TargetTemplate::LocalPodman { .. },
TargetLocator::LocalPodman {
container_id,
workspace_storage,
..
},
) => targets::TargetLocator::LocalPodman {
borrowed_from: None,
container_id: container_id.clone(),
workspace_storage: workspace_storage.into(),
},
(TargetTemplate::LocalDocker { .. }, TargetLocator::LocalDocker { container_id, .. }) => {
targets::TargetLocator::LocalDocker {
borrowed_from: None,
container_id: container_id.clone(),
}
}
(
TargetTemplate::AppleContainer { .. },
TargetLocator::AppleContainer { container_id, .. },
) => targets::TargetLocator::AppleContainer {
borrowed_from: None,
container_id: container_id.clone(),
},
(
TargetTemplate::SshPodman { ssh, .. },
TargetLocator::SshPodman {
container_id,
workspace_storage,
..
},
) => targets::TargetLocator::SshPodman {
borrowed_from: None,
ssh: SshTarget::from(ssh),
container_id: container_id.clone(),
workspace_storage: workspace_storage.into(),
},
(
TargetTemplate::SshDocker { ssh, .. },
TargetLocator::SshDocker {
host, container_id, ..
},
) => {
if host != &ssh.host {
bail!("recovery SSH Docker host does not match target template")
}
targets::TargetLocator::SshDocker {
borrowed_from: None,
ssh: SshTarget::from(ssh),
container_id: container_id.clone(),
}
}
(TargetTemplate::SshBare { ssh, .. }, TargetLocator::SshBare { workspace, .. }) => {
targets::TargetLocator::SshBare {
ssh: SshTarget::from(ssh),
workspace: workspace.to_string_lossy().into_owned(),
worker_id: None,
}
}
(
TargetTemplate::AwsEc2 {
aws_profile,
region,
ssh_user,
identity_file,
ssh_args,
..
},
TargetLocator::AwsEc2 {
instance_id,
address,
},
) => targets::TargetLocator::AwsEc2 {
profile: aws_profile.clone().unwrap_or_else(|| "default".into()),
region: region.clone(),
instance_id: instance_id.clone(),
ssh: SshTarget {
destination: format!(
"{ssh_user}@{}",
address.as_deref().unwrap_or("unavailable.invalid")
),
ssh_args: targets::ssh_args_with_identity(ssh_args, identity_file.as_deref()),
},
workspace: format!(".local/share/hel/workspaces/{session_id}"),
},
_ => bail!("recovery target locator does not match target template"),
})
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use crate::controller::test_support::{IsolatedTest, test_name};
use mj_core::config::{
AwsAddressSource, Config, ContainerTemplate as ConfigContainer, HarnessKind,
PodmanWorkspaceStorage, TargetTemplate,
};
use mj_core::state::{State, TargetLocator};
use crate::targets::ProcessExecutor;
use super::*;
const FAILED_ADOPTION_CHILD: &str = "MJ_TEST_FAILED_ADOPTION_CHILD";
#[tokio::test]
async fn a_failed_adoption_records_the_failure_and_stays_retryable() {
if std::env::var_os(FAILED_ADOPTION_CHILD).is_none() {
let directory = tempfile::tempdir().unwrap();
IsolatedTest::new(test_name(
module_path!(),
"a_failed_adoption_records_the_failure_and_stays_retryable",
))
.env(FAILED_ADOPTION_CHILD, "1")
.env("MJ_DATA_DIR", directory.path())
.run();
return;
}
let _writer = crate::database::install_isolated_test_writer();
let session_id = "0123456789abcdef0123456789abcdef";
let workers = tempfile::tempdir().unwrap();
let record = adopted_session_record(
session_id,
"local-bare",
"codex".into(),
HarnessKind::Codex,
"project".into(),
mj_core::workspace::DEFAULT_WORKSPACE_ID.to_owned(),
TargetLocator::LocalBare {
worker_root: workers.path().join(session_id),
},
);
assert!(
adoption_unfinished(&record, "local-bare"),
"the record adoption commits must be the record adoption can retry"
);
crate::database::save_session(&record).unwrap();
let mut config = Config::default();
config
.targets
.insert("local-bare".into(), TargetTemplate::LocalBare);
let mut state = State::default();
state.sessions.insert(session_id.to_owned(), record);
let mut controller = Controller { config, state };
let failure = controller
.adopt_orphan_worker(
session_id,
"local-bare",
None,
None,
false,
&ProcessExecutor,
)
.await
.expect_err("a worker root without a worker cannot complete the handshake");
assert!(
format!("{failure:#}").contains("orphan relay"),
"unexpected failure: {failure:#}"
);
let recorded = controller.state.sessions[session_id]
.last_error
.clone()
.expect("the failed handshake was recorded on the session");
assert!(
recorded.contains("orphan adoption failed"),
"unexpected recorded failure: {recorded}"
);
assert_eq!(
controller.state.sessions[session_id].state,
SessionState::Disconnected
);
let stored = crate::database::load_state().unwrap();
assert_eq!(
stored.sessions[session_id].last_error.as_deref(),
Some(recorded.as_str()),
"the adoption failure was not persisted"
);
let retry = controller
.adopt_orphan_worker(
session_id,
"local-bare",
None,
None,
false,
&ProcessExecutor,
)
.await
.expect_err("the worker is still unreachable");
let retry = format!("{retry:#}");
assert!(
retry.contains("orphan relay"),
"adoption did not retry the handshake: {retry}"
);
assert!(
!retry.contains("already tracked"),
"a session adoption never finished blocked its own retry: {retry}"
);
}
const RECOVERY_WORKSPACE_CHILD: &str = "MJ_TEST_RECOVERY_WORKSPACE_CHILD";
#[tokio::test]
async fn orphan_workspace_ids_are_reconciled_before_adoption_persistence() {
if std::env::var_os(RECOVERY_WORKSPACE_CHILD).is_none() {
let directory = tempfile::tempdir().unwrap();
IsolatedTest::new(test_name(
module_path!(),
"orphan_workspace_ids_are_reconciled_before_adoption_persistence",
))
.env(RECOVERY_WORKSPACE_CHILD, "1")
.env("MJ_DATA_DIR", directory.path())
.run();
return;
}
let _writer = crate::database::install_isolated_test_writer();
let default =
resolve_recovery_workspace_id(mj_core::workspace::DEFAULT_WORKSPACE_ID).unwrap();
assert_eq!(default, mj_core::workspace::DEFAULT_WORKSPACE_ID);
let known = crate::database::create_or_get_workspace("Known").unwrap();
assert_eq!(resolve_recovery_workspace_id(&known.id).unwrap(), known.id);
let recovered = resolve_recovery_workspace_id("workspace-from-old-controller").unwrap();
let repeated = resolve_recovery_workspace_id("another-old-workspace").unwrap();
assert_eq!(repeated, recovered);
assert_eq!(
crate::database::list_workspaces()
.unwrap()
.iter()
.filter(|workspace| workspace.name == "Recovered")
.count(),
1
);
let session_id = "0123456789abcdef0123456789abcdef";
let workers = tempfile::tempdir().unwrap();
let record = adopted_session_record(
session_id,
"local-bare",
"codex".into(),
HarnessKind::Codex,
"project".into(),
recovered.clone(),
TargetLocator::LocalBare {
worker_root: workers.path().join(session_id),
},
);
crate::database::save_session(&record).unwrap();
assert_eq!(
crate::database::load_state().unwrap().sessions[session_id].workspace_id,
recovered
);
let mut state = State::default();
state.sessions.insert(session_id.to_owned(), record);
let mut config = Config::default();
config
.targets
.insert("local-bare".into(), TargetTemplate::LocalBare);
let mut controller = Controller { config, state };
let failure = controller
.adopt_orphan_worker(
session_id,
"local-bare",
None,
None,
false,
&ProcessExecutor,
)
.await
.expect_err("a worker root without a worker cannot complete the handshake");
assert!(
format!("{failure:#}").contains("orphan relay"),
"unexpected failure: {failure:#}"
);
let stored = crate::database::load_state().unwrap();
assert_eq!(stored.sessions[session_id].workspace_id, recovered);
assert!(
stored.sessions[session_id]
.last_error
.as_deref()
.is_some_and(|error| error.contains("orphan adoption failed"))
);
}
#[test]
fn a_session_that_completed_its_handshake_is_not_adoptable_again() {
let mut record = adopted_session_record(
"0123456789abcdef0123456789abcdef",
"local-bare",
"codex".into(),
HarnessKind::Codex,
"project".into(),
mj_core::workspace::DEFAULT_WORKSPACE_ID.to_owned(),
TargetLocator::LocalBare {
worker_root: std::path::PathBuf::from("/workers/0123456789abcdef0123456789abcdef"),
},
);
record.native_session_id = Some("native-session".into());
assert!(!adoption_unfinished(&record, "local-bare"));
record.native_session_id = None;
assert!(
!adoption_unfinished(&record, "other-target"),
"a record adopted onto another target is not this target's retry"
);
}
#[test]
fn recovery_container_scan_requires_both_managed_and_session_labels() {
let template = TargetTemplate::LocalPodman {
container: ConfigContainer {
build_cache: None,
image: "ignored".into(),
pull_policy: Default::default(),
platform: None,
cpus: None,
memory: None,
environment: BTreeMap::new(),
workspace_storage: Default::default(),
},
};
let json = serde_json::json!([
{"Labels": {"dev.mj.managed": "true", "dev.mj.session": "0123456789abcdef0123456789abcdef", "dev.mj.instance": "qa0916"}},
{"Labels": {"dev.mj.managed": "false", "dev.mj.session": "not-owned"}},
{"configuration": {"labels": "dev.mj.managed=true,dev.mj.session=abcdef0123456789abcdef0123456789"}}
]);
let candidates = candidates_from_container_json(
"local",
&template,
serde_json::to_string(&json).unwrap().as_bytes(),
)
.unwrap();
assert_eq!(candidates.len(), 2);
assert_eq!(candidates[0].session_id, "0123456789abcdef0123456789abcdef");
assert_eq!(candidates[0].instance_id.as_deref(), Some("qa0916"));
assert_eq!(
candidates[1].instance_id, None,
"a container without an instance label is of unknown origin"
);
}
fn candidate(session_id: &str, instance_id: Option<&str>) -> RecoveryCandidate {
RecoveryCandidate {
session_id: session_id.to_owned(),
target_template_id: "local".to_owned(),
locator: TargetLocator::LocalDocker {
container_id: format!("mj-{session_id}"),
borrowed_from: None,
},
ownership: None,
instance_id: instance_id.map(str::to_owned),
}
}
#[test]
fn default_scan_scope_hides_other_and_unknown_instances() {
let mut scan = RecoveryScan {
candidates: vec![
candidate("mine", Some("qa")),
candidate("theirs", Some("prod")),
candidate("legacy", None),
],
instance_id: "qa".to_owned(),
..RecoveryScan::default()
};
restrict_to_instance(&mut scan);
assert_eq!(
scan.candidates
.iter()
.map(|candidate| candidate.session_id.as_str())
.collect::<Vec<_>>(),
["mine"]
);
assert_eq!(scan.hidden_other_instances, 2);
}
#[test]
fn acting_on_another_or_unknown_instance_requires_the_explicit_flag() {
require_instance_access(&candidate("mine", Some("qa")), "qa", false).unwrap();
let other = require_instance_access(&candidate("theirs", Some("prod")), "qa", false)
.expect_err("another instance's worker is refused by default");
assert!(
other.to_string().contains("belongs to instance \"prod\""),
"{other}"
);
require_instance_access(&candidate("theirs", Some("prod")), "qa", true).unwrap();
let unknown = require_instance_access(&candidate("legacy", None), "qa", false)
.expect_err("a worker without a stamp is refused by default");
assert!(
unknown.to_string().contains("no instance stamp"),
"{unknown}"
);
require_instance_access(&candidate("legacy", None), "qa", true).unwrap();
}
#[test]
fn a_podman_orphan_destroy_plan_removes_its_workspace_volume() {
let session = "0123456789abcdef0123456789abcdef";
let template = TargetTemplate::LocalPodman {
container: ConfigContainer {
image: "ignored".into(),
pull_policy: Default::default(),
platform: None,
cpus: None,
memory: None,
environment: BTreeMap::new(),
workspace_storage: PodmanWorkspaceStorage::PodmanVolume,
build_cache: None,
},
};
let json = serde_json::json!([
{"Labels": {"dev.mj.managed": "true", "dev.mj.session": session}}
]);
let candidates = candidates_from_container_json(
"local",
&template,
serde_json::to_string(&json).unwrap().as_bytes(),
)
.unwrap();
let [candidate] = candidates.as_slice() else {
panic!("expected one orphan candidate, got {candidates:?}");
};
let volume = format!("{}-workspace", targets::resource_name(session).unwrap());
assert!(
matches!(
&candidate.locator,
TargetLocator::LocalPodman {
workspace_storage: PodmanWorkspaceLocator::Volume { name },
..
} if name == &volume
),
"candidate locator lost the volume storage: {:?}",
candidate.locator
);
let backend = recovery_backend_locator(&template, &candidate.locator, session).unwrap();
let plan = targets::close_plan(&backend, session).unwrap();
assert!(
plan.commands.iter().any(|command| {
command.args.iter().any(|argument| argument == &volume)
&& command
.args
.iter()
.any(|argument| argument.contains("podman volume rm"))
}),
"destroy plan does not remove the workspace volume: {plan:?}"
);
}
#[test]
fn recovery_docker_scan_accepts_json_lines_and_builds_a_docker_locator() {
let template = TargetTemplate::LocalDocker {
container: ConfigContainer {
build_cache: None,
image: "ignored".into(),
pull_policy: Default::default(),
platform: None,
cpus: None,
memory: None,
environment: BTreeMap::new(),
workspace_storage: Default::default(),
},
};
let session = "0123456789abcdef0123456789abcdef";
let output = format!(
"{{\"Labels\":\"dev.mj.managed=true,dev.mj.session={session},dev.mj.instance=abc123\"}}\n{{\"Labels\":\"dev.mj.managed=false,dev.mj.session=ignored\"}}\n"
);
let candidates =
candidates_from_container_json("docker", &template, output.as_bytes()).unwrap();
assert_eq!(candidates.len(), 1);
assert_eq!(candidates[0].session_id, session);
assert_eq!(candidates[0].instance_id.as_deref(), Some("abc123"));
assert!(matches!(
&candidates[0].locator,
TargetLocator::LocalDocker { container_id, .. }
if container_id == &targets::resource_name(session).unwrap()
));
}
#[test]
fn recovery_aws_scan_uses_exact_tagged_instance_and_address() {
let json = serde_json::json!({"Reservations": [{"Instances": [{
"InstanceId": "i-exact",
"PrivateIpAddress": "10.0.0.7",
"Tags": [
{"Key": "dev.mj.managed", "Value": "true"},
{"Key": "dev.mj.session", "Value": "0123456789abcdef0123456789abcdef"},
{"Key": "dev.mj.instance", "Value": "qa0916"}
]
}]}]});
let candidates = candidates_from_aws_json(
"aws",
AwsAddressSource::PrivateIp,
serde_json::to_string(&json).unwrap().as_bytes(),
)
.unwrap();
assert_eq!(candidates[0].instance_id.as_deref(), Some("qa0916"));
assert!(matches!(
&candidates[0].locator,
TargetLocator::AwsEc2 { instance_id, address }
if instance_id == "i-exact" && address.as_deref() == Some("10.0.0.7")
));
}
}