From 5147cec1a95e32709d53fcb347aa98b08968a89d Mon Sep 17 00:00:00 2001 From: Will Killian Date: Wed, 15 Jul 2026 09:50:56 -0400 Subject: [PATCH 1/2] refactor: resolve CLI Sonar findings Signed-off-by: Will Killian --- crates/cli/src/commands/diagnostics.rs | 99 ++-- crates/cli/src/diagnostics/mod.rs | 170 +++--- crates/cli/src/hooks/delivery.rs | 71 +-- .../cli/src/installation/marketplace/mod.rs | 319 +++++++---- crates/cli/src/mcp_environment.rs | 47 +- .../cli/src/plugins/lifecycle/environment.rs | 123 +++-- crates/cli/src/plugins/lifecycle/mod.rs | 275 ++++++---- crates/cli/src/plugins/schema/secrets.rs | 108 ++-- crates/cli/src/process/detached.rs | 140 ++--- crates/cli/src/process/supervision/unix.rs | 29 +- crates/cli/src/process/supervision/windows.rs | 76 +-- crates/cli/src/server/mod.rs | 64 ++- crates/cli/src/sessions/idle.rs | 128 +++-- crates/cli/tests/cli_tests.rs | 517 ++++++++++-------- scripts/test-claude-plugin-e2e.sh | 4 + scripts/test-hermes-mcp-e2e.sh | 2 + scripts/test-support/codex_mock_provider.py | 2 + 17 files changed, 1318 insertions(+), 856 deletions(-) diff --git a/crates/cli/src/commands/diagnostics.rs b/crates/cli/src/commands/diagnostics.rs index bd689bb5c..b32fa7121 100644 --- a/crates/cli/src/commands/diagnostics.rs +++ b/crates/cli/src/commands/diagnostics.rs @@ -31,51 +31,64 @@ pub(crate) struct AgentsCommand { pub(super) async fn execute(command: DoctorCommand) -> Result { if let Some(plugin) = command.plugin { - let candidates = plugin.agents(); - let agents = if plugin.is_all() { - crate::agents::installed_integrations(&candidates, command.install_dir.as_deref()) - } else { - candidates - }; - if agents.is_empty() { - return Err(CliError::Install( - "no installed Claude Code, Codex, or Hermes integration state was found".into(), - )); - } - let options = crate::installation::marketplace::plugin_doctor_options(command.install_dir); - if command.json { - let reports = agents - .iter() - .copied() - .map(|agent| crate::agents::doctor_integration_report(agent, &options)) - .collect::, _>>()?; - let ready = reports - .iter() - .all(|report| report.get("ok").and_then(Value::as_bool) == Some(true)); - let output = if reports.len() > 1 { - json!({ "schema_version": 1, "plugins": reports }) - } else { - with_schema(reports.into_iter().next().expect("reports is not empty")) - }; - println!( - "{}", - serde_json::to_string_pretty(&output) - .map_err(|error| CliError::Install(error.to_string()))? - ); - Ok(if ready { - ExitCode::SUCCESS - } else { - ExitCode::FAILURE - }) - } else { - for agent in agents { - crate::agents::doctor_integration(agent, &options)?; - } - Ok(ExitCode::SUCCESS) - } + return execute_plugin_doctor(plugin, command.install_dir, command.json); + } + crate::diagnostics::run_doctor(command.agent.map(Into::into), command.json).await +} + +fn execute_plugin_doctor( + plugin: InstallTarget, + install_dir: Option, + json: bool, +) -> Result { + let candidates = plugin.agents(); + let agents = if plugin.is_all() { + crate::agents::installed_integrations(&candidates, install_dir.as_deref()) } else { - crate::diagnostics::run_doctor(command.agent.map(Into::into), command.json).await + candidates + }; + if agents.is_empty() { + return Err(CliError::Install( + "no installed Claude Code, Codex, or Hermes integration state was found".into(), + )); } + let options = crate::installation::marketplace::plugin_doctor_options(install_dir); + if !json { + for agent in agents { + crate::agents::doctor_integration(agent, &options)?; + } + return Ok(ExitCode::SUCCESS); + } + print_plugin_doctor_json(&agents, &options) +} + +fn print_plugin_doctor_json( + agents: &[crate::agents::CodingAgent], + options: &crate::installation::marketplace::state::PluginInstallOptions, +) -> Result { + let reports = agents + .iter() + .copied() + .map(|agent| crate::agents::doctor_integration_report(agent, options)) + .collect::, _>>()?; + let ready = reports + .iter() + .all(|report| report.get("ok").and_then(Value::as_bool) == Some(true)); + let output = if reports.len() > 1 { + json!({ "schema_version": 1, "plugins": reports }) + } else { + with_schema(reports.into_iter().next().expect("reports is not empty")) + }; + println!( + "{}", + serde_json::to_string_pretty(&output) + .map_err(|error| CliError::Install(error.to_string()))? + ); + Ok(if ready { + ExitCode::SUCCESS + } else { + ExitCode::FAILURE + }) } fn with_schema(mut value: Value) -> Value { diff --git a/crates/cli/src/diagnostics/mod.rs b/crates/cli/src/diagnostics/mod.rs index a5b7ab736..15d485d51 100644 --- a/crates/cli/src/diagnostics/mod.rs +++ b/crates/cli/src/diagnostics/mod.rs @@ -299,82 +299,108 @@ async fn collect_agents( if target_agent.is_some_and(|target| target != agent) { continue; } - let configured = agent_configured(agent, &resolved.agents); - let target_requested = target_agent == Some(agent); - let command = agent_command(agent, &resolved.agents); - let argv = crate::process::command_argv(&command); - let exec = argv.first().map(String::as_str).unwrap_or_default(); - let path = crate::process::resolve_executable(exec); - let version = match &path { - Some(_) => { - let probe = crate::process::version_probe_argv(agent, &argv); - probe_version(&probe).await - } - None => None, - }; - let mut status = agent_command_status(path.as_deref(), configured, target_requested); - let (hook_status, hook_details) = hook_status(agent, &resolved.agents); - status = combine_status(status, hook_status, configured || target_requested); - let mut details = Vec::new(); - details.push(if configured { - "configured".to_string() - } else if target_requested { - "not configured; first run will launch setup".to_string() - } else { - "not configured".to_string() - }); - if path.is_none() { - details.push(format!("command `{exec}` not found")); - } - if !hook_details.is_empty() { - details.push(hook_details); - } - let version_required = configured || target_requested; - match version.as_deref() { - Some(version) => { - if let Err(error) = agent.validate_version_output(version) { - status = combine_status( - status, - if version_required { - Status::Fail - } else { - Status::Warn - }, - true, - ); - details.push(error); - } - } - None if path.is_some() => { - status = combine_status( - status, - if version_required { - Status::Fail - } else { - Status::Warn - }, - true, - ); - details.push(format!( - "could not determine version; NeMo Relay requires {}", - agent.version_requirement() - )); - } - None => {} - } - out.push(AgentInfo { - name: agent.as_arg(), - status, - configured, - command, - path, - version, - annotation: details.join("; "), - }); + out.push(collect_agent(agent, target_agent == Some(agent), resolved).await); } out } +async fn collect_agent( + agent: CodingAgent, + target_requested: bool, + resolved: &ResolvedConfig, +) -> AgentInfo { + let configured = agent_configured(agent, &resolved.agents); + let command = agent_command(agent, &resolved.agents); + let argv = crate::process::command_argv(&command); + let exec = argv.first().map(String::as_str).unwrap_or_default(); + let path = crate::process::resolve_executable(exec); + let version = if path.is_some() { + probe_version(&crate::process::version_probe_argv(agent, &argv)).await + } else { + None + }; + let mut status = agent_command_status(path.as_deref(), configured, target_requested); + let (hook_status, hook_details) = hook_status(agent, &resolved.agents); + status = combine_status(status, hook_status, configured || target_requested); + let mut details = agent_details( + configured, + target_requested, + path.as_deref(), + exec, + hook_details, + ); + apply_agent_version_status( + agent, + version.as_deref(), + path.is_some(), + configured || target_requested, + &mut status, + &mut details, + ); + AgentInfo { + name: agent.as_arg(), + status, + configured, + command, + path, + version, + annotation: details.join("; "), + } +} + +fn agent_details( + configured: bool, + target_requested: bool, + path: Option<&Path>, + exec: &str, + hook_details: String, +) -> Vec { + let mut details = vec![if configured { + "configured".to_string() + } else if target_requested { + "not configured; first run will launch setup".to_string() + } else { + "not configured".to_string() + }]; + if path.is_none() { + details.push(format!("command `{exec}` not found")); + } + if !hook_details.is_empty() { + details.push(hook_details); + } + details +} + +fn apply_agent_version_status( + agent: CodingAgent, + version: Option<&str>, + executable_found: bool, + version_required: bool, + status: &mut Status, + details: &mut Vec, +) { + let problem = match version { + Some(version) => agent.validate_version_output(version).err(), + None if executable_found => Some(format!( + "could not determine version; NeMo Relay requires {}", + agent.version_requirement() + )), + None => None, + }; + if let Some(problem) = problem { + *status = combine_status( + *status, + if version_required { + Status::Fail + } else { + Status::Warn + }, + true, + ); + details.push(problem); + } +} + fn agent_command(agent: CodingAgent, agents: &AgentConfigs) -> String { configured_agent_command(agent, agents) .cloned() diff --git a/crates/cli/src/hooks/delivery.rs b/crates/cli/src/hooks/delivery.rs index d74939fb3..17bb60fff 100644 --- a/crates/cli/src/hooks/delivery.rs +++ b/crates/cli/src/hooks/delivery.rs @@ -8,7 +8,7 @@ use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue}; use serde_json::Value; use crate::error::CliError; -use crate::installation::generation::InstallGeneration; +use crate::installation::generation::{ActiveGenerationGuard, InstallGeneration}; use super::destination::{ HookGatewayLifecycle, hook_destination, recovery_plan, transparent_gateway_spec, @@ -32,10 +32,7 @@ pub(crate) async fn hook_forward(command: HookForwardRequest) -> Result<(), CliE let fail_closed = command.fail_closed || std::env::var("NEMO_RELAY_FAIL_CLOSED").ok().as_deref() == Some("1"); let destination = hook_destination(&command); - let persistent = match (destination.lifecycle != HookGatewayLifecycle::Transparent) - .then(|| recovery_plan(&destination.gateway_url)) - .transpose() - { + let persistent = match persistent_gateway(&destination) { Ok(persistent) => persistent, Err(error) => return handle_hook_error(error, fail_closed), }; @@ -47,34 +44,9 @@ pub(crate) async fn hook_forward(command: HookForwardRequest) -> Result<(), CliE Ok(gateway) => gateway, Err(error) => return handle_hook_error(error, fail_closed), }; - let _generation_guard = if destination.lifecycle == HookGatewayLifecycle::Existing - && !command.forward_only - { - let install_host = command.agent.install_arg(); - let Some(generation_file) = command.generation_file.clone() else { - return handle_hook_error( - CliError::Launch(format!( - "persistent {} hook is missing its install-generation fence; run `nemo-relay install {install_host} --force`", - command.agent.label() - )), - fail_closed, - ); - }; - let Some(generation_token) = command.generation_token.as_deref() else { - return handle_hook_error( - CliError::Launch(format!( - "persistent {} hook is missing its expected install-generation identity; run `nemo-relay install {install_host} --force`", - command.agent.label() - )), - fail_closed, - ); - }; - match InstallGeneration::capture_guarded_expected(generation_file, generation_token) { - Ok((_generation, guard)) => Some(guard), - Err(error) => return handle_hook_error(CliError::Launch(error), fail_closed), - } - } else { - None + let _generation_guard = match capture_generation_guard(&command, destination.lifecycle) { + Ok(guard) => guard, + Err(error) => return handle_hook_error(error, fail_closed), }; let input = match read_hook_payload(persistent.as_ref().map_or( crate::configuration::DEFAULT_MAX_HOOK_PAYLOAD_BYTES, @@ -126,6 +98,39 @@ pub(crate) async fn hook_forward(command: HookForwardRequest) -> Result<(), CliE handle_hook_forward_response(response, fail_closed).await } +fn persistent_gateway( + destination: &super::destination::HookDestination, +) -> Result, CliError> { + (destination.lifecycle != HookGatewayLifecycle::Transparent) + .then(|| recovery_plan(&destination.gateway_url)) + .transpose() +} + +fn capture_generation_guard( + command: &HookForwardRequest, + lifecycle: HookGatewayLifecycle, +) -> Result, CliError> { + if lifecycle != HookGatewayLifecycle::Existing || command.forward_only { + return Ok(None); + } + let install_host = command.agent.install_arg(); + let generation_file = command.generation_file.clone().ok_or_else(|| { + CliError::Launch(format!( + "persistent {} hook is missing its install-generation fence; run `nemo-relay install {install_host} --force`", + command.agent.label() + )) + })?; + let generation_token = command.generation_token.as_deref().ok_or_else(|| { + CliError::Launch(format!( + "persistent {} hook is missing its expected install-generation identity; run `nemo-relay install {install_host} --force`", + command.agent.label() + )) + })?; + InstallGeneration::capture_guarded_expected(generation_file, generation_token) + .map(|(_generation, guard)| Some(guard)) + .map_err(CliError::Launch) +} + fn handle_hook_error(error: CliError, fail_closed: bool) -> Result<(), CliError> { eprintln!("nemo-relay hook forward failed: {error}"); if fail_closed { Err(error) } else { Ok(()) } diff --git a/crates/cli/src/installation/marketplace/mod.rs b/crates/cli/src/installation/marketplace/mod.rs index 38d3daa44..4c25dc0dd 100644 --- a/crates/cli/src/installation/marketplace/mod.rs +++ b/crates/cli/src/installation/marketplace/mod.rs @@ -285,8 +285,64 @@ fn install_host_locked( require_host_cli(host, options, runner)?; host::validate_host_version(host, options, runner)?; let layout = PluginLayout::new(host, &options.install_dir); + let (staged, mut force_snapshot, mut replacement_generation_lock) = + prepare_install_transaction(host, &relay, &layout, options, runner, setup_runner)?; + install_marketplace_content( + host, + &relay, + &layout, + staged.as_ref(), + &mut force_snapshot, + &mut replacement_generation_lock, + options, + runner, + setup_runner, + )?; + write_install_state( + host, + &layout, + &mut force_snapshot, + &mut replacement_generation_lock, + options, + runner, + setup_runner, + )?; + finish_install_registration( + host, + &layout, + &mut force_snapshot, + &mut replacement_generation_lock, + options, + runner, + setup_runner, + )?; + if let Some(snapshot) = force_snapshot { + snapshot.commit(&layout.generation_lock); + } + println!( + "installed {} plugin marketplace at {}", + host.label(), + layout.marketplace_root.display() + ); + Ok(()) +} + +type PreparedInstall = ( + Option, + Option, + Option, +); + +fn prepare_install_transaction( + host: impl MarketplaceHost, + relay: &Path, + layout: &PluginLayout, + options: &PluginInstallOptions, + runner: &dyn CommandRunner, + setup_runner: &dyn PluginSetupRunner, +) -> Result { let plugin_preflight = if !options.dry_run { - Some(prepare_plugin_install(host, &layout, options, runner)?) + Some(prepare_plugin_install(host, layout, options, runner)?) } else { None }; @@ -306,7 +362,7 @@ fn install_host_locked( None => true, }; let staged = - stage_plugin_marketplace(host, &relay, &layout, initialize_generation_lock, options)?; + stage_plugin_marketplace(host, relay, layout, initialize_generation_lock, options)?; if initialize_generation_lock { replacement_generation_lock = match acquire_replacement_generation_lock( host, @@ -324,13 +380,13 @@ fn install_host_locked( } }; } - match begin_force_replacement(host, &layout, preflight, options, runner, setup_runner) { + match begin_force_replacement(host, layout, preflight, options, runner, setup_runner) { Ok(mut snapshot) => { if let Err(error) = setup_runner.refresh_gateway() { staged.cleanup(); return restore_force_replacement_after_error( host, - &layout, + layout, &mut snapshot, options, runner, @@ -349,15 +405,30 @@ fn install_host_locked( } else { None }; + Ok((staged, force_snapshot, replacement_generation_lock)) +} + +#[allow(clippy::too_many_arguments)] +fn install_marketplace_content( + host: impl MarketplaceHost, + relay: &Path, + layout: &PluginLayout, + staged: Option<&StagedPluginMarketplace>, + force_snapshot: &mut Option, + replacement_generation_lock: &mut Option, + options: &PluginInstallOptions, + runner: &dyn CommandRunner, + setup_runner: &dyn PluginSetupRunner, +) -> Result<(), String> { if options.force && staged.is_none() { - force_cleanup_existing_install(host, &layout, options, runner, setup_runner)?; + force_cleanup_existing_install(host, layout, options, runner, setup_runner)?; } - if let Some(staged) = staged.as_ref() { - if let Err(error) = staged.promote(&layout) { + if let Some(staged) = staged { + if let Err(error) = staged.promote(layout) { staged.cleanup(); return restore_force_replacement_after_error( host, - &layout, + layout, force_snapshot.as_mut().expect("force snapshot exists"), options, runner, @@ -375,7 +446,7 @@ fn install_host_locked( staged.cleanup(); return restore_force_replacement_after_error( host, - &layout, + layout, force_snapshot.as_mut().expect("force snapshot exists"), options, runner, @@ -387,7 +458,7 @@ fn install_host_locked( } else { let generation_lock_created = !options.dry_run && generation_lock_is_absent(&layout.generation_lock); - if let Err(error) = write_plugin_marketplace(host, &layout, &relay, options) { + if let Err(error) = write_plugin_marketplace(host, layout, relay, options) { let cleanup_error = (!options.dry_run) .then(|| remove_path(&layout.marketplace_root, options).err()) .flatten(); @@ -402,7 +473,7 @@ fn install_host_locked( }; } if !options.dry_run { - replacement_generation_lock = match acquire_replacement_generation_lock( + *replacement_generation_lock = match acquire_replacement_generation_lock( host, &layout.generation_fence, &layout.generation_lock, @@ -424,7 +495,20 @@ fn install_host_locked( }; } } - if let Err(error) = write_state(&layout, options) { + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn write_install_state( + host: impl MarketplaceHost, + layout: &PluginLayout, + force_snapshot: &mut Option, + replacement_generation_lock: &mut Option, + options: &PluginInstallOptions, + runner: &dyn CommandRunner, + setup_runner: &dyn PluginSetupRunner, +) -> Result<(), String> { + if let Err(error) = write_state(layout, options) { let _replacement_retirement = if force_snapshot.is_some() { let existing_retirement = replacement_generation_lock .as_mut() @@ -436,7 +520,7 @@ fn install_host_locked( }); match retire_replacement_before_rollback( host, - &layout, + layout, options, setup_runner, existing_retirement, @@ -453,7 +537,7 @@ fn install_host_locked( }; let cleanup_error = remove_path(&layout.marketplace_root, options).err(); let restore_error = force_snapshot.as_mut().and_then(|snapshot| { - restore_force_replacement(host, &layout, snapshot, options, runner, setup_runner).err() + restore_force_replacement(host, layout, snapshot, options, runner, setup_runner).err() }); let errors = [cleanup_error, restore_error] .into_iter() @@ -464,6 +548,19 @@ fn install_host_locked( } return Err(error); } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn finish_install_registration( + host: impl MarketplaceHost, + layout: &PluginLayout, + force_snapshot: &mut Option, + replacement_generation_lock: &mut Option, + options: &PluginInstallOptions, + runner: &dyn CommandRunner, + setup_runner: &dyn PluginSetupRunner, +) -> Result<(), String> { let mut registration = HostRegistrationProgress::default(); let mut registration_state_uncertain = false; let mut setup_installed = false; @@ -495,13 +592,13 @@ fn install_host_locked( } run_plugin_setup_with_generation( host, - &layout, + layout, options, setup_runner, generation_token.as_deref(), )?; setup_installed = true; - mark_plugin_setup_installed(host, &layout, options)?; + mark_plugin_setup_installed(host, layout, options)?; if !options.skip_doctor { run_plugin_doctor_with_generation( host, @@ -535,7 +632,7 @@ fn install_host_locked( }); match retire_replacement_before_rollback( host, - &layout, + layout, options, setup_runner, existing_retirement, @@ -552,7 +649,7 @@ fn install_host_locked( }; let rollback_error = rollback_install( host, - &layout, + layout, registration, setup_installed, options, @@ -561,7 +658,7 @@ fn install_host_locked( ) .err(); let restore_error = force_snapshot.as_mut().and_then(|snapshot| { - restore_force_replacement(host, &layout, snapshot, options, runner, setup_runner).err() + restore_force_replacement(host, layout, snapshot, options, runner, setup_runner).err() }); let rollback_errors = [ rollback_error.map(|error| format!("failed to roll back install: {error}")), @@ -578,14 +675,6 @@ fn install_host_locked( } return Err(error); } - if let Some(snapshot) = force_snapshot { - snapshot.commit(&layout.generation_lock); - } - println!( - "installed {} plugin marketplace at {}", - host.label(), - layout.marketplace_root.display() - ); Ok(()) } @@ -1779,98 +1868,141 @@ fn restore_force_replacement( setup_runner: &dyn PluginSetupRunner, ) -> Result<(), String> { let mut errors = Vec::new(); - if snapshot.replacement_promoted { - match host_registration_report(host, options, runner) { - Ok(report) => { - if report.host_plugin_registered - && let Err(error) = run_host_plugin_removal(host, options, runner) - { - errors.push(error); - } - if report.host_marketplace_registered - && let Err(error) = run_host_marketplace_removal(host, options, runner) - { - errors.push(error); - } + remove_promoted_replacement(host, layout, snapshot, options, runner, &mut errors); + restore_replaced_paths(snapshot, &mut errors); + if let Some(retirement) = snapshot.generation_retirement.as_mut() + && let Err(error) = retirement.restore_after_rollback() + { + errors.push(error); + } + reconcile_restored_registration(host, snapshot, options, runner, &mut errors); + restore_setup_snapshot(snapshot, setup_runner, &mut errors); + restore_install_state(layout, snapshot.state_bytes.as_deref(), &mut errors); + if errors.is_empty() { + Ok(()) + } else { + Err(errors.join("; ")) + } +} + +fn remove_promoted_replacement( + host: impl MarketplaceHost, + layout: &PluginLayout, + snapshot: &mut ForceInstallSnapshot, + options: &PluginInstallOptions, + runner: &dyn CommandRunner, + errors: &mut Vec, +) { + if !snapshot.replacement_promoted { + return; + } + match host_registration_report(host, options, runner) { + Ok(report) => { + if report.host_plugin_registered + && let Err(error) = run_host_plugin_removal(host, options, runner) + { + errors.push(error); + } + if report.host_marketplace_registered + && let Err(error) = run_host_marketplace_removal(host, options, runner) + { + errors.push(error); } - Err(error) => errors.push(error), - } - if let Err(error) = remove_path(&layout.marketplace_root, options) { - errors.push(error); } - snapshot.replacement_promoted = false; + Err(error) => errors.push(error), + } + if let Err(error) = remove_path(&layout.marketplace_root, options) { + errors.push(error); } + snapshot.replacement_promoted = false; +} + +fn restore_replaced_paths(snapshot: &mut ForceInstallSnapshot, errors: &mut Vec) { if snapshot.marketplace_moved { - if let Err(error) = fs::rename( + match fs::rename( &snapshot.backup_marketplace_root, &snapshot.original_marketplace_root, ) { - errors.push(format!( + Ok(()) => snapshot.marketplace_moved = false, + Err(error) => errors.push(format!( "failed to restore marketplace {}: {error}", snapshot.original_marketplace_root.display() - )); - } else { - snapshot.marketplace_moved = false; + )), } } if snapshot.plugin_moved && let Some(backup_plugin_root) = snapshot.backup_plugin_root.as_ref() { - if let Err(error) = fs::rename(backup_plugin_root, &snapshot.original_plugin_root) { - errors.push(format!( + match fs::rename(backup_plugin_root, &snapshot.original_plugin_root) { + Ok(()) => snapshot.plugin_moved = false, + Err(error) => errors.push(format!( "failed to restore plugin root {} containing generation marker {}: {error}", snapshot.original_plugin_root.display(), snapshot.original_generation_fence.display() - )); - } else { - snapshot.plugin_moved = false; + )), } } - if let Some(retirement) = snapshot.generation_retirement.as_mut() - && let Err(error) = retirement.restore_after_rollback() +} + +fn reconcile_restored_registration( + host: impl MarketplaceHost, + snapshot: &ForceInstallSnapshot, + options: &PluginInstallOptions, + runner: &dyn CommandRunner, + errors: &mut Vec, +) { + let report = match host_registration_report(host, options, runner) { + Ok(report) => report, + Err(error) => { + errors.push(error); + return; + } + }; + if report.host_plugin_registered + && !snapshot.plugin_registered + && let Err(error) = run_host_plugin_removal(host, options, runner) { errors.push(error); } - match host_registration_report(host, options, runner) { - Ok(report) => { - if report.host_plugin_registered - && !snapshot.plugin_registered - && let Err(error) = run_host_plugin_removal(host, options, runner) - { - errors.push(error); - } - if report.host_marketplace_registered - && !snapshot.marketplace_registered - && let Err(error) = run_host_marketplace_removal(host, options, runner) - { - errors.push(error); - } - if snapshot.marketplace_registered - && !report.host_marketplace_registered - && let Err(error) = run_host_marketplace_registration( - host, - &snapshot.original_marketplace_root, - options, - runner, - ) - { - errors.push(error); - } - if snapshot.plugin_registered - && !report.host_plugin_registered - && let Err(error) = run_host_plugin_registration(host, options, runner) - { - errors.push(error); - } - } - Err(error) => errors.push(error), + if report.host_marketplace_registered + && !snapshot.marketplace_registered + && let Err(error) = run_host_marketplace_removal(host, options, runner) + { + errors.push(error); + } + if snapshot.marketplace_registered + && !report.host_marketplace_registered + && let Err(error) = run_host_marketplace_registration( + host, + &snapshot.original_marketplace_root, + options, + runner, + ) + { + errors.push(error); + } + if snapshot.plugin_registered + && !report.host_plugin_registered + && let Err(error) = run_host_plugin_registration(host, options, runner) + { + errors.push(error); } +} + +fn restore_setup_snapshot( + snapshot: &ForceInstallSnapshot, + setup_runner: &dyn PluginSetupRunner, + errors: &mut Vec, +) { if let Some(setup_snapshot) = snapshot.setup_snapshot.as_ref() && let Err(error) = setup_runner.restore_snapshot(setup_snapshot) { errors.push(error); } - if let Some(bytes) = snapshot.state_bytes.as_deref() { +} + +fn restore_install_state(layout: &PluginLayout, bytes: Option<&[u8]>, errors: &mut Vec) { + if let Some(bytes) = bytes { if let Some(parent) = layout.state_path.parent() && let Err(error) = fs::create_dir_all(parent) { @@ -1890,11 +2022,6 @@ fn restore_force_replacement( layout.state_path.display() )); } - if errors.is_empty() { - Ok(()) - } else { - Err(errors.join("; ")) - } } fn force_cleanup_existing_install( diff --git a/crates/cli/src/mcp_environment.rs b/crates/cli/src/mcp_environment.rs index ef27f1890..8446d7cae 100644 --- a/crates/cli/src/mcp_environment.rs +++ b/crates/cli/src/mcp_environment.rs @@ -224,26 +224,7 @@ fn collect_config_names(value: &Value, names: &mut BTreeSet, windows: bo match value { Value::Object(object) => { for (key, value) in object { - match key.as_str() { - "header_env" => { - if let Some(headers) = value.as_object() { - for name in headers.values().filter_map(Value::as_str) { - if !name.is_empty() && !blocked(name) { - insert_name(names, name.to_owned(), windows); - } - } - } - } - "secret_access_key_var" | "session_token_var" => { - if let Some(name) = value.as_str() - && !name.is_empty() - && !blocked(name) - { - insert_name(names, name.to_owned(), windows); - } - } - _ => collect_config_names(value, names, windows), - } + collect_config_field(key, value, names, windows); } } Value::Array(values) => { @@ -254,3 +235,29 @@ fn collect_config_names(value: &Value, names: &mut BTreeSet, windows: bo _ => {} } } + +fn collect_config_field(key: &str, value: &Value, names: &mut BTreeSet, windows: bool) { + match key { + "header_env" => collect_header_env_names(value, names, windows), + "secret_access_key_var" | "session_token_var" => { + if let Some(name) = value.as_str() { + collect_config_name(name, names, windows); + } + } + _ => collect_config_names(value, names, windows), + } +} + +fn collect_header_env_names(value: &Value, names: &mut BTreeSet, windows: bool) { + if let Some(headers) = value.as_object() { + for name in headers.values().filter_map(Value::as_str) { + collect_config_name(name, names, windows); + } + } +} + +fn collect_config_name(name: &str, names: &mut BTreeSet, windows: bool) { + if !name.is_empty() && !blocked(name) { + insert_name(names, name.to_owned(), windows); + } +} diff --git a/crates/cli/src/plugins/lifecycle/environment.rs b/crates/cli/src/plugins/lifecycle/environment.rs index c8a0d572e..0298f25bf 100644 --- a/crates/cli/src/plugins/lifecycle/environment.rs +++ b/crates/cli/src/plugins/lifecycle/environment.rs @@ -429,60 +429,87 @@ fn digest_environment_directory( } children.sort_by_key(std::fs::DirEntry::file_name); for child in children { - let path = child.path(); - let relative = relative_directory.join(child.file_name()); - if relative == Path::new(ENVIRONMENT_ATTESTATION_FILE) - || path.file_name().and_then(|name| name.to_str()) == Some("__pycache__") - || path.extension().and_then(|extension| extension.to_str()) == Some("pyc") - { - continue; - } - let metadata = std::fs::symlink_metadata(&path) - .map_err(|error| format!("failed to inspect {}: {error}", path.display()))?; - let source = if metadata.file_type().is_symlink() { - std::fs::canonicalize(&path) - .map_err(|error| format!("failed to resolve {}: {error}", path.display()))? - } else { - path.clone() - }; - let source_metadata = std::fs::metadata(&source) - .map_err(|error| format!("failed to inspect {}: {error}", source.display()))?; - if source_metadata.is_dir() { - update_tree_digest(digest, b'd', &relative, &[]); - digest_environment_directory( - &source, - &relative, - ancestors, - digest, - total, - entries, - max_entries, - )?; - continue; - } - if !source_metadata.is_file() { - return Err(format!( - "managed Python environment entry {} must resolve to a regular file or directory", - path.display() - )); - } - let bytes = crate::filesystem::bounded::read_bounded_regular_file( - &source, - "managed Python environment file", + digest_environment_entry( + child, + relative_directory, + ancestors, + digest, + total, + entries, + max_entries, )?; - *total = total.saturating_add(bytes.len() as u64); - if *total > crate::filesystem::bounded::MAX_BOUNDED_FILE_BYTES { - return Err(format!( - "managed Python environment exceeds the {}-byte attestation budget", - crate::filesystem::bounded::MAX_BOUNDED_FILE_BYTES - )); - } - update_tree_digest(digest, b'f', &relative, &bytes); } ancestors.pop(); Ok(()) } +fn digest_environment_entry( + child: std::fs::DirEntry, + relative_directory: &Path, + ancestors: &mut Vec, + digest: &mut Sha256, + total: &mut u64, + entries: &mut usize, + max_entries: usize, +) -> Result<(), String> { + let path = child.path(); + let relative = relative_directory.join(child.file_name()); + if environment_entry_is_ignored(&path, &relative) { + return Ok(()); + } + let source = resolve_environment_entry(&path)?; + let metadata = std::fs::metadata(&source) + .map_err(|error| format!("failed to inspect {}: {error}", source.display()))?; + if metadata.is_dir() { + update_tree_digest(digest, b'd', &relative, &[]); + return digest_environment_directory( + &source, + &relative, + ancestors, + digest, + total, + entries, + max_entries, + ); + } + if !metadata.is_file() { + return Err(format!( + "managed Python environment entry {} must resolve to a regular file or directory", + path.display() + )); + } + let bytes = crate::filesystem::bounded::read_bounded_regular_file( + &source, + "managed Python environment file", + )?; + *total = total.saturating_add(bytes.len() as u64); + if *total > crate::filesystem::bounded::MAX_BOUNDED_FILE_BYTES { + return Err(format!( + "managed Python environment exceeds the {}-byte attestation budget", + crate::filesystem::bounded::MAX_BOUNDED_FILE_BYTES + )); + } + update_tree_digest(digest, b'f', &relative, &bytes); + Ok(()) +} + +fn environment_entry_is_ignored(path: &Path, relative: &Path) -> bool { + relative == Path::new(ENVIRONMENT_ATTESTATION_FILE) + || path.file_name().and_then(|name| name.to_str()) == Some("__pycache__") + || path.extension().and_then(|extension| extension.to_str()) == Some("pyc") +} + +fn resolve_environment_entry(path: &Path) -> Result { + let metadata = std::fs::symlink_metadata(path) + .map_err(|error| format!("failed to inspect {}: {error}", path.display()))?; + if metadata.file_type().is_symlink() { + std::fs::canonicalize(path) + .map_err(|error| format!("failed to resolve {}: {error}", path.display())) + } else { + Ok(path.to_path_buf()) + } +} + fn update_tree_digest(digest: &mut Sha256, entry_type: u8, path: &Path, payload: &[u8]) { let path = raw_path_bytes(path); digest.update([entry_type]); diff --git a/crates/cli/src/plugins/lifecycle/mod.rs b/crates/cli/src/plugins/lifecycle/mod.rs index af539db76..103d046d9 100644 --- a/crates/cli/src/plugins/lifecycle/mod.rs +++ b/crates/cli/src/plugins/lifecycle/mod.rs @@ -670,41 +670,14 @@ impl DynamicPluginActivationSnapshot { integrity.signature = Some(copied.to_string_lossy().into_owned()); } - let activation_environment_ref = if matches!( - &manifest.load, - DynamicPluginManifestLoad::Worker(load) - if load.runtime == Some(WorkerRuntime::Python) - ) { - let environment = environment_ref.ok_or_else(|| { - CliError::Config(format!( - "Python worker dynamic plugin '{expected_plugin_id}' has no managed environment" - )) - })?; - let source_artifact_sha256 = trusted_source_artifact_sha256(&manifest)?; - let environment = PathBuf::from(environment); - verify_environment_attestation(&environment, source_artifact_sha256) - .map_err(CliError::Config)?; - let environment_name = environment.file_name().ok_or_else(|| { - CliError::Config(format!( - "managed Python environment {} has no lifecycle environment name", - environment.display() - )) - })?; - let copied_environment = root.join(MANAGED_ENVIRONMENTS_DIR).join(environment_name); - copy_snapshot_directory( - &environment, - &copied_environment, - &mut copied_files, - &mut budget, - true, - &mut Vec::new(), - )?; - verify_environment_attestation(&copied_environment, source_artifact_sha256) - .map_err(CliError::Config)?; - Some(copied_environment.to_string_lossy().into_owned()) - } else { - None - }; + let activation_environment_ref = snapshot_python_environment( + &manifest, + environment_ref, + expected_plugin_id, + &root, + &mut copied_files, + &mut budget, + )?; let activation_manifest = runtime_root.join("relay-plugin.toml"); let rendered = toml::to_string(&manifest).map_err(|error| { @@ -796,6 +769,49 @@ impl DynamicPluginActivationSnapshot { } } +fn snapshot_python_environment( + manifest: &DynamicPluginManifest, + environment_ref: Option<&str>, + expected_plugin_id: &str, + root: &Path, + copied_files: &mut HashMap, + budget: &mut SnapshotBudget, +) -> Result, CliError> { + if !matches!( + &manifest.load, + DynamicPluginManifestLoad::Worker(load) if load.runtime == Some(WorkerRuntime::Python) + ) { + return Ok(None); + } + let environment = environment_ref.ok_or_else(|| { + CliError::Config(format!( + "Python worker dynamic plugin '{expected_plugin_id}' has no managed environment" + )) + })?; + let source_artifact_sha256 = trusted_source_artifact_sha256(manifest)?; + let environment = PathBuf::from(environment); + verify_environment_attestation(&environment, source_artifact_sha256) + .map_err(CliError::Config)?; + let environment_name = environment.file_name().ok_or_else(|| { + CliError::Config(format!( + "managed Python environment {} has no lifecycle environment name", + environment.display() + )) + })?; + let copied_environment = root.join(MANAGED_ENVIRONMENTS_DIR).join(environment_name); + copy_snapshot_directory( + &environment, + &copied_environment, + copied_files, + budget, + true, + &mut Vec::new(), + )?; + verify_environment_attestation(&copied_environment, source_artifact_sha256) + .map_err(CliError::Config)?; + Ok(Some(copied_environment.to_string_lossy().into_owned())) +} + struct SnapshotRootGuard(Option); impl Drop for SnapshotRootGuard { @@ -979,83 +995,128 @@ fn copy_snapshot_directory_contents( budget.record_entries(source, entries.len())?; entries.sort_by_key(fs::DirEntry::file_name); for entry in entries { - let source_path = entry.path(); - if skip_python_cache - && (entry.file_name() == "__pycache__" - || source_path.extension().and_then(|value| value.to_str()) == Some("pyc")) - { - continue; - } - let destination_path = destination.join(entry.file_name()); - let metadata = fs::symlink_metadata(&source_path) - .map_err(|error| CliError::Config(error.to_string()))?; - let resolved = if metadata.file_type().is_symlink() { - fs::canonicalize(&source_path).map_err(|error| { - CliError::Config(format!( - "failed to resolve dynamic plugin runtime symlink {}: {error}", - source_path.display() - )) - })? - } else { - source_path.clone() - }; - let resolved_metadata = - fs::metadata(&resolved).map_err(|error| CliError::Config(error.to_string()))?; - if resolved_metadata.is_dir() { - copy_snapshot_directory_contents( - &resolved, - &destination_path, - copied_files, - budget, - skip_python_cache, - ancestors, - )?; - } else if resolved_metadata.is_file() { - // A macOS venv's `bin/python` is normally an absolute symlink to the managed - // interpreter. Dereferencing that link while snapshotting turns the interpreter - // into a standalone file whose @rpath no longer points at libpython, so the worker - // exits before it can create its socket. Preserve only these launcher links; all - // other runtime symlinks remain dereferenced to keep the activation snapshot - // self-contained and deterministic. - #[cfg(unix)] - if metadata.file_type().is_symlink() && is_python_venv_launcher(&source_path) { - let target = fs::read_link(&source_path).map_err(|error| { - CliError::Config(format!( - "failed to read Python venv launcher symlink {}: {error}", - source_path.display() - )) - })?; - if let Some(parent) = destination_path.parent() { - fs::create_dir_all(parent) - .map_err(|error| CliError::Config(error.to_string()))?; - } - std::os::unix::fs::symlink(&target, &destination_path).map_err(|error| { - CliError::Config(format!( - "failed to preserve Python venv launcher symlink {}: {error}", - destination_path.display() - )) - })?; - copied_files.insert(resolved, destination_path); - continue; - } - copy_snapshot_regular_file( - &resolved, - &destination_path, - copied_files, - budget, - "runtime file", - )?; - } else { - return Err(CliError::Config(format!( - "dynamic plugin runtime entry {} must resolve to a regular file or directory", - source_path.display() - ))); - } + copy_snapshot_entry( + entry, + destination, + copied_files, + budget, + skip_python_cache, + ancestors, + )?; } ancestors.pop(); Ok(()) } +fn copy_snapshot_entry( + entry: fs::DirEntry, + destination: &Path, + copied_files: &mut HashMap, + budget: &mut SnapshotBudget, + skip_python_cache: bool, + ancestors: &mut Vec, +) -> Result<(), CliError> { + let source_path = entry.path(); + if skip_python_cache + && (entry.file_name() == "__pycache__" + || source_path.extension().and_then(|value| value.to_str()) == Some("pyc")) + { + return Ok(()); + } + let destination_path = destination.join(entry.file_name()); + let metadata = + fs::symlink_metadata(&source_path).map_err(|error| CliError::Config(error.to_string()))?; + let resolved = resolve_snapshot_entry(&source_path, &metadata)?; + let resolved_metadata = + fs::metadata(&resolved).map_err(|error| CliError::Config(error.to_string()))?; + if resolved_metadata.is_dir() { + return copy_snapshot_directory_contents( + &resolved, + &destination_path, + copied_files, + budget, + skip_python_cache, + ancestors, + ); + } + if !resolved_metadata.is_file() { + return Err(CliError::Config(format!( + "dynamic plugin runtime entry {} must resolve to a regular file or directory", + source_path.display() + ))); + } + if preserve_python_launcher( + &source_path, + &destination_path, + &resolved, + &metadata, + copied_files, + )? { + return Ok(()); + } + copy_snapshot_regular_file( + &resolved, + &destination_path, + copied_files, + budget, + "runtime file", + ) +} + +fn resolve_snapshot_entry(path: &Path, metadata: &fs::Metadata) -> Result { + if metadata.file_type().is_symlink() { + fs::canonicalize(path).map_err(|error| { + CliError::Config(format!( + "failed to resolve dynamic plugin runtime symlink {}: {error}", + path.display() + )) + }) + } else { + Ok(path.to_path_buf()) + } +} + +#[cfg(unix)] +fn preserve_python_launcher( + source: &Path, + destination: &Path, + resolved: &Path, + metadata: &fs::Metadata, + copied_files: &mut HashMap, +) -> Result { + if !metadata.file_type().is_symlink() || !is_python_venv_launcher(source) { + return Ok(false); + } + let target = fs::read_link(source).map_err(|error| { + CliError::Config(format!( + "failed to read Python venv launcher symlink {}: {error}", + source.display() + )) + })?; + if let Some(parent) = destination.parent() { + fs::create_dir_all(parent).map_err(|error| CliError::Config(error.to_string()))?; + } + std::os::unix::fs::symlink(&target, destination).map_err(|error| { + CliError::Config(format!( + "failed to preserve Python venv launcher symlink {}: {error}", + destination.display() + )) + })?; + copied_files.insert(resolved.to_path_buf(), destination.to_path_buf()); + Ok(true) +} + +#[cfg(not(unix))] +fn preserve_python_launcher( + _source: &Path, + _destination: &Path, + _resolved: &Path, + _metadata: &fs::Metadata, + _copied_files: &mut HashMap, +) -> Result { + Ok(false) +} + #[cfg(unix)] fn is_python_venv_launcher(path: &Path) -> bool { let Some(parent) = path.parent() else { diff --git a/crates/cli/src/plugins/schema/secrets.rs b/crates/cli/src/plugins/schema/secrets.rs index af608cadc..e54ea8a08 100644 --- a/crates/cli/src/plugins/schema/secrets.rs +++ b/crates/cli/src/plugins/schema/secrets.rs @@ -117,56 +117,92 @@ impl SecretPattern { } match &self.0[offset] { SecretSegment::Property(property) => { - if let Some(child) = value.get_mut(property) { - self.visit_matching_values(child, offset + 1, visit); - } + self.visit_property(value, property, offset, visit) } - SecretSegment::Any => match value { - Value::Object(object) => { - for child in object.values_mut() { - self.visit_matching_values(child, offset + 1, visit); - } - } - Value::Array(values) => { - for child in values { - self.visit_matching_values(child, offset + 1, visit); - } - } - _ => {} - }, + SecretSegment::Any => self.visit_any(value, offset, visit), SecretSegment::Pattern(pattern) => { - if let Value::Object(object) = value { - for (key, child) in object { - if pattern_matches(pattern, key) { - self.visit_matching_values(child, offset + 1, visit); - } - } - } + self.visit_object_matches(value, offset, visit, |key| { + pattern_matches(pattern, key) + }); } SecretSegment::UnmatchedProperties(selector) => { - if let Value::Object(object) = value { - for (key, child) in object { - if selector.matches(key) { - self.visit_matching_values(child, offset + 1, visit); - } - } + self.visit_object_matches(value, offset, visit, |key| selector.matches(key)); + } + SecretSegment::Index(index) => self.visit_index(value, *index, offset, visit), + SecretSegment::Tail(start) => self.visit_tail(value, *start, offset, visit), + } + } + + fn visit_property( + &self, + value: &mut Value, + property: &str, + offset: usize, + visit: &mut impl FnMut(&mut Value), + ) { + if let Some(child) = value.get_mut(property) { + self.visit_matching_values(child, offset + 1, visit); + } + } + + fn visit_any(&self, value: &mut Value, offset: usize, visit: &mut impl FnMut(&mut Value)) { + match value { + Value::Object(object) => { + for child in object.values_mut() { + self.visit_matching_values(child, offset + 1, visit); } } - SecretSegment::Index(index) => { - if let Some(child) = value.get_mut(*index) { + Value::Array(values) => { + for child in values { self.visit_matching_values(child, offset + 1, visit); } } - SecretSegment::Tail(start) => { - if let Value::Array(values) = value { - for child in values.iter_mut().skip(*start) { - self.visit_matching_values(child, offset + 1, visit); - } + _ => {} + } + } + + fn visit_object_matches( + &self, + value: &mut Value, + offset: usize, + visit: &mut impl FnMut(&mut Value), + matches: impl Fn(&str) -> bool, + ) { + if let Value::Object(object) = value { + for (key, child) in object { + if matches(key) { + self.visit_matching_values(child, offset + 1, visit); } } } } + fn visit_index( + &self, + value: &mut Value, + index: usize, + offset: usize, + visit: &mut impl FnMut(&mut Value), + ) { + if let Some(child) = value.get_mut(index) { + self.visit_matching_values(child, offset + 1, visit); + } + } + + fn visit_tail( + &self, + value: &mut Value, + start: usize, + offset: usize, + visit: &mut impl FnMut(&mut Value), + ) { + if let Value::Array(values) = value { + for child in values.iter_mut().skip(start) { + self.visit_matching_values(child, offset + 1, visit); + } + } + } + pub(super) fn applies_below(&self, path: &[String]) -> bool { self.0.len() >= path.len() && self diff --git a/crates/cli/src/process/detached.rs b/crates/cli/src/process/detached.rs index 9b2cc2ca3..c5b061fb3 100644 --- a/crates/cli/src/process/detached.rs +++ b/crates/cli/src/process/detached.rs @@ -14,6 +14,77 @@ use std::sync::Mutex; #[cfg(windows)] static SIDECAR_SPAWN_LOCK: Mutex<()> = Mutex::new(()); +#[cfg(windows)] +fn wide_nul(value: &std::ffi::OsStr) -> Vec { + use std::os::windows::ffi::OsStrExt; + value.encode_wide().chain(std::iter::once(0)).collect() +} + +#[cfg(windows)] +fn append_quoted(command_line: &mut Vec, value: &std::ffi::OsStr) { + use std::os::windows::ffi::OsStrExt; + + const BACKSLASH: u16 = b'\\' as u16; + const QUOTE: u16 = b'"' as u16; + command_line.push(QUOTE); + let mut slashes = 0; + for unit in value.encode_wide() { + if unit == BACKSLASH { + slashes += 1; + continue; + } + if unit == QUOTE { + command_line.extend(std::iter::repeat_n(BACKSLASH, slashes * 2 + 1)); + } else { + command_line.extend(std::iter::repeat_n(BACKSLASH, slashes)); + } + slashes = 0; + command_line.push(unit); + } + command_line.extend(std::iter::repeat_n(BACKSLASH, slashes * 2)); + command_line.push(QUOTE); +} + +#[cfg(windows)] +fn environment_block(command: &Command) -> Vec { + use std::collections::BTreeMap; + use std::ffi::OsString; + use std::os::windows::ffi::OsStrExt; + + let mut environment = BTreeMap::::new(); + for (name, value) in std::env::vars_os() { + environment.insert(name.to_string_lossy().to_uppercase(), (name, value)); + } + for (name, value) in command.get_envs() { + let key = name.to_string_lossy().to_uppercase(); + if let Some(value) = value { + environment.insert(key, (name.to_owned(), value.to_owned())); + } else { + environment.remove(&key); + } + } + let mut block = Vec::new(); + for (_, (name, value)) in environment { + block.extend(name.encode_wide()); + block.push(b'=' as u16); + block.extend(value.encode_wide()); + block.push(0); + } + block.push(0); + block +} + +#[cfg(windows)] +struct AttributeList(*mut std::ffi::c_void); + +#[cfg(windows)] +impl Drop for AttributeList { + fn drop(&mut self) { + // SAFETY: The list was initialized successfully and is still live. + unsafe { windows_sys::Win32::System::Threading::DeleteProcThreadAttributeList(self.0) }; + } +} + #[cfg(not(windows))] pub(crate) type DetachedChild = Child; @@ -83,77 +154,16 @@ impl Drop for DetachedChild { #[cfg(windows)] fn spawn_detached_with_handle_list(command: &Command) -> std::io::Result { - use std::collections::BTreeMap; - use std::ffi::{OsStr, OsString, c_void}; + use std::ffi::c_void; use std::fs::OpenOptions; - use std::os::windows::ffi::OsStrExt; use std::os::windows::io::AsRawHandle; use windows_sys::Win32::Foundation::{HANDLE, HANDLE_FLAG_INHERIT, SetHandleInformation}; use windows_sys::Win32::System::Threading::{ - CREATE_UNICODE_ENVIRONMENT, CreateProcessW, DeleteProcThreadAttributeList, - EXTENDED_STARTUPINFO_PRESENT, InitializeProcThreadAttributeList, - PROC_THREAD_ATTRIBUTE_HANDLE_LIST, PROCESS_INFORMATION, STARTF_USESTDHANDLES, - STARTUPINFOEXW, UpdateProcThreadAttribute, + CREATE_UNICODE_ENVIRONMENT, CreateProcessW, EXTENDED_STARTUPINFO_PRESENT, + InitializeProcThreadAttributeList, PROC_THREAD_ATTRIBUTE_HANDLE_LIST, PROCESS_INFORMATION, + STARTF_USESTDHANDLES, STARTUPINFOEXW, UpdateProcThreadAttribute, }; - fn wide_nul(value: &OsStr) -> Vec { - value.encode_wide().chain(std::iter::once(0)).collect() - } - - fn append_quoted(command_line: &mut Vec, value: &OsStr) { - const BACKSLASH: u16 = b'\\' as u16; - const QUOTE: u16 = b'"' as u16; - command_line.push(QUOTE); - let mut slashes = 0; - for unit in value.encode_wide() { - if unit == BACKSLASH { - slashes += 1; - continue; - } - if unit == QUOTE { - command_line.extend(std::iter::repeat_n(BACKSLASH, slashes * 2 + 1)); - } else { - command_line.extend(std::iter::repeat_n(BACKSLASH, slashes)); - } - slashes = 0; - command_line.push(unit); - } - command_line.extend(std::iter::repeat_n(BACKSLASH, slashes * 2)); - command_line.push(QUOTE); - } - - fn environment_block(command: &Command) -> Vec { - let mut environment = BTreeMap::::new(); - for (name, value) in std::env::vars_os() { - environment.insert(name.to_string_lossy().to_uppercase(), (name, value)); - } - for (name, value) in command.get_envs() { - let key = name.to_string_lossy().to_uppercase(); - if let Some(value) = value { - environment.insert(key, (name.to_owned(), value.to_owned())); - } else { - environment.remove(&key); - } - } - let mut block = Vec::new(); - for (_, (name, value)) in environment { - block.extend(name.encode_wide()); - block.push(b'=' as u16); - block.extend(value.encode_wide()); - block.push(0); - } - block.push(0); - block - } - - struct AttributeList(*mut c_void); - impl Drop for AttributeList { - fn drop(&mut self) { - // SAFETY: The list was initialized successfully and is still live. - unsafe { DeleteProcThreadAttributeList(self.0) }; - } - } - let stdin = OpenOptions::new().read(true).open(r"\\.\NUL")?; let stdout = OpenOptions::new().write(true).open(r"\\.\NUL")?; let handles = [ diff --git a/crates/cli/src/process/supervision/unix.rs b/crates/cli/src/process/supervision/unix.rs index e4101da18..6bd2f3137 100644 --- a/crates/cli/src/process/supervision/unix.rs +++ b/crates/cli/src/process/supervision/unix.rs @@ -70,15 +70,8 @@ pub(super) async fn wait( tokio::task::yield_now().await; continue; } - if tree.terminal.is_some() && child_is_stopped(tree.process_group)? { - tree.restore_terminal()?; - tree.stop_supervisor_group()?; - if let Some(status) = child.try_wait()? { - return Ok(status); - } - if let Some(terminal) = tree.terminal.as_mut() { - terminal.resume_after_supervisor()?; - } + if let Some(status) = handle_stopped_child(tree, child)? { + return Ok(status); } if termination_deadline.is_some_and(|deadline| tokio::time::Instant::now() >= deadline) { tree.terminate(child)?; @@ -99,6 +92,24 @@ pub(super) async fn wait( } } +fn handle_stopped_child( + tree: &mut ProcessTree, + child: &mut tokio::process::Child, +) -> std::io::Result> { + if tree.terminal.is_none() || !child_is_stopped(tree.process_group)? { + return Ok(None); + } + tree.restore_terminal()?; + tree.stop_supervisor_group()?; + if let Some(status) = child.try_wait()? { + return Ok(Some(status)); + } + if let Some(terminal) = tree.terminal.as_mut() { + terminal.resume_after_supervisor()?; + } + Ok(None) +} + impl ProcessTree { pub(super) fn restore_terminal(&mut self) -> std::io::Result<()> { self.terminal diff --git a/crates/cli/src/process/supervision/windows.rs b/crates/cli/src/process/supervision/windows.rs index 759d7a312..8f018d14f 100644 --- a/crates/cli/src/process/supervision/windows.rs +++ b/crates/cli/src/process/supervision/windows.rs @@ -135,7 +135,6 @@ fn resume_suspended_process(process_id: u32) -> std::io::Result<()> { use windows_sys::Win32::System::Diagnostics::ToolHelp::{ CreateToolhelp32Snapshot, TH32CS_SNAPTHREAD, THREADENTRY32, Thread32First, Thread32Next, }; - use windows_sys::Win32::System::Threading::{OpenThread, ResumeThread, THREAD_SUSPEND_RESUME}; // SAFETY: A system-wide thread snapshot does not borrow caller memory. let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0) }; @@ -152,42 +151,10 @@ fn resume_suspended_process(process_id: u32) -> std::io::Result<()> { let mut has_entry = unsafe { Thread32First(snapshot, &mut entry) } != 0; while has_entry { if entry.th32OwnerProcessID == process_id { - // SAFETY: The snapshot supplied this live thread identifier and only resume access is - // requested. - let thread = unsafe { OpenThread(THREAD_SUSPEND_RESUME, 0, entry.th32ThreadID) }; - if thread.is_null() { - let error = - last_windows_error("failed to open the suspended coding-agent primary thread"); - // SAFETY: `snapshot` is uniquely owned and closed exactly once on this path. - unsafe { CloseHandle(snapshot) }; - return Err(error); - } - // CREATE_SUSPENDED starts the primary thread with a suspend count of one. Resume until - // that count reaches zero, while rejecting a zero count that would imply the process - // had already run before Job Object assignment. - // SAFETY: `thread` is live and was opened with THREAD_SUSPEND_RESUME. - let mut previous_count = unsafe { ResumeThread(thread) }; - while previous_count > 1 && previous_count != u32::MAX { - // SAFETY: The same live thread handle remains owned by this function. - previous_count = unsafe { ResumeThread(thread) }; - } - let resume_error = if previous_count == u32::MAX { - Some(last_windows_error( - "failed to resume the Job-owned coding-agent process", - )) - } else if previous_count == 0 { - Some(std::io::Error::other( - "coding-agent primary thread was not suspended before Job Object assignment", - )) - } else { - None - }; - // SAFETY: Both handles are uniquely owned and closed exactly once on this path. - unsafe { - CloseHandle(thread); - CloseHandle(snapshot); - } - return resume_error.map_or(Ok(()), Err); + let result = resume_thread(entry.th32ThreadID); + // SAFETY: `snapshot` is uniquely owned and closed exactly once on this path. + unsafe { CloseHandle(snapshot) }; + return result; } // SAFETY: `snapshot` and `entry` remain valid for the next enumeration result. has_entry = unsafe { Thread32Next(snapshot, &mut entry) } != 0; @@ -199,6 +166,41 @@ fn resume_suspended_process(process_id: u32) -> std::io::Result<()> { ))) } +fn resume_thread(thread_id: u32) -> std::io::Result<()> { + use windows_sys::Win32::Foundation::CloseHandle; + use windows_sys::Win32::System::Threading::{OpenThread, ResumeThread, THREAD_SUSPEND_RESUME}; + + // SAFETY: The system snapshot supplied this live thread identifier and only resume access is + // requested. + let thread = unsafe { OpenThread(THREAD_SUSPEND_RESUME, 0, thread_id) }; + if thread.is_null() { + return Err(last_windows_error( + "failed to open the suspended coding-agent primary thread", + )); + } + // CREATE_SUSPENDED starts the primary thread with a suspend count of one. Resume until that + // count reaches zero, while rejecting a zero count that would imply the process had already + // run before Job Object assignment. + // SAFETY: `thread` is live and was opened with THREAD_SUSPEND_RESUME. + let mut previous_count = unsafe { ResumeThread(thread) }; + while previous_count > 1 && previous_count != u32::MAX { + // SAFETY: The same live thread handle remains owned by this function. + previous_count = unsafe { ResumeThread(thread) }; + } + let result = match previous_count { + u32::MAX => Err(last_windows_error( + "failed to resume the Job-owned coding-agent process", + )), + 0 => Err(std::io::Error::other( + "coding-agent primary thread was not suspended before Job Object assignment", + )), + _ => Ok(()), + }; + // SAFETY: `thread` is uniquely owned and closed exactly once. + unsafe { CloseHandle(thread) }; + result +} + fn last_windows_error(context: &str) -> std::io::Error { let source = std::io::Error::last_os_error(); std::io::Error::new(source.kind(), format!("{context}: {source}")) diff --git a/crates/cli/src/server/mod.rs b/crates/cli/src/server/mod.rs index 313f6ef65..2303a2de1 100644 --- a/crates/cli/src/server/mod.rs +++ b/crates/cli/src/server/mod.rs @@ -296,13 +296,36 @@ async fn serve_listener_with_dynamic_inner( if let Some(path) = ready_file { write_ready_file(path, local_address, &instance_id)?; } - let idle_shutdown = if matches!(&shutdown_mode, None | Some(ShutdownMode::ProcessSignal)) { - plugin_idle_timeout()? - .map(|timeout| idle_shutdown_future(last_activity, sessions.clone(), timeout)) - } else { - None + let idle_shutdown: Option = + if matches!(&shutdown_mode, None | Some(ShutdownMode::ProcessSignal)) { + plugin_idle_timeout()?.map(|timeout| { + Box::pin(idle_shutdown_future( + last_activity, + sessions.clone(), + timeout, + )) as ShutdownFuture + }) + } else { + None + }; + let shutdown = server_shutdown_future(shutdown_mode, idle_shutdown); + let shutdown = combine_shutdown_futures(shutdown, bootstrap_shutdown_rx); + let serve_result = match shutdown { + Some(shutdown) => { + axum::serve(listener, app) + .with_graceful_shutdown(shutdown) + .await + } + None => axum::serve(listener, app).await, }; - let shutdown: Option = match shutdown_mode { + finish_server_shutdown(serve_result, &sessions, plugin_activation).await +} + +fn server_shutdown_future( + shutdown_mode: Option, + idle_shutdown: Option, +) -> Option { + match shutdown_mode { Some(ShutdownMode::Receiver(receiver)) => Some(Box::pin(async move { let _ = receiver.await; })), @@ -316,9 +339,15 @@ async fn serve_listener_with_dynamic_inner( shutdown_signal().await; } })), - None => idle_shutdown.map(|idle| Box::pin(idle) as ShutdownFuture), - }; - let shutdown = match (shutdown, bootstrap_shutdown_rx) { + None => idle_shutdown, + } +} + +fn combine_shutdown_futures( + shutdown: Option, + bootstrap_shutdown_rx: Option>, +) -> Option { + match (shutdown, bootstrap_shutdown_rx) { (Some(shutdown), Some(receiver)) => Some(Box::pin(async move { tokio::select! { _ = shutdown => {} @@ -329,15 +358,14 @@ async fn serve_listener_with_dynamic_inner( let _ = receiver.await; }) as ShutdownFuture), (shutdown, None) => shutdown, - }; - let serve_result = match shutdown { - Some(shutdown) => { - axum::serve(listener, app) - .with_graceful_shutdown(shutdown) - .await - } - None => axum::serve(listener, app).await, - }; + } +} + +async fn finish_server_shutdown( + serve_result: std::io::Result<()>, + sessions: &SessionManager, + plugin_activation: Option, +) -> Result<(), CliError> { let close_result = sessions.close_all("gateway_shutdown").await; let flush_result = nemo_relay::api::runtime::flush_subscribers().map_err(CliError::from); let clear_result = plugin_activation diff --git a/crates/cli/src/sessions/idle.rs b/crates/cli/src/sessions/idle.rs index fc96b8ae1..70fe45fb8 100644 --- a/crates/cli/src/sessions/idle.rs +++ b/crates/cli/src/sessions/idle.rs @@ -40,41 +40,66 @@ pub(super) async fn close_idle_sessions_from_parts( timeout: Duration, reason: &str, ) -> Result { - let mut idle_sessions = Vec::new(); - { - let mut sessions = inner.lock().await; - let ids = sessions - .iter() - .filter_map(|(session_id, session)| { - session - .is_idle_for(now, timeout) - .then_some(session_id.clone()) - }) - .collect::>(); - for session_id in ids { - if let Some(session) = sessions.remove(&session_id) { - idle_sessions.push((session_id, session)); - } - } - } + let idle_sessions = take_idle_sessions(inner, now, timeout).await; if idle_sessions.is_empty() { return Ok(0); } + let (closed_turns, closed_subagents, retained_sessions, first_error) = + close_idle_turns(idle_sessions, reason).await; + let cleanup_sessions = + restore_retained_sessions(inner, retained_sessions, &closed_subagents).await; + clear_closed_subagents(alignment, closed_subagents, &cleanup_sessions).await; + first_error.map_or(Ok(closed_turns), Err) +} + +async fn take_idle_sessions( + inner: &Arc>>, + now: Instant, + timeout: Duration, +) -> Vec<(String, Session)> { + let mut sessions = inner.lock().await; + let ids = sessions + .iter() + .filter_map(|(session_id, session)| { + session + .is_idle_for(now, timeout) + .then_some(session_id.clone()) + }) + .collect::>(); + ids.into_iter() + .filter_map(|session_id| { + sessions + .remove(&session_id) + .map(|session| (session_id, session)) + }) + .collect() +} + +type ClosedIdleTurns = ( + usize, + Vec<(String, String)>, + Vec<(String, Session)>, + Option, +); + +async fn close_idle_turns(idle_sessions: Vec<(String, Session)>, reason: &str) -> ClosedIdleTurns { let mut closed_turns = 0; let mut closed_subagents = Vec::new(); let mut retained_sessions = Vec::new(); let mut first_error = None; for (session_id, mut session) in idle_sessions { let stack = session.scope_stack.clone(); - let result = TASK_SCOPE_STACK + match TASK_SCOPE_STACK .scope(stack, async { session.close_turn_for_reason(reason).await }) - .await; - match result { + .await + { Ok(subagent_ids) => { closed_turns += 1; - for subagent_id in subagent_ids { - closed_subagents.push((session_id.clone(), subagent_id)); - } + closed_subagents.extend( + subagent_ids + .into_iter() + .map(|subagent_id| (session_id.clone(), subagent_id)), + ); } Err(error) if first_error.is_none() => first_error = Some(error), Err(_) => {} @@ -83,28 +108,47 @@ pub(super) async fn close_idle_sessions_from_parts( retained_sessions.push((session_id, session)); } } - let mut alignment_cleanup_sessions = HashSet::new(); - { - let mut sessions = inner.lock().await; - for (session_id, session) in retained_sessions { - if let Entry::Vacant(entry) = sessions.entry(session_id.clone()) { - entry.insert(session); - alignment_cleanup_sessions.insert(session_id); - } + ( + closed_turns, + closed_subagents, + retained_sessions, + first_error, + ) +} + +async fn restore_retained_sessions( + inner: &Arc>>, + retained_sessions: Vec<(String, Session)>, + closed_subagents: &[(String, String)], +) -> HashSet { + let mut cleanup_sessions = HashSet::new(); + let mut sessions = inner.lock().await; + for (session_id, session) in retained_sessions { + if let Entry::Vacant(entry) = sessions.entry(session_id.clone()) { + entry.insert(session); + cleanup_sessions.insert(session_id); } - for (session_id, _) in &closed_subagents { - if !sessions.contains_key(session_id) { - alignment_cleanup_sessions.insert(session_id.clone()); - } + } + for (session_id, _) in closed_subagents { + if !sessions.contains_key(session_id) { + cleanup_sessions.insert(session_id.clone()); } } - if !closed_subagents.is_empty() { - let mut alignment_state = alignment.lock().await; - for (session_id, subagent_id) in closed_subagents { - if alignment_cleanup_sessions.contains(&session_id) { - alignment_state.clear_for_ended_subagent(&session_id, &subagent_id); - } + cleanup_sessions +} + +async fn clear_closed_subagents( + alignment: &Arc>, + closed_subagents: Vec<(String, String)>, + cleanup_sessions: &HashSet, +) { + if closed_subagents.is_empty() { + return; + } + let mut alignment_state = alignment.lock().await; + for (session_id, subagent_id) in closed_subagents { + if cleanup_sessions.contains(&session_id) { + alignment_state.clear_for_ended_subagent(&session_id, &subagent_id); } } - first_error.map_or(Ok(closed_turns), Err) } diff --git a/crates/cli/tests/cli_tests.rs b/crates/cli/tests/cli_tests.rs index f34986490..b76c6dd7f 100644 --- a/crates/cli/tests/cli_tests.rs +++ b/crates/cli/tests/cli_tests.rs @@ -771,116 +771,21 @@ fn run_fake_bootstrap_listener_with_options( let request = read_http_request(&mut stream); server_requests.lock().unwrap().push(request.clone()); if request.starts_with("GET /healthz ") { - let fingerprint = - bootstrap_request_header(&request, "x-nemo-relay-bootstrap-fingerprint") - .unwrap(); - let nonce = - bootstrap_request_header(&request, "x-nemo-relay-bootstrap-nonce").unwrap(); - let proof_header = match proof { - FakeBootstrapProof::Missing => String::new(), - FakeBootstrapProof::Wrong => { - "X-NeMo-Relay-Bootstrap-Proof: hmac-sha256:0000000000000000000000000000000000000000000000000000000000000000\r\n".into() - } - FakeBootstrapProof::Valid => format!( - "X-NeMo-Relay-Bootstrap-Proof: {}\r\n", - fake_bootstrap_proof( - &std::fs::read(&key_path).unwrap(), - fingerprint, - nonce - ) - ), - }; - let body = format!( - r#"{{"status":"ok","service":"nemo-relay","version":"{}","bootstrap_protocol":2,"instance_id":"test-instance"}}"#, - env!("CARGO_PKG_VERSION") - ); - stream - .write_all( - format!( - "HTTP/1.1 200 OK\r\n{proof_header}Content-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ) - .as_bytes(), - ) - .unwrap(); + write_fake_bootstrap_health(&mut stream, &request, proof, &key_path); continue; } if request.starts_with("GET /bootstrap/tunnel ") { - let fingerprint = - bootstrap_request_header(&request, "x-nemo-relay-bootstrap-fingerprint") - .unwrap(); - let nonce = - bootstrap_request_header(&request, "x-nemo-relay-bootstrap-nonce").unwrap(); - let proof_header = match proof { - FakeBootstrapProof::Missing => String::new(), - FakeBootstrapProof::Wrong => { - "X-NeMo-Relay-Bootstrap-Proof: hmac-sha256:0000000000000000000000000000000000000000000000000000000000000000\r\n".into() - } - FakeBootstrapProof::Valid => { - let key = std::fs::read(&key_path).unwrap(); - format!( - "X-NeMo-Relay-Bootstrap-Proof: {}\r\n", - fake_bootstrap_proof(&key, fingerprint, nonce) - ) - } - }; - stream - .write_all( - format!( - "HTTP/1.1 101 Switching Protocols\r\n{proof_header}Connection: upgrade\r\nUpgrade: nemo-relay-tls\r\nContent-Length: 0\r\n\r\n" - ) - .as_bytes(), - ) - .unwrap(); - if matches!(proof, FakeBootstrapProof::Valid) { - let connection = rustls::ServerConnection::new(tls.clone()).unwrap(); - let mut stream = rustls::StreamOwned::new(connection, stream); - let health = read_http_request(&mut stream); - server_requests.lock().unwrap().push(health); - let body = format!( - r#"{{"status":"ok","service":"nemo-relay","version":"{}","bootstrap_protocol":2,"instance_id":"test-instance"}}"#, - env!("CARGO_PKG_VERSION") - ); - stream - .write_all( - format!( - "HTTP/1.1 200 OK\r\nX-NeMo-Relay-Bootstrap-Proof: {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n{body}", - fake_bootstrap_proof( - &std::fs::read(&key_path).unwrap(), - fingerprint, - nonce - ), - body.len() - ) - .as_bytes(), - ) - .unwrap(); - let hook = read_http_request(&mut stream); - server_requests.lock().unwrap().push(hook); - if let Some(delay) = hook_delay { - thread::sleep(delay); - } - let body = r#"{"continue":true}"#; - let _ = stream.write_all( - format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ) - .as_bytes(), - ); - } - } else { - if let Some(delay) = hook_delay { - thread::sleep(delay); - } - let body = r#"{"continue":true}"#; - let _ = stream.write_all( - format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ) - .as_bytes(), + handle_fake_bootstrap_tunnel( + stream, + &request, + proof, + &key_path, + &tls, + &server_requests, + hook_delay, ); + } else { + write_fake_hook_response(&mut stream, hook_delay); } } }); @@ -925,6 +830,121 @@ fn run_fake_bootstrap_listener_with_options( (output, requests) } +fn fake_bootstrap_proof_header( + proof: FakeBootstrapProof, + key_path: &Path, + fingerprint: &str, + nonce: &str, +) -> String { + match proof { + FakeBootstrapProof::Missing => String::new(), + FakeBootstrapProof::Wrong => { + "X-NeMo-Relay-Bootstrap-Proof: hmac-sha256:0000000000000000000000000000000000000000000000000000000000000000\r\n".into() + } + FakeBootstrapProof::Valid => format!( + "X-NeMo-Relay-Bootstrap-Proof: {}\r\n", + fake_bootstrap_proof( + &std::fs::read(key_path).unwrap(), + fingerprint, + nonce + ) + ), + } +} + +fn write_fake_bootstrap_health( + stream: &mut std::net::TcpStream, + request: &str, + proof: FakeBootstrapProof, + key_path: &Path, +) { + let fingerprint = + bootstrap_request_header(request, "x-nemo-relay-bootstrap-fingerprint").unwrap(); + let nonce = bootstrap_request_header(request, "x-nemo-relay-bootstrap-nonce").unwrap(); + let proof_header = fake_bootstrap_proof_header(proof, key_path, fingerprint, nonce); + let body = format!( + r#"{{"status":"ok","service":"nemo-relay","version":"{}","bootstrap_protocol":2,"instance_id":"test-instance"}}"#, + env!("CARGO_PKG_VERSION") + ); + stream + .write_all( + format!( + "HTTP/1.1 200 OK\r\n{proof_header}Content-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .unwrap(); +} + +#[allow(clippy::too_many_arguments)] +fn handle_fake_bootstrap_tunnel( + mut stream: std::net::TcpStream, + request: &str, + proof: FakeBootstrapProof, + key_path: &Path, + tls: &Arc, + requests: &Arc>>, + hook_delay: Option, +) { + let fingerprint = + bootstrap_request_header(request, "x-nemo-relay-bootstrap-fingerprint").unwrap(); + let nonce = bootstrap_request_header(request, "x-nemo-relay-bootstrap-nonce").unwrap(); + let proof_header = fake_bootstrap_proof_header(proof, key_path, fingerprint, nonce); + stream + .write_all( + format!( + "HTTP/1.1 101 Switching Protocols\r\n{proof_header}Connection: upgrade\r\nUpgrade: nemo-relay-tls\r\nContent-Length: 0\r\n\r\n" + ) + .as_bytes(), + ) + .unwrap(); + if !matches!(proof, FakeBootstrapProof::Valid) { + return; + } + let connection = rustls::ServerConnection::new(tls.clone()).unwrap(); + let mut stream = rustls::StreamOwned::new(connection, stream); + let health = read_http_request(&mut stream); + requests.lock().unwrap().push(health); + let body = format!( + r#"{{"status":"ok","service":"nemo-relay","version":"{}","bootstrap_protocol":2,"instance_id":"test-instance"}}"#, + env!("CARGO_PKG_VERSION") + ); + stream + .write_all( + format!( + "HTTP/1.1 200 OK\r\nX-NeMo-Relay-Bootstrap-Proof: {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n{body}", + fake_bootstrap_proof( + &std::fs::read(key_path).unwrap(), + fingerprint, + nonce + ), + body.len() + ) + .as_bytes(), + ) + .unwrap(); + requests + .lock() + .unwrap() + .push(read_http_request(&mut stream)); + write_fake_hook_response(&mut stream, hook_delay); +} + +fn write_fake_hook_response(stream: &mut impl Write, hook_delay: Option) { + if let Some(delay) = hook_delay { + thread::sleep(delay); + } + let body = r#"{"continue":true}"#; + let _ = stream.write_all( + format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ); +} + #[test] fn cli_codex_hook_rejects_compatible_json_without_bootstrap_proof() { let (output, requests) = run_fake_bootstrap_listener(FakeBootstrapProof::Missing); @@ -3360,68 +3380,93 @@ while True: (libc::SIGQUIT, "SIGQUIT"), (libc::SIGTERM, "SIGTERM"), ] { - let pids = temp.path().join(format!("agent-pids-{signal_name}")); - let mut relay = Command::new(gateway_bin()) - .args([ - "--config", - config.to_str().unwrap(), - "run", - "--agent", - "codex", - ]) - .env("HOME", &home) - .env("XDG_CONFIG_HOME", &xdg) - .env("XDG_RUNTIME_DIR", &runtime) - .env("NEMO_RELAY_TEST_AGENT_PIDS", &pids) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::piped()) - .spawn() - .unwrap(); + assert_non_tty_signal_forwarding( + temp.path(), + &config, + &home, + &xdg, + &runtime, + signal, + signal_name, + ); + } +} - let deadline = Instant::now() + Duration::from_secs(10); - while !pids.is_file() { - if Instant::now() >= deadline { - // SAFETY: Relay's PID is live and owned by this test; SIGTERM exercises its - // registered cleanup path so any partially started descendants are reaped. - let _ = unsafe { libc::kill(relay.id() as i32, libc::SIGTERM) }; - let output = relay.wait_with_output().unwrap(); - panic!( - "agent PID file was not created for {signal_name}; stderr: {}", - String::from_utf8_lossy(&output.stderr) - ); - } - thread::sleep(Duration::from_millis(20)); +#[cfg(unix)] +#[allow(clippy::too_many_arguments)] +fn assert_non_tty_signal_forwarding( + root: &Path, + config: &Path, + home: &Path, + xdg: &Path, + runtime: &Path, + signal: i32, + signal_name: &str, +) { + let pids = root.join(format!("agent-pids-{signal_name}")); + let mut relay = Command::new(gateway_bin()) + .args([ + "--config", + config.to_str().unwrap(), + "run", + "--agent", + "codex", + ]) + .env("HOME", home) + .env("XDG_CONFIG_HOME", xdg) + .env("XDG_RUNTIME_DIR", runtime) + .env("NEMO_RELAY_TEST_AGENT_PIDS", &pids) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + wait_for_agent_pid_file(&mut relay, &pids, signal_name); + let supervised_pids = std::fs::read_to_string(&pids) + .unwrap() + .split_whitespace() + .map(|value| value.parse::().unwrap()) + .collect::>(); + assert_eq!(supervised_pids.len(), 2); + + // SAFETY: Relay's PID is live and owned by this test; each signal is caught and forwarded. + assert_eq!(unsafe { libc::kill(relay.id() as i32, signal) }, 0); + assert!(!wait_child(&mut relay).success()); + for pid in supervised_pids { + assert_process_exits(pid, signal_name); + } +} + +#[cfg(unix)] +fn wait_for_agent_pid_file(relay: &mut std::process::Child, pids: &Path, signal_name: &str) { + let deadline = Instant::now() + Duration::from_secs(10); + while !pids.is_file() { + if Instant::now() >= deadline { + // SAFETY: Relay's PID is live and owned by this test; SIGTERM exercises its registered + // cleanup path so any partially started descendants are reaped. + let _ = unsafe { libc::kill(relay.id() as i32, libc::SIGTERM) }; + let _ = relay.kill(); + let _ = relay.wait(); + panic!("agent PID file was not created for {signal_name}"); } - let supervised_pids = std::fs::read_to_string(&pids) - .unwrap() - .split_whitespace() - .map(|value| value.parse::().unwrap()) - .collect::>(); - assert_eq!(supervised_pids.len(), 2); - - // SAFETY: Relay's PID is live and owned by this test; each signal is caught and forwarded. - assert_eq!(unsafe { libc::kill(relay.id() as i32, signal) }, 0); - let status = wait_child(&mut relay); - assert!(!status.success()); - - for pid in supervised_pids { - let deadline = Instant::now() + Duration::from_secs(5); - loop { - // SAFETY: Signal zero is a read-only existence check for the recorded child PID. - let result = unsafe { libc::kill(pid, 0) }; - if result == -1 - && std::io::Error::last_os_error().raw_os_error() == Some(libc::ESRCH) - { - break; - } - assert!( - Instant::now() < deadline, - "coding-agent process {pid} survived Relay {signal_name}" - ); - thread::sleep(Duration::from_millis(20)); - } + thread::sleep(Duration::from_millis(20)); + } +} + +#[cfg(unix)] +fn assert_process_exits(pid: i32, signal_name: &str) { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + // SAFETY: Signal zero is a read-only existence check for the recorded child PID. + let result = unsafe { libc::kill(pid, 0) }; + if result == -1 && std::io::Error::last_os_error().raw_os_error() == Some(libc::ESRCH) { + return; } + assert!( + Instant::now() < deadline, + "coding-agent process {pid} survived Relay {signal_name}" + ); + thread::sleep(Duration::from_millis(20)); } } @@ -3511,22 +3556,7 @@ fn cli_forward_only_never_reconnects_payload_after_authenticated_connection_clos let server_requests = requests.clone(); let server = thread::spawn(move || -> Result { listener.set_nonblocking(true).unwrap(); - let deadline = Instant::now() + Duration::from_secs(4); - let (mut stream, _) = loop { - if server_stopped.load(Ordering::Relaxed) { - return Err("hook-forward exited before authenticated tunnel".into()); - } - match listener.accept() { - Ok(connection) => break connection, - Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { - if Instant::now() >= deadline { - return Err("timed out waiting for authenticated tunnel".into()); - } - thread::sleep(Duration::from_millis(5)); - } - Err(error) => return Err(format!("bootstrap listener failed: {error}")), - } - }; + let mut stream = accept_bootstrap_connection(&listener, &server_stopped)?; stream.set_nonblocking(false).unwrap(); stream .set_read_timeout(Some(Duration::from_secs(2))) @@ -3534,39 +3564,9 @@ fn cli_forward_only_never_reconnects_payload_after_authenticated_connection_clos let mut request = read_http_request(&mut stream); server_requests.lock().unwrap().push(request.clone()); if request.starts_with("GET /healthz ") { - let fingerprint = - bootstrap_request_header(&request, "x-nemo-relay-bootstrap-fingerprint") - .ok_or_else(|| "health probe omitted its fingerprint".to_string())?; - let nonce = bootstrap_request_header(&request, "x-nemo-relay-bootstrap-nonce") - .ok_or_else(|| "health probe omitted its nonce".to_string())?; - let proof = fake_bootstrap_proof(&key, fingerprint, nonce); - let body = format!( - r#"{{"status":"ok","service":"nemo-relay","version":"{}","bootstrap_protocol":2,"instance_id":"phase-health"}}"#, - env!("CARGO_PKG_VERSION") - ); - stream - .write_all( - format!( - "HTTP/1.1 200 OK\r\nX-NeMo-Relay-Bootstrap-Proof: {proof}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ) - .as_bytes(), - ) - .map_err(|error| format!("health response failed: {error}"))?; + write_phase_health_response(&mut stream, &request, &key)?; drop(stream); - let next_deadline = Instant::now() + Duration::from_secs(4); - stream = loop { - match listener.accept() { - Ok((stream, _)) => break stream, - Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { - if Instant::now() >= next_deadline { - return Err("timed out waiting for authenticated tunnel".into()); - } - thread::sleep(Duration::from_millis(5)); - } - Err(error) => return Err(format!("bootstrap listener failed: {error}")), - } - }; + stream = accept_bootstrap_connection(&listener, &server_stopped)?; stream.set_nonblocking(false).unwrap(); stream .set_read_timeout(Some(Duration::from_secs(2))) @@ -3593,23 +3593,7 @@ fn cli_forward_only_never_reconnects_payload_after_authenticated_connection_clos // Close after proving Relay identity but before completing TLS. A replacement listener on // the same port must never receive the lifecycle payload on a fresh connection. drop(stream); - let replacement_deadline = Instant::now() + Duration::from_secs(12); - while !server_stopped.load(Ordering::Relaxed) && Instant::now() < replacement_deadline { - match listener.accept() { - Ok((mut stream, _)) => { - stream.set_nonblocking(false).unwrap(); - stream - .set_read_timeout(Some(Duration::from_secs(2))) - .unwrap(); - let request = read_http_request(&mut stream); - server_requests.lock().unwrap().push(request); - } - Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { - thread::sleep(Duration::from_millis(5)); - } - Err(error) => panic!("replacement listener failed: {error}"), - } - } + collect_replacement_requests(&listener, &server_stopped, &server_requests); Ok(2) }); @@ -3654,6 +3638,79 @@ fn cli_forward_only_never_reconnects_payload_after_authenticated_connection_clos ); } +fn accept_bootstrap_connection( + listener: &TcpListener, + stopped: &AtomicBool, +) -> Result { + let deadline = Instant::now() + Duration::from_secs(4); + loop { + if stopped.load(Ordering::Relaxed) { + return Err("hook-forward exited before authenticated tunnel".into()); + } + match listener.accept() { + Ok((stream, _)) => return Ok(stream), + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + if Instant::now() >= deadline { + return Err("timed out waiting for authenticated tunnel".into()); + } + thread::sleep(Duration::from_millis(5)); + } + Err(error) => return Err(format!("bootstrap listener failed: {error}")), + } + } +} + +fn write_phase_health_response( + stream: &mut std::net::TcpStream, + request: &str, + key: &[u8], +) -> Result<(), String> { + let fingerprint = bootstrap_request_header(request, "x-nemo-relay-bootstrap-fingerprint") + .ok_or_else(|| "health probe omitted its fingerprint".to_string())?; + let nonce = bootstrap_request_header(request, "x-nemo-relay-bootstrap-nonce") + .ok_or_else(|| "health probe omitted its nonce".to_string())?; + let proof = fake_bootstrap_proof(key, fingerprint, nonce); + let body = format!( + r#"{{"status":"ok","service":"nemo-relay","version":"{}","bootstrap_protocol":2,"instance_id":"phase-health"}}"#, + env!("CARGO_PKG_VERSION") + ); + stream + .write_all( + format!( + "HTTP/1.1 200 OK\r\nX-NeMo-Relay-Bootstrap-Proof: {proof}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .map_err(|error| format!("health response failed: {error}")) +} + +fn collect_replacement_requests( + listener: &TcpListener, + stopped: &AtomicBool, + requests: &Mutex>, +) { + let deadline = Instant::now() + Duration::from_secs(12); + while !stopped.load(Ordering::Relaxed) && Instant::now() < deadline { + match listener.accept() { + Ok((mut stream, _)) => { + stream.set_nonblocking(false).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + requests + .lock() + .unwrap() + .push(read_http_request(&mut stream)); + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(5)); + } + Err(error) => panic!("replacement listener failed: {error}"), + } + } +} + #[test] fn cli_transparent_run_suppresses_persistent_hooks_and_rejects_a_foreign_gateway() { let persistent = Command::new(gateway_bin()) diff --git a/scripts/test-claude-plugin-e2e.sh b/scripts/test-claude-plugin-e2e.sh index 161759662..a3a28f156 100755 --- a/scripts/test-claude-plugin-e2e.sh +++ b/scripts/test-claude-plugin-e2e.sh @@ -35,6 +35,7 @@ cleanup() { else rm -rf "$work" fi + return 0 } trap cleanup EXIT @@ -143,6 +144,7 @@ while time.monotonic() < deadline: time.sleep(0.1) raise SystemExit("Relay port 47632 did not become free") PY + return 0 } run_claude() { @@ -176,6 +178,7 @@ assert log.count("SessionEnd:other") == 1, log assert log.count('MCP server "plugin:nemo-relay-plugin:nemo-relay": Successfully connected') == 1, log assert '"hasTools":false' in log, log PY + return 0 } run_transparent_claude() { @@ -211,6 +214,7 @@ assert 1 <= log.count('Hook Stop (Stop) success') <= 2, log assert 1 <= log.count("SessionEnd:other") <= 2, log assert log.count('MCP server "plugin:nemo-relay-plugin:nemo-relay": Successfully connected') == 1, log PY + return 0 } # The transparent wrapper preserves the explicit Claude settings source. The installed Relay MCP diff --git a/scripts/test-hermes-mcp-e2e.sh b/scripts/test-hermes-mcp-e2e.sh index 696d0a290..cd542d6b2 100755 --- a/scripts/test-hermes-mcp-e2e.sh +++ b/scripts/test-hermes-mcp-e2e.sh @@ -187,6 +187,7 @@ while time.monotonic() < deadline: time.sleep(0.1) raise SystemExit("Relay port 47632 did not become free") PY + return 0 } run_hermes() { @@ -204,6 +205,7 @@ from pathlib import Path output, stderr = map(Path, sys.argv[1:]) assert output.read_text().strip().lower() == "pong", (output.read_text(), stderr.read_text()) PY + return 0 } wait_for_relay_port_release diff --git a/scripts/test-support/codex_mock_provider.py b/scripts/test-support/codex_mock_provider.py index cf0b25b44..0868f5aa0 100644 --- a/scripts/test-support/codex_mock_provider.py +++ b/scripts/test-support/codex_mock_provider.py @@ -366,6 +366,8 @@ def main() -> None: args.log_file.parent.mkdir(parents=True, exist_ok=True) args.log_file.write_text("", encoding="utf-8") args.barrier_dir.mkdir(parents=True, exist_ok=True) + # This test-only HTTP provider is deliberately restricted to loopback and an ephemeral port; + # it never accepts remote traffic or production credentials. server = Provider(("127.0.0.1", 0), args.log_file, args.barrier_dir) temporary = args.ready_file.with_suffix(".tmp") temporary.write_text( From 4043f99b0f1be2a2fc02410a636dd5d3d6d45cfb Mon Sep 17 00:00:00 2001 From: Will Killian Date: Wed, 15 Jul 2026 11:12:40 -0400 Subject: [PATCH 2/2] refactor: model marketplace install transaction state Signed-off-by: Will Killian --- .../cli/src/installation/marketplace/mod.rs | 172 +++++++++--------- 1 file changed, 88 insertions(+), 84 deletions(-) diff --git a/crates/cli/src/installation/marketplace/mod.rs b/crates/cli/src/installation/marketplace/mod.rs index 4c25dc0dd..c8d8b07db 100644 --- a/crates/cli/src/installation/marketplace/mod.rs +++ b/crates/cli/src/installation/marketplace/mod.rs @@ -285,38 +285,18 @@ fn install_host_locked( require_host_cli(host, options, runner)?; host::validate_host_version(host, options, runner)?; let layout = PluginLayout::new(host, &options.install_dir); - let (staged, mut force_snapshot, mut replacement_generation_lock) = - prepare_install_transaction(host, &relay, &layout, options, runner, setup_runner)?; - install_marketplace_content( + let context = InstallPhaseContext { host, - &relay, - &layout, - staged.as_ref(), - &mut force_snapshot, - &mut replacement_generation_lock, - options, - runner, - setup_runner, - )?; - write_install_state( - host, - &layout, - &mut force_snapshot, - &mut replacement_generation_lock, - options, - runner, - setup_runner, - )?; - finish_install_registration( - host, - &layout, - &mut force_snapshot, - &mut replacement_generation_lock, options, runner, setup_runner, - )?; - if let Some(snapshot) = force_snapshot { + layout: &layout, + }; + let mut transaction = prepare_install_transaction(&context, &relay)?; + install_marketplace_content(&context, &relay, &mut transaction)?; + write_install_state(&context, &mut transaction)?; + finish_install_registration(&context, &mut transaction)?; + if let Some(snapshot) = transaction.force_snapshot { snapshot.commit(&layout.generation_lock); } println!( @@ -327,20 +307,29 @@ fn install_host_locked( Ok(()) } -type PreparedInstall = ( - Option, - Option, - Option, -); +struct InstallPhaseContext<'a, H: MarketplaceHost> { + host: H, + layout: &'a PluginLayout, + options: &'a PluginInstallOptions, + runner: &'a dyn CommandRunner, + setup_runner: &'a dyn PluginSetupRunner, +} -fn prepare_install_transaction( - host: impl MarketplaceHost, +struct InstallTransactionState { + staged: Option, + force_snapshot: Option, + replacement_generation_lock: Option, +} + +fn prepare_install_transaction( + context: &InstallPhaseContext<'_, H>, relay: &Path, - layout: &PluginLayout, - options: &PluginInstallOptions, - runner: &dyn CommandRunner, - setup_runner: &dyn PluginSetupRunner, -) -> Result { +) -> Result { + let host = context.host; + let layout = context.layout; + let options = context.options; + let runner = context.runner; + let setup_runner = context.setup_runner; let plugin_preflight = if !options.dry_run { Some(prepare_plugin_install(host, layout, options, runner)?) } else { @@ -405,21 +394,24 @@ fn prepare_install_transaction( } else { None }; - Ok((staged, force_snapshot, replacement_generation_lock)) + Ok(InstallTransactionState { + staged, + force_snapshot, + replacement_generation_lock, + }) } -#[allow(clippy::too_many_arguments)] -fn install_marketplace_content( - host: impl MarketplaceHost, +fn install_marketplace_content( + context: &InstallPhaseContext<'_, H>, relay: &Path, - layout: &PluginLayout, - staged: Option<&StagedPluginMarketplace>, - force_snapshot: &mut Option, - replacement_generation_lock: &mut Option, - options: &PluginInstallOptions, - runner: &dyn CommandRunner, - setup_runner: &dyn PluginSetupRunner, + transaction: &mut InstallTransactionState, ) -> Result<(), String> { + let host = context.host; + let layout = context.layout; + let options = context.options; + let runner = context.runner; + let setup_runner = context.setup_runner; + let staged = transaction.staged.as_ref(); if options.force && staged.is_none() { force_cleanup_existing_install(host, layout, options, runner, setup_runner)?; } @@ -429,25 +421,32 @@ fn install_marketplace_content( return restore_force_replacement_after_error( host, layout, - force_snapshot.as_mut().expect("force snapshot exists"), + transaction + .force_snapshot + .as_mut() + .expect("force snapshot exists"), options, runner, setup_runner, error, ); } - force_snapshot + transaction + .force_snapshot .as_mut() .expect("force snapshot exists") .replacement_promoted = true; - if let Some(lock) = replacement_generation_lock.as_mut() + if let Some(lock) = transaction.replacement_generation_lock.as_mut() && let Err(error) = lock.retarget_promoted_marker(&layout.generation_fence) { staged.cleanup(); return restore_force_replacement_after_error( host, layout, - force_snapshot.as_mut().expect("force snapshot exists"), + transaction + .force_snapshot + .as_mut() + .expect("force snapshot exists"), options, runner, setup_runner, @@ -473,7 +472,7 @@ fn install_marketplace_content( }; } if !options.dry_run { - *replacement_generation_lock = match acquire_replacement_generation_lock( + transaction.replacement_generation_lock = match acquire_replacement_generation_lock( host, &layout.generation_fence, &layout.generation_lock, @@ -498,23 +497,24 @@ fn install_marketplace_content( Ok(()) } -#[allow(clippy::too_many_arguments)] -fn write_install_state( - host: impl MarketplaceHost, - layout: &PluginLayout, - force_snapshot: &mut Option, - replacement_generation_lock: &mut Option, - options: &PluginInstallOptions, - runner: &dyn CommandRunner, - setup_runner: &dyn PluginSetupRunner, +fn write_install_state( + context: &InstallPhaseContext<'_, H>, + transaction: &mut InstallTransactionState, ) -> Result<(), String> { + let host = context.host; + let layout = context.layout; + let options = context.options; + let runner = context.runner; + let setup_runner = context.setup_runner; if let Err(error) = write_state(layout, options) { - let _replacement_retirement = if force_snapshot.is_some() { - let existing_retirement = replacement_generation_lock + let _replacement_retirement = if transaction.force_snapshot.is_some() { + let existing_retirement = transaction + .replacement_generation_lock .as_mut() .map(ReplacementGenerationLock::retirement_mut) .or_else(|| { - force_snapshot + transaction + .force_snapshot .as_mut() .and_then(|snapshot| snapshot.generation_retirement.as_mut()) }); @@ -536,7 +536,7 @@ fn write_install_state( None }; let cleanup_error = remove_path(&layout.marketplace_root, options).err(); - let restore_error = force_snapshot.as_mut().and_then(|snapshot| { + let restore_error = transaction.force_snapshot.as_mut().and_then(|snapshot| { restore_force_replacement(host, layout, snapshot, options, runner, setup_runner).err() }); let errors = [cleanup_error, restore_error] @@ -551,25 +551,26 @@ fn write_install_state( Ok(()) } -#[allow(clippy::too_many_arguments)] -fn finish_install_registration( - host: impl MarketplaceHost, - layout: &PluginLayout, - force_snapshot: &mut Option, - replacement_generation_lock: &mut Option, - options: &PluginInstallOptions, - runner: &dyn CommandRunner, - setup_runner: &dyn PluginSetupRunner, +fn finish_install_registration( + context: &InstallPhaseContext<'_, H>, + transaction: &mut InstallTransactionState, ) -> Result<(), String> { + let host = context.host; + let layout = context.layout; + let options = context.options; + let runner = context.runner; + let setup_runner = context.setup_runner; let mut registration = HostRegistrationProgress::default(); let mut registration_state_uncertain = false; let mut setup_installed = false; let result = (|| { - let generation_token = replacement_generation_lock + let generation_token = transaction + .replacement_generation_lock .as_ref() .map(ReplacementGenerationLock::retirement) .or_else(|| { - force_snapshot + transaction + .force_snapshot .as_ref() .and_then(|snapshot| snapshot.generation_retirement.as_ref()) }) @@ -620,13 +621,16 @@ fn finish_install_registration( registration.host_plugin_added |= observed.host_plugin_registered; registration.host_marketplace_added |= observed.host_marketplace_registered; } - let replacement_may_be_live = force_snapshot.is_some() || registration.host_plugin_added; + let replacement_may_be_live = + transaction.force_snapshot.is_some() || registration.host_plugin_added; let _replacement_retirement = if replacement_may_be_live { - let existing_retirement = replacement_generation_lock + let existing_retirement = transaction + .replacement_generation_lock .as_mut() .map(ReplacementGenerationLock::retirement_mut) .or_else(|| { - force_snapshot + transaction + .force_snapshot .as_mut() .and_then(|snapshot| snapshot.generation_retirement.as_mut()) }); @@ -657,7 +661,7 @@ fn finish_install_registration( setup_runner, ) .err(); - let restore_error = force_snapshot.as_mut().and_then(|snapshot| { + let restore_error = transaction.force_snapshot.as_mut().and_then(|snapshot| { restore_force_replacement(host, layout, snapshot, options, runner, setup_runner).err() }); let rollback_errors = [