hyperhive/hive-agent/src/swarm_queue.rs

//! This agent's swarm-queue connection, made with the agent's own credential.
//!
//! The credential is minted per agent at swarm level. This process reads it
//! from the swarm secret store under the agent's own certificate and holds it
//! in memory only ([`AgentSecret`]), and publishes on the agent's own
//! hive-free subjects.
//!
//! **The connection is required.** [`init`] fails when the queue's address or
//! the store's coordinates are missing, and [`client`] retries the first
//! connect on [`CONNECT_BACKOFF`] and then exits the process, so the harness
//! unit fails and systemd restarts it.

use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;

use anyhow::{Context as _, anyhow};
use tokio::sync::OnceCell;

use swarm_secret_client::{
    SecretStore,
    client::{DEFAULT_CERT_MOUNT, ENV_ADDR, ENV_CACERT, Settings},
    policy, queue,
};

/// Where the swarm queue listens, as this container reaches it.
const ENV_NATS_URL: &str = "HIVE_AGENT_NATS_URL";

/// Names the agent whose queue secret this process reads. The cert-auth role
/// and the store path are both built from it by `swarm_secret_client`.
const ENV_AGENT_NAME: &str = "HIVE_AGENT_NAME";

/// How long one read of the store may take. The read runs inside a queue
/// connection attempt, and a store that never answers would otherwise stall
/// every reconnect behind it.
const STORE_READ_TIMEOUT: Duration = Duration::from_secs(10);

/// The waits between first-connect attempts: 5s doubling, about five minutes
/// in all. After the last attempt fails the process exits.
const CONNECT_BACKOFF: [Duration; 6] = [
    Duration::from_secs(5),
    Duration::from_secs(10),
    Duration::from_secs(20),
    Duration::from_secs(40),
    Duration::from_secs(80),
    Duration::from_secs(160),
];

/// Where to connect, resolved once at boot by [`init`].
static AGENT: OnceLock<AgentPath> = OnceLock::new();

/// The one connection every publisher in this process shares. Separate from
/// [`AGENT`] because resolving the coordinates is synchronous boot work and
/// connecting is not — see [`client`].
static CLIENT: OnceCell<async_nats::Client> = OnceCell::const_new();

/// What connecting with this agent's own credential needs.
#[derive(Debug, PartialEq, Eq)]
struct AgentPath {
    url: String,
    store: Arc<StoreSource>,
}

/// This agent's queue secret as the store holds it: where the store is, the
/// role this agent logs in as, and the path it reads.
#[derive(Debug, PartialEq, Eq)]
struct StoreSource {
    settings: Settings,
    role: String,
    path: String,
}

/// Where [`AgentSecret`] reads the secret from. A trait so the retry rules can
/// be exercised without a store.
trait SecretSource: Send + Sync + 'static {
    /// The store path, for log lines and errors. Never the value.
    fn origin(&self) -> &str;

    /// The value stored now, or `None` when nothing is stored.
    fn fetch(&self) -> impl Future<Output = anyhow::Result<Option<String>>> + Send;
}

impl SecretSource for Arc<StoreSource> {
    fn origin(&self) -> &str {
        &self.path
    }

    /// A fresh login per read: reads happen once per connection attempt, and a
    /// token held between them would be one more secret in memory to expire.
    async fn fetch(&self) -> anyhow::Result<Option<String>> {
        let store = SecretStore::connect(&self.settings, &self.role, DEFAULT_CERT_MOUNT)
            .await
            .context("logging in to the swarm secret store as this agent")?;
        // `read_optional`: only a 404 is "nothing stored". A 403 stays an
        // error, so a policy that stopped covering the path is not reported
        // as a secret that was deleted.
        let credential: Option<queue::AgentCredential> = store
            .read_optional(&self.path)
            .await
            .with_context(|| format!("reading {} from the store", self.path))?;
        Ok(credential.map(|c| c.value))
    }
}

/// This agent's queue secret, held in memory between connection attempts.
///
/// No `Debug`: [`Held`] carries the secret.
struct AgentSecret<S> {
    source: S,
    timeout: Duration,
    held: Mutex<Held>,
}

#[derive(Default)]
struct Held {
    value: Option<String>,
    /// `value` was read by [`AgentSecret::prime`] and no attempt has
    /// presented it yet.
    unpresented: bool,
}

impl<S: SecretSource> AgentSecret<S> {
    fn new(source: S) -> Self {
        Self {
            source,
            timeout: STORE_READ_TIMEOUT,
            held: Mutex::default(),
        }
    }

    fn held(&self) -> std::sync::MutexGuard<'_, Held> {
        // A poisoned lock only means another attempt panicked mid-update; the
        // worst it left is a stale value, which the next read replaces.
        self.held
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// Read the secret before the first connection attempt. An `Err` here
    /// fails that attempt.
    async fn prime(&self) -> anyhow::Result<()> {
        let value = self
            .read()
            .await?
            .with_context(|| format!("nothing is stored at {}", self.source.origin()))?;
        *self.held() = Held {
            value: Some(value),
            unpresented: true,
        };
        Ok(())
    }

    /// The secret to present on one connection attempt.
    ///
    /// The value [`prime`](Self::prime) read serves the first attempt. Every
    /// later attempt reads the store again, so after the queue refuses a
    /// re-minted secret the next attempt presents the current one. A store
    /// that cannot be read leaves the held value in use; a store that holds
    /// nothing drops it. Either way the attempt is one read, and the next one
    /// waits out the connection's reconnect backoff.
    async fn for_attempt(&self) -> anyhow::Result<String> {
        {
            let mut held = self.held();
            if std::mem::take(&mut held.unpresented)
                && let Some(value) = &held.value
            {
                return Ok(value.clone());
            }
        }
        match self.read().await {
            Ok(Some(value)) => {
                self.held().value = Some(value.clone());
                Ok(value)
            }
            Ok(None) => {
                self.held().value = None;
                tracing::warn!(
                    path = %self.source.origin(),
                    "nothing is stored at this agent's queue credential path; its queue \
                     connection retries under backoff"
                );
                Err(anyhow!("nothing is stored at {}", self.source.origin()))
            }
            Err(e) => {
                let held = self.held().value.clone();
                tracing::warn!(
                    error = format!("{e:#}"),
                    presenting_held = held.is_some(),
                    "re-reading this agent's queue credential from the store failed"
                );
                held.ok_or(e)
            }
        }
    }

    /// One bounded read. An empty stored value is no value.
    async fn read(&self) -> anyhow::Result<Option<String>> {
        let value = tokio::time::timeout(self.timeout, self.source.fetch())
            .await
            .map_err(|_| anyhow!("the store did not answer within {:?}", self.timeout))??;
        Ok(value.map(|v| v.trim().to_owned()).filter(|v| !v.is_empty()))
    }
}

/// Where this agent reads its own queue secret from.
///
/// Taken through a lookup so it can be asserted without mutating the process
/// environment. The shape follows `hive_matrix_mcp::credential`'s
/// `store_settings`, which reads the same variables for the same identity.
///
/// # Errors
/// No store address, or one without the agent name or the identity beside it.
fn store_source(get: impl Fn(&str) -> Option<String>) -> anyhow::Result<StoreSource> {
    if get(ENV_ADDR).is_none_or(|v| v.is_empty()) {
        anyhow::bail!(
            "{ENV_ADDR} is unset or empty, so there is no secret store to read this agent's \
             queue credential from"
        );
    }
    let agent = get(ENV_AGENT_NAME)
        .filter(|v| !v.is_empty())
        .with_context(|| format!("{ENV_AGENT_NAME} is unset or empty"))?;
    let settings = Settings::from_lookup(|k| get(k).filter(|v| k != ENV_CACERT || ca_is_usable(v)))
        .context("reading the swarm secret store's coordinates from the environment")?;
    Ok(StoreSource {
        settings,
        role: policy::agent_object_name(&agent)?,
        path: queue::agent_queue_path(&agent)?,
    })
}

/// Whether the CA bundle at `path` is a file with bytes in it. `BAO_CACERT`
/// names a systemd credential that may not have been delivered; absent means
/// the container's own trust store.
fn ca_is_usable(path: &str) -> bool {
    std::fs::metadata(path).is_ok_and(|m| m.len() > 0)
}

/// What connecting needs: the queue's address and a store to read this agent's
/// secret from.
///
/// # Errors
/// Either one missing.
fn agent_path(get: impl Fn(&str) -> Option<String>) -> anyhow::Result<AgentPath> {
    let url = get(ENV_NATS_URL)
        .filter(|v| !v.is_empty())
        .with_context(|| format!("{ENV_NATS_URL} is unset or empty"))?;
    Ok(AgentPath {
        url,
        store: Arc::new(store_source(&get)?),
    })
}

/// Resolve this agent's queue coordinates and record them for [`client`].
///
/// # Errors
/// A coordinate is missing. The harness exits on it: an agent without a queue
/// connection is invisible to the operator at swarm level.
pub fn init() -> anyhow::Result<()> {
    let agent = agent_path(|k| std::env::var(k).ok())
        .context("resolving this agent's swarm queue coordinates")?;
    tracing::info!(
        url = %agent.url,
        path = %agent.store.path,
        "this agent connects to the swarm queue with its own credential from the secret store"
    );
    let _ = AGENT.set(agent);
    Ok(())
}

/// The shared queue connection, made on first call and memoized for the rest
/// of the process.
///
/// One connection per process, not per publisher: the agent authenticates as
/// one client, so a second `connect` would be a second live connection for the
/// same identity rather than a second credential.
///
/// The first call retries the connect on [`CONNECT_BACKOFF`], and every other
/// caller waits on it. When no attempt succeeds the process exits with status
/// 1, so this never returns without a client.
pub async fn client() -> async_nats::Client {
    CLIENT.get_or_init(connect_or_exit).await.clone()
}

async fn connect_or_exit() -> async_nats::Client {
    let connected = match AGENT.get() {
        Some(agent) => retry_first_connect(|| connect_once(agent), tokio::time::sleep).await,
        None => Err(anyhow!("swarm_queue::init has not run")),
    };
    connected.unwrap_or_else(|e| {
        tracing::error!(
            error = format!("{e:#}"),
            "no swarm queue connection; exiting, since this agent cannot run without one"
        );
        std::process::exit(1)
    })
}

/// One connection attempt with this agent's own credential.
async fn connect_once(agent: &AgentPath) -> anyhow::Result<async_nats::Client> {
    let secret = AgentSecret::new(Arc::clone(&agent.store));
    let name = crate::identity::label();
    let client = connect_as_agent(&agent.url, name, secret).await?;
    tracing::info!(
        url = %agent.url,
        "connected to the swarm queue with this agent's own credential"
    );
    Ok(client)
}

/// Run `connect` until it succeeds, sleeping each wait in [`CONNECT_BACKOFF`]
/// between attempts, and return the last error once every attempt has failed.
async fn retry_first_connect<T, C, CFut, S, SFut>(mut connect: C, mut sleep: S) -> anyhow::Result<T>
where
    C: FnMut() -> CFut,
    CFut: Future<Output = anyhow::Result<T>>,
    S: FnMut(Duration) -> SFut,
    SFut: Future<Output = ()>,
{
    let mut waits = CONNECT_BACKOFF.iter();
    let mut attempt = 1;
    loop {
        let e = match connect().await {
            Ok(client) => return Ok(client),
            Err(e) => e,
        };
        let Some(&wait) = waits.next() else {
            return Err(e.context(format!("{attempt} connection attempts failed")));
        };
        tracing::warn!(
            error = format!("{e:#}"),
            attempt,
            retry_in = ?wait,
            "connecting to the swarm queue failed"
        );
        sleep(wait).await;
        attempt += 1;
    }
}

/// Connect presenting this agent's own token as `name`.
///
/// The secret is read before the first attempt, and any failure of that read
/// or of the first attempt, a refusal included, is returned for
/// [`retry_first_connect`] to retry. Once connected, each reconnect attempt
/// takes its secret from [`AgentSecret::for_attempt`].
async fn connect_as_agent<S: SecretSource>(
    url: &str,
    name: String,
    secret: AgentSecret<S>,
) -> anyhow::Result<async_nats::Client> {
    secret.prime().await?;
    let secret = Arc::new(secret);
    let agent = name.clone();
    // No extra trust anchor: the container trusts the store's root, which
    // `hive_c0re::meta` embeds at build time.
    swarm_queue_client::connect_with_token(url, None, &agent, move || {
        let secret = Arc::clone(&secret);
        let name = name.clone();
        // Spawned so the future handed back is only a `JoinHandle`, which is
        // `Sync` as async-nats's auth callback requires; the store read is not.
        let attempt = tokio::spawn(async move {
            let value = secret.for_attempt().await.map_err(|e| format!("{e:#}"))?;
            swarm_queue_client::agent_token::format_agent_token(&name, &value)
                .map_err(|e| e.to_string())
        });
        async move { attempt.await.map_err(|e| e.to_string())? }
    })
    .await
    .map_err(|e| anyhow!(swarm_queue_client::chain(&e)))
}

#[cfg(test)]
mod tests {
    use std::collections::VecDeque;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::{Arc, Mutex};
    use std::time::Duration;

    use anyhow::anyhow;

    use super::{
        AgentPath, AgentSecret, CONNECT_BACKOFF, SecretSource, StoreSource, agent_path,
        connect_as_agent, retry_first_connect, store_source,
    };

    /// A lookup standing in for a container that was handed a store and the
    /// queue's address.
    fn with_store(k: &str) -> Option<String> {
        match k {
            "BAO_ADDR" => Some("https://bao.t.local:8200".to_owned()),
            "BAO_CLIENT_CERT" => Some("/run/credentials/x/cert".to_owned()),
            "BAO_CLIENT_KEY" => Some("/run/credentials/x/key".to_owned()),
            "HIVE_AGENT_NAME" => Some("a1".to_owned()),
            "HIVE_AGENT_NATS_URL" => Some("nats://10.42.0.1:4222".to_owned()),
            _ => None,
        }
    }

    fn source() -> StoreSource {
        store_source(with_store).expect("a complete store environment")
    }

    #[test]
    fn a_store_address_reads_the_agents_own_path_as_its_own_role() {
        let s = source();
        assert_eq!(s.path, "swarm/agents/a1/queue");
        assert_eq!(s.role, "hive-agent-a1");
    }

    #[test]
    fn a_missing_store_address_is_an_error() {
        for addr in [None, Some(String::new())] {
            let e = store_source(|k| {
                if k == "BAO_ADDR" {
                    addr.clone()
                } else {
                    with_store(k)
                }
            })
            .expect_err("no store to read the credential from");
            assert!(format!("{e:#}").contains("BAO_ADDR"), "{e:#}");
        }
    }

    #[test]
    fn a_store_address_without_the_agent_or_its_identity_is_an_error() {
        for var in ["HIVE_AGENT_NAME", "BAO_CLIENT_CERT", "BAO_CLIENT_KEY"] {
            let e = store_source(|k| if k == var { None } else { with_store(k) })
                .expect_err("a half-delivered store");
            assert!(format!("{e:#}").contains(var), "dropping {var}: {e:#}");
        }
    }

    /// `BAO_CACERT` is always named and may point at nothing; that is the
    /// container's own trust store, not a CA path that fails the handshake.
    #[test]
    fn an_undelivered_ca_is_dropped() {
        let with_missing_ca = store_source(|k| {
            if k == "BAO_CACERT" {
                Some("/nonexistent/ca".to_owned())
            } else {
                with_store(k)
            }
        })
        .expect("complete");
        assert_eq!(with_missing_ca, source());
    }

    #[test]
    fn a_store_and_the_queue_address_are_all_an_agent_connects_with() {
        assert_eq!(
            agent_path(with_store).expect("complete"),
            AgentPath {
                url: "nats://10.42.0.1:4222".to_owned(),
                store: Arc::new(source()),
            }
        );
    }

    #[test]
    fn without_a_store_or_the_queue_address_an_agent_has_nothing_to_connect_with() {
        for var in ["HIVE_AGENT_NATS_URL", "BAO_ADDR"] {
            let e = agent_path(|k| if k == var { None } else { with_store(k) })
                .expect_err("a coordinate is missing");
            assert!(format!("{e:#}").contains(var), "dropping {var}: {e:#}");
        }
    }

    /// Runs `retry_first_connect` against a connect that fails `failures`
    /// times and then succeeds, recording each wait instead of sleeping it.
    async fn retry_with(failures: usize) -> (anyhow::Result<usize>, usize, Vec<Duration>) {
        let mut calls = 0;
        let mut waits = Vec::new();
        let result = retry_first_connect(
            || {
                calls += 1;
                let n = calls;
                async move {
                    if n > failures {
                        Ok(n)
                    } else {
                        Err(anyhow!("attempt {n} refused"))
                    }
                }
            },
            |wait| {
                waits.push(wait);
                async {}
            },
        )
        .await;
        (result, calls, waits)
    }

    #[tokio::test]
    async fn a_first_connect_that_succeeds_waits_for_nothing() {
        let (result, calls, waits) = retry_with(0).await;
        assert_eq!(result.expect("connected"), 1);
        assert_eq!(calls, 1);
        assert!(waits.is_empty(), "{waits:?}");
    }

    /// 5s doubling, one wait per failed attempt, and the attempt after the
    /// last wait is the one that connects.
    #[tokio::test]
    async fn a_failed_connect_is_retried_on_a_doubling_backoff() {
        let (result, calls, waits) = retry_with(3).await;
        assert_eq!(result.expect("connected on the fourth attempt"), 4);
        assert_eq!(calls, 4);
        assert_eq!(
            waits,
            [5, 10, 20].map(Duration::from_secs),
            "each wait doubles the one before"
        );
    }

    /// About five minutes of waits, then the last attempt's error rather than
    /// another wait.
    #[tokio::test]
    async fn the_retry_gives_up_with_the_last_error_after_about_five_minutes() {
        let (result, calls, waits) = retry_with(usize::MAX).await;
        let e = result.expect_err("every attempt failed");
        assert_eq!(calls, CONNECT_BACKOFF.len() + 1);
        assert_eq!(waits, CONNECT_BACKOFF);
        let total: Duration = waits.iter().sum();
        assert!(
            (Duration::from_mins(4)..=Duration::from_mins(6)).contains(&total),
            "{total:?}"
        );
        assert!(
            format!("{e:#}").contains(&format!("attempt {calls} refused")),
            "{e:#}"
        );
    }

    enum Answer {
        Value(&'static str),
        Absent,
        Fails,
        Hangs,
    }

    /// A store that answers each read with the next scripted answer and
    /// counts the reads.
    struct Scripted {
        answers: Mutex<VecDeque<Answer>>,
        reads: AtomicUsize,
    }

    fn scripted(answers: impl IntoIterator<Item = Answer>) -> Arc<Scripted> {
        Arc::new(Scripted {
            answers: Mutex::new(answers.into_iter().collect()),
            reads: AtomicUsize::new(0),
        })
    }

    impl SecretSource for Arc<Scripted> {
        fn origin(&self) -> &'static str {
            "swarm/agents/a1/queue"
        }

        async fn fetch(&self) -> anyhow::Result<Option<String>> {
            self.reads.fetch_add(1, Ordering::SeqCst);
            let next = self.answers.lock().expect("lock").pop_front();
            match next.expect("a read nobody scripted") {
                Answer::Value(v) => Ok(Some(v.to_owned())),
                Answer::Absent => Ok(None),
                Answer::Fails => Err(anyhow!("the store is unreachable")),
                Answer::Hangs => std::future::pending().await,
            }
        }
    }

    fn reads(s: &Arc<Scripted>) -> usize {
        s.reads.load(Ordering::SeqCst)
    }

    /// (a) The first attempt presents what was just read, and does not read
    /// the store a second time for it.
    #[tokio::test]
    async fn the_first_attempt_presents_the_value_read_before_it() {
        let store = scripted([Answer::Value("v1\n")]);
        let secret = AgentSecret::new(Arc::clone(&store));
        secret.prime().await.expect("the store holds a secret");
        assert_eq!(secret.for_attempt().await.expect("primed"), "v1");
        assert_eq!(reads(&store), 1);
    }

    /// (b) After the queue refuses a re-minted secret, the next attempt reads
    /// the store and presents what it holds now: one read per attempt.
    #[tokio::test]
    async fn each_later_attempt_presents_what_the_store_holds_now() {
        let store = scripted([
            Answer::Value("v1"),
            Answer::Value("v1"),
            Answer::Value("v2"),
        ]);
        let secret = AgentSecret::new(Arc::clone(&store));
        secret.prime().await.expect("the store holds a secret");
        assert_eq!(secret.for_attempt().await.expect("primed"), "v1");
        assert_eq!(secret.for_attempt().await.expect("re-read"), "v1");
        assert_eq!(secret.for_attempt().await.expect("re-read"), "v2");
        assert_eq!(reads(&store), 3);
    }

    /// (b) Bounded: a store that never answers costs one attempt the read
    /// timeout, and the attempt goes ahead with the value it already holds.
    #[tokio::test]
    async fn a_store_that_does_not_answer_is_given_up_on() {
        let store = scripted([Answer::Value("v1"), Answer::Hangs, Answer::Hangs]);
        let mut secret = AgentSecret::new(Arc::clone(&store));
        secret.timeout = Duration::from_millis(20);
        secret.prime().await.expect("the store holds a secret");
        assert_eq!(secret.for_attempt().await.expect("primed"), "v1");
        assert_eq!(
            secret
                .for_attempt()
                .await
                .expect("the held value stands in"),
            "v1"
        );
        assert_eq!(reads(&store), 2);

        let cold = AgentSecret::new(scripted([Answer::Hangs]));
        let mut cold = cold;
        cold.timeout = Duration::from_millis(20);
        let e = cold.prime().await.expect_err("nothing to stand in");
        assert!(format!("{e:#}").contains("did not answer"), "{e:#}");
    }

    /// (b) A secret the store no longer holds (the agent was destroyed) is not
    /// presented again, even when a later read fails.
    #[tokio::test]
    async fn a_secret_removed_from_the_store_is_dropped() {
        let store = scripted([Answer::Value("v1"), Answer::Absent, Answer::Fails]);
        let secret = AgentSecret::new(Arc::clone(&store));
        secret.prime().await.expect("the store holds a secret");
        assert_eq!(secret.for_attempt().await.expect("primed"), "v1");
        let gone = secret.for_attempt().await.expect_err("nothing stored");
        assert!(
            format!("{gone:#}").contains("swarm/agents/a1/queue"),
            "{gone:#}"
        );
        secret
            .for_attempt()
            .await
            .expect_err("the dropped value does not come back");
    }

    /// (c) Nothing to connect with before the first attempt fails that attempt.
    /// The queue is never dialled: nothing listens at the address given.
    #[tokio::test]
    async fn no_secret_before_the_first_connect_fails_the_attempt() {
        for (answer, says) in [
            (Answer::Fails, "unreachable"),
            (Answer::Absent, "nothing is stored"),
            (Answer::Value("  \n"), "nothing is stored"),
        ] {
            let secret = AgentSecret::new(scripted([answer]));
            let e = connect_as_agent("nats://127.0.0.1:9", "a1".to_owned(), secret)
                .await
                .expect_err("no secret to present");
            assert!(format!("{e:#}").contains(says), "{e:#}");
        }
    }

    /// Serve one NATS handshake on `sock`: send INFO, read CONNECT up to the
    /// PING, and return the token presented.
    async fn handshake(sock: &mut tokio::net::TcpStream, port: u16) -> String {
        use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _};
        let info = format!(
            "INFO {{\"server_id\":\"fake\",\"version\":\"2.10.0\",\"go\":\"go\",\
             \"host\":\"127.0.0.1\",\"port\":{port},\"headers\":true,\
             \"max_payload\":1048576,\"proto\":1,\"auth_required\":true,\
             \"nonce\":\"n0nce\"}}\r\n"
        );
        sock.write_all(info.as_bytes()).await.expect("send INFO");
        let mut lines = tokio::io::BufReader::new(sock).lines();
        let mut token = String::new();
        while let Some(line) = lines.next_line().await.expect("read") {
            if let Some(json) = line.strip_prefix("CONNECT ") {
                let v: serde_json::Value = serde_json::from_str(json).expect("CONNECT is JSON");
                token = v["auth_token"].as_str().unwrap_or_default().to_owned();
            } else if line.starts_with("PING") {
                break;
            }
        }
        token
    }

    /// (b) end to end through async-nats: connected with `v1`, the connection
    /// drops, the reconnect presenting `v1` is refused, and the next attempt
    /// reads `v2` and is accepted.
    #[tokio::test]
    async fn a_refused_reconnect_reconnects_with_the_secret_read_after_it() {
        use tokio::io::AsyncWriteExt as _;

        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
            .await
            .expect("bind");
        let port = listener.local_addr().expect("addr").port();
        let (seen_tx, mut seen_rx) = tokio::sync::mpsc::unbounded_channel();
        tokio::spawn(async move {
            let mut kept = Vec::new();
            for n in 0.. {
                let (mut sock, _) = listener.accept().await.expect("accept");
                let token = handshake(&mut sock, port).await;
                let accept = n == 0 || token.rsplit_once('.').is_some_and(|(_, s)| s == "v2");
                let reply: &[u8] = if accept {
                    b"PONG\r\n"
                } else {
                    b"-ERR 'Authorization Violation'\r\n"
                };
                let _ = sock.write_all(reply).await;
                let _ = seen_tx.send(token);
                // The first connection is dropped to force a reconnect; the
                // one accepted after the refusal stays open.
                if accept && n > 0 {
                    kept.push(sock);
                }
            }
        });

        let store = scripted([
            Answer::Value("v1"),
            Answer::Value("v1"),
            Answer::Value("v2"),
        ]);
        let secret = AgentSecret::new(Arc::clone(&store));
        let url = format!("nats://127.0.0.1:{port}");
        let _client = connect_as_agent(&url, "a1".to_owned(), secret)
            .await
            .expect("the first connect is accepted");

        let mut seen = Vec::new();
        while seen.len() < 3 {
            let token = tokio::time::timeout(Duration::from_secs(20), seen_rx.recv())
                .await
                .expect("the client keeps reconnecting")
                .expect("the server is running");
            seen.push(token);
        }
        let secrets: Vec<&str> = seen
            .iter()
            .map(|t| t.rsplit_once('.').map_or("", |(_, s)| s))
            .collect();
        assert_eq!(secrets, ["v1", "v1", "v2"]);
        assert_eq!(reads(&store), 3);
    }
}