diff --git a/bt-daemon/Cargo.lock b/bt-daemon/Cargo.lock index cd61f55..6ef96e3 100644 --- a/bt-daemon/Cargo.lock +++ b/bt-daemon/Cargo.lock @@ -272,6 +272,7 @@ dependencies = [ "fs2", "regex", "reqwest", + "semver", "serde", "serde_json", "sha2", @@ -1594,6 +1595,12 @@ dependencies = [ "libc", ] +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + [[package]] name = "serde" version = "1.0.229" diff --git a/bt-daemon/Cargo.toml b/bt-daemon/Cargo.toml index 41845ca..bea64a9 100644 --- a/bt-daemon/Cargo.toml +++ b/bt-daemon/Cargo.toml @@ -27,6 +27,7 @@ fs2 = "0.4" regex = "1" serde = { version = "1", features = ["derive"] } serde_json = "1" +semver = "1" sha2 = "0.10" sysinfo = { version = "0.33.1", default-features = false, features = ["system"] } thiserror = "2" diff --git a/bt-daemon/docs/protocol.md b/bt-daemon/docs/protocol.md index d5e0dc5..09f92f4 100644 --- a/bt-daemon/docs/protocol.md +++ b/bt-daemon/docs/protocol.md @@ -82,7 +82,12 @@ Params: ```json { "protocol_version": 1, - "client": { "source": "codex", "plugin_version": "1.2.3", "pid": 12345 } + "client": { + "source": "codex", + "daemon_version": "0.20.0", + "plugin_version": "1.2.3", + "pid": 12345 + } } ``` Result: @@ -93,11 +98,11 @@ Result: "capabilities": { "sources": ["codex", "claude-code", "opencode", "pi", "debug"] } } ``` -If `protocol_version` is incompatible the daemon returns an application error; -the client decides whether to drop events or (if the client is newer) trigger a -version handover (`daemon.shutdown` → respawn). +If `protocol_version` is incompatible the daemon returns an application error. +Daemon versions move forward: an older client uses a newer compatible daemon, +while a newer client drains and replaces an older daemon once. -### `event.log` (request or notification) +### `event.log` (request) The hot path. Params are the **Envelope** (see below). Request result: ```json @@ -112,6 +117,11 @@ reconciled by the daemon. On restart, uncheckpointed journal entries are queued again automatically. Explicit status and flush requests act as daemon-worker barriers, but hook capture never does. +During version handover, the draining daemon returns `{ "accepted": false }`. +The client waits for the replacement daemon, initializes a new connection, and +retries the same envelope. Capture adapters must use requests because a JSON-RPC +notification has no acknowledgement and therefore cannot make this retry safe. + ### `session.flush` (request) Block until every route's spans for the session are delivered, or `timeout_ms` @@ -173,9 +183,12 @@ widget. ### `daemon.shutdown` (request) -Graceful: stop accepting new events, drain all session queues, flush sinks, -release the local endpoint, exit. Result `{ "ok": true }` is sent before exit. -Used for version handover and by tests. +Graceful: reject new events, wait for captures already being journaled, drain all +session queues, and flush sinks. Result `{ "ok": true }` is sent only after the +drain completes; the daemon then releases the local endpoint and exits. Used +for version handover and explicit stop commands. A client whose declared daemon +version is older than the running daemon receives `{ "ok": false }` and cannot +downgrade it. ## Envelope (`event.log` params) @@ -208,6 +221,10 @@ Field notes: - **`source`** selects the daemon-side translator. `debug` is a built-in pass-through translator used by the prototype and tests. + +The Braintrust sink keeps the capture package version in +`context.span_origin.version` and adds `metadata.bt_daemon_version` to every +span so the daemon build that performed translation can be queried separately. - **`session_id`** identifies the source agent session. Combined with `route` it forms the queue + state key (see "Multiple routes per session" below). The shim extracts it from the payload (default JSON field `session_id`, @@ -288,10 +305,12 @@ journal, logs, status, or RPC response. Envelopes journal only their non-secret event rebuilds it from the journal, and deterministic span ids merge the re-emitted rows. This matters because the idle exit above requires *every* session to be quiet, which for a continuously active user never happens. -- **Version handover.** `initialize` compares versions. A newer client sends - `daemon.shutdown`, waits until the endpoint no longer accepts connections, - and spawns its own daemon. In-flight session state is rebuilt from the - journal. +- **Version handover.** `initialize` compares versions. Older clients use a + newer compatible daemon. A newer client sends `daemon.shutdown`; the old + daemon quiesces and drains before acknowledging, then the client spawns its + daemon. In-flight session state is rebuilt from the journal. This makes + upgrades monotonic instead of allowing mixed-version clients to repeatedly + replace each other. ## Durability & idempotence diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index d3124a1..20ff3e2 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -460,11 +460,12 @@ pub(crate) fn apply_tags(route: &mut SessionRoute, tags: &[String]) -> anyhow::R Ok(()) } -fn initialize_params(env: &Envelope) -> serde_json::Value { +fn initialize_params(env: &Envelope, daemon_version: &str) -> serde_json::Value { serde_json::json!({ "protocol_version": PROTOCOL_VERSION, "client": { "source": env.source, + "daemon_version": daemon_version, "plugin_version": env.plugin_version, "pid": std::process::id() } @@ -480,36 +481,103 @@ pub async fn forward_envelope( host: &HostInfo, no_spawn: bool, ) -> anyhow::Result<()> { - let stream = client::ensure_daemon(socket, host, no_spawn).await?; - let mut conn = client::Conn::new(stream); - let initialized = conn - .request(method::INITIALIZE, initialize_params(env)) - .await?; - let initialized: wire::InitializeResult = serde_json::from_value(initialized)?; - if initialized.daemon_version != host.version { - if no_spawn { + const MAX_HANDOVER_ATTEMPTS: usize = 3; + let mut handover_attempts = 0; + let mut capture_retry_deadline = None; + + loop { + if capture_retry_deadline.is_some_and(|deadline| tokio::time::Instant::now() >= deadline) { anyhow::bail!( - "daemon version {} does not match client {} and --no-spawn is set", - initialized.daemon_version, - host.version + "replacement daemon did not accept the event at {}", + socket.display() ); } - conn.request(method::DAEMON_SHUTDOWN, serde_json::json!({})) - .await?; - drop(conn); - for _ in 0..100 { - if client::connect(socket).await.is_err() { - break; + + let stream = match client::connect(socket).await { + Ok(stream) => stream, + Err(_) if capture_retry_deadline.is_some() && no_spawn => { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + continue; + } + Err(_) => client::ensure_daemon(socket, host, no_spawn).await?, + }; + let mut conn = client::Conn::new(stream); + let initialized = match conn + .request(method::INITIALIZE, initialize_params(env, &host.version)) + .await + { + Ok(initialized) => initialized, + Err(_) if capture_retry_deadline.is_some() => { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + continue; + } + Err(error) => return Err(error), + }; + let initialized: wire::InitializeResult = serde_json::from_value(initialized)?; + if daemon_needs_upgrade(&initialized.daemon_version, &host.version) { + if no_spawn { + anyhow::bail!( + "daemon version {} is older than client {} and --no-spawn is set", + initialized.daemon_version, + host.version + ); + } + if handover_attempts == MAX_HANDOVER_ATTEMPTS { + anyhow::bail!( + "daemon version {} is still older than client {} after {MAX_HANDOVER_ATTEMPTS} handover attempts", + initialized.daemon_version, + host.version + ); + } + conn.request(method::DAEMON_SHUTDOWN, serde_json::json!({})) + .await?; + handover_attempts += 1; + drop(conn); + wait_for_daemon_exit(socket).await; + continue; + } + let result = match conn.request(method::EVENT_LOG, env).await { + Ok(result) => result, + Err(_) if capture_retry_deadline.is_some() => { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + continue; } - tokio::time::sleep(std::time::Duration::from_millis(20)).await; + Err(error) => return Err(error), + }; + let result: wire::EventLogResult = serde_json::from_value(result)?; + if result.accepted { + return Ok(()); } - let stream = client::ensure_daemon(socket, host, false).await?; - conn = client::Conn::new(stream); - conn.request(method::INITIALIZE, initialize_params(env)) - .await?; + capture_retry_deadline + .get_or_insert_with(|| tokio::time::Instant::now() + std::time::Duration::from_secs(5)); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; } - conn.request(method::EVENT_LOG, env).await?; - Ok(()) +} + +async fn wait_for_daemon_exit(socket: &std::path::Path) { + for _ in 0..100 { + if client::connect(socket).await.is_err() { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } +} + +/// A shared daemon only moves forward. Released versions use semver; opaque +/// development versions retain the old exact-match handover behavior. +fn daemon_needs_upgrade(daemon_version: &str, client_version: &str) -> bool { + if daemon_version == client_version { + return false; + } + compare_daemon_versions(daemon_version, client_version).is_none_or(|ordering| ordering.is_lt()) +} + +pub(crate) fn compare_daemon_versions(left: &str, right: &str) -> Option { + Some( + semver::Version::parse(left) + .ok()? + .cmp(&semver::Version::parse(right).ok()?), + ) } /// Ask the daemon to flush a session, bounded by `timeout_ms`. A reliable @@ -1451,8 +1519,9 @@ mod tests { // Both the initial connection and the post-restart retry use this // shared parameter builder. - let initialize = initialize_params(&env); + let initialize = initialize_params(&env, "1.0.13"); assert_eq!(initialize["client"]["source"], "grok"); + assert_eq!(initialize["client"]["daemon_version"], "1.0.13"); assert_eq!(initialize["client"]["plugin_version"], "0.1.0"); assert_ne!(initialize["client"]["plugin_version"], "1.0.13"); } diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs index 3148b63..ef75bea 100644 --- a/bt-daemon/src/server.rs +++ b/bt-daemon/src/server.rs @@ -8,9 +8,10 @@ use crate::sink::SinkFactory; use crate::translate::Registry; use crate::transport::{self, Listener, ServerStream}; use crate::wire::{ - error_code, method, Capabilities, Envelope, EventLogResult, FlushParams, FlushResult, - InitializeParams, InitializeResult, ManagedRunFlushParams, Message, Request, Response, - RpcError, SessionStatus, ShutdownResult, StatusParams, StatusResult, PROTOCOL_VERSION, + error_code, method, Capabilities, ClientInfo, Envelope, EventLogResult, FlushParams, + FlushResult, InitializeParams, InitializeResult, ManagedRunFlushParams, Message, Request, + Response, RpcError, SessionStatus, ShutdownResult, StatusParams, StatusResult, + PROTOCOL_VERSION, }; use crate::wire::{AuthSelection, BackendAuth, SessionRoute}; use crate::{paths, ServeArgs}; @@ -190,6 +191,11 @@ pub struct Daemon { ingress_dispatched: Mutex>, started: Instant, last_activity: Mutex, + /// Prevents a shutdown drain from racing capture that has passed the + /// quiescing check but has not yet appended and queued its event. + capture_gate: tokio::sync::RwLock<()>, + quiescing: AtomicBool, + drained: tokio::sync::Mutex, shutting_down: AtomicBool, shutdown: Notify, } @@ -222,6 +228,9 @@ impl Daemon { ingress_dispatched: Mutex::new(HashMap::new()), started: Instant::now(), last_activity: Mutex::new(Instant::now()), + capture_gate: tokio::sync::RwLock::new(()), + quiescing: AtomicBool::new(false), + drained: tokio::sync::Mutex::new(false), shutting_down: AtomicBool::new(false), shutdown: Notify::new(), }); @@ -583,7 +592,11 @@ impl Daemon { .sum() } - async fn capture_event(&self, mut env: Envelope) -> Result<(), String> { + async fn capture_event(&self, mut env: Envelope) -> Result { + let _capture_guard = self.capture_gate.read().await; + if self.quiescing.load(Ordering::SeqCst) { + return Ok(false); + } env.source = self .translators .canonical_source(&env.source) @@ -612,7 +625,7 @@ impl Daemon { tracing::warn!("journaled event will be recovered after daemon restart"); } } - Ok(()) + Ok(true) } async fn ingress_barrier(&self) { @@ -796,9 +809,14 @@ impl Daemon { } fn trigger_shutdown(&self) { + self.quiescing.store(true, Ordering::SeqCst); self.shutting_down.store(true, Ordering::SeqCst); self.shutdown.notify_waiters(); } + + fn begin_quiesce(&self) { + self.quiescing.store(true, Ordering::SeqCst); + } } fn lease_is_expiring(lease: &AuthLease) -> bool { @@ -1083,6 +1101,7 @@ async fn serve_connection(daemon: Arc, stream: ServerStream) -> anyhow:: if line.trim().is_empty() { continue; } + let mut shutdown_after_response = false; let response = match Message::from_line(&line) { Ok(Message::Request(req)) => { let request_id = req.id.clone(); @@ -1093,6 +1112,14 @@ async fn serve_connection(daemon: Arc, stream: ServerStream) -> anyhow:: "request received" ); let response = handle_request(&daemon, req, &mut client).await; + shutdown_after_response = method == method::DAEMON_SHUTDOWN + && response.error.is_none() + && response + .result + .as_ref() + .and_then(|result| result.get("ok")) + .and_then(Value::as_bool) + == Some(true); if let Some(error) = &response.error { tracing::warn!( request_id = ?request_id, @@ -1112,13 +1139,22 @@ async fn serve_connection(daemon: Arc, stream: ServerStream) -> anyhow:: } Ok(Message::Notification(note)) => { tracing::info!(method = %note.method, "notification received"); - // Hot-path notifications (in-process clients): process, no reply. + // Legacy best-effort notifications cannot participate in the + // acknowledged handover retry used by current clients. if note.method == method::EVENT_LOG { if let Some(params) = note.params { match serde_json::from_value::(params) { Ok(mut env) => { attach_process_capture(&mut env, client.as_ref()); - let _ = daemon.capture_event(env).await; + match daemon.capture_event(env).await { + Ok(true) => {} + Ok(false) => tracing::warn!( + "event notification arrived while daemon was shutting down" + ), + Err(error) => { + tracing::warn!(%error, "event notification was not captured") + } + } } Err(error) => tracing::warn!( method = %note.method, @@ -1143,8 +1179,15 @@ async fn serve_connection(daemon: Arc, stream: ServerStream) -> anyhow:: if let Some(resp) = response { let mut buf = Message::Response(resp).to_line()?; buf.push('\n'); - write_half.write_all(buf.as_bytes()).await?; - write_half.flush().await?; + let write_result = async { + write_half.write_all(buf.as_bytes()).await?; + write_half.flush().await + } + .await; + if shutdown_after_response { + daemon.trigger_shutdown(); + } + write_result?; } } Ok(()) @@ -1860,9 +1903,9 @@ async fn handle_request( let mut env = parse!(Envelope); attach_process_capture(&mut env, client.as_ref()); match daemon.capture_event(env).await { - Ok(()) => Response::ok( + Ok(accepted) => Response::ok( id, - serde_json::to_value(EventLogResult { accepted: true }).unwrap(), + serde_json::to_value(EventLogResult { accepted }).unwrap(), ), Err(error) => Response::err(id, RpcError::new(error_code::INTERNAL, error)), } @@ -1923,12 +1966,17 @@ async fn handle_request( Response::ok(id, serde_json::to_value(daemon.status(p)).unwrap()) } method::DAEMON_SHUTDOWN => { - let resp = Response::ok( + if !client_may_shutdown(client.as_ref(), &daemon.version) { + return Response::ok( + id, + serde_json::to_value(ShutdownResult { ok: false }).unwrap(), + ); + } + drain_all(daemon).await; + Response::ok( id, serde_json::to_value(ShutdownResult { ok: true }).unwrap(), - ); - daemon.trigger_shutdown(); - resp + ) } other => Response::err( id, @@ -1940,6 +1988,17 @@ async fn handle_request( } } +fn client_may_shutdown(client: Option<&ClientInfo>, daemon_version: &str) -> bool { + let Some(client) = client else { + return true; // Explicit stop commands do not initialize first. + }; + let Some(client_version) = client.daemon_version.as_deref() else { + return false; // Legacy initialized hooks must not downgrade the daemon. + }; + crate::compare_daemon_versions(client_version, daemon_version) + .is_none_or(|ordering| !ordering.is_lt()) +} + impl Daemon { fn status(&self, p: StatusParams) -> StatusResult { let map = self.sessions.lock().unwrap(); @@ -2083,11 +2142,18 @@ fn spawn_idle_watchdog(daemon: Arc, idle_timeout: Duration) { } async fn drain_all(daemon: &Arc) { + let mut drained = daemon.drained.lock().await; + if *drained { + return; + } + daemon.begin_quiesce(); + let _capture_guard = daemon.capture_gate.write().await; daemon.settle_ingress().await; let sessions: Vec> = daemon.sessions.lock().unwrap().values().cloned().collect(); for s in sessions { s.shutdown().await; } + *drained = true; } /// Is a live daemon answering at the endpoint? Connect and expect any line @@ -2118,3 +2184,35 @@ async fn probe_alive(endpoint: &std::path::Path) -> bool { Ok(Ok(Some(_))) ) } + +#[cfg(test)] +mod tests { + use super::{client_may_shutdown, ClientInfo}; + + fn initialized(version: Option<&str>) -> ClientInfo { + ClientInfo { + source: "codex".into(), + daemon_version: version.map(str::to_string), + plugin_version: None, + pid: None, + } + } + + #[test] + fn legacy_initialized_clients_cannot_downgrade_the_daemon() { + assert!(client_may_shutdown(None, "0.20.0")); + assert!(!client_may_shutdown(Some(&initialized(None)), "0.20.0")); + assert!(!client_may_shutdown( + Some(&initialized(Some("0.19.3"))), + "0.20.0" + )); + assert!(client_may_shutdown( + Some(&initialized(Some("0.20.0"))), + "0.20.0" + )); + assert!(client_may_shutdown( + Some(&initialized(Some("0.21.0"))), + "0.20.0" + )); + } +} diff --git a/bt-daemon/src/sink/braintrust.rs b/bt-daemon/src/sink/braintrust.rs index 99180f5..cfad709 100644 --- a/bt-daemon/src/sink/braintrust.rs +++ b/bt-daemon/src/sink/braintrust.rs @@ -173,6 +173,7 @@ impl SinkFactory for BraintrustSinkFactory { default_api_url: self.default_api_url.clone(), default_app_url: self.default_app_url.clone(), version: plugin_version.unwrap_or(&self.version).to_string(), + daemon_version: self.version.clone(), source: source.to_string(), creds: None, urls: None, @@ -205,6 +206,7 @@ struct BraintrustSink { default_api_url: Option, default_app_url: Option, version: String, + daemon_version: String, source: String, creds: Option, /// Resolved `(api_url, app_url)` for this session, from its config. @@ -300,7 +302,7 @@ impl BraintrustSink { fn update_open(&mut self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { self.ensure_handle(client, row)?; let handle = self.open.get(&row.span_id).expect("just inserted"); - handle.log(build_log(row)?); + handle.log(build_log(row, &self.daemon_version)?); if let Some(end) = row.end_ms { handle.end_with_time(ms_to_secs(end)); // SpanHandle retains the complete accumulated input/output. Once a @@ -323,7 +325,7 @@ impl BraintrustSink { creds.token.clone(), creds.org_id.clone(), &components, - build_log(row)?, + build_log(row, &self.daemon_version)?, ) .map_err(|error| anyhow::anyhow!("braintrust span merge failed: {error}")) } @@ -558,7 +560,7 @@ fn ms_to_secs(ms: i64) -> f64 { ms as f64 / 1000.0 } -fn build_log(row: &SpanRow) -> anyhow::Result { +fn build_log(row: &SpanRow, daemon_version: &str) -> anyhow::Result { // The span's display name is carried on the log event, not the builder. // An empty name means "unchanged" (many merge ops use `..Default::default()` // and don't rename the span) — omitting `.name()` avoids overwriting the @@ -573,9 +575,19 @@ fn build_log(row: &SpanRow) -> anyhow::Result { if let Some(output) = &row.output { lb = lb.output(output.clone()); } - if let Some(Value::Object(md)) = &row.metadata { - lb = lb.metadata(md.clone()); - } + let mut metadata = row + .metadata + .as_ref() + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + // The plugin version remains in span_origin for backwards compatibility; + // this records the daemon build that actually translated the event. + metadata.insert( + "bt_daemon_version".into(), + Value::String(daemon_version.to_string()), + ); + lb = lb.metadata(metadata); let mut metrics = row .metrics .as_ref() diff --git a/bt-daemon/src/wire/methods.rs b/bt-daemon/src/wire/methods.rs index 143b1ea..2690e95 100644 --- a/bt-daemon/src/wire/methods.rs +++ b/bt-daemon/src/wire/methods.rs @@ -21,6 +21,10 @@ pub struct InitializeParams { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ClientInfo { pub source: String, + /// Version of the executable that would host a replacement daemon. + /// Distinct from the instrumentation plugin version below. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub daemon_version: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub plugin_version: Option, #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index d9aafed..492fe0f 100644 --- a/bt-daemon/tests/braintrust_sink.rs +++ b/bt-daemon/tests/braintrust_sink.rs @@ -512,6 +512,10 @@ async fn braintrust_sink_delivers_spans_to_collector() { bodies.contains("0.9.0"), "plugin version missing from shared span origin" ); + assert!( + bodies.contains("bt_daemon_version"), + "daemon version metadata missing from spans: {bodies}" + ); // Project registration happened (org_name path, no login). assert!( diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 7be859c..525eb2d 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -1530,18 +1530,93 @@ async fn no_spawn_errors_when_daemon_absent() { } #[tokio::test] -async fn no_spawn_rejects_a_mismatched_daemon_version() { - let (_data_dir, socket, handle, _tmp) = start_daemon().await; - let host = HostInfo { +async fn daemon_version_handover_only_moves_forward() { + let (older_socket, older_handle, _flushes, _tmp) = start_tracking_daemon("1.0.0").await; + let newer_host = HostInfo { + serve_argv: vec![OsString::from("unused")], + version: "2.0.0".into(), + }; + let error = forward_envelope( + &envelope("newer-client", "SessionStart", 1), + &older_socket, + &newer_host, + true, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("is older than client")); + shutdown(&older_socket).await; + older_handle.await.unwrap(); + + let (newer_socket, newer_handle, _flushes, _tmp) = start_tracking_daemon("2.0.0").await; + let older_host = HostInfo { serve_argv: vec![OsString::from("unused")], - version: "newer-client".into(), + version: "1.0.0".into(), }; - let err = forward_envelope(&envelope("x", "y", 1), &socket, &host, true) + forward_envelope( + &envelope("older-client", "SessionStart", 1), + &newer_socket, + &older_host, + true, + ) + .await + .unwrap(); + shutdown(&newer_socket).await; + newer_handle.await.unwrap(); +} + +#[tokio::test] +async fn shutdown_drains_then_retries_late_capture_on_replacement() { + let gate = Arc::new(tokio::sync::Notify::new()); + let emitted = Arc::new(std::sync::atomic::AtomicU64::new(0)); + let (data_dir, socket, handle, _tmp) = start_gated_daemon(gate.clone(), emitted.clone()).await; + + forward_envelope( + &envelope("drained-shutdown", "SessionStart", 1), + &socket, + &dummy_host(), + false, + ) + .await + .unwrap(); + + let shutdown_socket = socket.clone(); + let shutdown_task = tokio::spawn(async move { shutdown_daemon(&shutdown_socket).await }); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + !shutdown_task.is_finished(), + "shutdown returned before the sink drained" + ); + + let late_socket = socket.clone(); + let late_event = tokio::spawn(async move { + forward_envelope( + &envelope("drained-shutdown", "PostToolUse", 2), + &late_socket, + &dummy_host(), + true, + ) .await - .unwrap_err(); - assert!(err.to_string().contains("does not match client")); + }); + + shutdown_task.abort(); + let _ = shutdown_task.await; + gate.notify_one(); + tokio::time::timeout(Duration::from_secs(2), handle) + .await + .expect("daemon stayed alive after the shutdown client disconnected") + .unwrap(); + let replacement = start_daemon_at(data_dir.clone(), socket.clone()).await; + late_event.await.unwrap().unwrap(); shutdown(&socket).await; - handle.await.unwrap(); + replacement.await.unwrap(); + let journal = + std::fs::read_to_string(source_journal_path(&data_dir, "debug", "drained-shutdown")) + .unwrap(); + assert!( + journal.contains(r#""n":2"#), + "replacement daemon did not journal the retried event: {journal}" + ); } #[cfg(feature = "cli")]