mirror of
https://github.com/akitaonrails/ai-memory.git
synced 2026-10-02 03:24:46 +08:00
fix(serve): stop on SIGINT and SIGTERM on both transports (#699)
The stdio transport listened for no signal at all, and the HTTP transport listened for SIGINT only, so SIGTERM — what `docker stop`, `docker compose down` and `systemctl stop` send — reached no handler on either. What that cost depended on whether the server was PID 1. In the image it is: the ENTRYPOINT is exec form with no init shim, and for PID 1 the kernel discards a signal whose handler is not installed, so the signal was not merely unhandled, it was invisible. `docker stop` sat out its whole grace period and ended in SIGKILL, `docker kill` was the only way out, and Ctrl-C on stdio did nothing at all. Everywhere else — under the native systemd unit, or a plain `ai-memory serve` in a terminal — the process is not PID 1, so the same signal fell through to the kernel's default disposition and killed the process instantly instead. That stop was fast but unclean: no drain, and the durable SessionEnd consolidation worker cut off mid-flight. For those operators this change makes stopping slower, up to the five-second bound, and correct. Handling the signal was still not enough on stdio. The MCP transport reads stdin from a tokio blocking thread, a blocking read cannot be cancelled, and dropping the runtime waits for in-flight blocking work forever — so a handled Ctrl-C left the process parked on that read. The runtime is now built by hand and shut down without waiting for blocking work nothing is reading the result of any more. Both transports install their listeners before the transport starts, so a signal arriving during a slow boot (migrations, the pre-migration archive) is handled rather than lost, and every wait on the shutdown path — the connection drain and the consolidation worker's join — is bounded at five seconds, so a stateful or SSE client holding a connection open cannot stall the exit.
This commit is contained in:
viniciusdsandrade
@@ -51,6 +51,28 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
and Podman. Additionally, `emit_docker_run_script` now preserves volume mount
|
||||
modes (such as `:Z` on SELinux/Podman environments) and filters transient
|
||||
runtime environment variables (`HOSTNAME`, `container=podman`). (#673)
|
||||
- `ai-memory serve` now stops on Ctrl-C and on SIGTERM, on both transports.
|
||||
The stdio transport listened for no signal at all, and the HTTP transport
|
||||
listened for SIGINT alone — so SIGTERM, what `docker stop`, `docker compose
|
||||
down` and `systemctl stop` send, reached no handler on either. What that
|
||||
cost depended on whether the server was PID 1. In the container it is (the
|
||||
image's ENTRYPOINT is exec form, with no init shim), and for PID 1 the
|
||||
kernel discards a signal whose handler is not installed: the signal was not
|
||||
merely unhandled, it was invisible, so `docker stop` sat out its whole grace
|
||||
period and ended in SIGKILL, `docker kill` was the only way out, and Ctrl-C
|
||||
on stdio did nothing at all. Everywhere else — under the native systemd
|
||||
unit, or a plain `ai-memory serve` in a terminal — the process is not PID 1,
|
||||
so the same signal fell through to the kernel's default disposition and
|
||||
killed it instantly instead, with no drain at all: the durable SessionEnd
|
||||
consolidation worker was cut off mid-flight rather than drained. Both
|
||||
transports now listen for SIGINT and SIGTERM, log which one arrived, and
|
||||
bound the drain at five seconds so a stateful or SSE MCP client holding a
|
||||
connection open cannot stall the exit — a stop that used to be instant and
|
||||
unclean now takes up to those five seconds and drains. The listeners are
|
||||
installed before the transport starts, so a signal arriving during a slow
|
||||
boot — migrations, the pre-migration archive — is handled rather than lost,
|
||||
and no container init shim (`tini`, `docker run --init`) is needed for the
|
||||
server to stop as PID 1 (#699).
|
||||
|
||||
## [2.1.1] - 2026-09-07
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! `ai-memory serve` — MCP server with optional filesystem watcher.
|
||||
|
||||
use std::future::Future;
|
||||
use std::future::{Future, IntoFuture};
|
||||
use std::net::SocketAddr;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
@@ -70,6 +70,98 @@ const SESSION_CONSOLIDATION_LEASE: Duration = Duration::from_secs(10 * 60);
|
||||
/// Lock file guarding a data dir against a second `ai-memory serve` (#563).
|
||||
const SERVE_LOCK_FILE: &str = ".serve.lock";
|
||||
|
||||
/// How long a wait on the shutdown path may run before the drain it is
|
||||
/// waiting for is abandoned. axum's graceful shutdown waits for every
|
||||
/// in-flight connection and a stateful or SSE MCP client can hold one open
|
||||
/// indefinitely, so an unbounded drain is indistinguishable from ignoring the
|
||||
/// signal: `docker stop` and `systemctl stop` would still burn their own
|
||||
/// grace period and finish with SIGKILL (#699). Each wait is bounded on its
|
||||
/// own, so a stop can take a small multiple of this.
|
||||
const SHUTDOWN_GRACE: Duration = Duration::from_secs(5);
|
||||
|
||||
/// The signals that stop a running server, listened for on both transports.
|
||||
///
|
||||
/// Installed before the transport starts. The server runs as PID 1 under
|
||||
/// `docker run` (no init shim), and for PID 1 the kernel discards any signal
|
||||
/// whose handler is not installed — so a SIGTERM arriving during a slow boot
|
||||
/// (the pre-migration archive, wiki migrations) must already have a listener
|
||||
/// waiting for it. A tokio `Signal` queues a signal received before the first
|
||||
/// `recv`, so registering early loses nothing (#699).
|
||||
struct ShutdownSignals {
|
||||
#[cfg(unix)]
|
||||
interrupt: Option<tokio::signal::unix::Signal>,
|
||||
#[cfg(unix)]
|
||||
terminate: Option<tokio::signal::unix::Signal>,
|
||||
}
|
||||
|
||||
impl ShutdownSignals {
|
||||
/// Install the listeners.
|
||||
///
|
||||
/// A listener that cannot be registered degrades to the remaining one with
|
||||
/// a warning: losing one way to stop the server is bad, refusing to start
|
||||
/// over it is worse.
|
||||
#[cfg(unix)]
|
||||
fn install() -> Self {
|
||||
use tokio::signal::unix::{SignalKind, signal};
|
||||
|
||||
fn listen(kind: SignalKind, name: &str) -> Option<tokio::signal::unix::Signal> {
|
||||
match signal(kind) {
|
||||
Ok(stream) => Some(stream),
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
%error,
|
||||
signal = name,
|
||||
"cannot listen for this shutdown signal; the server will not stop on it"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Self {
|
||||
interrupt: listen(SignalKind::interrupt(), "SIGINT"),
|
||||
terminate: listen(SignalKind::terminate(), "SIGTERM"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Install the listeners. Non-unix has only ctrl-c.
|
||||
#[cfg(not(unix))]
|
||||
fn install() -> Self {
|
||||
Self {}
|
||||
}
|
||||
|
||||
/// Resolve with the name of the first shutdown signal to arrive.
|
||||
#[cfg(unix)]
|
||||
async fn recv(&mut self) -> &'static str {
|
||||
match (self.interrupt.as_mut(), self.terminate.as_mut()) {
|
||||
(Some(interrupt), Some(terminate)) => tokio::select! {
|
||||
_ = interrupt.recv() => "SIGINT",
|
||||
_ = terminate.recv() => "SIGTERM",
|
||||
},
|
||||
(Some(interrupt), None) => {
|
||||
interrupt.recv().await;
|
||||
"SIGINT"
|
||||
}
|
||||
(None, Some(terminate)) => {
|
||||
terminate.recv().await;
|
||||
"SIGTERM"
|
||||
}
|
||||
// Both registrations failed: there is nothing left to wait for.
|
||||
(None, None) => std::future::pending().await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve with the name of the first shutdown signal to arrive.
|
||||
#[cfg(not(unix))]
|
||||
async fn recv(&mut self) -> &'static str {
|
||||
if let Err(error) = tokio::signal::ctrl_c().await {
|
||||
tracing::warn!(%error, "ctrl-c listener failed; the server will not stop on it");
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
"ctrl-c"
|
||||
}
|
||||
}
|
||||
|
||||
/// The single-instance guard for `ai-memory serve`: an exclusive `flock` on
|
||||
/// `<data-dir>/.serve.lock` held for the process lifetime. The OS releases it
|
||||
/// when the process exits, so a crashed server never locks the operator out of
|
||||
@@ -707,6 +799,11 @@ async fn run_session_consolidation_worker(
|
||||
/// Returns an error if the store cannot be opened, the watcher cannot
|
||||
/// install, or the transport setup fails.
|
||||
pub async fn run(config: &Config, args: ServeArgs) -> Result<()> {
|
||||
// Before anything slow: boot takes the pre-migration archive and runs the
|
||||
// wiki migrations, and a signal arriving in that window has to be caught
|
||||
// rather than fall through to the default disposition (#699).
|
||||
let mut shutdown = ShutdownSignals::install();
|
||||
|
||||
validate_web_ui_args(args.enable_web, args.web_ui_dir.as_deref())?;
|
||||
|
||||
// Merge config + CLI CORS origins (config first, CLI adds new entries).
|
||||
@@ -905,8 +1002,50 @@ pub async fn run(config: &Config, args: ServeArgs) -> Result<()> {
|
||||
match args.transport {
|
||||
TransportKind::Stdio => {
|
||||
info!("MCP server ready on stdio (Ctrl-C to stop)");
|
||||
let service = server.serve(stdio()).await?;
|
||||
service.waiting().await?;
|
||||
// `serve` resolves only once a client has completed the MCP
|
||||
// `initialize` handshake, so the signal races the handshake as
|
||||
// well as the session that follows it: a Ctrl-C before any client
|
||||
// connected is the exact state a launched-but-unused server sits
|
||||
// in, and until #699 nothing here listened for one at all.
|
||||
let service = tokio::select! {
|
||||
service = server.serve(stdio()) => Some(service?),
|
||||
signal = shutdown.recv() => {
|
||||
info!(signal, "shutdown signal received before a client connected; stopping");
|
||||
None
|
||||
}
|
||||
};
|
||||
if let Some(service) = service {
|
||||
// Take the token before `waiting` consumes the service:
|
||||
// stopping the transport is the only way out of that await.
|
||||
let stop = service.cancellation_token();
|
||||
let mut waiting = std::pin::pin!(service.waiting());
|
||||
let signal = tokio::select! {
|
||||
result = &mut waiting => {
|
||||
result?;
|
||||
None
|
||||
}
|
||||
signal = shutdown.recv() => Some(signal),
|
||||
};
|
||||
if let Some(signal) = signal {
|
||||
info!(
|
||||
signal,
|
||||
"shutdown signal received; stopping the stdio transport"
|
||||
);
|
||||
stop.cancel();
|
||||
// A signal is a normal stop, so nothing below turns it
|
||||
// into a failing exit — but neither does it wait forever.
|
||||
match tokio::time::timeout(SHUTDOWN_GRACE, waiting).await {
|
||||
Ok(Ok(_)) => {}
|
||||
Ok(Err(error)) => {
|
||||
tracing::warn!(%error, "stdio transport ended abnormally during shutdown");
|
||||
}
|
||||
Err(_) => tracing::warn!(
|
||||
grace_secs = SHUTDOWN_GRACE.as_secs(),
|
||||
"stdio transport did not stop within the shutdown grace period; exiting anyway"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
TransportKind::Http => {
|
||||
let bind = args.bind.unwrap_or_else(|| config.bind.clone());
|
||||
@@ -1286,23 +1425,58 @@ pub async fn run(config: &Config, args: ServeArgs) -> Result<()> {
|
||||
);
|
||||
}
|
||||
let shutdown_cancel = cancel.clone();
|
||||
let serve_result = axum::serve(
|
||||
listener,
|
||||
router.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(async move {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
info!("ctrl-c received; shutting down");
|
||||
shutdown_cancel.cancel();
|
||||
})
|
||||
.await;
|
||||
let serve_result = {
|
||||
let serve = axum::serve(
|
||||
listener,
|
||||
router.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(async move {
|
||||
let signal = shutdown.recv().await;
|
||||
info!(signal, "shutdown signal received; draining");
|
||||
shutdown_cancel.cancel();
|
||||
})
|
||||
.into_future();
|
||||
let mut serve = std::pin::pin!(serve);
|
||||
// Bound the drain. axum waits for every in-flight connection
|
||||
// to close, and a stateful or SSE MCP client never closes one
|
||||
// on its own — without this the shutdown outlives the
|
||||
// supervisor's own patience and ends in SIGKILL (#699).
|
||||
tokio::select! {
|
||||
result = &mut serve => Some(result),
|
||||
() = async {
|
||||
cancel.cancelled().await;
|
||||
tokio::time::sleep(SHUTDOWN_GRACE).await;
|
||||
} => {
|
||||
tracing::warn!(
|
||||
grace_secs = SHUTDOWN_GRACE.as_secs(),
|
||||
"connections still open past the shutdown grace period; exiting anyway"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
};
|
||||
cancel.cancel();
|
||||
if let Some(task) = session_consolidation_task
|
||||
&& let Err(error) = task.await
|
||||
{
|
||||
tracing::warn!(%error, "SessionEnd consolidation worker join failed");
|
||||
if let Some(task) = session_consolidation_task {
|
||||
// Bounded like the drain above. The worker can be parked in
|
||||
// `claim_session_consolidation` or `release_session_consolidation`,
|
||||
// neither of which races the cancellation, behind the
|
||||
// single-writer actor's queue — an unbounded join here would
|
||||
// sit outside the shutdown bound entirely (#699).
|
||||
match tokio::time::timeout(SHUTDOWN_GRACE, task).await {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(error)) => {
|
||||
tracing::warn!(%error, "SessionEnd consolidation worker join failed");
|
||||
}
|
||||
Err(_) => tracing::warn!(
|
||||
grace_secs = SHUTDOWN_GRACE.as_secs(),
|
||||
"SessionEnd consolidation worker did not stop within the shutdown \
|
||||
grace period; exiting anyway"
|
||||
),
|
||||
}
|
||||
}
|
||||
if let Some(serve_result) = serve_result {
|
||||
serve_result?;
|
||||
}
|
||||
serve_result?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
|
||||
@@ -5,9 +5,26 @@
|
||||
|
||||
#![doc(html_no_source)]
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
ai_memory_cli::run().await
|
||||
fn main() -> Result<()> {
|
||||
// The runtime is built by hand rather than with `#[tokio::main]` for the
|
||||
// sake of the `shutdown_timeout` below; the builder settings are the ones
|
||||
// that attribute would have used.
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()?;
|
||||
let result = runtime.block_on(ai_memory_cli::run());
|
||||
// `run` has returned, so every guard it owns — the log-flush guard, the
|
||||
// store, the wiki watcher, the serve lock — has already been dropped in
|
||||
// order. What can still be parked is the MCP stdio transport's read of
|
||||
// stdin: tokio serves it from a blocking thread, a blocking read cannot be
|
||||
// cancelled, and dropping a runtime waits for in-flight blocking work
|
||||
// forever. That await is what kept `serve --transport stdio` alive after a
|
||||
// handled Ctrl-C (#699). Nothing is waiting on that read's result any more,
|
||||
// so stop waiting for it and let the process exit.
|
||||
runtime.shutdown_timeout(Duration::ZERO);
|
||||
result
|
||||
}
|
||||
|
||||
@@ -15,3 +15,4 @@ mod removal;
|
||||
mod repo_layout;
|
||||
mod routing_instructions;
|
||||
mod routing_skills;
|
||||
mod shutdown_signals;
|
||||
|
||||
@@ -0,0 +1,403 @@
|
||||
//! Subprocess smoke tests for signal shutdown of `ai-memory serve` (#699).
|
||||
//!
|
||||
//! Both transports used to be unstoppable in their own way: stdio listened
|
||||
//! for no signal at all, and http listened only for SIGINT — so SIGTERM,
|
||||
//! what `docker stop` and `systemctl stop` send, reached no handler. As PID 1
|
||||
//! in the container the kernel discarded it outright and `docker stop` waited
|
||||
//! out its whole grace period before SIGKILL; anywhere else it fell through
|
||||
//! to the kernel's default disposition and killed the process instantly, with
|
||||
//! no drain. Only a real spawned process can prove the difference: signal
|
||||
//! disposition is a property of the process, not of any function these tests
|
||||
//! could call in-crate.
|
||||
//!
|
||||
//! Each test spawns the built binary against a temp data dir, waits for the
|
||||
//! transport's own ready line (so a child that died during boot can never
|
||||
//! pass vacuously), sends the signal with `kill`, and requires the child to
|
||||
//! be gone inside a deadline. `kill` is shelled out to deliberately: sending
|
||||
//! a signal from Rust needs either `libc` (a new dependency) or `unsafe`
|
||||
//! (forbidden by the workspace lint). The third test completes the MCP
|
||||
//! `initialize` handshake before signalling, because the stdio arm takes a
|
||||
//! different branch before and after one.
|
||||
//!
|
||||
//! The exit status is the load-bearing half of each assertion. A test child is
|
||||
//! not PID 1, so an unhandled signal still ends it — just by the kernel's
|
||||
//! default disposition, which reports as "killed by signal N" rather than a
|
||||
//! clean code. Requiring `success()` is therefore what separates a handled
|
||||
//! shutdown from no handler at all; the deadline covers the other failure
|
||||
//! shape, a handler that runs but leaves the process alive.
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::process::{Child, ChildStdin, Command, ExitStatus, Stdio};
|
||||
use std::sync::mpsc::{self, Receiver};
|
||||
use std::thread::JoinHandle;
|
||||
use std::time::{Duration, Instant};
|
||||
use tempfile::TempDir;
|
||||
|
||||
/// Boot does migrations, the pre-migration archive and index work; the
|
||||
/// autoscope suite budgets 8s for the same startup, and this one runs three
|
||||
/// servers in parallel with the rest of the suite. Widened from 20s after one
|
||||
/// loaded machine spent 11s of wall clock on what these tests normally finish
|
||||
/// in about 1s: a 10x resource-pressure factor, not a logic race. Being
|
||||
/// generous is free, because elapsed time is not what any test here
|
||||
/// discriminates on — see [`EXIT_TIMEOUT`].
|
||||
const READY_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
|
||||
/// The shutdown itself is bounded by `SHUTDOWN_GRACE` (5s) in `serve.rs`.
|
||||
/// Anything past that plus process teardown means the signal was ignored,
|
||||
/// which is the bug.
|
||||
///
|
||||
/// 30s is 6x that bound, and widening it from 10s weakens no assertion: the
|
||||
/// discriminator is the exit status, not the clock. If the handler regresses
|
||||
/// away, the child falls back to the kernel's default disposition, dies at
|
||||
/// once, and fails `status.success()` immediately however long this deadline
|
||||
/// is. The deadline covers only the other failure shape — a handler that runs
|
||||
/// but leaves the process alive, #699's PID-1 symptom — and any value past
|
||||
/// the grace period catches that one. A tight deadline therefore buys no
|
||||
/// sensitivity and costs flakiness on a loaded runner.
|
||||
///
|
||||
/// The ceiling for both constants is nextest's everyday profile, which kills
|
||||
/// a test at 120s (`slow-timeout` 5s x 24): 60 + 10 + 30 keeps even the
|
||||
/// handshake test's worst case under it, so a wedged child still fails with
|
||||
/// this file's own diagnostic instead of a bare timeout. The tests stay in
|
||||
/// that tier — they cost about 0.25s when nothing is wrong.
|
||||
const EXIT_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
/// The handshake is one round trip over a pipe against an already-booted
|
||||
/// server; this is slack for a loaded box, not a real budget.
|
||||
const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
const HTTP_READY: &str = "MCP HTTP server ready";
|
||||
const STDIO_READY: &str = "MCP server ready on stdio";
|
||||
/// The stdio arm's post-handshake shutdown log — the branch only
|
||||
/// [`stdio_server_exits_on_sigterm_after_initialize`] reaches.
|
||||
const STDIO_STOP_AFTER_HANDSHAKE: &str = "stopping the stdio transport";
|
||||
/// And the pre-handshake one, from the select that runs before a client
|
||||
/// connects. Distinguishing the two is what keeps that test honest.
|
||||
const STDIO_STOP_BEFORE_CLIENT: &str = "shutdown signal received before a client connected";
|
||||
|
||||
fn bin() -> &'static str {
|
||||
env!("CARGO_BIN_EXE_ai-memory")
|
||||
}
|
||||
|
||||
/// Pump one of the child's output streams into a channel from a background
|
||||
/// thread, so waiting for a line has a wall-clock timeout rather than blocking
|
||||
/// on a read syscall. The handle yields everything read, for failure messages.
|
||||
fn pump_lines<R>(reader: R, tx: mpsc::Sender<String>) -> JoinHandle<String>
|
||||
where
|
||||
R: std::io::Read + Send + 'static,
|
||||
{
|
||||
std::thread::spawn(move || {
|
||||
let mut all = String::new();
|
||||
for line in BufReader::new(reader).lines() {
|
||||
let Ok(line) = line else { break };
|
||||
all.push_str(&line);
|
||||
all.push('\n');
|
||||
if tx.send(line).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
all
|
||||
})
|
||||
}
|
||||
|
||||
/// A spawned `ai-memory serve` whose stderr is pumped into a channel by a
|
||||
/// background thread, so waiting for a log line has a wall-clock timeout
|
||||
/// rather than blocking on a read syscall.
|
||||
struct Server {
|
||||
child: Child,
|
||||
lines: Receiver<String>,
|
||||
pump: Option<JoinHandle<String>>,
|
||||
/// Held open for stdio: closing it is EOF, which stops the transport on
|
||||
/// its own and would let the test pass without any signal being handled.
|
||||
stdin: Option<ChildStdin>,
|
||||
/// JSON-RPC frames, present only when stdout was piped: the handshake is
|
||||
/// the one thing here that needs to read the transport itself.
|
||||
frames: Option<Receiver<String>>,
|
||||
seen: String,
|
||||
_data_dir: TempDir,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
/// stdout at `/dev/null`: a test that only signals never reads it.
|
||||
fn spawn(transport: &str) -> Self {
|
||||
Self::spawn_inner(transport, Stdio::null())
|
||||
}
|
||||
|
||||
/// stdout piped, so [`Server::initialize`] can read the response frame.
|
||||
fn spawn_with_stdout(transport: &str) -> Self {
|
||||
Self::spawn_inner(transport, Stdio::piped())
|
||||
}
|
||||
|
||||
fn spawn_inner(transport: &str, stdout: Stdio) -> Self {
|
||||
let data_dir = TempDir::new().expect("tempdir for serve");
|
||||
let mut cmd = Command::new(bin());
|
||||
cmd.args(["serve", "--transport", transport]);
|
||||
if transport == "http" {
|
||||
// Any free port: nothing here talks to the listener, it only has
|
||||
// to exist so the server reaches its ready line.
|
||||
cmd.args(["--bind", "127.0.0.1:0"]);
|
||||
}
|
||||
cmd.arg("--data-dir")
|
||||
.arg(data_dir.path())
|
||||
.env("AI_MEMORY_DATA_DIR", data_dir.path())
|
||||
// Hermetic: the default provider would start a model download in
|
||||
// the background of every spawned server.
|
||||
.env("AI_MEMORY_EMBEDDING_PROVIDER", "none")
|
||||
.env("RUST_LOG", "info")
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(stdout)
|
||||
.stderr(Stdio::piped());
|
||||
let mut child = cmd.spawn().expect("spawn ai-memory serve");
|
||||
|
||||
let stderr = child.stderr.take().expect("stderr piped");
|
||||
let stdin = child.stdin.take();
|
||||
let (tx, lines) = mpsc::channel::<String>();
|
||||
let pump = pump_lines(stderr, tx);
|
||||
let frames = child.stdout.take().map(|stdout| {
|
||||
let (tx, frames) = mpsc::channel::<String>();
|
||||
// Detached: only stderr's accumulated copy is used, and the thread
|
||||
// ends at EOF when the child exits.
|
||||
drop(pump_lines(stdout, tx));
|
||||
frames
|
||||
});
|
||||
|
||||
Self {
|
||||
child,
|
||||
lines,
|
||||
pump: Some(pump),
|
||||
stdin,
|
||||
frames,
|
||||
seen: String::new(),
|
||||
_data_dir: data_dir,
|
||||
}
|
||||
}
|
||||
|
||||
/// Drive the MCP `initialize` handshake to completion, returning the
|
||||
/// server's response frame.
|
||||
///
|
||||
/// `serve` resolves only once a client has initialized, so this is what
|
||||
/// moves the stdio arm out of its pre-handshake select and into the branch
|
||||
/// that owns the service — where every real harness ends up.
|
||||
fn initialize(&mut self, timeout: Duration) -> String {
|
||||
// `2025-11-25` is `ProtocolVersion::LATEST` in rmcp 1.7 (its
|
||||
// `model.rs`); the server answers with its own version when the
|
||||
// client's is newer, so this only has to parse. `capabilities` and
|
||||
// `clientInfo` are the request's required fields.
|
||||
const INITIALIZE: &str = concat!(
|
||||
r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"#,
|
||||
r#""protocolVersion":"2025-11-25","capabilities":{},"#,
|
||||
r#""clientInfo":{"name":"shutdown-signals-test","version":"0"}}}"#,
|
||||
"\n",
|
||||
);
|
||||
let stdin = self.stdin.as_mut().expect("stdin piped");
|
||||
stdin
|
||||
.write_all(INITIALIZE.as_bytes())
|
||||
.expect("write the initialize request");
|
||||
stdin.flush().expect("flush the initialize request");
|
||||
let response = self
|
||||
.frames
|
||||
.as_ref()
|
||||
.expect("stdout piped")
|
||||
.recv_timeout(timeout);
|
||||
match response {
|
||||
Ok(frame) => frame,
|
||||
// Disconnected also lands here: stdout closed means the child died.
|
||||
Err(_) => panic!(
|
||||
"no initialize response on stdout within {timeout:?}.\nstderr:\n{}",
|
||||
self.stderr()
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Block until a stderr line contains `needle`, or the timeout expires.
|
||||
fn wait_for(&mut self, needle: &str, timeout: Duration) -> bool {
|
||||
let deadline = Instant::now() + timeout;
|
||||
loop {
|
||||
let remaining = deadline.saturating_duration_since(Instant::now());
|
||||
if remaining.is_zero() {
|
||||
return false;
|
||||
}
|
||||
match self.lines.recv_timeout(remaining) {
|
||||
Ok(line) => {
|
||||
let hit = line.contains(needle);
|
||||
self.seen.push_str(&line);
|
||||
self.seen.push('\n');
|
||||
if hit {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
// Disconnected means stderr closed: the child died during boot.
|
||||
Err(_) => return false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn signal(&self, name: &str) {
|
||||
let pid = self.child.id().to_string();
|
||||
let status = Command::new("kill")
|
||||
.arg(format!("-{name}"))
|
||||
.arg(&pid)
|
||||
.status()
|
||||
.expect("run kill");
|
||||
assert!(status.success(), "kill -{name} {pid} failed: {status}");
|
||||
}
|
||||
|
||||
/// Poll for the child's exit, returning `None` if it outlives `timeout`.
|
||||
fn wait_for_exit(&mut self, timeout: Duration) -> Option<ExitStatus> {
|
||||
let deadline = Instant::now() + timeout;
|
||||
loop {
|
||||
match self.child.try_wait() {
|
||||
Ok(Some(status)) => return Some(status),
|
||||
Ok(None) => {
|
||||
if Instant::now() >= deadline {
|
||||
return None;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(25));
|
||||
}
|
||||
Err(error) => panic!("try_wait on the serve child failed: {error}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Everything the child wrote to stderr, for a failure message. Kills the
|
||||
/// child first, so the pump thread's read reaches EOF; call it last.
|
||||
fn stderr(&mut self) -> String {
|
||||
let _ = self.child.kill();
|
||||
let _ = self.child.wait();
|
||||
self.stdin.take();
|
||||
match self.pump.take() {
|
||||
Some(pump) => pump.join().unwrap_or_else(|_| self.seen.clone()),
|
||||
None => self.seen.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Server {
|
||||
fn drop(&mut self) {
|
||||
// A panicking test must not leave a server holding the temp data
|
||||
// dir's serve lock.
|
||||
let _ = self.child.kill();
|
||||
let _ = self.child.wait();
|
||||
}
|
||||
}
|
||||
|
||||
/// `docker stop` and `systemctl stop` both send SIGTERM, which the http arm
|
||||
/// never listened for: it handled ctrl-c only, so a stop either burned the
|
||||
/// full grace period and ended in SIGKILL (as PID 1, where the kernel
|
||||
/// discards an unhandled signal) or killed the server outright, undrained
|
||||
/// (#699).
|
||||
#[test]
|
||||
fn http_server_exits_on_sigterm() {
|
||||
let mut server = Server::spawn("http");
|
||||
if !server.wait_for(HTTP_READY, READY_TIMEOUT) {
|
||||
panic!(
|
||||
"http server never logged {HTTP_READY:?} within {READY_TIMEOUT:?}; \
|
||||
the signal was never sent.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
}
|
||||
|
||||
server.signal("TERM");
|
||||
|
||||
let Some(status) = server.wait_for_exit(EXIT_TIMEOUT) else {
|
||||
panic!(
|
||||
"http server still running {EXIT_TIMEOUT:?} after SIGTERM.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
};
|
||||
assert!(
|
||||
status.success(),
|
||||
"SIGTERM is a normal stop, so the exit must be clean. Got {status}.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
}
|
||||
|
||||
/// The stdio arm awaited the transport and nothing else, so ctrl-c could not
|
||||
/// stop it at all — the reported symptom of #699.
|
||||
#[test]
|
||||
fn stdio_server_exits_on_sigint() {
|
||||
let mut server = Server::spawn("stdio");
|
||||
if !server.wait_for(STDIO_READY, READY_TIMEOUT) {
|
||||
panic!(
|
||||
"stdio server never logged {STDIO_READY:?} within {READY_TIMEOUT:?}; \
|
||||
the signal was never sent.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
}
|
||||
// Precondition for a non-vacuous result: stdin is still open, so the
|
||||
// transport has no reason of its own to stop.
|
||||
assert!(
|
||||
server.stdin.is_some(),
|
||||
"stdin must stay piped and open, or EOF — not the signal — ends the transport"
|
||||
);
|
||||
|
||||
server.signal("INT");
|
||||
|
||||
let Some(status) = server.wait_for_exit(EXIT_TIMEOUT) else {
|
||||
panic!(
|
||||
"stdio server still running {EXIT_TIMEOUT:?} after SIGINT.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
};
|
||||
assert!(
|
||||
status.success(),
|
||||
"a signal-triggered stop is a normal shutdown. Got {status}.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
}
|
||||
|
||||
/// Both tests above signal before any client connects, so on stdio they
|
||||
/// resolve in the pre-handshake select and return early. A real harness
|
||||
/// completes the MCP `initialize` handshake first, which lands the arm in the
|
||||
/// branch that takes the service's cancellation token, cancels it and waits
|
||||
/// under the grace period — the code that actually runs in production, and
|
||||
/// which nothing else here reaches (#699).
|
||||
#[test]
|
||||
fn stdio_server_exits_on_sigterm_after_initialize() {
|
||||
let mut server = Server::spawn_with_stdout("stdio");
|
||||
if !server.wait_for(STDIO_READY, READY_TIMEOUT) {
|
||||
panic!(
|
||||
"stdio server never logged {STDIO_READY:?} within {READY_TIMEOUT:?}; \
|
||||
the handshake was never started.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
}
|
||||
|
||||
// Precondition for a non-vacuous result: until the handshake completes,
|
||||
// the signal lands in the pre-handshake select the other tests cover.
|
||||
let response = server.initialize(HANDSHAKE_TIMEOUT);
|
||||
assert!(
|
||||
response.contains(r#""result""#) && response.contains("serverInfo"),
|
||||
"the initialize response must carry a result, or the handshake did not \
|
||||
complete: {response}.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
|
||||
server.signal("TERM");
|
||||
|
||||
let Some(status) = server.wait_for_exit(EXIT_TIMEOUT) else {
|
||||
panic!(
|
||||
"stdio server still running {EXIT_TIMEOUT:?} after SIGTERM.\nstderr:\n{}",
|
||||
server.stderr()
|
||||
);
|
||||
};
|
||||
let stderr = server.stderr();
|
||||
assert!(
|
||||
status.success(),
|
||||
"a signal-triggered stop is a normal shutdown. Got {status}.\nstderr:\n{stderr}"
|
||||
);
|
||||
// Which line was logged says which branch ran, so this test cannot drift
|
||||
// back into covering the early path it exists to complement.
|
||||
assert!(
|
||||
stderr.contains(STDIO_STOP_AFTER_HANDSHAKE),
|
||||
"expected the post-handshake shutdown log {STDIO_STOP_AFTER_HANDSHAKE:?}.\
|
||||
\nstderr:\n{stderr}"
|
||||
);
|
||||
assert!(
|
||||
!stderr.contains(STDIO_STOP_BEFORE_CLIENT),
|
||||
"the pre-handshake branch handled the signal, so this test covers the \
|
||||
same path as the others.\nstderr:\n{stderr}"
|
||||
);
|
||||
}
|
||||
@@ -138,6 +138,25 @@ homelab, and restarts. The compose file + env file on the homelab are
|
||||
unchanged between deploys; if you ever need to change them, scp the
|
||||
new copy + re-run `bin/deploy`.
|
||||
|
||||
The restart step stops the running container with SIGTERM. The server
|
||||
handles SIGINT and SIGTERM on both transports and bounds each wait in
|
||||
its shutdown path at five seconds, so a restart — or a plain `docker
|
||||
stop` or `docker compose down` — drains and exits in a few seconds
|
||||
instead of waiting out the supervisor's grace period and ending in
|
||||
SIGKILL. Before those handlers existed the two deployment shapes failed
|
||||
differently. In the container the server is PID 1, and for PID 1 the
|
||||
kernel discards a signal whose handler is not installed, so `docker
|
||||
stop` burned its full grace period and `docker kill` was the only way
|
||||
out. Under the native systemd unit the server is not PID 1, so
|
||||
`systemctl stop` fell through to the kernel's default disposition and
|
||||
killed it instantly instead — fast, but with no drain and the durable
|
||||
SessionEnd consolidation worker cut off mid-flight. There the stop is
|
||||
now slower and clean. The five-second bound is fixed and not
|
||||
configurable. `docker kill` remains the way to stop the server without
|
||||
waiting for the drain. No container init shim is required: the binary
|
||||
installs its own signal handlers, so it stops correctly as PID 1 and you
|
||||
do not need `tini` or `docker run --init`.
|
||||
|
||||
## Updating API keys
|
||||
|
||||
```bash
|
||||
|
||||
Reference in New Issue
Block a user