diff --git a/crates/cli/src/bootstrap/mod.rs b/crates/cli/src/bootstrap/mod.rs index a7cd3d61e..126ce61be 100644 --- a/crates/cli/src/bootstrap/mod.rs +++ b/crates/cli/src/bootstrap/mod.rs @@ -571,8 +571,23 @@ pub(crate) fn plugin_idle_timeout() -> Result { Ok(Duration::from_secs(seconds)) } +/// Returns the persistent MCP gateway heartbeat interval. pub(crate) fn plugin_heartbeat_interval() -> Result { - Ok((plugin_idle_timeout()? / 3).clamp(Duration::from_millis(100), Duration::from_secs(30))) + let raw = env::var(crate::configuration::PLUGIN_HEARTBEAT_INTERVAL_ENV) + .unwrap_or_else(|_| "3".into()); + let seconds = raw.parse::().map_err(|error| { + format!( + "{} must be a positive integer: {error}", + crate::configuration::PLUGIN_HEARTBEAT_INTERVAL_ENV + ) + })?; + if seconds == 0 { + return Err(format!( + "{} must be greater than 0", + crate::configuration::PLUGIN_HEARTBEAT_INTERVAL_ENV + )); + } + Ok(Duration::from_secs(seconds)) } #[cfg(test)] diff --git a/crates/cli/src/bootstrap/state.rs b/crates/cli/src/bootstrap/state.rs index 920b01102..7e6f483f4 100644 --- a/crates/cli/src/bootstrap/state.rs +++ b/crates/cli/src/bootstrap/state.rs @@ -16,7 +16,9 @@ use serde::{Deserialize, Serialize}; use crate::filesystem::{LockAttempt, atomic_write, try_lock_exclusive}; use super::{BOOTSTRAP_LOCK_TIMEOUT, BOOTSTRAP_PROTOCOL_VERSION}; -use crate::gateway::client::{RelayHealth, probe, request_shutdown}; +use crate::gateway::client::{ + RelayHealth, probe, probe_with_instance, request_lifecycle_shutdown, request_shutdown, +}; pub(crate) const BOOTSTRAP_STATE_DIR_ENV: &str = "NEMO_RELAY_BOOTSTRAP_STATE_DIR"; pub(crate) const BOOTSTRAP_SHUTDOWN_TOKEN_ENV: &str = "NEMO_RELAY_BOOTSTRAP_SHUTDOWN_TOKEN"; @@ -225,7 +227,11 @@ pub(crate) fn stop_owned_and_reset(url: &str) -> Result<(), String> { return Ok(()); } let _lock = lock_endpoint(&state, url)?; - let path = owner_path(&state, url); + stop_owned_and_reset_locked(&state, url) +} + +fn stop_owned_and_reset_locked(state: &Path, url: &str) -> Result<(), String> { + let path = owner_path(state, url); let Some(owner) = read_owner_record(&path)? else { return Ok(()); }; @@ -275,6 +281,66 @@ pub(crate) fn stop_owned_and_reset(url: &str) -> Result<(), String> { remove_if_matches(&path, &owner) } +pub(crate) fn stop_gateway_and_reset(url: &str) -> Result<(), String> { + let state = state_dir()?; + if !state.exists() { + return stop_gateway_without_owner(url); + } + let _lock = lock_endpoint(&state, url)?; + let path = owner_path(&state, url); + let previous_owner = read_optional_bytes(&path)?; + let owner_stop = stop_owned_and_reset_locked(&state, url); + if owner_stop.is_ok() && probe(url, None) == RelayHealth::Unavailable { + return Ok(()); + } + stop_gateway_without_owner(url)?; + if let Some(previous_owner) = previous_owner { + remove_if_bytes_match(&path, &previous_owner)?; + } + Ok(()) +} + +fn stop_gateway_without_owner(url: &str) -> Result<(), String> { + let (health, _) = probe_with_instance(url, None); + match health { + RelayHealth::Unavailable => return Ok(()), + RelayHealth::Compatible => {} + RelayHealth::Incompatible => { + return Err(format!( + "Relay gateway at {url} uses an incompatible lifecycle protocol" + )); + } + RelayHealth::Foreign => { + return Err(format!( + "refusing to stop an unverified process at gateway URL {url}" + )); + } + } + let expected_instance = request_lifecycle_shutdown(url)?; + let deadline = Instant::now() + SHUTDOWN_TIMEOUT; + loop { + match probe_with_instance(url, None) { + (RelayHealth::Unavailable, _) => return Ok(()), + (RelayHealth::Compatible, Some(instance)) + if instance == expected_instance && Instant::now() < deadline => + { + thread::sleep(Duration::from_millis(50)); + } + (RelayHealth::Compatible, Some(instance)) if instance == expected_instance => { + return Err(format!("Relay gateway at {url} did not stop")); + } + (RelayHealth::Foreign, _) if Instant::now() < deadline => { + thread::sleep(Duration::from_millis(50)); + } + _ => { + return Err(format!( + "a different process replaced the Relay gateway at {url} during shutdown" + )); + } + } + } +} + fn write_owner_record(path: &Path, record: &OwnerRecord) -> Result<(), String> { let bytes = serde_json::to_vec(record) .map_err(|error| format!("failed to encode gateway ownership: {error}"))?; @@ -282,13 +348,20 @@ fn write_owner_record(path: &Path, record: &OwnerRecord) -> Result<(), String> { } pub(super) fn read_owner_record(path: &Path) -> Result, String> { - match fs::read(path) { - Ok(bytes) => serde_json::from_slice(&bytes).map(Some).map_err(|error| { + match read_optional_bytes(path)? { + Some(bytes) => serde_json::from_slice(&bytes).map(Some).map_err(|error| { format!( "failed to parse gateway ownership {}: {error}", path.display() ) }), + None => Ok(None), + } +} + +fn read_optional_bytes(path: &Path) -> Result>, String> { + match fs::read(path) { + Ok(bytes) => Ok(Some(bytes)), Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), Err(error) => Err(format!( "failed to read gateway ownership {}: {error}", @@ -311,6 +384,20 @@ fn remove_if_matches(path: &Path, expected: &OwnerRecord) -> Result<(), String> } } +fn remove_if_bytes_match(path: &Path, expected: &[u8]) -> Result<(), String> { + if read_optional_bytes(path)?.as_deref() != Some(expected) { + return Ok(()); + } + match fs::remove_file(path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(format!( + "failed to remove gateway ownership {}: {error}", + path.display() + )), + } +} + pub(crate) fn lock_name(url: &str) -> String { let raw = Url::parse(url) .ok() diff --git a/crates/cli/src/commands/gateway.rs b/crates/cli/src/commands/gateway.rs new file mode 100644 index 000000000..3e36fa3fa --- /dev/null +++ b/crates/cli/src/commands/gateway.rs @@ -0,0 +1,50 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::net::SocketAddr; +use std::process::ExitCode; + +use clap::{Args, Subcommand}; + +use crate::error::CliError; + +use super::serve::ServerArgs; + +pub(super) fn stop_bind(server: &ServerArgs) -> SocketAddr { + server + .bind + .unwrap_or_else(|| crate::configuration::GatewayConfig::default().bind) +} + +#[derive(Debug, Clone, Args)] +pub(crate) struct GatewayCommand { + #[command(subcommand)] + command: GatewaySubcommand, +} + +#[derive(Debug, Clone, Subcommand)] +enum GatewaySubcommand { + /// Start the gateway with the same server configuration as a bare daemon invocation. + Start, + /// Stop the Relay gateway at the configured loopback endpoint. + Stop, +} + +impl GatewayCommand { + /// Returns whether this command only stops an existing gateway. + pub(crate) fn is_stop(&self) -> bool { + matches!(self.command, GatewaySubcommand::Stop) + } +} + +/// Executes a gateway lifecycle command. +pub(crate) async fn execute( + command: GatewayCommand, + server: &ServerArgs, + bootstrap_shutdown_token: Option, +) -> Result { + match command.command { + GatewaySubcommand::Start => super::serve_gateway(server, bootstrap_shutdown_token).await, + GatewaySubcommand::Stop => crate::mcp::stop(stop_bind(server)), + } +} diff --git a/crates/cli/src/commands/mod.rs b/crates/cli/src/commands/mod.rs index 3c2153019..a23a8fd64 100644 --- a/crates/cli/src/commands/mod.rs +++ b/crates/cli/src/commands/mod.rs @@ -6,6 +6,7 @@ mod completions; mod configure; mod diagnostics; +mod gateway; mod hook_forward; mod install; mod logging; @@ -60,6 +61,10 @@ struct LoggingSetup { fn configure_logging(cli: &Cli) -> Result { let initialize = match cli.command.as_ref() { + Some(Command::Gateway(command)) if !command.is_stop() => { + cli.server.to_runtime().requested_daemon_mode() + || runtime_configuration::any_config_file_exists() + } Some(command) => !command.skips_logging(), None => { cli.server.to_runtime().requested_daemon_mode() @@ -78,6 +83,9 @@ fn configure_logging(cli: &Cli) -> Result { // Uninstall uses persisted integration state, not Relay runtime configuration. Preserve // direct logging settings while avoiding ambient config discovery that could block cleanup. Some(Command::Uninstall(_)) => cli.logging.resolve_without_ambient_config(), + Some(Command::Gateway(command)) if command.is_stop() => { + cli.logging.resolve_without_ambient_config() + } Some(Command::Mcp) => cli.logging.resolve(None), Some(Command::Run(command)) => cli .logging @@ -125,7 +133,15 @@ async fn dispatch(bootstrap_shutdown_token: Option) -> Result run_command(command, &cli.server, logging.fallback_error.as_ref()).await, + Some(command) => { + run_command( + command, + &cli.server, + logging.fallback_error.as_ref(), + bootstrap_shutdown_token, + ) + .await + } None => run_default(&cli.server, bootstrap_shutdown_token).await, }; match &result { @@ -165,6 +181,7 @@ async fn run_command( command: Command, server: &ServerArgs, logging_fallback_error: Option<&error::CliError>, + bootstrap_shutdown_token: Option, ) -> Result { match command { Command::HookForward(command) => { @@ -177,6 +194,9 @@ async fn run_command( Command::Claude(command) => run::easy_path(CodingAgent::ClaudeCode, command, server).await, Command::Codex(command) => run::easy_path(CodingAgent::Codex, command, server).await, Command::Mcp => mcp::execute(server).await, + Command::Gateway(command) => { + gateway::execute(command, server, bootstrap_shutdown_token).await + } Command::Config(command) => configure::execute(command, server).await, Command::Plugins(command) => plugins::execute(command, server), Command::ModelPricing(command) => model_pricing::execute(command), @@ -210,29 +230,7 @@ async fn run_default( // exists. Once configured, bare `nemo-relay` becomes a quick health check; explicit // `nemo-relay config` remains the reconfiguration path. if runtime_args.requested_daemon_mode() { - let resolved = runtime_configuration::resolve_server_config(&runtime_args)?; - let explicit_plugin_config = crate::configuration::explicit_plugin_config_path( - runtime_args.config.as_ref(), - runtime_args.plugin_config_path.as_ref(), - ); - let dynamic_plugins = crate::plugins::lifecycle::active_dynamic_plugin_components( - explicit_plugin_config.as_ref(), - &resolved, - )?; - let managed_bootstrap = runtime_configuration::managed_bootstrap_identity( - &runtime_args, - &resolved, - &dynamic_plugins, - )?; - server::serve_with_dynamic( - resolved.gateway, - dynamic_plugins, - managed_bootstrap, - runtime_args.ready_file.as_deref(), - bootstrap_shutdown_token, - ) - .await?; - Ok(ExitCode::SUCCESS) + serve_gateway(server_args, bootstrap_shutdown_token).await } else if runtime_configuration::any_config_file_exists() { runtime_diagnostics::run_doctor( None, @@ -248,6 +246,37 @@ async fn run_default( } } +/// Resolves and serves the configured gateway until it shuts down. +async fn serve_gateway( + server_args: &ServerArgs, + bootstrap_shutdown_token: Option, +) -> Result { + let runtime_args = server_args.to_runtime(); + let resolved = runtime_configuration::resolve_server_config(&runtime_args)?; + let explicit_plugin_config = crate::configuration::explicit_plugin_config_path( + runtime_args.config.as_ref(), + runtime_args.plugin_config_path.as_ref(), + ); + let dynamic_plugins = crate::plugins::lifecycle::active_dynamic_plugin_components( + explicit_plugin_config.as_ref(), + &resolved, + )?; + let managed_bootstrap = runtime_configuration::managed_bootstrap_identity( + &runtime_args, + &resolved, + &dynamic_plugins, + )?; + server::serve_with_dynamic( + resolved.gateway, + dynamic_plugins, + managed_bootstrap, + runtime_args.ready_file.as_deref(), + bootstrap_shutdown_token, + ) + .await?; + Ok(ExitCode::SUCCESS) +} + #[cfg(test)] fn run_completions(command: CompletionsCommand) -> Result { completions::execute(command) diff --git a/crates/cli/src/commands/root.rs b/crates/cli/src/commands/root.rs index 79e9ffaf4..c2db9d706 100644 --- a/crates/cli/src/commands/root.rs +++ b/crates/cli/src/commands/root.rs @@ -6,6 +6,7 @@ use clap::{Parser, Subcommand, ValueEnum}; use super::completions::CompletionsCommand; use super::configure::ConfigCommand; use super::diagnostics::{AgentsCommand, DoctorCommand}; +use super::gateway::GatewayCommand; use super::hook_forward::HookForwardCommand; use super::install::{InstallCommand, UninstallCommand}; use super::logging::LoggingArgs; @@ -88,6 +89,8 @@ pub(crate) enum Command { nemo-relay --bind 127.0.0.1:4041 mcp # explicit standalone/test bind" )] Mcp, + /// Manage the persistent shared Relay gateway. + Gateway(GatewayCommand), /// Run the interactive setup (writes the XDG user `config.toml`) Config(ConfigCommand), /// Create or edit plugin configuration (writes `plugins.toml`) @@ -117,6 +120,7 @@ impl Command { Self::Claude(_) => "claude", Self::Codex(_) => "codex", Self::Mcp => "mcp", + Self::Gateway(_) => "gateway", Self::Config(_) => "config", Self::Plugins(_) => "plugins", Self::Install(_) => "install", @@ -134,6 +138,7 @@ impl Command { /// invalid, so users can repair their configuration. pub(crate) fn skips_logging(&self) -> bool { matches!(self, Self::Config(_)) + || matches!(self, Self::Gateway(command) if command.is_stop()) || matches!(self, Self::Plugins(command) if command.is_edit()) || matches!(self, Self::HookForward(command) if transparent_hook_is_inert(command)) } diff --git a/crates/cli/src/configuration/mod.rs b/crates/cli/src/configuration/mod.rs index 677a8d5bd..1ffd51094 100644 --- a/crates/cli/src/configuration/mod.rs +++ b/crates/cli/src/configuration/mod.rs @@ -41,6 +41,7 @@ use crate::server::GatewayOverrides; pub(crate) const BOOTSTRAP_FINGERPRINT_ENV: &str = "NEMO_RELAY_BOOTSTRAP_FINGERPRINT"; pub(crate) const PLUGIN_IDLE_TIMEOUT_ENV: &str = "NEMO_RELAY_PLUGIN_IDLE_TIMEOUT_SECS"; +pub(crate) const PLUGIN_HEARTBEAT_INTERVAL_ENV: &str = "NEMO_RELAY_PLUGIN_HEARTBEAT_INTERVAL_SECS"; pub(crate) const RELAY_PLUGIN_ID: &str = "nemo-relay-plugin@nemo-relay-local"; pub(crate) const RELAY_SOURCE_PLUGIN_ID: &str = "nemo-relay-plugin@nemo-relay"; pub(crate) const DEFAULT_MAX_HOOK_PAYLOAD_BYTES: usize = 20 * 1024 * 1024; @@ -269,7 +270,7 @@ fn persistent_bootstrap_fingerprint( let environment = env::vars_os().filter_map(|(name, _)| name.into_string().ok()); for name in crate::mcp_environment::forwarded_names(environment, gateway.plugin_config.as_ref()) { - if name == PLUGIN_IDLE_TIMEOUT_ENV { + if name == PLUGIN_IDLE_TIMEOUT_ENV || name == PLUGIN_HEARTBEAT_INTERVAL_ENV { continue; } digest.update(&[0]); @@ -476,6 +477,11 @@ const BOOTSTRAP_HMAC_LOCK_TIMEOUT: Duration = Duration::from_secs(5); const BOOTSTRAP_CHALLENGE_DOMAIN: &[u8] = b"nemo-relay/bootstrap-health/v1\0"; const BOOTSTRAP_CLIENT_TOKEN_DOMAIN: &[u8] = b"nemo-relay/bootstrap-client/v1\0"; const HOOK_CLIENT_TOKEN_DOMAIN: &[u8] = b"nemo-relay/hook-client/v1\0"; +const GATEWAY_LIFECYCLE_DOMAIN: &[u8] = b"nemo-relay/gateway-lifecycle/v1\0"; +// Health proofs are visible to any loopback client. A separate action domain prevents a health +// challenge response from being replayed as shutdown authorization. +const GATEWAY_LIFECYCLE_HEALTH_ACTION: &[u8] = b"health"; +const GATEWAY_LIFECYCLE_SHUTDOWN_ACTION: &[u8] = b"shutdown"; const TRANSPARENT_GATEWAY_DOMAIN: &[u8] = b"nemo-relay/transparent-gateway/v1\0"; const PYTHON_ENVIRONMENT_ATTESTATION_DOMAIN: &[u8] = b"nemo-relay/python-environment-attestation/v1\0"; @@ -483,6 +489,8 @@ const PYTHON_ENVIRONMENT_ATTESTATION_DOMAIN: &[u8] = /// Private proof installed into supported coding-agent provider configuration. pub(crate) const BOOTSTRAP_CLIENT_TOKEN_HEADER: &str = "x-nemo-relay-client-token"; pub(crate) const HOOK_CLIENT_TOKEN_HEADER: &str = "x-nemo-relay-hook-client"; +pub(crate) const GATEWAY_LIFECYCLE_NONCE_HEADER: &str = "x-nemo-relay-lifecycle-nonce"; +pub(crate) const GATEWAY_LIFECYCLE_PROOF_HEADER: &str = "x-nemo-relay-lifecycle-proof"; /// Stable health-proof context shared by a transparent wrapper and plugin-owned MCP client. pub(crate) fn transparent_gateway_fingerprint(gateway_url: &str) -> String { @@ -543,6 +551,113 @@ impl BootstrapChallengeKey { hmac::verify(&self.0, &message, &tag).is_ok() } + pub(crate) fn gateway_lifecycle_health_proof( + &self, + instance_id: &str, + address: &str, + nonce: &str, + ) -> String { + self.gateway_lifecycle_proof(GATEWAY_LIFECYCLE_HEALTH_ACTION, instance_id, address, nonce) + } + + pub(crate) fn verify_gateway_lifecycle_health_proof( + &self, + instance_id: &str, + address: &str, + nonce: &str, + proof: &str, + ) -> bool { + self.verify_gateway_lifecycle_proof( + GATEWAY_LIFECYCLE_HEALTH_ACTION, + instance_id, + address, + nonce, + proof, + ) + } + + pub(crate) fn gateway_lifecycle_shutdown_proof( + &self, + instance_id: &str, + address: &str, + nonce: &str, + ) -> String { + self.gateway_lifecycle_proof( + GATEWAY_LIFECYCLE_SHUTDOWN_ACTION, + instance_id, + address, + nonce, + ) + } + + pub(crate) fn verify_gateway_lifecycle_shutdown_proof( + &self, + instance_id: &str, + address: &str, + nonce: &str, + proof: &str, + ) -> bool { + self.verify_gateway_lifecycle_proof( + GATEWAY_LIFECYCLE_SHUTDOWN_ACTION, + instance_id, + address, + nonce, + proof, + ) + } + + fn gateway_lifecycle_proof( + &self, + action: &[u8], + instance_id: &str, + address: &str, + nonce: &str, + ) -> String { + let mut context = hmac::Context::with_key(&self.0); + context.update(GATEWAY_LIFECYCLE_DOMAIN); + context.update(action); + context.update(&[0]); + context.update(instance_id.as_bytes()); + context.update(&[0]); + context.update(address.as_bytes()); + context.update(&[0]); + context.update(nonce.as_bytes()); + encode_hmac_tag(context.sign()) + } + + fn verify_gateway_lifecycle_proof( + &self, + action: &[u8], + instance_id: &str, + address: &str, + nonce: &str, + proof: &str, + ) -> bool { + let Some(encoded) = proof.strip_prefix("hmac-sha256:") else { + return false; + }; + let Some(tag) = decode_fixed_hex::<32>(encoded) else { + return false; + }; + let mut message = Vec::with_capacity( + GATEWAY_LIFECYCLE_DOMAIN.len() + + action.len() + + instance_id.len() + + address.len() + + nonce.len() + + 3, + ); + message.extend_from_slice(GATEWAY_LIFECYCLE_DOMAIN); + message.extend_from_slice(action); + message.push(0); + message.extend_from_slice(instance_id.as_bytes()); + message.push(0); + message.extend_from_slice(address.as_bytes()); + message.push(0); + message.extend_from_slice(nonce.as_bytes()); + hmac::verify(&self.0, &message, &tag).is_ok() + } + /// Returns a stable, per-user proof that authorizes use of credentials forwarded to a /// managed sidecar. The HMAC key remains in Relay's private bootstrap state; coding-agent /// configuration stores only this domain-separated proof. diff --git a/crates/cli/src/gateway/client.rs b/crates/cli/src/gateway/client.rs index dae33aca1..074eb68d3 100644 --- a/crates/cli/src/gateway/client.rs +++ b/crates/cli/src/gateway/client.rs @@ -14,8 +14,10 @@ use reqwest::Url; use ring::rand::{SecureRandom, SystemRandom}; use serde_json::Value; -use crate::configuration::BOOTSTRAP_CLIENT_TOKEN_HEADER; -use crate::configuration::BootstrapChallengeKey; +use crate::configuration::{ + BOOTSTRAP_CLIENT_TOKEN_HEADER, BootstrapChallengeKey, GATEWAY_LIFECYCLE_NONCE_HEADER, + GATEWAY_LIFECYCLE_PROOF_HEADER, +}; use crate::bootstrap::{BOOTSTRAP_PROTOCOL_VERSION, HEALTHZ_TIMEOUT}; @@ -356,14 +358,7 @@ pub(crate) fn request_shutdown( .map_err(|error| format!("failed to configure sidecar shutdown write timeout: {error}"))?; let key = cached_bootstrap_challenge_key() .map_err(|error| format!("failed to load the Relay bootstrap challenge key: {error}"))?; - let mut nonce = [0_u8; 32]; - SystemRandom::new() - .fill(&mut nonce) - .map_err(|_| "failed to generate a Relay bootstrap shutdown challenge".to_string())?; - let nonce = nonce - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::(); + let nonce = random_challenge_nonce()?; let authority = loopback_authority(&host, port); let challenge = format!( "GET /healthz HTTP/1.1\r\nHost: {authority}\r\nX-NeMo-Relay-Bootstrap-Fingerprint: {bootstrap_fingerprint}\r\nX-NeMo-Relay-Bootstrap-Nonce: {nonce}\r\nConnection: keep-alive\r\n\r\n" @@ -415,6 +410,89 @@ pub(crate) fn request_shutdown( } } +/// Authenticates the Relay instance at `url` and requests shutdown on the same connection. +pub(crate) fn request_lifecycle_shutdown(url: &str) -> Result { + let (host, port) = parse_loopback_url(url)?; + let addresses = (host.as_str(), port) + .to_socket_addrs() + .map_err(|error| format!("failed to resolve Relay gateway {url}: {error}"))?; + let mut stream = connect_loopback(addresses, HEALTHZ_TIMEOUT) + .map_err(|error| format!("failed to connect to Relay gateway {url}: {error}"))?; + stream + .set_read_timeout(Some(HEALTHZ_TIMEOUT)) + .map_err(|error| format!("failed to configure gateway shutdown read timeout: {error}"))?; + stream + .set_write_timeout(Some(HEALTHZ_TIMEOUT)) + .map_err(|error| format!("failed to configure gateway shutdown write timeout: {error}"))?; + let address = stream + .peer_addr() + .map_err(|error| format!("failed to identify Relay gateway address: {error}"))? + .to_string(); + let key = cached_bootstrap_challenge_key() + .map_err(|error| format!("failed to load the Relay gateway lifecycle key: {error}"))?; + let nonce = random_challenge_nonce()?; + let authority = loopback_authority(&host, port); + let health_request = format!( + "GET /healthz HTTP/1.1\r\nHost: {authority}\r\n{GATEWAY_LIFECYCLE_NONCE_HEADER}: {nonce}\r\nConnection: keep-alive\r\n\r\n" + ); + stream + .write_all(health_request.as_bytes()) + .map_err(|error| format!("failed to verify Relay gateway before shutdown: {error}"))?; + let (health_headers, health_body) = read_http_message(&mut stream, 16 * 1024) + .map_err(|error| format!("failed to read Relay gateway lifecycle proof: {error}"))?; + let (health, instance_id) = classify_health_response(&health_headers, &health_body, None); + if health != RelayHealth::Compatible { + return Err("shutdown target did not verify as a compatible Relay gateway".into()); + } + let instance_id = instance_id.expect("compatible health responses include an instance ID"); + let proof_valid = + http_header(&health_headers, GATEWAY_LIFECYCLE_PROOF_HEADER).is_some_and(|proof| { + key.verify_gateway_lifecycle_health_proof(&instance_id, &address, &nonce, proof) + }); + if !proof_valid { + return Err("Relay gateway did not authenticate lifecycle shutdown".into()); + } + if http_header(&health_headers, "connection") + .is_some_and(|value| value.eq_ignore_ascii_case("close")) + { + return Err("Relay gateway closed the connection before the shutdown request".into()); + } + let shutdown_proof = key.gateway_lifecycle_shutdown_proof(&instance_id, &address, &nonce); + let request = format!( + "POST /bootstrap/shutdown HTTP/1.1\r\nHost: {authority}\r\n{GATEWAY_LIFECYCLE_NONCE_HEADER}: {nonce}\r\n{GATEWAY_LIFECYCLE_PROOF_HEADER}: {shutdown_proof}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ); + stream + .write_all(request.as_bytes()) + .map_err(|error| format!("failed to request Relay gateway shutdown: {error}"))?; + let mut response = Vec::new(); + stream + .take(16 * 1024) + .read_to_end(&mut response) + .map_err(|error| format!("failed to read Relay gateway shutdown response: {error}"))?; + let Some((headers, _)) = split_http_response(&response) else { + return Err("Relay gateway returned a malformed shutdown response".into()); + }; + if headers.starts_with(b"HTTP/1.1 204") || headers.starts_with(b"HTTP/1.0 204") { + Ok(instance_id) + } else { + Err(format!( + "Relay gateway rejected shutdown: {}", + String::from_utf8_lossy(headers) + .lines() + .next() + .unwrap_or("unknown response") + )) + } +} + +fn random_challenge_nonce() -> Result { + let mut nonce = [0_u8; 32]; + SystemRandom::new() + .fill(&mut nonce) + .map_err(|_| "failed to generate a Relay gateway lifecycle challenge".to_string())?; + Ok(nonce.iter().map(|byte| format!("{byte:02x}")).collect()) +} + fn cached_bootstrap_challenge_key() -> Result, String> { let state = crate::bootstrap::state::state_dir()?; let cache = CHALLENGE_KEY_CACHE.get_or_init(|| Mutex::new(HashMap::new())); diff --git a/crates/cli/src/mcp/mod.rs b/crates/cli/src/mcp/mod.rs index 3b512349e..4fd14debb 100644 --- a/crates/cli/src/mcp/mod.rs +++ b/crates/cli/src/mcp/mod.rs @@ -103,6 +103,18 @@ pub(crate) async fn run(server_args: &GatewayOverrides) -> Result Result { + if bind.port() == 0 { + return Err(CliError::Config( + "gateway stop requires a concrete nonzero port".into(), + )); + } + let url = format!("http://{bind}"); + crate::bootstrap::state::stop_gateway_and_reset(&url).map_err(CliError::Launch)?; + Ok(ExitCode::SUCCESS) +} + /// Builds the host-independent persistent MCP launch contract. /// /// Host adapters add only schema-specific activation and environment-forwarding fields. Keeping diff --git a/crates/cli/src/mcp_environment.rs b/crates/cli/src/mcp_environment.rs index d5f522fda..5f4cc1e5f 100644 --- a/crates/cli/src/mcp_environment.rs +++ b/crates/cli/src/mcp_environment.rs @@ -45,6 +45,7 @@ const BASE_MCP_ENV_VARS: &[&str] = &[ "NEMO_RELAY_MAX_PASSTHROUGH_BODY_BYTES", "NEMO_RELAY_OPENAI_AUTH_HEADER", "NEMO_RELAY_OPENAI_BASE_URL", + "NEMO_RELAY_PLUGIN_HEARTBEAT_INTERVAL_SECS", "NEMO_RELAY_PLUGIN_IDLE_TIMEOUT_SECS", "NEMO_RELAY_PYTHON", "NEMO_RELAY_TRANSPARENT_RUN", diff --git a/crates/cli/src/server/mod.rs b/crates/cli/src/server/mod.rs index 3dfef0f2e..dc605e9d9 100644 --- a/crates/cli/src/server/mod.rs +++ b/crates/cli/src/server/mod.rs @@ -37,7 +37,8 @@ use tokio::sync::oneshot; use crate::agents::shared::adapters::{claude_code, codex}; use crate::configuration::{ - BOOTSTRAP_CLIENT_TOKEN_HEADER, BootstrapChallengeKey, GatewayConfig, HOOK_CLIENT_TOKEN_HEADER, + BOOTSTRAP_CLIENT_TOKEN_HEADER, BootstrapChallengeKey, GATEWAY_LIFECYCLE_NONCE_HEADER, + GATEWAY_LIFECYCLE_PROOF_HEADER, GatewayConfig, HOOK_CLIENT_TOKEN_HEADER, ManagedBootstrapIdentity, }; use crate::error::CliError; @@ -68,7 +69,7 @@ pub(crate) struct AppState { #[derive(Clone)] pub(crate) struct BootstrapShutdown { - token: String, + token: Option, sender: Arc>>>, } @@ -78,6 +79,7 @@ struct BootstrapServeOptions<'a> { identity: Option, ready_file: Option<&'a Path>, shutdown_token: Option, + lifecycle_shutdown: bool, transparent_proxy_credential: Option, } @@ -117,6 +119,7 @@ pub(crate) async fn serve_with_dynamic( identity: managed_bootstrap, ready_file, shutdown_token: bootstrap_shutdown_token, + lifecycle_shutdown: true, ..BootstrapServeOptions::default() }, ) @@ -270,6 +273,7 @@ async fn serve_listener_with_dynamic_inner( identity: managed_bootstrap, ready_file, shutdown_token: bootstrap_shutdown_token, + lifecycle_shutdown, transparent_proxy_credential, } = bootstrap; let bootstrap_challenge_key = Some(BootstrapChallengeKey::load()?); @@ -286,7 +290,7 @@ async fn serve_listener_with_dynamic_inner( let plugin_activation = initialize_plugin_host(config.plugin_config.clone(), dynamic_plugins).await?; let (bootstrap_shutdown, bootstrap_shutdown_rx) = - bootstrap_shutdown_channel(bootstrap_shutdown_token.clone()); + bootstrap_shutdown_channel(bootstrap_shutdown_token.clone(), lifecycle_shutdown); let mut state = AppState::new_with_bootstrap( config, bootstrap_fingerprint, @@ -717,10 +721,11 @@ async fn bootstrap_tls_tunnel( fn bootstrap_shutdown_channel( token: Option, + lifecycle_enabled: bool, ) -> (Option, Option>) { - let Some(token) = token else { + if token.is_none() && !lifecycle_enabled { return (None, None); - }; + } let (sender, receiver) = oneshot::channel(); ( Some(BootstrapShutdown { @@ -738,11 +743,14 @@ async fn shutdown_bootstrap_sidecar( let Some(shutdown) = state.bootstrap_shutdown.as_ref() else { return StatusCode::NOT_FOUND; }; - if headers + let owner_token_matches = headers .get("x-nemo-relay-bootstrap-token") .and_then(|value| value.to_str().ok()) - != Some(shutdown.token.as_str()) - { + .zip(shutdown.token.as_deref()) + .is_some_and(|(presented, expected)| { + bool::from(presented.as_bytes().ct_eq(expected.as_bytes())) + }); + if !owner_token_matches && !valid_gateway_lifecycle_shutdown(&state, &headers) { return StatusCode::FORBIDDEN; } let Ok(mut sender) = shutdown.sender.lock() else { @@ -755,6 +763,36 @@ async fn shutdown_bootstrap_sidecar( StatusCode::NO_CONTENT } +fn valid_gateway_lifecycle_shutdown(state: &AppState, headers: &HeaderMap) -> bool { + let Some(nonce) = headers + .get(GATEWAY_LIFECYCLE_NONCE_HEADER) + .and_then(|value| value.to_str().ok()) + .filter(|nonce| valid_gateway_lifecycle_nonce(nonce)) + else { + return false; + }; + let Some(proof) = headers + .get(GATEWAY_LIFECYCLE_PROOF_HEADER) + .and_then(|value| value.to_str().ok()) + else { + return false; + }; + let (Some(key), Some(address)) = (state.bootstrap_challenge_key.as_ref(), state.local_address) + else { + return false; + }; + key.verify_gateway_lifecycle_shutdown_proof( + &state.instance_id, + &address.to_string(), + nonce, + proof, + ) +} + +fn valid_gateway_lifecycle_nonce(nonce: &str) -> bool { + nonce.len() == 64 && nonce.bytes().all(|byte| byte.is_ascii_hexdigit()) +} + async fn healthz(State(state): State, headers: HeaderMap) -> Response { let presented_fingerprint = headers .get("x-nemo-relay-bootstrap-fingerprint") @@ -791,6 +829,21 @@ async fn healthz(State(state): State, headers: HeaderMap) -> Response } } }; + if compatible + && let Some(nonce) = headers + .get(GATEWAY_LIFECYCLE_NONCE_HEADER) + .and_then(|value| value.to_str().ok()) + .filter(|nonce| valid_gateway_lifecycle_nonce(nonce)) + && let (Some(key), Some(address)) = + (state.bootstrap_challenge_key.as_ref(), state.local_address) + { + let proof = + key.gateway_lifecycle_health_proof(&state.instance_id, &address.to_string(), nonce); + response_headers.insert( + GATEWAY_LIFECYCLE_PROOF_HEADER, + HeaderValue::from_str(&proof).expect("gateway lifecycle proof is an ASCII value"), + ); + } ( if compatible { StatusCode::OK diff --git a/crates/cli/tests/cli_tests.rs b/crates/cli/tests/cli_tests.rs index 41610fbe8..326426e6a 100644 --- a/crates/cli/tests/cli_tests.rs +++ b/crates/cli/tests/cli_tests.rs @@ -729,7 +729,8 @@ fn start_mcp_client_with_generation( .env("XDG_CONFIG_HOME", temp.join("xdg")) .env("XDG_RUNTIME_DIR", temp.join("runtime")) .env("TMPDIR", temp) - .env("NEMO_RELAY_PLUGIN_IDLE_TIMEOUT_SECS", idle_timeout_secs); + .env("NEMO_RELAY_PLUGIN_IDLE_TIMEOUT_SECS", idle_timeout_secs) + .env("NEMO_RELAY_PLUGIN_HEARTBEAT_INTERVAL_SECS", "1"); if let Some(generation) = generation { let token = std::fs::read_to_string(generation).unwrap(); command @@ -1377,6 +1378,11 @@ impl ChildGuard { let _ = child.kill(); wait_child_with_output(child) } + + fn wait(mut self) -> ExitStatus { + let mut child = self.0.take().unwrap(); + wait_child(&mut child) + } } impl Drop for ChildGuard { @@ -1572,6 +1578,61 @@ fn relay_health(address: SocketAddr) -> serde_json::Value { serde_json::from_str(response.split("\r\n\r\n").nth(1).unwrap()).unwrap() } +fn wait_for_relay_health(address: SocketAddr) -> serde_json::Value { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if TcpStream::connect_timeout(&address, Duration::from_millis(100)).is_ok() { + return relay_health(address); + } + assert!( + Instant::now() < deadline, + "Relay gateway did not become ready at {address}" + ); + thread::sleep(Duration::from_millis(20)); + } +} + +#[test] +fn cli_gateway_stop_stops_explicit_and_legacy_gateway_starts() { + for explicit_start in [true, false] { + let temp = tempfile::tempdir().unwrap(); + let probe = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = probe.local_addr().unwrap(); + drop(probe); + let xdg = temp.path().join("xdg"); + + let mut start = Command::new(gateway_bin()); + start.args(["--bind", &address.to_string()]); + if explicit_start { + start.args(["gateway", "start"]); + } + start + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", &xdg) + .env("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", "1") + .stdout(Stdio::null()) + .stderr(Stdio::null()); + let gateway = ChildGuard::new(start.spawn().unwrap()); + + let health = wait_for_relay_health(address); + assert_eq!(health["service"], "nemo-relay"); + let stop = Command::new(gateway_bin()) + .args(["--bind", &address.to_string(), "gateway", "stop"]) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", &xdg) + .env("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", "1") + .output() + .unwrap(); + assert!( + stop.status.success(), + "gateway stop failed: {}", + String::from_utf8_lossy(&stop.stderr) + ); + assert!(gateway.wait().success()); + wait_for_port_closed(address); + } +} + #[test] fn cli_mcp_clients_share_gateway_until_final_idle_shutdown() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/cli/tests/coverage/agents/plugin_host_tests.rs b/crates/cli/tests/coverage/agents/plugin_host_tests.rs index de331e658..cba9f5990 100644 --- a/crates/cli/tests/coverage/agents/plugin_host_tests.rs +++ b/crates/cli/tests/coverage/agents/plugin_host_tests.rs @@ -3887,13 +3887,11 @@ fn shared_defaults_cover_idle_lifecycle_and_lock_names() { let _plugin_url = EnvVarGuard::remove("NEMO_RELAY_PLUGIN_GATEWAY_URL"); let _claude_url = EnvVarGuard::remove("NEMO_RELAY_GATEWAY_URL"); let _timeout = EnvVarGuard::remove("NEMO_RELAY_PLUGIN_IDLE_TIMEOUT_SECS"); + let _heartbeat = EnvVarGuard::remove("NEMO_RELAY_PLUGIN_HEARTBEAT_INTERVAL_SECS"); let _fail_closed = EnvVarGuard::remove("NEMO_RELAY_FAIL_CLOSED"); assert_eq!(plugin_idle_timeout().unwrap(), Duration::from_secs(300)); - assert_eq!( - plugin_heartbeat_interval().unwrap(), - Duration::from_secs(30) - ); + assert_eq!(plugin_heartbeat_interval().unwrap(), Duration::from_secs(3)); assert_eq!(bootstrap_lock_name(""), "unknown"); } diff --git a/crates/cli/tests/coverage/commands/main_tests.rs b/crates/cli/tests/coverage/commands/main_tests.rs index 95525e17b..fe8c1a8a5 100644 --- a/crates/cli/tests/coverage/commands/main_tests.rs +++ b/crates/cli/tests/coverage/commands/main_tests.rs @@ -3,6 +3,8 @@ use clap::Parser; use std::ffi::OsString; +use std::io::{Read, Write}; +use std::net::TcpListener; use std::path::PathBuf; use super::completions::CompletionsCommand; @@ -83,12 +85,117 @@ fn operational_command_names_cover_logging_exempt_commands() { for (args, expected) in [ (vec!["nemo-relay", "codex"], "codex"), (vec!["nemo-relay", "config"], "config"), + (vec!["nemo-relay", "gateway", "start"], "gateway"), + (vec!["nemo-relay", "gateway", "stop"], "gateway"), ] { let cli = Cli::try_parse_from(args).unwrap(); assert_eq!(cli.command.unwrap().log_name(), expected); } } +#[test] +fn gateway_stop_uses_the_daemon_default_unless_bind_is_explicit() { + let default = Cli::try_parse_from(["nemo-relay", "gateway", "stop"]).unwrap(); + assert_eq!( + gateway::stop_bind(&default.server), + "127.0.0.1:4040".parse().unwrap() + ); + + let explicit = + Cli::try_parse_from(["nemo-relay", "--bind", "127.0.0.1:47632", "gateway", "stop"]) + .unwrap(); + assert_eq!( + gateway::stop_bind(&explicit.server), + "127.0.0.1:47632".parse().unwrap() + ); + + let error = crate::mcp::stop("127.0.0.1:0".parse().unwrap()).unwrap_err(); + assert!(error.to_string().contains("nonzero port"), "{error}"); +} + +#[tokio::test] +async fn gateway_lifecycle_commands_use_the_requested_bind_and_refuse_foreign_listeners() { + let temp = tempfile::tempdir().unwrap(); + let xdg = temp.path().join("xdg"); + std::fs::create_dir_all(&xdg).unwrap(); + let bootstrap_state = xdg.join("nemo-relay/bootstrap"); + let _environment = EnvScope::set_with_cwd_guard( + &[ + ("HOME", Some(temp.path().as_os_str())), + ("XDG_CONFIG_HOME", Some(xdg.as_os_str())), + ( + crate::bootstrap::state::BOOTSTRAP_STATE_DIR_ENV, + Some(bootstrap_state.as_os_str()), + ), + ( + crate::configuration::BOOTSTRAP_FINGERPRINT_ENV, + Some(std::ffi::OsStr::new("test-fingerprint")), + ), + ], + Some(crate::test_support::CwdTestScope::locked()), + ); + + let occupied = TcpListener::bind("127.0.0.1:0").unwrap(); + let occupied_address = occupied.local_addr().unwrap(); + let start = Cli::try_parse_from([ + "nemo-relay", + "--bind", + &occupied_address.to_string(), + "gateway", + "start", + ]) + .unwrap(); + assert!( + run_command(start.command.unwrap(), &start.server, None, None) + .await + .is_err(), + "gateway start should attempt to serve on the requested bind" + ); + drop(occupied); + + let foreign = TcpListener::bind("127.0.0.1:0").unwrap(); + foreign.set_nonblocking(true).unwrap(); + let foreign_address = foreign.local_addr().unwrap(); + let running = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)); + let server_running = running.clone(); + let foreign_server = std::thread::spawn(move || { + while server_running.load(std::sync::atomic::Ordering::Relaxed) { + match foreign.accept() { + Ok((mut stream, _)) => { + stream.set_nonblocking(false).unwrap(); + let mut request = [0_u8; 1024]; + let _ = stream.read(&mut request); + let _ = stream.write_all( + b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}", + ); + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + std::thread::sleep(std::time::Duration::from_millis(5)); + } + Err(error) => panic!("foreign listener failed: {error}"), + } + } + }); + let stop = Cli::try_parse_from([ + "nemo-relay", + "--bind", + &foreign_address.to_string(), + "gateway", + "stop", + ]) + .unwrap(); + let error = run_command(stop.command.unwrap(), &stop.server, None, None) + .await + .unwrap_err() + .to_string(); + assert!( + error.contains("refusing to stop an unverified process"), + "{error}" + ); + running.store(false, std::sync::atomic::Ordering::Relaxed); + foreign_server.join().unwrap(); +} + #[test] fn bootstrap_shutdown_token_is_removed_before_runtime_startup() { let _environment = crate::test_support::EnvScope::set(&[( @@ -580,7 +687,7 @@ async fn run_command_dispatches_safe_plugin_and_install_paths() { ]) .unwrap(); assert_eq!( - run_command(cli.command.unwrap(), &cli.server, None) + run_command(cli.command.unwrap(), &cli.server, None, None) .await .unwrap(), ExitCode::SUCCESS @@ -596,7 +703,7 @@ async fn run_command_dispatches_safe_plugin_and_install_paths() { ]) .unwrap(); assert_eq!( - run_command(cli.command.unwrap(), &cli.server, None) + run_command(cli.command.unwrap(), &cli.server, None, None) .await .unwrap(), ExitCode::SUCCESS @@ -624,7 +731,7 @@ async fn run_command_install_requires_a_valid_bootstrap_key_except_dry_runs() { "--skip-doctor", ]) .unwrap(); - let error = run_command(cli.command.unwrap(), &cli.server, None) + let error = run_command(cli.command.unwrap(), &cli.server, None, None) .await .unwrap_err() .to_string(); @@ -643,7 +750,7 @@ async fn run_command_install_requires_a_valid_bootstrap_key_except_dry_runs() { ]) .unwrap(); assert_eq!( - run_command(cli.command.unwrap(), &cli.server, None) + run_command(cli.command.unwrap(), &cli.server, None, None) .await .unwrap(), ExitCode::SUCCESS diff --git a/crates/cli/tests/coverage/shared/bootstrap_state_tests.rs b/crates/cli/tests/coverage/shared/bootstrap_state_tests.rs index c94580871..2014d6c63 100644 --- a/crates/cli/tests/coverage/shared/bootstrap_state_tests.rs +++ b/crates/cli/tests/coverage/shared/bootstrap_state_tests.rs @@ -6,6 +6,7 @@ use crate::test_support::{EnvScope, accept_bounded, header, read_headers}; use std::ffi::OsStr; use std::io::Write; use std::net::TcpListener; +use std::sync::mpsc; #[test] fn owner_records_are_versioned_endpoint_scoped_and_round_trip() { @@ -219,3 +220,168 @@ fn authenticated_owned_gateway_is_shut_down_and_cleaned_up() { server.join().unwrap(); assert!(!path.exists()); } + +#[test] +fn ownerless_shutdown_ignores_malformed_owner_and_holds_the_startup_lock() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config"); + let _scope = EnvScope::set(&[ + ("XDG_CONFIG_HOME", Some(config.as_os_str())), + ("HOME", Some(dir.path().as_os_str())), + ("USERPROFILE", None), + ]); + let key = crate::configuration::BootstrapChallengeKey::load().unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let url = format!("http://{address}"); + let state = state_dir().unwrap(); + create_private_dir(&state).unwrap(); + let path = owner_path(&state, &url); + std::fs::write(&path, b"{not-valid-json").unwrap(); + let replacement = OwnerRecord::new(84, &url, "replacement-token", Some("replacement")); + let (shutdown_seen, shutdown_received) = mpsc::channel(); + let (lock_attempted, lock_attempt) = mpsc::channel(); + + let publisher_state = state.clone(); + let publisher_url = url.clone(); + let publisher = std::thread::spawn(move || { + shutdown_received + .recv_timeout(Duration::from_secs(2)) + .unwrap(); + let result = lock_endpoint_for(&publisher_state, &publisher_url, Duration::from_millis(25)) + .map(drop); + lock_attempted.send(result).unwrap(); + }); + + let server = std::thread::spawn(move || { + let body = format!( + "{{\"status\":\"ok\",\"service\":\"nemo-relay\",\"version\":\"{}\",\"bootstrap_protocol\":{},\"instance_id\":\"ownerless-instance\"}}", + env!("CARGO_PKG_VERSION"), + BOOTSTRAP_PROTOCOL_VERSION + ); + + let mut probe = accept_bounded(&listener); + read_headers(&mut probe); + probe + .write_all( + format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .unwrap(); + + let mut lifecycle = accept_bounded(&listener); + let challenge = read_headers(&mut lifecycle); + let nonce = header( + &challenge, + crate::configuration::GATEWAY_LIFECYCLE_NONCE_HEADER, + ); + let proof = + key.gateway_lifecycle_health_proof("ownerless-instance", &address.to_string(), &nonce); + lifecycle + .write_all( + format!( + "HTTP/1.1 200 OK\r\n{}: {proof}\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n{body}", + crate::configuration::GATEWAY_LIFECYCLE_PROOF_HEADER, + body.len() + ) + .as_bytes(), + ) + .unwrap(); + let shutdown = read_headers(&mut lifecycle); + assert!(shutdown.starts_with("POST /bootstrap/shutdown HTTP/1.1")); + shutdown_seen.send(()).unwrap(); + let error = lock_attempt + .recv_timeout(Duration::from_secs(2)) + .unwrap() + .unwrap_err(); + assert!(error.contains("timed out waiting"), "{error}"); + drop(listener); + lifecycle + .write_all(b"HTTP/1.1 204 No Content\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + .unwrap(); + }); + + stop_gateway_and_reset(&url).unwrap(); + server.join().unwrap(); + publisher.join().unwrap(); + assert!(!path.exists()); + + let _lock = lock_endpoint(&state, &url).unwrap(); + write_owner_record(&path, &replacement).unwrap(); + assert_eq!(read_owner_record(&path).unwrap(), Some(replacement)); +} + +#[test] +fn ownerless_shutdown_retries_a_transient_probe_failure() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config"); + let _scope = EnvScope::set(&[ + ("XDG_CONFIG_HOME", Some(config.as_os_str())), + ("HOME", Some(dir.path().as_os_str())), + ("USERPROFILE", None), + ]); + let key = crate::configuration::BootstrapChallengeKey::load().unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let url = format!("http://{address}"); + let state = state_dir().unwrap(); + create_private_dir(&state).unwrap(); + let path = owner_path(&state, &url); + std::fs::write(&path, b"{not-valid-json").unwrap(); + + let server = std::thread::spawn(move || { + let body = format!( + "{{\"status\":\"ok\",\"service\":\"nemo-relay\",\"version\":\"{}\",\"bootstrap_protocol\":{},\"instance_id\":\"ownerless-instance\"}}", + env!("CARGO_PKG_VERSION"), + BOOTSTRAP_PROTOCOL_VERSION + ); + + let mut probe = accept_bounded(&listener); + read_headers(&mut probe); + probe + .write_all( + format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .unwrap(); + + let mut lifecycle = accept_bounded(&listener); + let challenge = read_headers(&mut lifecycle); + let nonce = header( + &challenge, + crate::configuration::GATEWAY_LIFECYCLE_NONCE_HEADER, + ); + let proof = + key.gateway_lifecycle_health_proof("ownerless-instance", &address.to_string(), &nonce); + lifecycle + .write_all( + format!( + "HTTP/1.1 200 OK\r\n{}: {proof}\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n{body}", + crate::configuration::GATEWAY_LIFECYCLE_PROOF_HEADER, + body.len() + ) + .as_bytes(), + ) + .unwrap(); + let shutdown = read_headers(&mut lifecycle); + assert!(shutdown.starts_with("POST /bootstrap/shutdown HTTP/1.1")); + lifecycle + .write_all(b"HTTP/1.1 204 No Content\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + .unwrap(); + drop(lifecycle); + + let transient_probe = accept_bounded(&listener); + drop(listener); + drop(transient_probe); + }); + + stop_gateway_and_reset(&url).unwrap(); + server.join().unwrap(); + assert!(!path.exists()); +} diff --git a/crates/cli/tests/coverage/shared/bootstrap_tests.rs b/crates/cli/tests/coverage/shared/bootstrap_tests.rs index c5fed8d32..ff2aca6d2 100644 --- a/crates/cli/tests/coverage/shared/bootstrap_tests.rs +++ b/crates/cli/tests/coverage/shared/bootstrap_tests.rs @@ -193,7 +193,7 @@ fn persistent_gateway_resolution_keeps_server_configuration_in_one_spec() { } #[test] -fn idle_timeout_drives_heartbeat_and_rejects_invalid_values() { +fn heartbeat_interval_is_independent_of_idle_timeout_and_rejects_invalid_values() { let _environment = EnvScope::set(&[( crate::configuration::PLUGIN_IDLE_TIMEOUT_ENV, Some(OsStr::new("9")), @@ -202,6 +202,20 @@ fn idle_timeout_drives_heartbeat_and_rejects_invalid_values() { assert_eq!(plugin_heartbeat_interval().unwrap(), Duration::from_secs(3)); drop(_environment); + let _environment = EnvScope::set(&[ + ( + crate::configuration::PLUGIN_IDLE_TIMEOUT_ENV, + Some(OsStr::new("300")), + ), + ( + crate::configuration::PLUGIN_HEARTBEAT_INTERVAL_ENV, + Some(OsStr::new("7")), + ), + ]); + assert_eq!(plugin_idle_timeout().unwrap(), Duration::from_secs(300)); + assert_eq!(plugin_heartbeat_interval().unwrap(), Duration::from_secs(7)); + drop(_environment); + let _environment = EnvScope::set(&[( crate::configuration::PLUGIN_IDLE_TIMEOUT_ENV, Some(OsStr::new("0")), @@ -211,6 +225,17 @@ fn idle_timeout_drives_heartbeat_and_rejects_invalid_values() { .unwrap_err() .contains("greater than 0") ); + drop(_environment); + + let _environment = EnvScope::set(&[( + crate::configuration::PLUGIN_HEARTBEAT_INTERVAL_ENV, + Some(OsStr::new("0")), + )]); + assert!( + plugin_heartbeat_interval() + .unwrap_err() + .contains("greater than 0") + ); } #[test] diff --git a/crates/cli/tests/coverage/shared/config_tests.rs b/crates/cli/tests/coverage/shared/config_tests.rs index 54d63b2d7..82d01d9d4 100644 --- a/crates/cli/tests/coverage/shared/config_tests.rs +++ b/crates/cli/tests/coverage/shared/config_tests.rs @@ -229,6 +229,7 @@ struct PluginConfigDiscoveryScope { previous_anthropic_auth_header: Option, previous_bootstrap_fingerprint: Option, previous_plugin_idle_timeout: Option, + previous_plugin_heartbeat_interval: Option, previous_test_skip_implicit_config: Option, } @@ -247,6 +248,7 @@ impl PluginConfigDiscoveryScope { let previous_anthropic_auth_header = std::env::var_os("NEMO_RELAY_ANTHROPIC_AUTH_HEADER"); let previous_bootstrap_fingerprint = std::env::var_os(BOOTSTRAP_FINGERPRINT_ENV); let previous_plugin_idle_timeout = std::env::var_os(PLUGIN_IDLE_TIMEOUT_ENV); + let previous_plugin_heartbeat_interval = std::env::var_os(PLUGIN_HEARTBEAT_INTERVAL_ENV); let previous_test_skip_implicit_config = std::env::var_os("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG"); unsafe { @@ -258,6 +260,7 @@ impl PluginConfigDiscoveryScope { std::env::remove_var("NEMO_RELAY_ANTHROPIC_AUTH_HEADER"); std::env::remove_var(BOOTSTRAP_FINGERPRINT_ENV); std::env::remove_var(PLUGIN_IDLE_TIMEOUT_ENV); + std::env::remove_var(PLUGIN_HEARTBEAT_INTERVAL_ENV); std::env::remove_var("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG"); } std::env::set_current_dir(cwd).unwrap(); @@ -273,6 +276,7 @@ impl PluginConfigDiscoveryScope { previous_anthropic_auth_header, previous_bootstrap_fingerprint, previous_plugin_idle_timeout, + previous_plugin_heartbeat_interval, previous_test_skip_implicit_config, } } @@ -345,6 +349,10 @@ impl Drop for PluginConfigDiscoveryScope { Some(value) => std::env::set_var(PLUGIN_IDLE_TIMEOUT_ENV, value), None => std::env::remove_var(PLUGIN_IDLE_TIMEOUT_ENV), } + match self.previous_plugin_heartbeat_interval.take() { + Some(value) => std::env::set_var(PLUGIN_HEARTBEAT_INTERVAL_ENV, value), + None => std::env::remove_var(PLUGIN_HEARTBEAT_INTERVAL_ENV), + } match self.previous_test_skip_implicit_config.take() { Some(value) => std::env::set_var("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", value), None => std::env::remove_var("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG"), diff --git a/crates/cli/tests/coverage/shared/gateway_client_tests.rs b/crates/cli/tests/coverage/shared/gateway_client_tests.rs index 6d7b3284d..f3b48905c 100644 --- a/crates/cli/tests/coverage/shared/gateway_client_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_client_tests.rs @@ -79,6 +79,63 @@ fn serve_verified_shutdown( (url, receiver, server) } +fn serve_lifecycle_shutdown( + key: crate::configuration::BootstrapChallengeKey, +) -> (String, mpsc::Receiver, thread::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap().to_string(); + let url = format!("http://{address}"); + let (sender, receiver) = mpsc::channel(); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + let challenge = read_headers(&mut stream); + let nonce = header( + &challenge, + crate::configuration::GATEWAY_LIFECYCLE_NONCE_HEADER, + ); + let instance_id = "lifecycle-instance"; + let proof = key.gateway_lifecycle_health_proof(instance_id, &address, &nonce); + let body = format!( + "{{\"status\":\"ok\",\"service\":\"nemo-relay\",\"version\":\"{}\",\"bootstrap_protocol\":{},\"instance_id\":\"{instance_id}\"}}", + env!("CARGO_PKG_VERSION"), + BOOTSTRAP_PROTOCOL_VERSION + ); + stream + .write_all( + format!( + "HTTP/1.1 200 OK\r\n{}: {proof}\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n{body}", + crate::configuration::GATEWAY_LIFECYCLE_PROOF_HEADER, + body.len() + ) + .as_bytes(), + ) + .unwrap(); + let request = read_headers(&mut stream); + assert_eq!( + header( + &request, + crate::configuration::GATEWAY_LIFECYCLE_NONCE_HEADER + ), + nonce + ); + assert_eq!( + header( + &request, + crate::configuration::GATEWAY_LIFECYCLE_PROOF_HEADER + ), + key.gateway_lifecycle_shutdown_proof(instance_id, &address, &nonce) + ); + let _ = sender.send(request); + stream + .write_all(b"HTTP/1.1 204 No Content\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + .unwrap(); + }); + (url, receiver, server) +} + #[test] fn shutdown_request_sends_the_private_token_and_accepts_no_content() { let temp = tempfile::tempdir().unwrap(); @@ -106,6 +163,28 @@ fn shutdown_request_sends_the_private_token_and_accepts_no_content() { server.join().unwrap(); } +#[test] +fn lifecycle_shutdown_authenticates_the_bind_without_an_owner_record() { + let temp = tempfile::tempdir().unwrap(); + let _environment = EnvScope::set(&[ + ("XDG_CONFIG_HOME", Some(temp.path().as_os_str())), + ("HOME", Some(temp.path().as_os_str())), + ]); + let key = crate::configuration::BootstrapChallengeKey::load().unwrap(); + let (url, request, server) = serve_lifecycle_shutdown(key); + + assert_eq!( + request_lifecycle_shutdown(&url).unwrap(), + "lifecycle-instance" + ); + let request = request.recv_timeout(Duration::from_secs(2)).unwrap(); + assert!( + request.starts_with("POST /bootstrap/shutdown HTTP/1.1"), + "{request}" + ); + server.join().unwrap(); +} + #[test] fn shutdown_request_reports_rejection_without_hiding_the_status() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index 65968230f..be163d6b3 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -3,6 +3,7 @@ use std::ffi::OsString; use std::future::Future; +use std::net::SocketAddr; use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -707,7 +708,7 @@ async fn bootstrap_shutdown_requires_the_private_owner_token() { None, false, Some(BootstrapShutdown { - token: "private-token".into(), + token: Some("private-token".into()), sender: Arc::new(std::sync::Mutex::new(Some(sender))), }), None, @@ -744,6 +745,89 @@ async fn bootstrap_shutdown_requires_the_private_owner_token() { .unwrap(); } +#[tokio::test] +async fn gateway_lifecycle_shutdown_requires_a_domain_separated_proof() { + let key = BootstrapChallengeKey::from_bytes(b"test lifecycle key"); + let address: SocketAddr = "127.0.0.1:4040".parse().unwrap(); + let nonce = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + let (sender, receiver) = oneshot::channel(); + let mut state = AppState::new_with_bootstrap( + test_config(), + None, + Some(key.clone()), + false, + Some(BootstrapShutdown { + token: None, + sender: Arc::new(std::sync::Mutex::new(Some(sender))), + }), + None, + ); + state.local_address = Some(address); + let instance_id = state.instance_id.clone(); + let app = router_with_state(state); + + let health = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/healthz") + .header(crate::configuration::GATEWAY_LIFECYCLE_NONCE_HEADER, nonce) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(health.status(), StatusCode::OK); + assert_eq!( + health + .headers() + .get(crate::configuration::GATEWAY_LIFECYCLE_PROOF_HEADER) + .unwrap(), + key.gateway_lifecycle_health_proof(&instance_id, &address.to_string(), nonce) + .as_str() + ); + + let rejected = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/bootstrap/shutdown") + .header(crate::configuration::GATEWAY_LIFECYCLE_NONCE_HEADER, nonce) + .header( + crate::configuration::GATEWAY_LIFECYCLE_PROOF_HEADER, + key.gateway_lifecycle_health_proof(&instance_id, &address.to_string(), nonce), + ) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(rejected.status(), StatusCode::FORBIDDEN); + + let accepted = app + .oneshot( + Request::builder() + .method("POST") + .uri("/bootstrap/shutdown") + .header(crate::configuration::GATEWAY_LIFECYCLE_NONCE_HEADER, nonce) + .header( + crate::configuration::GATEWAY_LIFECYCLE_PROOF_HEADER, + key.gateway_lifecycle_shutdown_proof(&instance_id, &address.to_string(), nonce), + ) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(accepted.status(), StatusCode::NO_CONTENT); + tokio::time::timeout(std::time::Duration::from_secs(1), receiver) + .await + .expect("shutdown signal was not delivered") + .unwrap(); +} + #[test] fn readiness_file_is_published_atomically_with_gateway_identity() { let directory = tempfile::tempdir().unwrap(); diff --git a/docs/nemo-relay-cli/claude-code.mdx b/docs/nemo-relay-cli/claude-code.mdx index e5ac10ae7..aaa5fa183 100644 --- a/docs/nemo-relay-cli/claude-code.mdx +++ b/docs/nemo-relay-cli/claude-code.mdx @@ -68,8 +68,11 @@ The plugin starts `nemo-relay mcp`, a lightweight Rust lifecycle client that starts or reuses the shared gateway on `127.0.0.1:47632` immediately when the MCP process launches. The client verifies the gateway identity and effective persistent configuration, heartbeats it while MCP stdio remains open, and -performs one coordinated restart if the gateway becomes unhealthy. Claude -Code and Codex MCP clients share a compatible gateway, and the gateway +performs one coordinated restart if the gateway becomes unhealthy. The +heartbeat defaults to 3 seconds; set +`NEMO_RELAY_PLUGIN_HEARTBEAT_INTERVAL_SECS` to a positive integer to adjust it +without changing the gateway's idle timeout. Claude Code and Codex MCP clients +share a compatible gateway, and the gateway exits after the final client's idle timeout. The MCP server advertises no tools. The generated entry sets `alwaysLoad: true`, so Claude Code 2.1.121 or newer @@ -84,6 +87,22 @@ wrapper for project-specific `.nemo-relay` configuration. Run `nemo-relay install claude-code --force` to replace an existing generation-fenced installation safely. +To explicitly stop the managed shared gateway, first close active MCP clients, +then run: + +```bash +nemo-relay --bind 127.0.0.1:47632 gateway stop +``` + +The command authenticates the Relay process at the configured endpoint and +never stops a foreign process using the same endpoint. An open MCP client can +restart the managed gateway as part of its one allowed recovery. + +`nemo-relay gateway start` starts the gateway with the same server +configuration as a bare daemon invocation. `nemo-relay gateway stop` stops that +gateway at the default daemon endpoint, `127.0.0.1:4040`. Pass the same +`--bind` value to both commands when using a custom endpoint. + Check or remove the installed plugin with: ```bash diff --git a/docs/nemo-relay-cli/codex.mdx b/docs/nemo-relay-cli/codex.mdx index d36d2e363..6988fc1b7 100644 --- a/docs/nemo-relay-cli/codex.mdx +++ b/docs/nemo-relay-cli/codex.mdx @@ -158,10 +158,28 @@ snapshot budget error, remove unrelated files from the manifest or load-target directory, flatten deeply nested directories, or reduce the managed Python environment before retrying. Concurrent Codex and Claude Code processes can share the sidecar. Each MCP client sends a heartbeat -while its stdio connection is open, and the sidecar exits after 300 seconds +while its stdio connection is open. The heartbeat defaults to 3 seconds; set +`NEMO_RELAY_PLUGIN_HEARTBEAT_INTERVAL_SECS` to a positive integer to adjust it +without changing the gateway's idle timeout. The sidecar exits after 300 seconds without activity by default. Relay does not install a wrapper, launch agent, system user service, scheduled task, login item, or persistent supervisor. +To explicitly stop the managed shared gateway, first close active MCP clients, +then run: + +```bash +nemo-relay --bind 127.0.0.1:47632 gateway stop +``` + +The command authenticates the Relay process at the configured endpoint and +never stops a foreign process using the same endpoint. An open MCP client can +restart the managed gateway as part of its one allowed recovery. + +`nemo-relay gateway start` starts the gateway with the same server +configuration as a bare daemon invocation. `nemo-relay gateway stop` stops that +gateway at the default daemon endpoint, `127.0.0.1:4040`. Pass the same +`--bind` value to both commands when using a custom endpoint. + On Windows, Relay requests Job Object breakaway when the host job permits it. When breakaway is unavailable and the host permits nested assignment, the gateway remains scoped to the host job and the usual 300-second idle reuse diff --git a/docs/nemo-relay-cli/plugin-installation.mdx b/docs/nemo-relay-cli/plugin-installation.mdx index 562ce5661..34e43340a 100644 --- a/docs/nemo-relay-cli/plugin-installation.mdx +++ b/docs/nemo-relay-cli/plugin-installation.mdx @@ -132,7 +132,9 @@ Claude Code and Codex processes use the same `nemo-relay mcp` lifecycle client. Before reading MCP protocol frames, it starts or reuses a detached native sidecar on `127.0.0.1:47632`. The client verifies Relay identity, version, protocol readiness, and effective user-level -configuration, then heartbeats the sidecar every 30 seconds. Concurrent +configuration, then heartbeats the sidecar every 3 seconds by default. Set +`NEMO_RELAY_PLUGIN_HEARTBEAT_INTERVAL_SECS` to a positive integer to adjust it +without changing the gateway's idle timeout. Concurrent processes from any host share the gateway. After the final MCP client closes, the sidecar exits after 300 idle seconds by default. If the gateway exits while MCP stdio is open, the client performs one coordinated restart and fails if @@ -156,6 +158,22 @@ payload once. They never start or recover the gateway, and Relay does not retry after payload transmission begins. The MCP server advertises no tools in any host. +To explicitly stop the managed shared gateway, first close active MCP clients, +then run: + +```bash +nemo-relay --bind 127.0.0.1:47632 gateway stop +``` + +The command authenticates the Relay process at the configured endpoint and +never stops a foreign process using the same endpoint. An open MCP client can +restart the managed gateway as part of its one allowed recovery. + +`nemo-relay gateway start` starts the gateway with the same server +configuration as a bare daemon invocation. `nemo-relay gateway stop` stops that +gateway at the default daemon endpoint, `127.0.0.1:4040`. Pass the same +`--bind` value to both commands when using a custom endpoint. + ### Hook Delivery and Upgrade Safety Before contacting the gateway, Relay rejects a fixed-endpoint `hook-forward` diff --git a/integrations/coding-agents/codex/.mcp.json b/integrations/coding-agents/codex/.mcp.json index c45737935..bd116f0af 100644 --- a/integrations/coding-agents/codex/.mcp.json +++ b/integrations/coding-agents/codex/.mcp.json @@ -43,6 +43,7 @@ "NEMO_RELAY_MAX_PASSTHROUGH_BODY_BYTES", "NEMO_RELAY_OPENAI_AUTH_HEADER", "NEMO_RELAY_OPENAI_BASE_URL", + "NEMO_RELAY_PLUGIN_HEARTBEAT_INTERVAL_SECS", "NEMO_RELAY_PLUGIN_IDLE_TIMEOUT_SECS", "NEMO_RELAY_PYTHON", "NEMO_RELAY_TRANSPARENT_RUN",