Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion crates/cli/src/bootstrap/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -571,8 +571,23 @@ pub(crate) fn plugin_idle_timeout() -> Result<Duration, String> {
Ok(Duration::from_secs(seconds))
}

/// Returns the persistent MCP gateway heartbeat interval.
pub(crate) fn plugin_heartbeat_interval() -> Result<Duration, String> {
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::<u64>().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)]
Expand Down
95 changes: 91 additions & 4 deletions crates/cli/src/bootstrap/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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(());
};
Expand Down Expand Up @@ -275,20 +281,87 @@ 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"
));
}
Comment thread
mnajafian-nv marked this conversation as resolved.
}
}
}

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}"))?;
atomic_write(path, &bytes)
}

pub(super) fn read_owner_record(path: &Path) -> Result<Option<OwnerRecord>, 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<Option<Vec<u8>>, 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}",
Expand All @@ -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()
Expand Down
50 changes: 50 additions & 0 deletions crates/cli/src/commands/gateway.rs
Original file line number Diff line number Diff line change
@@ -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<String>,
) -> Result<ExitCode, CliError> {
match command.command {
GatewaySubcommand::Start => super::serve_gateway(server, bootstrap_shutdown_token).await,
Comment thread
mnajafian-nv marked this conversation as resolved.
GatewaySubcommand::Stop => crate::mcp::stop(stop_bind(server)),
}
}
77 changes: 53 additions & 24 deletions crates/cli/src/commands/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
mod completions;
mod configure;
mod diagnostics;
mod gateway;
mod hook_forward;
mod install;
mod logging;
Expand Down Expand Up @@ -60,6 +61,10 @@ struct LoggingSetup {

fn configure_logging(cli: &Cli) -> Result<LoggingSetup, error::CliError> {
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()
Expand All @@ -78,6 +83,9 @@ fn configure_logging(cli: &Cli) -> Result<LoggingSetup, error::CliError> {
// 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
Expand Down Expand Up @@ -125,7 +133,15 @@ async fn dispatch(bootstrap_shutdown_token: Option<String>) -> Result<ExitCode,
);

let result = match cli.command {
Some(command) => 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 {
Expand Down Expand Up @@ -165,6 +181,7 @@ async fn run_command(
command: Command,
server: &ServerArgs,
logging_fallback_error: Option<&error::CliError>,
bootstrap_shutdown_token: Option<String>,
) -> Result<ExitCode, error::CliError> {
match command {
Command::HookForward(command) => {
Expand All @@ -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),
Expand Down Expand Up @@ -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,
Expand All @@ -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<String>,
) -> Result<ExitCode, error::CliError> {
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<ExitCode, error::CliError> {
completions::execute(command)
Expand Down
5 changes: 5 additions & 0 deletions crates/cli/src/commands/root.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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`)
Expand Down Expand Up @@ -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",
Expand All @@ -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))
}
Expand Down
Loading
Loading