klieo-core 0.41.3

//! Per-replica leader claim + liveness probe for orphan detection.
//!
//! Wraps a [`KvStore`] bucket. [`LeaderRegistry::claim`] writes a
//! [`LeaderEntry`] envelope and returns a [`LeaderHandle`] that spawns
//! an internal heartbeat task re-putting the key at half-TTL
//! intervals. `Drop` on the handle aborts the heartbeat and
//! best-effort deletes the key. Followers call
//! [`LeaderRegistry::is_alive`] on resume entry; `false` means orphan
//! -> write terminal "leader died" frame OR re-invoke from the
//! cached payload (cluster 0.24, see ADR-024).

use crate::bus::{KvStore, Revision};
use crate::error::BusError;
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;
use tokio::task::JoinHandle;
use tracing::Instrument as _;

/// Hard cap on the number of failover re-invoke attempts per stream.
///
/// After this many followers have re-invoked + died, the next
/// follower writes the terminal "leader died" frame instead of
/// re-invoking. Prevents bad-state tools from draining the cluster
/// via infinite failover. See ADR-024.
pub const FAILOVER_ATTEMPT_CAP: u32 = 3;

/// Persisted leader claim envelope.
///
/// Cluster 0.24 extends cluster-0.20's raw-replica-id value to carry
/// the original invoke payload + principal + attempt counter so
/// followers can re-invoke idempotent handlers on orphan detection.
///
/// Backward compat: cluster-0.20 entries (raw `replica_id` string)
/// parse via [`LeaderEntry::from_bytes`]'s fallback path. Pre-0.24
/// entries deserialise with `payload = None`, `attempt = 0`,
/// `principal = None`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LeaderEntry {
    /// Identifier of the replica that holds (or last held) the claim.
    pub replica_id: String,
    /// Original invoke payload (serialised JSON params). `None` for
    /// non-idempotent handlers or pre-0.24 entries. JSON-serialises
    /// as a byte-array via the `bytes_opt_serde` adapter.
    #[serde(with = "bytes_opt_serde", default)]
    pub payload: Option<Bytes>,
    /// Re-invoke attempt counter. 0 = original; cap = [`FAILOVER_ATTEMPT_CAP`].
    #[serde(default)]
    pub attempt: u32,
    /// Original authenticated principal (cluster 0.22). Used by the
    /// follower to fabricate a `RequestContext` for re-invoke that
    /// respects cross-tenant ownership gates. `None` = anonymous or
    /// no auth wired at original invoke.
    #[serde(default)]
    pub principal: Option<String>,
}

impl LeaderEntry {
    /// Parse from KV value bytes. JSON [`LeaderEntry`] first; falls
    /// back to pre-0.24 raw-string interpretation when JSON fails
    /// (treats bytes as `replica_id` only).
    pub fn from_bytes(bytes: &[u8]) -> Self {
        if let Ok(entry) = serde_json::from_slice::<LeaderEntry>(bytes) {
            return entry;
        }
        let replica_id = std::str::from_utf8(bytes).unwrap_or("").to_string();
        Self {
            replica_id,
            payload: None,
            attempt: 0,
            principal: None,
        }
    }

    /// Serialise to KV value bytes. Errors as [`BusError::Permanent`]
    /// because a serialisation failure here indicates a programmer
    /// bug (the envelope's fields are all serde-friendly), not a
    /// transient bus condition.
    pub fn to_bytes(&self) -> Result<Bytes, BusError> {
        let v = serde_json::to_vec(self)
            .map_err(|e| BusError::Permanent(format!("leader-entry encode: {e}")))?;
        Ok(Bytes::from(v))
    }
}

/// Serde adapter for `Option<bytes::Bytes>`. JSON encodes as
/// `Option<Vec<u8>>` (a byte array); deserialise rebuilds the `Bytes`
/// from the vector. Bus already enforces a 1 MiB payload cap upstream
/// so the byte-array encoding is acceptable.
mod bytes_opt_serde {
    use bytes::Bytes;
    use serde::{Deserialize, Deserializer, Serializer};

    pub fn serialize<S: Serializer>(v: &Option<Bytes>, s: S) -> Result<S::Ok, S::Error> {
        match v {
            Some(b) => s.serialize_some(b.as_ref()),
            None => s.serialize_none(),
        }
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Option<Bytes>, D::Error> {
        let v: Option<Vec<u8>> = Option::deserialize(d)?;
        Ok(v.map(Bytes::from))
    }
}

/// Per-replica leader claim handle. Backed by a [`KvStore`] bucket
/// dedicated to leader entries; both transports share one bucket
/// keyed by `{kind}.{id}` (e.g. `a2a.t-1`, `mcp.tok-1`).
#[derive(Clone)]
pub struct LeaderRegistry {
    kv: Arc<dyn KvStore>,
    bucket: String,
    replica_id: String,
}

impl LeaderRegistry {
    /// Construct a registry bound to `bucket` on `kv`. `replica_id`
    /// is recorded in the [`LeaderEntry`] envelope; followers only
    /// check existence + freshness, the identity is informational.
    pub fn new(kv: Arc<dyn KvStore>, bucket: String, replica_id: String) -> Self {
        Self {
            kv,
            bucket,
            replica_id,
        }
    }

    /// Borrow the backing [`KvStore`]. Cluster 0.25 reaper wiring
    /// uses this to spawn its scan task against the same bucket the
    /// registry writes to.
    pub fn kv(&self) -> &Arc<dyn KvStore> {
        &self.kv
    }

    /// Name of the KV bucket this registry writes leader entries to
    /// (typically `klieo-leaders`).
    pub fn bucket(&self) -> &str {
        &self.bucket
    }

    /// Claim leadership over `key` with `ttl`. Writes a fresh
    /// [`LeaderEntry`] envelope (attempt = 0) and spawns a heartbeat
    /// task that re-puts the key every `ttl / 2`. Returned
    /// [`LeaderHandle`]'s `Drop` aborts the heartbeat + best-effort
    /// deletes the key.
    ///
    /// `payload` carries the original invoke payload for cluster-0.24
    /// idempotent re-invoke; pass `None` from non-idempotent paths.
    /// `principal` records the original authenticated subject so a
    /// re-invoking follower can fabricate the same cross-tenant
    /// authorization context.
    ///
    /// Note: [`KvStore::put`] does not currently take an explicit
    /// TTL parameter — the bucket-level TTL configured by the
    /// operator is what enforces eviction. The `ttl` parameter
    /// here is used only to compute the heartbeat interval
    /// (`ttl / 2`). For production NATS deployments, configure the
    /// `klieo-leaders` bucket with `max_age = 5s` to match the
    /// default ttl.
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.leader.key = %key,
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "put",
        ),
        err,
    )]
    pub async fn claim(
        &self,
        key: String,
        ttl: Duration,
        payload: Option<Bytes>,
        principal: Option<String>,
    ) -> Result<LeaderHandle, BusError> {
        self.claim_with_heartbeat(key, ttl, ttl / 2, payload, principal)
            .await
    }

    /// Claim leadership with an explicit `heartbeat` interval. Otherwise
    /// identical to [`Self::claim`], which uses `ttl / 2`. Cluster 0.25
    /// callers that opt into tuning the leader TTL use this variant so
    /// the heartbeat cadence stays independent of the TTL.
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.leader.key = %key,
            klieo.leader.ttl_ms = %ttl.as_millis(),
            klieo.leader.heartbeat_ms = %heartbeat.as_millis(),
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "put",
        ),
        err,
    )]
    pub async fn claim_with_heartbeat(
        &self,
        key: String,
        ttl: Duration,
        heartbeat: Duration,
        payload: Option<Bytes>,
        principal: Option<String>,
    ) -> Result<LeaderHandle, BusError> {
        let _ = ttl; // `ttl` retained for callers + tracing; eviction is bucket-level.
        let entry = LeaderEntry {
            replica_id: self.replica_id.clone(),
            payload,
            attempt: 0,
            principal,
        };
        let bytes = entry.to_bytes()?;
        self.kv.put(&self.bucket, &key, bytes.clone()).await?;
        let heartbeat_task = self.spawn_heartbeat_with_interval(key.clone(), heartbeat, bytes);
        Ok(LeaderHandle {
            kv: self.kv.clone(),
            bucket: self.bucket.clone(),
            key,
            replica_id: self.replica_id.clone(),
            heartbeat: Some(heartbeat_task),
        })
    }

    /// CAS-claim leadership atomically based on the expired entry's
    /// revision. Increments `attempt` from `prior` and preserves
    /// `payload` + `principal`. Returns [`BusError::CasConflict`]
    /// when another follower won the race.
    ///
    /// Caller MUST gate this on `prior.attempt < FAILOVER_ATTEMPT_CAP`;
    /// the registry does not enforce the cap so the resume-gate site
    /// (T3/T5) can emit a distinct "cap reached" terminal frame.
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.leader.key = %key,
            klieo.leader.expected_rev = %expected_rev,
            klieo.leader.prior_attempt = %prior.attempt,
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "cas",
        ),
        err,
    )]
    pub async fn claim_with_attempt_cas(
        &self,
        key: String,
        ttl: Duration,
        expected_rev: Revision,
        prior: &LeaderEntry,
    ) -> Result<LeaderHandle, BusError> {
        self.claim_with_attempt_cas_and_heartbeat(key, ttl, ttl / 2, expected_rev, prior)
            .await
    }

    /// CAS-claim leadership with an explicit heartbeat interval. Same
    /// semantics as [`Self::claim_with_attempt_cas`] but lets cluster-0.25
    /// callers override the heartbeat cadence independently of `ttl`.
    ///
    /// [`Self::claim_with_attempt_cas`]: Self::claim_with_attempt_cas
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.leader.key = %key,
            klieo.leader.ttl_ms = %ttl.as_millis(),
            klieo.leader.heartbeat_ms = %heartbeat.as_millis(),
            klieo.leader.expected_rev = %expected_rev,
            klieo.leader.prior_attempt = %prior.attempt,
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "cas",
        ),
        err,
    )]
    pub async fn claim_with_attempt_cas_and_heartbeat(
        &self,
        key: String,
        ttl: Duration,
        heartbeat: Duration,
        expected_rev: Revision,
        prior: &LeaderEntry,
    ) -> Result<LeaderHandle, BusError> {
        let _ = ttl; // see claim_with_heartbeat — TTL is informational here.
        let new_entry = LeaderEntry {
            replica_id: self.replica_id.clone(),
            payload: prior.payload.clone(),
            attempt: prior.attempt + 1,
            principal: prior.principal.clone(),
        };
        let bytes = new_entry.to_bytes()?;
        self.kv
            .cas(&self.bucket, &key, bytes.clone(), Some(expected_rev))
            .await?;
        let heartbeat_task = self.spawn_heartbeat_with_interval(key.clone(), heartbeat, bytes);
        Ok(LeaderHandle {
            kv: self.kv.clone(),
            bucket: self.bucket.clone(),
            key,
            replica_id: self.replica_id.clone(),
            heartbeat: Some(heartbeat_task),
        })
    }

    /// Spawn the periodic heartbeat task that re-puts the same
    /// envelope every `heartbeat`. Extracted so the claim variants
    /// share the loop.
    fn spawn_heartbeat_with_interval(
        &self,
        key: String,
        heartbeat: Duration,
        entry_bytes: Bytes,
    ) -> JoinHandle<()> {
        let kv = self.kv.clone();
        let bucket = self.bucket.clone();
        let key_for_task = key.clone();
        tokio::spawn(
            async move {
                loop {
                    tokio::time::sleep(heartbeat).await;
                    if let Err(e) = kv.put(&bucket, &key_for_task, entry_bytes.clone()).await {
                        tracing::warn!(
                            target: "klieo.leader",
                            bucket = %bucket,
                            key = %key_for_task,
                            error = %e,
                            "leader heartbeat failed; leadership will lapse",
                        );
                        return;
                    }
                    tracing::debug!(
                        target: "klieo.leader",
                        bucket = %bucket,
                        key = %key_for_task,
                        "heartbeat",
                    );
                }
            }
            .instrument(tracing::info_span!(
                "leader_heartbeat_task",
                klieo.leader.key = %key,
                klieo.leader.heartbeat_interval_ms = %heartbeat.as_millis(),
            )),
        )
    }

    /// True when an entry exists for `key`. KV-layer TTL eviction
    /// turns dead leaders into `Ok(false)` automatically.
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.leader.key = %key,
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "get",
        ),
        err,
        level = "debug",
    )]
    pub async fn is_alive(&self, key: &str) -> Result<bool, BusError> {
        Ok(self.kv.get(&self.bucket, key).await?.is_some())
    }

    /// Look up the leader entry (alive or recently expired). Returns
    /// `None` when KV has no entry. Used by cluster-0.24 followers to
    /// read the original payload + attempt count during
    /// orphan-detection re-invoke decisions.
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.leader.key = %key,
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "get",
        ),
        err,
        level = "debug",
    )]
    pub async fn lookup_entry(&self, key: &str) -> Result<Option<LeaderEntry>, BusError> {
        match self.kv.get(&self.bucket, key).await? {
            Some(e) => Ok(Some(LeaderEntry::from_bytes(&e.value))),
            None => Ok(None),
        }
    }

    /// Look up the leader entry together with its current KV
    /// revision. Followers performing CAS-based failover re-invoke
    /// (cluster 0.24, see [`claim_with_attempt_cas`]) need both the
    /// payload from the prior entry and the revision so the CAS can
    /// detect a concurrent winning follower.
    ///
    /// [`claim_with_attempt_cas`]: Self::claim_with_attempt_cas
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.leader.key = %key,
            db.system = "klieo-kv",
            db.namespace = %self.bucket,
            db.operation = "get",
        ),
        err,
        level = "debug",
    )]
    pub async fn lookup_entry_with_revision(
        &self,
        key: &str,
    ) -> Result<Option<(LeaderEntry, Revision)>, BusError> {
        match self.kv.get(&self.bucket, key).await? {
            Some(e) => Ok(Some((LeaderEntry::from_bytes(&e.value), e.revision))),
            None => Ok(None),
        }
    }
}

/// Drop-guard for a leader claim. Aborts the internal heartbeat
/// task + spawns a best-effort delete of the KV key. Operators
/// that need synchronous cleanup should drop early and let the
/// spawned delete settle (a few KV round-trips at most).
pub struct LeaderHandle {
    kv: Arc<dyn KvStore>,
    bucket: String,
    key: String,
    replica_id: String,
    heartbeat: Option<JoinHandle<()>>,
}

impl LeaderHandle {
    /// Identifier of the replica that holds this leader claim.
    ///
    /// Cluster 0.24 failover marker frames record this so resume
    /// clients can correlate "the leader changed mid-stream" with the
    /// new owner in subsequent observability traces. See ADR-024.
    pub fn replica_id(&self) -> &str {
        &self.replica_id
    }
}

impl Drop for LeaderHandle {
    fn drop(&mut self) {
        if let Some(task) = self.heartbeat.take() {
            task.abort();
        }
        let kv = self.kv.clone();
        let bucket = std::mem::take(&mut self.bucket);
        let key = std::mem::take(&mut self.key);
        match tokio::runtime::Handle::try_current() {
            Ok(handle) => {
                handle.spawn(async move {
                    if let Err(e) = kv.delete(&bucket, &key).await {
                        tracing::warn!(
                            target: "klieo.leader",
                            bucket = %bucket,
                            key = %key,
                            error = %e,
                            "leader delete failed; lapse to TTL",
                        );
                    }
                });
            }
            Err(_) => {
                tracing::warn!(
                    target: "klieo.leader",
                    bucket = %bucket,
                    key = %key,
                    "leader delete skipped: no tokio runtime active during Drop",
                );
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bus::{KvEntry, Lease, Revision};
    use crate::test_utils::fake_kv;
    use async_trait::async_trait;
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::time::Duration;

    const BUCKET: &str = "klieo-leaders";

    /// Spy [`KvStore`] that counts `put` calls, delegating to an inner store.
    struct CountingKvStore {
        inner: Arc<dyn KvStore>,
        put_count: Arc<AtomicU64>,
    }

    impl CountingKvStore {
        fn new(inner: Arc<dyn KvStore>) -> (Arc<Self>, Arc<AtomicU64>) {
            let put_count = Arc::new(AtomicU64::new(0));
            (
                Arc::new(Self {
                    inner,
                    put_count: put_count.clone(),
                }),
                put_count,
            )
        }
    }

    #[async_trait]
    impl KvStore for CountingKvStore {
        async fn get(&self, bucket: &str, key: &str) -> Result<Option<KvEntry>, BusError> {
            self.inner.get(bucket, key).await
        }

        async fn put(&self, bucket: &str, key: &str, value: Bytes) -> Result<Revision, BusError> {
            self.put_count.fetch_add(1, Ordering::SeqCst);
            self.inner.put(bucket, key, value).await
        }

        async fn cas(
            &self,
            bucket: &str,
            key: &str,
            value: Bytes,
            expected: Option<Revision>,
        ) -> Result<Revision, BusError> {
            self.inner.cas(bucket, key, value, expected).await
        }

        async fn delete(&self, bucket: &str, key: &str) -> Result<(), BusError> {
            self.inner.delete(bucket, key).await
        }

        async fn lease(
            &self,
            bucket: &str,
            key: &str,
            ttl: std::time::Duration,
        ) -> Result<Lease, BusError> {
            self.inner.lease(bucket, key, ttl).await
        }

        async fn keys(&self, bucket: &str) -> Result<Vec<String>, BusError> {
            self.inner.keys(bucket).await
        }
    }

    #[tokio::test]
    async fn claim_writes_entry_and_is_alive_returns_true() {
        let kv = fake_kv();
        let reg = LeaderRegistry::new(kv.clone(), BUCKET.into(), "replica-a".into());
        let _handle = reg
            .claim("a2a.t-1".into(), Duration::from_secs(5), None, None)
            .await
            .unwrap();
        assert!(reg.is_alive("a2a.t-1").await.unwrap());
    }

    #[tokio::test]
    async fn drop_handle_deletes_entry() {
        let kv = fake_kv();
        let reg = LeaderRegistry::new(kv.clone(), BUCKET.into(), "replica-a".into());
        {
            let _handle = reg
                .claim("a2a.t-2".into(), Duration::from_secs(5), None, None)
                .await
                .unwrap();
            assert!(reg.is_alive("a2a.t-2").await.unwrap());
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(!reg.is_alive("a2a.t-2").await.unwrap());
    }

    #[tokio::test]
    async fn is_alive_false_on_unknown_key() {
        let kv = fake_kv();
        let reg = LeaderRegistry::new(kv, BUCKET.into(), "replica-a".into());
        assert!(!reg.is_alive("never").await.unwrap());
    }

    #[test]
    fn leader_entry_round_trip_json() {
        let entry = LeaderEntry {
            replica_id: "replica-a".into(),
            payload: Some(Bytes::from_static(b"{\"params\":1}")),
            attempt: 2,
            principal: Some("alice".into()),
        };
        let bytes = entry.to_bytes().unwrap();
        let back = LeaderEntry::from_bytes(&bytes);
        assert_eq!(back.replica_id, "replica-a");
        assert_eq!(back.payload.as_deref().unwrap(), b"{\"params\":1}");
        assert_eq!(back.attempt, 2);
        assert_eq!(back.principal.as_deref(), Some("alice"));
    }

    #[test]
    fn leader_entry_pre_024_raw_string_fallback() {
        // Cluster-0.20 entries stored just the replica_id bytes.
        let bytes = b"replica-old";
        let entry = LeaderEntry::from_bytes(bytes);
        assert_eq!(entry.replica_id, "replica-old");
        assert!(entry.payload.is_none());
        assert_eq!(entry.attempt, 0);
        assert!(entry.principal.is_none());
    }

    #[tokio::test]
    async fn claim_writes_envelope_with_payload_and_principal() {
        let kv = fake_kv();
        let reg = LeaderRegistry::new(kv.clone(), BUCKET.into(), "replica-a".into());
        let _handle = reg
            .claim(
                "a2a.t-payload".into(),
                Duration::from_secs(5),
                Some(Bytes::from_static(b"{\"params\":42}")),
                Some("alice".into()),
            )
            .await
            .unwrap();

        let entry = reg.lookup_entry("a2a.t-payload").await.unwrap().unwrap();
        assert_eq!(entry.replica_id, "replica-a");
        assert_eq!(entry.payload.as_deref().unwrap(), b"{\"params\":42}");
        assert_eq!(entry.attempt, 0);
        assert_eq!(entry.principal.as_deref(), Some("alice"));
    }

    #[tokio::test]
    async fn claim_with_attempt_cas_increments_attempt() {
        let kv = fake_kv();
        let reg_a = LeaderRegistry::new(kv.clone(), BUCKET.into(), "replica-a".into());
        let reg_b = LeaderRegistry::new(kv.clone(), BUCKET.into(), "replica-b".into());
        let original = reg_a
            .claim(
                "a2a.t-cas".into(),
                Duration::from_secs(5),
                Some(Bytes::from_static(b"{}")),
                Some("alice".into()),
            )
            .await
            .unwrap();
        // Drop the original claim's heartbeat task — keep the KV
        // entry in place by reading it (don't Drop the handle yet,
        // that would best-effort-delete the key).
        let (prior, rev) = reg_b
            .lookup_entry_with_revision("a2a.t-cas")
            .await
            .unwrap()
            .unwrap();
        let _new = reg_b
            .claim_with_attempt_cas("a2a.t-cas".into(), Duration::from_secs(5), rev, &prior)
            .await
            .unwrap();
        std::mem::forget(original);

        let updated = reg_b.lookup_entry("a2a.t-cas").await.unwrap().unwrap();
        assert_eq!(updated.replica_id, "replica-b");
        assert_eq!(updated.attempt, 1);
        assert_eq!(updated.principal.as_deref(), Some("alice"));
    }

    #[tokio::test]
    async fn claim_with_attempt_cas_rejects_on_revision_mismatch() {
        let kv = fake_kv();
        let reg_a = LeaderRegistry::new(kv.clone(), BUCKET.into(), "replica-a".into());
        let reg_b = LeaderRegistry::new(kv.clone(), BUCKET.into(), "replica-b".into());
        let original = reg_a
            .claim("a2a.t-stale".into(), Duration::from_secs(5), None, None)
            .await
            .unwrap();
        let stale_rev: Revision = 99_999;
        let prior = LeaderEntry {
            replica_id: "replica-a".into(),
            payload: None,
            attempt: 0,
            principal: None,
        };
        let result = reg_b
            .claim_with_attempt_cas(
                "a2a.t-stale".into(),
                Duration::from_secs(5),
                stale_rev,
                &prior,
            )
            .await;
        match result {
            Err(BusError::CasConflict { .. }) => {}
            Err(other) => panic!("expected CasConflict, got {other:?}"),
            Ok(_handle) => panic!("expected CasConflict, claim succeeded"),
        }
        std::mem::forget(original);
    }

    #[tokio::test(start_paused = true)]
    async fn heartbeat_renews_lease_before_expiry() {
        let base_kv = fake_kv();
        let (spy_kv, put_count) = CountingKvStore::new(base_kv);
        let reg = LeaderRegistry::new(spy_kv as Arc<dyn KvStore>, BUCKET.into(), "replica-a".into());
        let heartbeat_interval = Duration::from_millis(100);
        let _handle = reg
            .claim_with_heartbeat(
                "a2a.heartbeat-test".into(),
                Duration::from_millis(500),
                heartbeat_interval,
                None,
                None,
            )
            .await
            .unwrap();

        let puts_after_claim = put_count.load(Ordering::SeqCst);
        assert!(puts_after_claim >= 1, "claim must have issued at least one put");

        // Advance time past two heartbeat intervals so the heartbeat task fires at least once.
        // Yield before the advance to ensure the spawned task has entered its first sleep.
        tokio::task::yield_now().await;
        tokio::time::advance(heartbeat_interval * 3).await;
        // Multiple yields give the spawned heartbeat task opportunity to wake and run.
        for _ in 0..10 {
            tokio::task::yield_now().await;
        }

        let puts_after_heartbeat = put_count.load(Ordering::SeqCst);
        assert!(
            puts_after_heartbeat > puts_after_claim,
            "heartbeat must have issued additional puts after time advance; \
             puts after claim={puts_after_claim}, puts after advance={puts_after_heartbeat}"
        );
    }

    #[tokio::test]
    async fn lookup_entry_with_revision_returns_nonzero_revision() {
        let kv = fake_kv();
        let reg = LeaderRegistry::new(kv.clone(), BUCKET.into(), "replica-a".into());
        let _handle = reg
            .claim("a2a.rev-test".into(), Duration::from_secs(5), None, None)
            .await
            .unwrap();
        let result = reg
            .lookup_entry_with_revision("a2a.rev-test")
            .await
            .unwrap();
        let (entry, revision) = result.expect("entry must be present after claim");
        assert_eq!(entry.replica_id, "replica-a");
        assert!(revision > 0, "revision must be non-zero after put");
    }
}