From 1a78217091f562b1a01d81fc3cbe9598ef97d9e2 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 16 Sep 2026 14:26:31 -0400 Subject: [PATCH] fix(daemon): make version handover forward-only Prevent mixed-version clients from repeatedly replacing the shared daemon. Shutdown now quiesces capture and drains pending delivery before the replacement starts. Before: 0.19 -> stop 0.20 -> start 0.19 -> stop 0.19 -> start 0.20 After: 0.19 -> reuse 0.20 0.21 -> drain 0.20 -> start 0.21 Concurrent upgrades retry up to three handovers: 0.21 + 0.22 -> drain 0.20 0.21 wins socket 0.22 -> drain 0.21 -> start 0.22 Captures rejected during a drain reconnect to the replacement, or restore the daemon on demand after an explicit stop. Record the translating daemon version separately from the capture plugin version so production traces retain both provenance signals. Fixes #88 --- bt-daemon/Cargo.lock | 7 ++ bt-daemon/Cargo.toml | 1 + bt-daemon/docs/protocol.md | 43 +++++++--- bt-daemon/src/lib.rs | 121 +++++++++++++++++++++------ bt-daemon/src/server.rs | 128 +++++++++++++++++++++++++---- bt-daemon/src/sink/braintrust.rs | 24 ++++-- bt-daemon/src/wire/methods.rs | 4 + bt-daemon/tests/braintrust_sink.rs | 4 + bt-daemon/tests/pipeline.rs | 91 ++++++++++++++++++-- 9 files changed, 356 insertions(+), 67 deletions(-) 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")]