dynamo-llm 1.5.1

// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::collections::HashMap;
use std::future::Future;
use std::sync::{Arc, Weak};
use std::time::Duration;

use dynamo_kv_router::identity::PoolId;
use dynamo_kv_router::indexer::cuckoo::{CkfBuildError, CkfConfig, DcCkfState, ProducerIdentity};
use dynamo_kv_router::protocols::{ActiveLoad, WorkerId};
use dynamo_runtime::protocols::EndpointId;
use parking_lot::Mutex;
use tokio::sync::{Notify, OwnedSemaphorePermit, TryAcquireError};
use tokio::sync::{Semaphore, mpsc, oneshot, watch};
use tokio_util::sync::CancellationToken;

use super::actor::{ActorFault, KvDcRelayHandle, StreamScope};
use super::host::KvDcRelayError;
use super::identity::{
    CanonicalModelId, CanonicalModelRegistration, DcPoolCatalog, DcPoolDescriptor, DcRelayIdentity,
    KvQuerySemantics, ModelAlias, WorkerRole,
};
use super::load::{LoadObservationOutcome, PoolLoadSnapshot, PoolLoadState};
use super::publication::{
    PublicationHub, PublicationHubConfig, PublicationHubError, PublicationHubSubscription,
    TerminalFailure, publication_lease,
};
use crate::local_model::runtime_config::ModelRuntimeConfig;

const DEFAULT_CKF_ALLOCATION_CONCURRENCY: usize = 2;
const DEFAULT_INITIALIZED_POOL_HUBS: usize = 64;

#[derive(Debug, Clone, Copy)]
pub(super) struct PoolActorConfig {
    pub(super) expected_unique_blocks: usize,
    pub(super) publication_threshold: usize,
    pub(super) publication_delay: Duration,
}

#[derive(Debug)]
pub(super) struct PoolAttachRequest {
    pub(super) pool_id: PoolId,
    pub(super) endpoint: EndpointId,
    pub(super) registrations: Vec<CanonicalModelRegistration>,
    pub(super) query_semantics: KvQuerySemantics,
    pub(super) roles: Vec<WorkerRole>,
    pub(super) serving_facts: Option<PoolServingFacts>,
}

#[derive(Debug, Clone)]
pub(super) struct PoolServingFacts {
    pub(super) runtime_configs: HashMap<WorkerId, ModelRuntimeConfig>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum PoolRetirementMode {
    Graceful,
    Fenced,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PoolEntryState {
    Active,
    Withdrawn,
    Fenced,
}

struct PoolEntry {
    endpoint: EndpointId,
    handle: KvDcRelayHandle,
    identity: ProducerIdentity,
    layout_generation: u64,
    registrations: Arc<[CanonicalModelRegistration]>,
    query_semantics: KvQuerySemantics,
    roles: Arc<[WorkerRole]>,
    cancel: CancellationToken,
    state: PoolEntryState,
    serving: Option<PoolServingState>,
    hub: Arc<PoolHubSlot>,
}

struct PoolServingState {
    load: PoolLoadState,
}

#[derive(Clone)]
pub(super) struct PoolPublicationConfig {
    pub(super) hub: PublicationHubConfig,
    pub(super) max_initialized_pool_hubs: usize,
    #[cfg(test)]
    pub(super) eviction_gate: Option<Arc<Semaphore>>,
}

impl Default for PoolPublicationConfig {
    fn default() -> Self {
        Self {
            hub: PublicationHubConfig::default(),
            max_initialized_pool_hubs: DEFAULT_INITIALIZED_POOL_HUBS,
            #[cfg(test)]
            eviction_gate: None,
        }
    }
}

enum PoolHubState {
    Vacant,
    Initializing,
    Ready(InitializedPublicationHub),
    Evicting,
    Failed(String),
    Retired,
}

struct InitializedPublicationHub {
    hub: PublicationHub,
    permit: OwnedSemaphorePermit,
}

struct PoolHubSlotState {
    phase: PoolHubState,
    pending_admissions: usize,
    retirement_requested: bool,
}

struct PoolHubSlot {
    state: Mutex<PoolHubSlotState>,
    changed: Notify,
}

impl PoolHubSlot {
    fn new() -> Self {
        Self {
            state: Mutex::new(PoolHubSlotState {
                phase: PoolHubState::Vacant,
                pending_admissions: 0,
                retirement_requested: false,
            }),
            changed: Notify::new(),
        }
    }

    fn begin_admission(self: &Arc<Self>) -> Result<PoolHubAdmission, PublicationHubError> {
        let mut state = self.state.lock();
        if state.retirement_requested || matches!(state.phase, PoolHubState::Retired) {
            return Err(PublicationHubError::Unavailable(
                "pool generation retired".to_string(),
            ));
        }
        let Some(pending_admissions) = state.pending_admissions.checked_add(1) else {
            return Err(PublicationHubError::Unavailable(
                "pool publication admission counter exhausted".to_string(),
            ));
        };
        state.pending_admissions = pending_admissions;
        Ok(PoolHubAdmission { slot: self.clone() })
    }

    async fn get_or_start(
        self: &Arc<Self>,
        registry: &Arc<PoolRegistry>,
        actor: KvDcRelayHandle,
        generation_cancel: CancellationToken,
        terminal_failure: TerminalFailure,
    ) -> Result<PublicationHub, PublicationHubError> {
        let config = registry.publication_config.clone();
        let initialization_permits = registry.publication_hub_permits.clone();
        let max_initialized_hubs = registry.max_initialized_pool_hubs;
        let registry = Arc::downgrade(registry);
        let mut attempted_initialization = false;
        loop {
            let changed = self.changed.notified();
            tokio::pin!(changed);
            changed.as_mut().enable();
            let initialization_permit = {
                let mut state = self.state.lock();
                match &state.phase {
                    PoolHubState::Ready(initialized) => return Ok(initialized.hub.clone()),
                    PoolHubState::Failed(reason) => {
                        return Err(PublicationHubError::Unavailable(reason.clone()));
                    }
                    PoolHubState::Retired => {
                        return Err(PublicationHubError::Unavailable(
                            "pool generation retired".to_string(),
                        ));
                    }
                    PoolHubState::Initializing | PoolHubState::Evicting => None,
                    PoolHubState::Vacant => {
                        if attempted_initialization {
                            return Err(PublicationHubError::InitializedHubLimit {
                                limit: max_initialized_hubs,
                            });
                        }
                        let permit = initialization_permits.clone().try_acquire_owned().ok();
                        state.phase = PoolHubState::Initializing;
                        attempted_initialization = true;
                        Some(permit)
                    }
                }
            };
            if let Some(initialization_permit) = initialization_permit {
                let slot = self.clone();
                let registry = registry.clone();
                let cancel = generation_cancel.clone();
                let failure = terminal_failure.clone();
                let lease = publication_lease(actor.identity());
                let actor = actor.clone();
                let cleanup_actor = actor.clone();
                let config = config.clone();
                tokio::spawn(async move {
                    let initialization_permit = match initialization_permit {
                        Some(permit) => permit,
                        None => {
                            let Some(registry) = registry.upgrade() else {
                                slot.state.lock().phase = PoolHubState::Retired;
                                slot.changed.notify_waiters();
                                return;
                            };
                            match registry
                                .reclaim_publication_hub_permit(&slot, &cancel)
                                .await
                            {
                                Ok(permit) => permit,
                                Err(PublicationHubError::InitializedHubLimit { .. }) => {
                                    let mut state = slot.state.lock();
                                    if matches!(state.phase, PoolHubState::Initializing) {
                                        state.phase = if state.retirement_requested
                                            || cancel.is_cancelled()
                                        {
                                            PoolHubState::Retired
                                        } else {
                                            PoolHubState::Vacant
                                        };
                                    }
                                    drop(state);
                                    slot.changed.notify_waiters();
                                    return;
                                }
                                Err(_error) if cancel.is_cancelled() => {
                                    slot.state.lock().phase = PoolHubState::Retired;
                                    slot.changed.notify_waiters();
                                    return;
                                }
                                Err(error) => {
                                    let reason = format!(
                                        "failed to reserve publication hub capacity: {error}"
                                    );
                                    slot.state.lock().phase = PoolHubState::Failed(reason.clone());
                                    failure(reason);
                                    slot.changed.notify_waiters();
                                    return;
                                }
                            }
                        }
                    };
                    let mut start =
                        tokio::spawn(PublicationHub::start(actor, lease, config, failure.clone()));
                    let result = tokio::select! {
                        result = &mut start => match result {
                            Ok(Ok(hub)) => Ok(hub),
                            Ok(Err(error)) => Err(format!(
                                "failed to initialize publication hub: {error}"
                            )),
                            Err(error) => Err(format!(
                                "publication hub initialization task failed: {error}"
                            )),
                        },
                        _ = cancel.cancelled() => {
                            start.abort();
                            if let Ok(Ok(hub)) = start.await {
                                hub.shutdown().await;
                            }
                            if let Err(error) = cleanup_actor.retire_publication_lease(lease).await {
                                tracing::debug!(
                                    pool_id = %cleanup_actor.identity().pool_id(),
                                    %error,
                                    "KV DC Relay actor stopped before its cancelled publication lease retired"
                                );
                            }
                            slot.state.lock().phase = PoolHubState::Retired;
                            slot.changed.notify_waiters();
                            return;
                        }
                    };
                    if cancel.is_cancelled() {
                        if let Ok(hub) = result {
                            hub.shutdown().await;
                        }
                        if let Err(error) = cleanup_actor.retire_publication_lease(lease).await {
                            tracing::debug!(
                                pool_id = %cleanup_actor.identity().pool_id(),
                                %error,
                                "KV DC Relay actor stopped before its cancelled publication lease retired"
                            );
                        }
                        slot.state.lock().phase = PoolHubState::Retired;
                        slot.changed.notify_waiters();
                        return;
                    }
                    match result {
                        Ok(hub) => {
                            let hub = {
                                let mut state = slot.state.lock();
                                if state.retirement_requested || cancel.is_cancelled() {
                                    state.phase = PoolHubState::Retired;
                                    Some(hub)
                                } else {
                                    state.phase = PoolHubState::Ready(InitializedPublicationHub {
                                        hub,
                                        permit: initialization_permit,
                                    });
                                    None
                                }
                            };
                            if let Some(hub) = hub {
                                hub.shutdown().await;
                                if let Err(error) =
                                    cleanup_actor.retire_publication_lease(lease).await
                                {
                                    tracing::debug!(
                                        pool_id = %cleanup_actor.identity().pool_id(),
                                        %error,
                                        "KV DC Relay actor stopped before its cancelled publication lease retired"
                                    );
                                }
                            }
                        }
                        Err(reason) => {
                            let should_fail = {
                                let mut state = slot.state.lock();
                                if state.retirement_requested || cancel.is_cancelled() {
                                    state.phase = PoolHubState::Retired;
                                    false
                                } else {
                                    state.phase = PoolHubState::Failed(reason.clone());
                                    true
                                }
                            };
                            if should_fail {
                                failure(reason);
                            }
                        }
                    }
                    slot.changed.notify_waiters();
                });
            }
            tokio::select! {
                _ = &mut changed => {}
                _ = generation_cancel.cancelled() => {
                    return Err(PublicationHubError::Unavailable(
                        "pool generation retired".to_string(),
                    ));
                }
            }
        }
    }

    async fn shutdown(&self) {
        loop {
            let changed = self.changed.notified();
            tokio::pin!(changed);
            changed.as_mut().enable();
            let hub = {
                let mut state = self.state.lock();
                state.retirement_requested = true;
                match std::mem::replace(&mut state.phase, PoolHubState::Retired) {
                    PoolHubState::Ready(initialized) => Some(initialized),
                    PoolHubState::Initializing => {
                        state.phase = PoolHubState::Initializing;
                        None
                    }
                    PoolHubState::Evicting => {
                        state.phase = PoolHubState::Evicting;
                        None
                    }
                    PoolHubState::Vacant | PoolHubState::Failed(_) | PoolHubState::Retired => {
                        return;
                    }
                }
            };
            if let Some(initialized) = hub {
                initialized.hub.shutdown().await;
                self.changed.notify_waiters();
                return;
            }
            changed.await;
        }
    }

    #[cfg(test)]
    fn phase(&self) -> &'static str {
        match &self.state.lock().phase {
            PoolHubState::Vacant => "vacant",
            PoolHubState::Initializing => "initializing",
            PoolHubState::Ready(_) => "ready",
            PoolHubState::Evicting => "evicting",
            PoolHubState::Failed(_) => "failed",
            PoolHubState::Retired => "retired",
        }
    }

    fn retire(&self) {
        let hub = {
            let mut state = self.state.lock();
            state.retirement_requested = true;
            match &state.phase {
                PoolHubState::Ready(initialized) => Some(initialized.hub.clone()),
                PoolHubState::Vacant | PoolHubState::Failed(_) => {
                    state.phase = PoolHubState::Retired;
                    None
                }
                PoolHubState::Initializing | PoolHubState::Evicting | PoolHubState::Retired => None,
            }
        };
        if let Some(hub) = hub {
            hub.retire();
        }
    }

    fn try_claim_idle(
        self: &Arc<Self>,
        identity: ProducerIdentity,
        generation_cancel: CancellationToken,
        registry: Weak<PoolRegistry>,
    ) -> Option<IdlePublicationHub> {
        if generation_cancel.is_cancelled() {
            return None;
        }
        let initialized = {
            let mut state = self.state.lock();
            if state.retirement_requested || state.pending_admissions != 0 {
                return None;
            }
            let PoolHubState::Ready(initialized) = &state.phase else {
                return None;
            };
            if !initialized.hub.try_begin_idle_eviction() {
                return None;
            }
            match std::mem::replace(&mut state.phase, PoolHubState::Evicting) {
                PoolHubState::Ready(initialized) => initialized,
                phase => {
                    state.phase = phase;
                    return None;
                }
            }
        };
        Some(IdlePublicationHub {
            guard: PoolHubEvictionGuard {
                slot: self.clone(),
                identity,
                generation_cancel,
                registry,
                is_finished: false,
            },
            initialized,
        })
    }

    fn finish_eviction(&self, is_retired: bool) {
        let mut state = self.state.lock();
        if matches!(state.phase, PoolHubState::Evicting) {
            state.phase = if is_retired || state.retirement_requested {
                PoolHubState::Retired
            } else {
                PoolHubState::Vacant
            };
        }
        drop(state);
        self.changed.notify_waiters();
    }
}

struct PoolHubAdmission {
    slot: Arc<PoolHubSlot>,
}

impl Drop for PoolHubAdmission {
    fn drop(&mut self) {
        let mut state = self.slot.state.lock();
        debug_assert_ne!(state.pending_admissions, 0);
        state.pending_admissions -= 1;
    }
}

struct IdlePublicationHub {
    guard: PoolHubEvictionGuard,
    initialized: InitializedPublicationHub,
}

impl IdlePublicationHub {
    async fn evict(self) -> (OwnedSemaphorePermit, Result<(), String>) {
        let Self {
            mut guard,
            initialized,
        } = self;
        let InitializedPublicationHub { hub, permit } = initialized;
        let result = hub.evict_idle().await;
        drop(hub);
        if let Err(reason) = &result {
            guard.fence_if_active(reason);
        }
        guard.finish();
        (permit, result)
    }
}

struct PoolHubEvictionGuard {
    slot: Arc<PoolHubSlot>,
    identity: ProducerIdentity,
    generation_cancel: CancellationToken,
    registry: Weak<PoolRegistry>,
    is_finished: bool,
}

struct PoolHubEvictionCandidate {
    identity: ProducerIdentity,
    generation_cancel: CancellationToken,
    slot: Arc<PoolHubSlot>,
}

impl PoolHubEvictionGuard {
    fn fence_if_active(&self, reason: &str) {
        if self.generation_cancel.is_cancelled() {
            return;
        }
        let Some(registry) = self.registry.upgrade() else {
            return;
        };
        let reason = format!("idle publication hub eviction failed: {reason}");
        registry.fence_generation(
            self.identity.pool_id(),
            self.identity.layout_generation(),
            &reason,
        );
    }

    fn finish(&mut self) {
        self.slot
            .finish_eviction(self.generation_cancel.is_cancelled());
        self.is_finished = true;
    }
}

impl Drop for PoolHubEvictionGuard {
    fn drop(&mut self) {
        if self.is_finished {
            return;
        }
        self.fence_if_active("cleanup task was abandoned");
        self.slot.finish_eviction(true);
    }
}

struct PoolReservation {
    endpoint: EndpointId,
    layout_generation: u64,
}

struct PoolReservationGuard<'a> {
    state: &'a Mutex<PoolRegistryState>,
    pool_id: PoolId,
    layout_generation: u64,
    is_armed: bool,
}

impl<'a> PoolReservationGuard<'a> {
    fn new(state: &'a Mutex<PoolRegistryState>, pool_id: PoolId, layout_generation: u64) -> Self {
        Self {
            state,
            pool_id,
            layout_generation,
            is_armed: true,
        }
    }

    fn disarm(&mut self) {
        self.is_armed = false;
    }
}

impl Drop for PoolReservationGuard<'_> {
    fn drop(&mut self) {
        if self.is_armed {
            rollback_reservation(&mut self.state.lock(), self.pool_id, self.layout_generation);
        }
    }
}

struct PoolRegistryState {
    pools: HashMap<PoolId, PoolEntry>,
    reservations: HashMap<PoolId, PoolReservation>,
    next_layout_generation: u64,
    catalog_revision: u64,
    accepting: bool,
}

impl Default for PoolRegistryState {
    fn default() -> Self {
        Self {
            pools: HashMap::new(),
            reservations: HashMap::new(),
            next_layout_generation: 1,
            catalog_revision: 0,
            accepting: true,
        }
    }
}

pub(super) struct PoolAttachment {
    pub(super) pool_id: PoolId,
    pub(super) layout_generation: u64,
    pub(super) handle: KvDcRelayHandle,
    registrations: Arc<[CanonicalModelRegistration]>,
    roles: Arc<[WorkerRole]>,
    pub(super) faults: mpsc::Receiver<ActorFault>,
    pub(super) pool_cancel: CancellationToken,
}

pub(super) struct PoolRegistry {
    relay_identity: DcRelayIdentity,
    actor_config: PoolActorConfig,
    ckf_allocation_permits: Arc<Semaphore>,
    state: Mutex<PoolRegistryState>,
    catalog_tx: watch::Sender<DcPoolCatalog>,
    load_tx: watch::Sender<Vec<PoolLoadSnapshot>>,
    publication_config: PublicationHubConfig,
    publication_hub_permits: Arc<Semaphore>,
    publication_hub_eviction_permits: Arc<Semaphore>,
    max_initialized_pool_hubs: usize,
    #[cfg(test)]
    publication_eviction_gate: Option<Arc<Semaphore>>,
}

impl PoolRegistry {
    #[cfg(test)]
    pub(super) fn new(relay_identity: DcRelayIdentity, actor_config: PoolActorConfig) -> Self {
        Self::new_inner(
            relay_identity,
            actor_config,
            PoolPublicationConfig::default(),
        )
    }

    pub(super) fn new_with_publication_config(
        relay_identity: DcRelayIdentity,
        actor_config: PoolActorConfig,
        publication_config: PoolPublicationConfig,
    ) -> Self {
        Self::new_inner(relay_identity, actor_config, publication_config)
    }

    fn new_inner(
        relay_identity: DcRelayIdentity,
        actor_config: PoolActorConfig,
        publication_config: PoolPublicationConfig,
    ) -> Self {
        debug_assert_ne!(publication_config.max_initialized_pool_hubs, 0);
        let (catalog_tx, _) = watch::channel(DcPoolCatalog::new(relay_identity, 0, Vec::new()));
        let (load_tx, _) = watch::channel(Vec::new());
        Self {
            relay_identity,
            actor_config,
            ckf_allocation_permits: Arc::new(Semaphore::new(DEFAULT_CKF_ALLOCATION_CONCURRENCY)),
            state: Mutex::new(PoolRegistryState::default()),
            catalog_tx,
            load_tx,
            publication_config: publication_config.hub,
            publication_hub_permits: Arc::new(Semaphore::new(
                publication_config.max_initialized_pool_hubs,
            )),
            publication_hub_eviction_permits: Arc::new(Semaphore::new(1)),
            max_initialized_pool_hubs: publication_config.max_initialized_pool_hubs,
            #[cfg(test)]
            publication_eviction_gate: publication_config.eviction_gate,
        }
    }

    pub(super) async fn attach(
        &self,
        request: PoolAttachRequest,
    ) -> anyhow::Result<PoolAttachment> {
        self.attach_with_builder(request, DcCkfState::new).await
    }

    async fn attach_with_builder<Builder>(
        &self,
        request: PoolAttachRequest,
        builder: Builder,
    ) -> anyhow::Result<PoolAttachment>
    where
        Builder: FnOnce(CkfConfig) -> Result<DcCkfState, CkfBuildError> + Send + 'static,
    {
        anyhow::ensure!(
            !request.registrations.is_empty(),
            "pool {} requires at least one canonical model binding",
            request.pool_id
        );
        validate_registrations(&request.registrations)?;
        validate_roles(&request.roles)?;
        let serving = request.serving_facts.as_ref().map(|facts| {
            let load = match PoolLoadState::from_runtime_configs(&facts.runtime_configs) {
                Ok(load) => load,
                Err(error) => {
                    // Load telemetry is supplemental to CKF materialization. Attach
                    // the pool with an explicitly degraded snapshot rather than make
                    // malformed capacity metadata block KV evidence publication.
                    tracing::warn!(
                        pool_id = %request.pool_id,
                        endpoint = %request.endpoint,
                        %error,
                        "Ignoring invalid KV DC Relay pool load capacity during attach"
                    );
                    PoolLoadState::default()
                }
            };
            PoolServingState { load }
        });

        let layout_generation = {
            let mut state = self.state.lock();
            anyhow::ensure!(
                state.accepting,
                "KV DC Relay pool registry is shutting down"
            );
            if let Some(endpoint) = pool_owner(&state, request.pool_id) {
                anyhow::bail!(
                    "pool {} is already owned by endpoint {} and cannot also attach endpoint {}",
                    request.pool_id,
                    endpoint,
                    request.endpoint
                );
            }
            if let Some(pool_id) = endpoint_pool(&state, &request.endpoint) {
                anyhow::bail!(
                    "endpoint {} is already assigned to pool {} and cannot attach pool {}",
                    request.endpoint,
                    pool_id,
                    request.pool_id
                );
            }
            let layout_generation = allocate_layout_generation(&mut state)?;
            state.reservations.insert(
                request.pool_id,
                PoolReservation {
                    endpoint: request.endpoint.clone(),
                    layout_generation,
                },
            );
            layout_generation
        };
        let mut reservation =
            PoolReservationGuard::new(&self.state, request.pool_id, layout_generation);

        let mut config = CkfConfig::new(self.actor_config.expected_unique_blocks);
        config.publish_every_n_events = self.actor_config.publication_threshold;
        let permit = self
            .ckf_allocation_permits
            .clone()
            .acquire_owned()
            .await
            .map_err(|_| anyhow::anyhow!("KV DC Relay pool registry is shutting down"))?;
        let ckf_state = tokio::task::spawn_blocking(move || {
            let result = builder(config);
            drop(permit);
            result
        })
        .await
        .map_err(|error| anyhow::anyhow!("KV DC Relay CKF allocation task failed: {error}"))??;

        let registrations: Arc<[CanonicalModelRegistration]> = request.registrations.into();
        let roles: Arc<[WorkerRole]> = request.roles.into();
        let cancel = CancellationToken::new();

        let mut state = self.state.lock();
        let reservation_matches = state
            .reservations
            .get(&request.pool_id)
            .is_some_and(|reservation| reservation.layout_generation == layout_generation);
        anyhow::ensure!(
            state.accepting && reservation_matches,
            "pool {} generation {} reservation was retired before commit",
            request.pool_id,
            layout_generation
        );
        let (handle, faults) = KvDcRelayHandle::spawn_with_state_and_publication_delay(
            ckf_state,
            StreamScope {
                relay_incarnation: self.relay_identity.relay_incarnation(),
                layout_generation,
                pool_id: request.pool_id,
            },
            self.actor_config.publication_delay,
        );
        let identity = handle.identity();
        let descriptor = DcPoolDescriptor::new(
            identity,
            request.endpoint.clone(),
            registrations.clone(),
            request.query_semantics,
            roles.clone(),
        );
        state.reservations.remove(&request.pool_id);
        debug_assert!(!state.pools.contains_key(&request.pool_id));
        state.pools.insert(
            request.pool_id,
            PoolEntry {
                endpoint: request.endpoint.clone(),
                handle: handle.clone(),
                identity,
                layout_generation,
                registrations: registrations.clone(),
                query_semantics: request.query_semantics,
                roles: roles.clone(),
                cancel: cancel.clone(),
                state: PoolEntryState::Active,
                serving,
                hub: Arc::new(PoolHubSlot::new()),
            },
        );
        publish_catalog_upsert(&mut state, &self.catalog_tx, descriptor);
        {
            publish_load_if_changed(&state, &self.load_tx, request.pool_id);
        }
        reservation.disarm();

        Ok(PoolAttachment {
            pool_id: request.pool_id,
            layout_generation,
            handle,
            registrations,
            roles,
            faults,
            pool_cancel: cancel,
        })
    }

    pub(super) async fn detach(&self, attachment: PoolAttachment) -> Result<(), KvDcRelayError> {
        let PoolAttachment {
            pool_id,
            layout_generation,
            handle,
            mut faults,
            ..
        } = attachment;
        self.withdraw(pool_id, layout_generation, PoolRetirementMode::Graceful)
            .await;
        let result = drain_faults_while(pool_id, &mut faults, handle.shutdown()).await;
        self.remove(pool_id, layout_generation).await;
        result
    }

    pub(super) async fn replace_registrations(
        &self,
        attachment: &mut PoolAttachment,
        registrations: Vec<CanonicalModelRegistration>,
    ) -> anyhow::Result<()> {
        anyhow::ensure!(
            !registrations.is_empty(),
            "pool {} requires at least one canonical model binding",
            attachment.pool_id
        );
        if attachment.registrations.as_ref() == registrations.as_slice() {
            return Ok(());
        }

        validate_registrations(&registrations)?;
        let registrations: Arc<[CanonicalModelRegistration]> = registrations.into();
        let mut state = self.state.lock();
        let entry = state
            .pools
            .get_mut(&attachment.pool_id)
            .ok_or_else(|| anyhow::anyhow!("pool {} is not attached", attachment.pool_id))?;
        anyhow::ensure!(
            entry.layout_generation == attachment.layout_generation
                && entry.state == PoolEntryState::Active,
            "pool {} generation {} is no longer active",
            attachment.pool_id,
            attachment.layout_generation
        );
        entry.registrations = registrations.clone();
        let descriptor = DcPoolDescriptor::new(
            entry.identity,
            entry.endpoint.clone(),
            registrations.clone(),
            entry.query_semantics,
            entry.roles.clone(),
        );
        attachment.registrations = registrations;
        publish_catalog_upsert(&mut state, &self.catalog_tx, descriptor);
        Ok(())
    }

    pub(super) fn replace_roles(
        &self,
        attachment: &mut PoolAttachment,
        roles: Vec<WorkerRole>,
    ) -> anyhow::Result<()> {
        validate_roles(&roles)?;
        if attachment.roles.as_ref() == roles.as_slice() {
            return Ok(());
        }

        let roles: Arc<[WorkerRole]> = roles.into();
        let mut state = self.state.lock();
        let entry = state
            .pools
            .get_mut(&attachment.pool_id)
            .ok_or_else(|| anyhow::anyhow!("pool {} is not attached", attachment.pool_id))?;
        anyhow::ensure!(
            entry.layout_generation == attachment.layout_generation
                && entry.state == PoolEntryState::Active,
            "pool {} generation {} is no longer active",
            attachment.pool_id,
            attachment.layout_generation
        );
        entry.roles = roles.clone();
        let descriptor = DcPoolDescriptor::new(
            entry.identity,
            entry.endpoint.clone(),
            entry.registrations.clone(),
            entry.query_semantics,
            roles.clone(),
        );
        attachment.roles = roles;
        publish_catalog_upsert(&mut state, &self.catalog_tx, descriptor);
        Ok(())
    }

    pub(super) fn replace_load_capacity(
        &self,
        pool_id: PoolId,
        layout_generation: u64,
        runtime_configs: &HashMap<WorkerId, ModelRuntimeConfig>,
    ) -> anyhow::Result<bool> {
        let mut state = self.state.lock();
        let Some(entry) = state.pools.get_mut(&pool_id) else {
            return Ok(false);
        };
        if entry.layout_generation != layout_generation || entry.state != PoolEntryState::Active {
            return Ok(false);
        }
        let Some(serving) = entry.serving.as_mut() else {
            return Ok(false);
        };
        let capacity_update = serving.load.replace_capacity(runtime_configs);
        match capacity_update {
            Ok(changed) => {
                if changed {
                    publish_load_if_changed(&state, &self.load_tx, pool_id);
                }
                // The caller uses true to record that this active generation accepted
                // the full runtime config, including fields outside the load contract.
                Ok(true)
            }
            Err(error) => {
                // replace_capacity invalidates state before returning an error. Publish
                // that degraded state immediately so the previous complete snapshot
                // cannot remain authoritative while the caller retries metadata refresh.
                publish_load_if_changed(&state, &self.load_tx, pool_id);
                Err(error)
            }
        }
    }

    pub(super) fn observe_load(
        &self,
        pool_id: PoolId,
        layout_generation: u64,
        load: ActiveLoad,
    ) -> bool {
        let mut state = self.state.lock();
        let Some(entry) = state.pools.get_mut(&pool_id) else {
            return false;
        };
        if entry.layout_generation != layout_generation || entry.state != PoolEntryState::Active {
            return false;
        }
        let Some(serving) = entry.serving.as_mut() else {
            return false;
        };
        match serving.load.observe(load) {
            LoadObservationOutcome::UnknownRank => false,
            LoadObservationOutcome::IgnoredAdvisory => true,
            LoadObservationOutcome::Updated => {
                publish_load_if_changed(&state, &self.load_tx, pool_id);
                true
            }
        }
    }

    pub(super) fn clear_load_observations(&self, pool_id: PoolId, layout_generation: u64) -> bool {
        let mut state = self.state.lock();
        let Some(entry) = state.pools.get_mut(&pool_id) else {
            return false;
        };
        if entry.layout_generation != layout_generation || entry.state != PoolEntryState::Active {
            return false;
        }
        let Some(serving) = entry.serving.as_mut() else {
            return false;
        };
        if serving.load.clear_observations() {
            publish_load_if_changed(&state, &self.load_tx, pool_id);
        }
        true
    }

    pub(super) async fn withdraw(
        &self,
        pool_id: PoolId,
        layout_generation: u64,
        mode: PoolRetirementMode,
    ) -> bool {
        let Some((hub, _)) = self.withdraw_publication_visibility(pool_id, layout_generation, mode)
        else {
            return false;
        };
        hub.shutdown().await;
        true
    }

    fn withdraw_publication_visibility(
        &self,
        pool_id: PoolId,
        layout_generation: u64,
        mode: PoolRetirementMode,
    ) -> Option<(Arc<PoolHubSlot>, bool)> {
        let mut state = self.state.lock();
        let entry = state.pools.get_mut(&pool_id)?;
        if entry.layout_generation != layout_generation {
            return None;
        }
        let was_active = entry.state == PoolEntryState::Active;
        entry.state = match (entry.state, mode) {
            (PoolEntryState::Fenced, _) | (_, PoolRetirementMode::Fenced) => PoolEntryState::Fenced,
            _ => PoolEntryState::Withdrawn,
        };
        entry.cancel.cancel();
        let hub = entry.hub.clone();
        if was_active {
            publish_catalog_remove(&mut state, &self.catalog_tx, pool_id);
            publish_load_if_changed(&state, &self.load_tx, pool_id);
        }
        Some((hub, was_active))
    }

    pub(super) fn fence_generation(
        &self,
        pool_id: PoolId,
        layout_generation: u64,
        reason: &str,
    ) -> bool {
        let retired = self.withdraw_publication_visibility(
            pool_id,
            layout_generation,
            PoolRetirementMode::Fenced,
        );
        let Some((hub, was_active)) = retired else {
            return false;
        };
        hub.retire();
        if was_active {
            tracing::error!(%pool_id, layout_generation, %reason, "Fencing KV DC Relay pool after terminal publication failure");
        }
        was_active
    }

    pub(super) async fn remove(&self, pool_id: PoolId, layout_generation: u64) -> bool {
        let entry = {
            let mut state = self.state.lock();
            let Some(entry) = state.pools.get(&pool_id) else {
                return false;
            };
            if entry.layout_generation != layout_generation {
                return false;
            }
            let was_active = entry.state == PoolEntryState::Active;
            let Some(entry) = state.pools.remove(&pool_id) else {
                return false;
            };
            entry.cancel.cancel();
            if was_active {
                publish_catalog_remove(&mut state, &self.catalog_tx, pool_id);
            }
            {
                publish_load_if_changed(&state, &self.load_tx, pool_id);
            }
            entry
        };
        entry.hub.shutdown().await;
        true
    }

    async fn reclaim_publication_hub_permit(
        self: &Arc<Self>,
        target_slot: &Arc<PoolHubSlot>,
        target_cancel: &CancellationToken,
    ) -> Result<OwnedSemaphorePermit, PublicationHubError> {
        let pressure_permit = tokio::select! {
            biased;
            _ = target_cancel.cancelled() => {
                return Err(PublicationHubError::Unavailable(
                    "pool generation retired while waiting for publication hub capacity".to_string(),
                ));
            }
            permit = self.publication_hub_eviction_permits.clone().acquire_owned() => {
                permit.map_err(|_| PublicationHubError::Unavailable(
                    "publication hub eviction coordinator is shutting down".to_string(),
                ))?
            }
        };

        match self.publication_hub_permits.clone().try_acquire_owned() {
            Ok(permit) => return Ok(permit),
            Err(TryAcquireError::Closed) => {
                return Err(PublicationHubError::Unavailable(
                    "publication hub capacity is shutting down".to_string(),
                ));
            }
            Err(TryAcquireError::NoPermits) => {}
        }

        let candidates = {
            let state = self.state.lock();
            state
                .pools
                .values()
                .filter(|entry| {
                    entry.state == PoolEntryState::Active && !Arc::ptr_eq(&entry.hub, target_slot)
                })
                .map(|entry| PoolHubEvictionCandidate {
                    identity: entry.identity,
                    generation_cancel: entry.cancel.child_token(),
                    slot: entry.hub.clone(),
                })
                .collect::<Vec<_>>()
        };
        let registry = Arc::downgrade(self);
        for candidate in candidates {
            let Some(victim) = candidate.slot.try_claim_idle(
                candidate.identity,
                candidate.generation_cancel,
                registry.clone(),
            ) else {
                continue;
            };
            let (permit_tx, permit_rx) = oneshot::channel();
            #[cfg(test)]
            let eviction_gate = self.publication_eviction_gate.clone();
            // Cleanup outlives the requesting admission so cancellation cannot strand the victim
            // in Evicting or leak its permit.
            tokio::spawn(async move {
                let _pressure_permit = pressure_permit;
                #[cfg(test)]
                if let Some(gate) = eviction_gate {
                    let _gate_permit = gate.acquire().await;
                }
                let (permit, result) = victim.evict().await;
                if let Err(error) = result {
                    tracing::debug!(%error, "Idle KV DC Relay publication hub eviction failed");
                }
                let _ = permit_tx.send(permit);
            });
            return tokio::select! {
                biased;
                _ = target_cancel.cancelled() => Err(PublicationHubError::Unavailable(
                    "pool generation retired while reclaiming publication hub capacity".to_string(),
                )),
                permit = permit_rx => permit.map_err(|_| PublicationHubError::InitializedHubLimit {
                    limit: self.max_initialized_pool_hubs,
                }),
            };
        }

        Err(PublicationHubError::InitializedHubLimit {
            limit: self.max_initialized_pool_hubs,
        })
    }

    pub(super) fn validate_active_producer(
        &self,
        expected: ProducerIdentity,
    ) -> Result<CancellationToken, PublicationHubError> {
        let pool_id = expected.pool_id();
        let state = self.state.lock();
        let entry = state
            .pools
            .get(&pool_id)
            .filter(|entry| entry.state == PoolEntryState::Active)
            .ok_or(PublicationHubError::UnknownPool(pool_id))?;
        if entry.identity != expected {
            return Err(PublicationHubError::ProducerMismatch(pool_id));
        }
        // A child observes generation retirement without giving admission code authority to retire
        // the generation itself.
        Ok(entry.cancel.child_token())
    }

    pub(super) async fn subscribe_pool(
        self: &Arc<Self>,
        expected: ProducerIdentity,
    ) -> Result<PublicationHubSubscription, PublicationHubError> {
        let pool_id = expected.pool_id();
        let (identity, actor, generation_cancel, hub_slot) = {
            let state = self.state.lock();
            let entry = state
                .pools
                .get(&pool_id)
                .filter(|entry| entry.state == PoolEntryState::Active)
                .ok_or(PublicationHubError::UnknownPool(pool_id))?;
            (
                entry.identity,
                entry.handle.clone(),
                entry.cancel.clone(),
                entry.hub.clone(),
            )
        };
        if identity != expected {
            return Err(PublicationHubError::ProducerMismatch(pool_id));
        }
        let admission = hub_slot.begin_admission()?;
        let failure_registry = Arc::downgrade(self);
        let terminal_failure: TerminalFailure = Arc::new(move |reason| {
            let Some(registry) = failure_registry.upgrade() else {
                return;
            };
            let reason = format!("publication hub: {reason}");
            registry.fence_generation(pool_id, identity.layout_generation(), &reason);
        });
        let hub = hub_slot
            .get_or_start(self, actor, generation_cancel, terminal_failure)
            .await?;
        let (current_identity, same_hub) = {
            let state = self.state.lock();
            let entry = state
                .pools
                .get(&pool_id)
                .filter(|entry| entry.state == PoolEntryState::Active)
                .ok_or(PublicationHubError::UnknownPool(pool_id))?;
            (entry.identity, Arc::ptr_eq(&entry.hub, &hub_slot))
        };
        if current_identity != expected {
            return Err(PublicationHubError::ProducerMismatch(pool_id));
        }
        if current_identity != identity || !same_hub {
            return Err(PublicationHubError::Unavailable(
                "pool publication generation changed".to_string(),
            ));
        }
        let subscription = hub.subscribe();
        drop(admission);
        subscription
    }

    pub(super) fn catalog(&self) -> DcPoolCatalog {
        self.catalog_tx.borrow().clone()
    }

    pub(super) fn watch_catalog(&self) -> watch::Receiver<DcPoolCatalog> {
        self.catalog_tx.subscribe()
    }

    pub(super) fn load_snapshots(&self) -> Vec<PoolLoadSnapshot> {
        self.load_tx.borrow().clone()
    }

    pub(super) fn watch_load(&self) -> watch::Receiver<Vec<PoolLoadSnapshot>> {
        self.load_tx.subscribe()
    }

    pub(super) async fn shutdown(&self) {
        let entries = {
            let mut state = self.state.lock();
            state.accepting = false;
            self.ckf_allocation_permits.close();
            state.reservations.clear();
            let entries = state.pools.drain().collect::<Vec<_>>();
            publish_catalog_clear(&mut state, &self.catalog_tx);
            {
                self.load_tx.send_if_modified(|snapshots| {
                    if snapshots.is_empty() {
                        return false;
                    }
                    snapshots.clear();
                    true
                });
            }
            entries
        };
        for (pool_id, entry) in entries {
            entry.cancel.cancel();
            entry.hub.shutdown().await;
            if let Err(error) = entry.handle.fence().await {
                tracing::warn!(%pool_id, %error, "KV Relay pool actor failed to fence during registry shutdown");
            }
        }
    }

    #[cfg(test)]
    pub(super) async fn pool_count(&self) -> usize {
        self.state.lock().pools.len()
    }
}

pub(super) async fn drain_faults_while<T>(
    pool_id: PoolId,
    faults: &mut mpsc::Receiver<ActorFault>,
    future: impl Future<Output = T>,
) -> T {
    tokio::pin!(future);
    loop {
        tokio::select! {
            result = &mut future => return result,
            fault = faults.recv() => match fault {
                Some(fault) => tracing::debug!(
                    %pool_id,
                    worker_id = fault.worker_id,
                    dp_rank = fault.dp_rank,
                    category = ?fault.category,
                    error = %fault.message,
                    "Draining KV DC Relay actor fault while retiring its pool"
                ),
                None => return future.await,
            },
        }
    }
}

fn allocate_layout_generation(state: &mut PoolRegistryState) -> anyhow::Result<u64> {
    let generation = state.next_layout_generation;
    state.next_layout_generation = generation
        .checked_add(1)
        .ok_or_else(|| anyhow::anyhow!("KV DC Relay layout generation space exhausted"))?;
    Ok(generation)
}

fn pool_owner(state: &PoolRegistryState, pool_id: PoolId) -> Option<&EndpointId> {
    state
        .pools
        .get(&pool_id)
        .map(|entry| &entry.endpoint)
        .or_else(|| {
            state
                .reservations
                .get(&pool_id)
                .map(|reservation| &reservation.endpoint)
        })
}

fn endpoint_pool(state: &PoolRegistryState, endpoint: &EndpointId) -> Option<PoolId> {
    state
        .pools
        .iter()
        .find_map(|(&pool_id, entry)| (&entry.endpoint == endpoint).then_some(pool_id))
        .or_else(|| {
            state
                .reservations
                .iter()
                .find_map(|(&pool_id, reservation)| {
                    (&reservation.endpoint == endpoint).then_some(pool_id)
                })
        })
}

fn rollback_reservation(state: &mut PoolRegistryState, pool_id: PoolId, layout_generation: u64) {
    if state
        .reservations
        .get(&pool_id)
        .is_some_and(|reservation| reservation.layout_generation == layout_generation)
    {
        state.reservations.remove(&pool_id);
    }
}

fn advance_catalog_revision(state: &mut PoolRegistryState) -> u64 {
    state.catalog_revision = state.catalog_revision.saturating_add(1);
    state.catalog_revision
}

fn publish_catalog_upsert(
    state: &mut PoolRegistryState,
    sender: &watch::Sender<DcPoolCatalog>,
    descriptor: DcPoolDescriptor,
) {
    let revision = advance_catalog_revision(state);
    sender.send_modify(|catalog| catalog.upsert(revision, descriptor));
}

fn publish_catalog_remove(
    state: &mut PoolRegistryState,
    sender: &watch::Sender<DcPoolCatalog>,
    pool_id: PoolId,
) {
    let revision = advance_catalog_revision(state);
    sender.send_modify(|catalog| catalog.remove(revision, pool_id));
}

fn publish_catalog_clear(state: &mut PoolRegistryState, sender: &watch::Sender<DcPoolCatalog>) {
    let revision = advance_catalog_revision(state);
    sender.send_modify(|catalog| catalog.clear(revision));
}

fn publish_load_if_changed(
    state: &PoolRegistryState,
    sender: &watch::Sender<Vec<PoolLoadSnapshot>>,
    pool_id: PoolId,
) {
    let snapshot = state
        .pools
        .get(&pool_id)
        .filter(|entry| entry.state == PoolEntryState::Active)
        .and_then(|entry| {
            entry
                .serving
                .as_ref()
                .map(|serving| serving.load.snapshot(entry.identity))
        });
    sender.send_if_modified(|snapshots| {
        match (
            snapshots.binary_search_by_key(&pool_id, |item| item.producer.pool_id()),
            snapshot,
        ) {
            (Ok(index), Some(snapshot)) if snapshots[index] != snapshot => {
                snapshots[index] = snapshot;
                true
            }
            (Ok(index), None) => {
                snapshots.remove(index);
                true
            }
            (Err(index), Some(snapshot)) => {
                snapshots.insert(index, snapshot);
                true
            }
            _ => false,
        }
    });
}

fn validate_registrations(registrations: &[CanonicalModelRegistration]) -> anyhow::Result<()> {
    let mut request_models = HashMap::with_capacity(registrations.len());
    for registration in registrations {
        if let Some(previous) =
            request_models.insert(registration.model().clone(), registration.target().clone())
        {
            anyhow::ensure!(
                previous == *registration.target(),
                "canonical model {} resolves to conflicting targets in the same pool",
                registration.model()
            );
            anyhow::bail!(
                "duplicate canonical model registration {}",
                registration.model()
            );
        }
    }

    let mut request_aliases = HashMap::<ModelAlias, CanonicalModelId>::new();
    for registration in registrations {
        for alias in registration.aliases() {
            let alias_as_model = CanonicalModelId::new(alias.as_str().to_string())?;
            anyhow::ensure!(
                !request_models.contains_key(&alias_as_model),
                "model alias {} conflicts with canonical model {} in the same pool",
                alias,
                alias_as_model
            );
            if let Some(owner) = request_aliases.insert(alias.clone(), registration.model().clone())
            {
                anyhow::ensure!(
                    owner == *registration.model(),
                    "model alias {} is claimed by both {} and {} in the same pool",
                    alias,
                    owner,
                    registration.model()
                );
            }
        }
    }
    Ok(())
}

fn validate_roles(roles: &[WorkerRole]) -> anyhow::Result<()> {
    anyhow::ensure!(
        !roles.is_empty(),
        "pool requires at least one declared worker role"
    );
    let mut unique = std::collections::HashSet::with_capacity(roles.len());
    for role in roles {
        anyhow::ensure!(
            unique.insert(*role),
            "pool repeats declared worker role {role:?}"
        );
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, mpsc as std_mpsc};

    use dynamo_kv_router::identity::{
        CacheSemanticsId, DcId, IdentitySource, IndexerDomainId, RoutingScopeId,
    };
    use dynamo_kv_router::protocols::{
        ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheStoreData,
        KvCacheStoredBlockData, LocalBlockHash, RouterEvent,
    };
    use futures::StreamExt;
    use tokio_util::sync::CancellationToken;

    use super::*;
    use crate::kv_dc_relay::discovery::DcMembershipView;
    use crate::kv_dc_relay::identity::{KvQueryHashFormat, ModelTarget};
    use crate::kv_dc_relay::publication::{
        PublicationErrorKind, PublicationFrameKind, RegistryPublicationSource,
        RelayPublicationSource,
    };
    use crate::kv_dc_relay::topology::TopologyPublisher;

    type TestCkfBuilder =
        Box<dyn FnOnce(CkfConfig) -> Result<DcCkfState, CkfBuildError> + Send + 'static>;

    struct GatedBuilder {
        builder: TestCkfBuilder,
        started: tokio::sync::oneshot::Receiver<()>,
        release: std_mpsc::Sender<()>,
        finished: tokio::sync::oneshot::Receiver<()>,
    }

    fn pool(seed: u8) -> PoolId {
        PoolId::new(
            IndexerDomainId::new(
                CacheSemanticsId::new([seed; 16], IdentitySource::Explicit),
                RoutingScopeId::new([seed.wrapping_add(1); 16], IdentitySource::Explicit),
            ),
            DcId::new(3),
        )
    }

    fn config() -> PoolActorConfig {
        PoolActorConfig {
            expected_unique_blocks: 32,
            publication_threshold: 1,
            publication_delay: Duration::from_millis(1),
        }
    }

    fn relay_identity() -> DcRelayIdentity {
        DcRelayIdentity::new(11, 7)
    }

    fn registration(model: &str) -> CanonicalModelRegistration {
        CanonicalModelRegistration::new(
            CanonicalModelId::new(model).unwrap(),
            vec![ModelAlias::new(format!("{model}-alias")).unwrap()],
        )
    }

    fn query_semantics() -> KvQuerySemantics {
        KvQuerySemantics::new(64, KvQueryHashFormat::DynamoStandardV1).unwrap()
    }

    fn request(pool_id: PoolId, endpoint: &str, model: &str) -> PoolAttachRequest {
        PoolAttachRequest {
            pool_id,
            endpoint: EndpointId::from(endpoint),
            registrations: vec![registration(model)],
            query_semantics: query_semantics(),
            roles: vec![WorkerRole::Aggregated],
            serving_facts: Some(PoolServingFacts {
                runtime_configs: HashMap::new(),
            }),
        }
    }

    fn stored_event(event_id: u64) -> RouterEvent {
        RouterEvent::new(
            1,
            KvCacheEvent {
                event_id,
                data: KvCacheEventData::Stored(KvCacheStoreData {
                    parent_hash: None,
                    start_position: None,
                    blocks: vec![KvCacheStoredBlockData {
                        block_hash: ExternalSequenceBlockHash(event_id),
                        tokens_hash: LocalBlockHash(event_id),
                        mm_extra_info: None,
                    }],
                }),
                dp_rank: 0,
            },
        )
    }

    fn gated_builder() -> GatedBuilder {
        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
        let (release_tx, release_rx) = std_mpsc::channel();
        let (finished_tx, finished_rx) = tokio::sync::oneshot::channel();
        let builder = move |config| {
            let _ = started_tx.send(());
            release_rx
                .recv_timeout(Duration::from_secs(1))
                .expect("test must release the CKF builder");
            let result = DcCkfState::new(config);
            let _ = finished_tx.send(());
            result
        };
        GatedBuilder {
            builder: Box::new(builder),
            started: started_rx,
            release: release_tx,
            finished: finished_rx,
        }
    }

    fn descriptor(catalog: &DcPoolCatalog, pool_id: PoolId) -> &DcPoolDescriptor {
        catalog
            .pools()
            .iter()
            .find(|descriptor| descriptor.pool_id() == pool_id)
            .unwrap()
    }

    async fn retire(registry: &PoolRegistry, attachment: PoolAttachment, mode: PoolRetirementMode) {
        let PoolAttachment {
            pool_id,
            layout_generation,
            handle,
            mut faults,
            ..
        } = attachment;
        assert!(registry.withdraw(pool_id, layout_generation, mode).await);
        let result = match mode {
            PoolRetirementMode::Graceful => {
                drain_faults_while(pool_id, &mut faults, handle.shutdown()).await
            }
            PoolRetirementMode::Fenced => {
                drain_faults_while(pool_id, &mut faults, handle.fence()).await
            }
        };
        result.unwrap();
        assert!(registry.remove(pool_id, layout_generation).await);
    }

    fn hub_slot(registry: &PoolRegistry, pool_id: PoolId) -> Arc<PoolHubSlot> {
        registry
            .state
            .lock()
            .pools
            .get(&pool_id)
            .expect("pool must be attached")
            .hub
            .clone()
    }

    async fn subscribe(
        registry: &Arc<PoolRegistry>,
        producer: ProducerIdentity,
    ) -> Result<PublicationHubSubscription, PublicationHubError> {
        registry.subscribe_pool(producer).await
    }

    fn publication_source(
        registry: Arc<PoolRegistry>,
        lifecycle: CancellationToken,
        max_active_streams: usize,
    ) -> Arc<RegistryPublicationSource> {
        let topology = Arc::new(TopologyPublisher::new(
            DcMembershipView::default(),
            &registry.catalog(),
        ));
        Arc::new(RegistryPublicationSource::new(
            registry,
            topology,
            relay_identity(),
            lifecycle,
            Arc::new(Semaphore::new(2)),
            max_active_streams,
            Duration::from_secs(1),
        ))
    }

    #[tokio::test]
    async fn one_model_binds_to_independent_pools() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let first = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let second = registry
            .attach(request(pool(2), "slow.router.generate", "llama"))
            .await
            .unwrap();

        assert_eq!(registry.pool_count().await, 2);
        let catalog = registry.catalog();
        assert_eq!(catalog.drt_instance_id(), 11);
        assert_eq!(catalog.relay_incarnation(), 7);
        assert_eq!(catalog.pools().len(), 2);
        assert_eq!(
            descriptor(&catalog, pool(1)).serving_endpoint(),
            &EndpointId::from("fast.router.generate")
        );
        assert!(
            catalog
                .pools()
                .iter()
                .all(|descriptor| { descriptor.registrations()[0].model().as_str() == "llama" })
        );

        registry.detach(first).await.unwrap();
        registry.detach(second).await.unwrap();
    }

    #[tokio::test]
    async fn default_registry_materializes_publication_images_eagerly() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();

        attachment
            .handle
            .admit_event(
                0,
                RouterEvent::new(
                    1,
                    KvCacheEvent {
                        event_id: 1,
                        data: KvCacheEventData::Stored(KvCacheStoreData {
                            parent_hash: None,
                            start_position: None,
                            blocks: vec![KvCacheStoredBlockData {
                                block_hash: ExternalSequenceBlockHash(1),
                                tokens_hash: LocalBlockHash(1),
                                mm_extra_info: None,
                            }],
                        }),
                        dp_rank: 0,
                    },
                ),
            )
            .await
            .unwrap();
        attachment.handle.flush().await.unwrap();
        assert!(
            attachment
                .handle
                .state_stats()
                .await
                .unwrap()
                .0
                .publication()
                .emitted_images()
                > 0
        );

        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn burst_attach_defers_catalog_materialization_until_observed() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let catalog_rx = registry.watch_catalog();
        let mut attachments = Vec::new();

        for seed in 1..=32 {
            attachments.push(
                registry
                    .attach(request(
                        pool(seed),
                        &format!("pool-{seed}.router.generate"),
                        &format!("model-{seed}"),
                    ))
                    .await
                    .unwrap(),
            );
        }

        assert_eq!(catalog_rx.borrow().revision(), 32);
        assert!(!catalog_rx.borrow().is_materialized());

        let catalog = registry.catalog();
        let pool_ids: Vec<_> = catalog
            .pools()
            .iter()
            .map(DcPoolDescriptor::pool_id)
            .collect();
        assert!(catalog.is_materialized());
        assert!(pool_ids.windows(2).all(|pair| pair[0] < pair[1]));
        let serialized = serde_json::to_value(&catalog).unwrap();
        assert_eq!(serialized["drt_instance_id"], 11);
        assert_eq!(serialized["relay_incarnation"], 7);
        assert!(serialized.get("process_incarnation").is_none());
        assert_eq!(serialized["revision"], 32);
        assert_eq!(serialized["pools"].as_array().unwrap().len(), 32);

        registry.shutdown().await;
        assert_eq!(catalog.pools().len(), attachments.len());
        assert!(registry.catalog().pools().is_empty());
    }

    #[tokio::test]
    async fn catalog_publishes_and_preserves_generation_query_semantics() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let expected = KvQuerySemantics::new(128, KvQueryHashFormat::DynamoEagleV1).unwrap();
        let mut attach_request = request(pool(1), "fast.router.generate", "llama");
        attach_request.query_semantics = expected;
        let mut attachment = registry.attach(attach_request).await.unwrap();
        let producer = attachment.handle.identity();

        let catalog = registry.catalog();
        let initial_descriptor = descriptor(&catalog, pool(1));
        assert_eq!(initial_descriptor.producer(), producer);
        assert_eq!(initial_descriptor.query_semantics(), expected);

        registry
            .replace_registrations(
                &mut attachment,
                vec![CanonicalModelRegistration::new(
                    CanonicalModelId::new("llama").unwrap(),
                    vec![ModelAlias::new("chat").unwrap()],
                )],
            )
            .await
            .unwrap();
        let catalog = registry.catalog();
        let updated_descriptor = descriptor(&catalog, pool(1));
        assert_eq!(updated_descriptor.producer(), producer);
        assert_eq!(updated_descriptor.query_semantics(), expected);

        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn registration_updates_remain_pool_local() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let mut first = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let mut second = registry
            .attach(request(pool(2), "slow.router.generate", "mistral"))
            .await
            .unwrap();

        registry
            .replace_registrations(
                &mut first,
                vec![CanonicalModelRegistration::new(
                    CanonicalModelId::new("llama").unwrap(),
                    vec![ModelAlias::new("mistral-alias").unwrap()],
                )],
            )
            .await
            .unwrap();
        let catalog = registry.catalog();
        assert!(catalog.pools().iter().all(|descriptor| {
            descriptor.registrations()[0].aliases() == [ModelAlias::new("mistral-alias").unwrap()]
        }));

        registry
            .replace_registrations(
                &mut second,
                vec![CanonicalModelRegistration::new(
                    CanonicalModelId::new("mistral").unwrap(),
                    vec![ModelAlias::new("llama-alias").unwrap()],
                )],
            )
            .await
            .unwrap();
        let catalog = registry.catalog();
        assert_eq!(
            descriptor(&catalog, pool(1)).registrations()[0].aliases(),
            [ModelAlias::new("mistral-alias").unwrap()]
        );
        assert_eq!(
            descriptor(&catalog, pool(2)).registrations()[0].aliases(),
            [ModelAlias::new("llama-alias").unwrap()]
        );

        registry.detach(first).await.unwrap();
        registry.detach(second).await.unwrap();
    }

    #[tokio::test]
    async fn one_pool_cannot_be_owned_by_two_endpoints() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let pool_id = pool(1);
        let attachment = registry
            .attach(request(pool_id, "first.router.generate", "llama"))
            .await
            .unwrap();

        let error = registry
            .attach(request(pool_id, "second.router.generate", "llama"))
            .await
            .err()
            .unwrap();
        assert!(error.to_string().contains("already owned"));

        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn endpoint_reassignment_waits_for_prior_pool_removal() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let endpoint = "fast.router.generate";
        let first = registry
            .attach(request(pool(1), endpoint, "llama"))
            .await
            .unwrap();

        assert!(
            registry
                .withdraw(
                    first.pool_id,
                    first.layout_generation,
                    PoolRetirementMode::Graceful,
                )
                .await
        );
        assert!(registry.catalog().pools().is_empty());

        let error = registry
            .attach(request(pool(2), endpoint, "llama"))
            .await
            .err()
            .unwrap();
        assert!(error.to_string().contains("already assigned"));
        assert!(registry.catalog().pools().is_empty());

        registry.detach(first).await.unwrap();
        let second = registry
            .attach(request(pool(2), endpoint, "llama"))
            .await
            .unwrap();
        let catalog = registry.catalog();
        assert_eq!(catalog.pools().len(), 1);
        assert_eq!(catalog.pools()[0].pool_id(), pool(2));
        assert_eq!(
            catalog.pools()[0].serving_endpoint(),
            &EndpointId::from(endpoint)
        );

        registry.detach(second).await.unwrap();
    }

    #[tokio::test]
    async fn cancelled_allocation_rolls_back_and_allows_reattach() {
        let registry = Arc::new(PoolRegistry::new(relay_identity(), config()));
        let pool_id = pool(1);
        let GatedBuilder {
            builder,
            started: started_rx,
            release: release_tx,
            finished: finished_rx,
        } = gated_builder();
        let task_registry = registry.clone();
        let attach = tokio::spawn(async move {
            task_registry
                .attach_with_builder(request(pool_id, "first.router.generate", "llama"), builder)
                .await
        });

        started_rx.await.unwrap();
        attach.abort();
        assert!(matches!(attach.await, Err(error) if error.is_cancelled()));
        assert!(registry.state.lock().reservations.is_empty());

        release_tx.send(()).unwrap();
        finished_rx.await.unwrap();
        let attachment = registry
            .attach(request(pool_id, "second.router.generate", "llama"))
            .await
            .unwrap();
        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn shutdown_during_allocation_never_publishes_the_pool() {
        let registry = Arc::new(PoolRegistry::new(relay_identity(), config()));
        let GatedBuilder {
            builder,
            started: started_rx,
            release: release_tx,
            finished: finished_rx,
        } = gated_builder();
        let task_registry = registry.clone();
        let attach = tokio::spawn(async move {
            task_registry
                .attach_with_builder(request(pool(1), "first.router.generate", "llama"), builder)
                .await
        });

        started_rx.await.unwrap();
        registry.shutdown().await;
        assert!(registry.state.lock().reservations.is_empty());
        assert!(registry.catalog().pools().is_empty());

        release_tx.send(()).unwrap();
        finished_rx.await.unwrap();
        let Err(error) = attach.await.unwrap() else {
            panic!("pool attached after registry shutdown");
        };
        assert!(error.to_string().contains("retired before commit"));
        assert_eq!(registry.pool_count().await, 0);
        assert!(registry.catalog().pools().is_empty());
    }

    #[tokio::test(flavor = "current_thread")]
    async fn ckf_allocation_does_not_block_the_async_executor() {
        let registry = Arc::new(PoolRegistry::new(relay_identity(), config()));
        let GatedBuilder {
            builder,
            started: started_rx,
            release: release_tx,
            finished: finished_rx,
        } = gated_builder();
        let task_registry = registry.clone();
        let attach = tokio::spawn(async move {
            task_registry
                .attach_with_builder(request(pool(1), "first.router.generate", "llama"), builder)
                .await
        });

        tokio::time::timeout(Duration::from_millis(500), started_rx)
            .await
            .expect("blocking CKF allocation did not start")
            .unwrap();
        tokio::time::timeout(
            Duration::from_millis(500),
            tokio::time::sleep(Duration::from_millis(1)),
        )
        .await
        .expect("CKF allocation blocked the async executor");

        release_tx.send(()).unwrap();
        finished_rx.await.unwrap();
        let attachment = attach.await.unwrap().unwrap();
        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn failed_actor_build_rolls_back_the_pool_reservation() {
        let registry = PoolRegistry::new(
            relay_identity(),
            PoolActorConfig {
                expected_unique_blocks: 0,
                publication_threshold: 1,
                publication_delay: Duration::from_millis(1),
            },
        );
        let pool_id = pool(1);

        let first_error = registry
            .attach(request(pool_id, "first.router.generate", "llama"))
            .await
            .err()
            .unwrap();
        let second_error = registry
            .attach(request(pool_id, "second.router.generate", "llama"))
            .await
            .err()
            .unwrap();

        assert!(first_error.to_string().contains("greater than zero"));
        assert!(second_error.to_string().contains("greater than zero"));
        assert_eq!(registry.pool_count().await, 0);
        assert_eq!(registry.catalog().revision(), 0);
    }

    #[tokio::test]
    async fn reattaching_a_pool_allocates_a_new_layout_generation() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let pool_id = pool(1);
        let first = registry
            .attach(request(pool_id, "fast.router.generate", "llama"))
            .await
            .unwrap();
        let first_generation = first.layout_generation;
        registry.detach(first).await.unwrap();

        let replacement = registry
            .attach(request(pool_id, "fast.router.generate", "llama"))
            .await
            .unwrap();
        assert!(replacement.layout_generation > first_generation);

        registry.detach(replacement).await.unwrap();
    }

    #[tokio::test]
    async fn relay_incarnation_fences_an_identical_pool_layout() {
        let first_registry = PoolRegistry::new(DcRelayIdentity::new(11, 7), config());
        let second_registry = PoolRegistry::new(DcRelayIdentity::new(11, 8), config());
        let pool_id = pool(1);
        let first = first_registry
            .attach(request(pool_id, "fast.router.generate", "llama"))
            .await
            .unwrap();
        let second = second_registry
            .attach(request(pool_id, "fast.router.generate", "llama"))
            .await
            .unwrap();

        assert_ne!(first.handle.identity(), second.handle.identity());
        assert_eq!(first.layout_generation, second.layout_generation);

        first_registry.detach(first).await.unwrap();
        second_registry.detach(second).await.unwrap();
    }

    #[tokio::test]
    async fn withdraw_removes_catalog_before_actor_retirement() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();

        assert!(
            registry
                .withdraw(
                    attachment.pool_id,
                    attachment.layout_generation,
                    PoolRetirementMode::Graceful,
                )
                .await
        );
        assert!(registry.catalog().pools().is_empty());
        attachment.handle.state_stats().await.unwrap();

        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn adapter_registration_changes_without_replacing_pool_generation() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let mut attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let generation = attachment.layout_generation;
        let base = CanonicalModelId::new("llama").unwrap();
        let adapter = CanonicalModelId::new("tenant-a").unwrap();
        registry
            .replace_registrations(
                &mut attachment,
                vec![
                    CanonicalModelRegistration::new(base.clone(), Vec::new()),
                    CanonicalModelRegistration::with_target(
                        adapter.clone(),
                        ModelTarget::Lora {
                            base_model: base,
                            adapter: adapter.clone(),
                        },
                        Vec::new(),
                    ),
                ],
            )
            .await
            .unwrap();

        assert_eq!(attachment.layout_generation, generation);
        let catalog = registry.catalog();
        assert!(
            descriptor(&catalog, pool(1))
                .registrations()
                .iter()
                .any(|registration| registration.model() == &adapter)
        );
        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn declared_role_change_updates_catalog_without_replacing_generation() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let mut attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let producer = attachment.handle.identity();

        registry
            .replace_roles(&mut attachment, vec![WorkerRole::Decode])
            .unwrap();
        let catalog = registry.catalog();
        assert_eq!(descriptor(&catalog, pool(1)).producer(), producer);
        assert_eq!(
            descriptor(&catalog, pool(1)).pool_roles(),
            [WorkerRole::Decode]
        );

        assert!(registry.replace_roles(&mut attachment, Vec::new()).is_err());
        assert_eq!(
            descriptor(&registry.catalog(), pool(1)).pool_roles(),
            [WorkerRole::Decode]
        );
        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn one_lora_target_binds_to_independent_pools() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let base = CanonicalModelId::new("llama").unwrap();
        let adapter = CanonicalModelId::new("tenant-a").unwrap();
        let registrations = || {
            vec![
                CanonicalModelRegistration::new(base.clone(), Vec::new()),
                CanonicalModelRegistration::with_target(
                    adapter.clone(),
                    ModelTarget::Lora {
                        base_model: base.clone(),
                        adapter: adapter.clone(),
                    },
                    Vec::new(),
                ),
            ]
        };
        let first = registry
            .attach(PoolAttachRequest {
                pool_id: pool(1),
                endpoint: EndpointId::from("fast.router.generate"),
                registrations: registrations(),
                query_semantics: query_semantics(),
                roles: vec![WorkerRole::Aggregated],
                serving_facts: Some(PoolServingFacts {
                    runtime_configs: HashMap::new(),
                }),
            })
            .await
            .unwrap();
        let second = registry
            .attach(PoolAttachRequest {
                pool_id: pool(2),
                endpoint: EndpointId::from("slow.router.generate"),
                registrations: registrations(),
                query_semantics: query_semantics(),
                roles: vec![WorkerRole::Aggregated],
                serving_facts: Some(PoolServingFacts {
                    runtime_configs: HashMap::new(),
                }),
            })
            .await
            .unwrap();

        let catalog = registry.catalog();
        assert_eq!(catalog.pools().len(), 2);
        assert!(catalog.pools().iter().all(|descriptor| {
            descriptor.registrations().iter().any(|registration| {
                registration.target()
                    == &ModelTarget::Lora {
                        base_model: base.clone(),
                        adapter: adapter.clone(),
                    }
            })
        }));

        registry.detach(first).await.unwrap();
        registry.detach(second).await.unwrap();
    }

    #[tokio::test]
    async fn fencing_withdraws_pool_from_catalog() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        retire(&registry, attachment, PoolRetirementMode::Fenced).await;
        assert!(registry.watch_catalog().borrow().pools().is_empty());
    }

    #[tokio::test]
    async fn fencing_withdraws_only_the_target_pool_descriptor() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let with_alias = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let without_alias = registry
            .attach(PoolAttachRequest {
                pool_id: pool(2),
                endpoint: EndpointId::from("slow.router.generate"),
                registrations: vec![CanonicalModelRegistration::new(
                    CanonicalModelId::new("llama").unwrap(),
                    Vec::new(),
                )],
                query_semantics: query_semantics(),
                roles: vec![WorkerRole::Aggregated],
                serving_facts: Some(PoolServingFacts {
                    runtime_configs: HashMap::new(),
                }),
            })
            .await
            .unwrap();

        let catalog = registry.catalog();
        assert_eq!(
            descriptor(&catalog, pool(1)).registrations()[0]
                .aliases()
                .len(),
            1
        );
        assert!(
            descriptor(&catalog, pool(2)).registrations()[0]
                .aliases()
                .is_empty()
        );
        retire(&registry, with_alias, PoolRetirementMode::Fenced).await;
        let catalog = registry.catalog();
        assert_eq!(catalog.pools().len(), 1);
        assert_eq!(catalog.pools()[0].pool_id(), pool(2));
        assert!(catalog.pools()[0].registrations()[0].aliases().is_empty());

        registry.detach(without_alias).await.unwrap();
    }

    #[tokio::test]
    async fn publication_hub_is_lazy_and_shared_by_concurrent_subscribers() {
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig::default(),
        ));
        let attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let producer = attachment.handle.identity();
        let hub = hub_slot(&registry, attachment.pool_id);
        assert_eq!(hub.phase(), "vacant");

        let mut tasks = Vec::new();
        for _ in 0..8 {
            let registry = registry.clone();
            tasks.push(tokio::spawn(
                async move { subscribe(&registry, producer).await },
            ));
        }
        let mut subscriptions = Vec::new();
        for task in tasks {
            let subscription = task.await.unwrap().unwrap();
            assert_eq!(subscription.snapshot().unwrap().identity(), producer);
            subscriptions.push(subscription);
        }
        assert_eq!(hub.phase(), "ready");

        drop(subscriptions);
        registry.detach(attachment).await.unwrap();
        assert_eq!(hub.phase(), "retired");
    }

    #[tokio::test]
    async fn initialized_hub_limit_protects_active_hubs_and_reclaims_idle_hubs() {
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig {
                max_initialized_pool_hubs: 1,
                ..PoolPublicationConfig::default()
            },
        ));
        let first = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let second = registry
            .attach(request(pool(2), "slow.router.generate", "llama"))
            .await
            .unwrap();
        let first_hub = hub_slot(&registry, first.pool_id);
        let second_hub = hub_slot(&registry, second.pool_id);

        let first_subscription = subscribe(&registry, first.handle.identity()).await.unwrap();
        let error = subscribe(&registry, second.handle.identity())
            .await
            .err()
            .expect("second pool must hit the global initialized-hub limit");
        assert_eq!(error, PublicationHubError::InitializedHubLimit { limit: 1 });

        drop(first_subscription);
        let second_subscription = subscribe(&registry, second.handle.identity())
            .await
            .unwrap();
        assert_eq!(first_hub.phase(), "vacant");
        assert_eq!(second_hub.phase(), "ready");
        let catalog = registry.catalog();
        assert_eq!(catalog.pools().len(), 2);
        assert!(
            catalog
                .pools()
                .iter()
                .any(|descriptor| descriptor.pool_id() == first.pool_id)
        );
        assert!(
            catalog
                .pools()
                .iter()
                .any(|descriptor| descriptor.pool_id() == second.pool_id)
        );

        drop(second_subscription);
        let reopened_first = subscribe(&registry, first.handle.identity()).await.unwrap();
        assert_eq!(first_hub.phase(), "ready");
        assert_eq!(second_hub.phase(), "vacant");

        drop(reopened_first);
        registry.detach(first).await.unwrap();
        registry.detach(second).await.unwrap();
    }

    #[tokio::test]
    async fn pending_admission_prevents_idle_hub_eviction() {
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig {
                max_initialized_pool_hubs: 1,
                ..PoolPublicationConfig::default()
            },
        ));
        let first = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let second = registry
            .attach(request(pool(2), "slow.router.generate", "llama"))
            .await
            .unwrap();
        let first_hub = hub_slot(&registry, first.pool_id);

        let first_subscription = subscribe(&registry, first.handle.identity()).await.unwrap();
        drop(first_subscription);
        let pending_admission = first_hub.begin_admission().unwrap();

        let error = subscribe(&registry, second.handle.identity())
            .await
            .err()
            .expect("pending admission must protect the idle hub");
        assert_eq!(error, PublicationHubError::InitializedHubLimit { limit: 1 });
        assert_eq!(first_hub.phase(), "ready");

        drop(pending_admission);
        let second_subscription = subscribe(&registry, second.handle.identity())
            .await
            .unwrap();
        assert_eq!(first_hub.phase(), "vacant");

        drop(second_subscription);
        registry.detach(first).await.unwrap();
        registry.detach(second).await.unwrap();
    }

    #[tokio::test]
    async fn concurrent_pressure_claims_one_idle_victim() {
        let eviction_gate = Arc::new(Semaphore::new(0));
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig {
                max_initialized_pool_hubs: 1,
                eviction_gate: Some(eviction_gate.clone()),
                ..PoolPublicationConfig::default()
            },
        ));
        let first = registry
            .attach(request(pool(1), "first.router.generate", "llama"))
            .await
            .unwrap();
        let second = registry
            .attach(request(pool(2), "second.router.generate", "llama"))
            .await
            .unwrap();
        let third = registry
            .attach(request(pool(3), "third.router.generate", "llama"))
            .await
            .unwrap();
        let first_hub = hub_slot(&registry, first.pool_id);

        let subscription = subscribe(&registry, first.handle.identity()).await.unwrap();
        drop(subscription);
        let second_subscribe = tokio::spawn({
            let registry = registry.clone();
            let producer = second.handle.identity();
            async move { subscribe(&registry, producer).await }
        });
        tokio::time::timeout(Duration::from_secs(1), async {
            while first_hub.phase() != "evicting" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("first pressure request must claim the idle hub");

        let third_subscribe = tokio::spawn({
            let registry = registry.clone();
            let producer = third.handle.identity();
            async move { subscribe(&registry, producer).await }
        });
        eviction_gate.add_permits(1);

        let second_subscription = second_subscribe.await.unwrap().unwrap();
        let third_error = third_subscribe
            .await
            .unwrap()
            .err()
            .expect("one idle victim cannot satisfy two pressure requests");
        assert_eq!(
            third_error,
            PublicationHubError::InitializedHubLimit { limit: 1 }
        );
        assert_eq!(first_hub.phase(), "vacant");
        assert_eq!(registry.publication_hub_permits.available_permits(), 0);

        drop(second_subscription);
        registry.detach(first).await.unwrap();
        registry.detach(second).await.unwrap();
        registry.detach(third).await.unwrap();
        assert_eq!(registry.publication_hub_permits.available_permits(), 1);
    }

    #[tokio::test]
    async fn generation_retirement_wins_idle_eviction_race() {
        let eviction_gate = Arc::new(Semaphore::new(0));
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig {
                max_initialized_pool_hubs: 1,
                eviction_gate: Some(eviction_gate.clone()),
                ..PoolPublicationConfig::default()
            },
        ));
        let first = registry
            .attach(request(pool(1), "first.router.generate", "llama"))
            .await
            .unwrap();
        let second = registry
            .attach(request(pool(2), "second.router.generate", "llama"))
            .await
            .unwrap();
        let first_hub = hub_slot(&registry, first.pool_id);

        let subscription = subscribe(&registry, first.handle.identity()).await.unwrap();
        drop(subscription);
        let second_subscribe = tokio::spawn({
            let registry = registry.clone();
            let producer = second.handle.identity();
            async move { subscribe(&registry, producer).await }
        });
        tokio::time::timeout(Duration::from_secs(1), async {
            while first_hub.phase() != "evicting" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("pressure request must claim the idle hub");

        assert!(registry.fence_generation(
            first.pool_id,
            first.layout_generation,
            "test retirement during idle eviction",
        ));
        eviction_gate.add_permits(1);

        let second_subscription = second_subscribe.await.unwrap().unwrap();
        tokio::time::timeout(Duration::from_secs(1), async {
            while first_hub.phase() != "retired" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("retired generation must not return to the idle cache");

        drop(second_subscription);
        registry.detach(first).await.unwrap();
        registry.detach(second).await.unwrap();
        assert_eq!(registry.publication_hub_permits.available_permits(), 1);
    }

    #[tokio::test]
    async fn requester_retirement_does_not_abandon_claimed_idle_eviction() {
        let eviction_gate = Arc::new(Semaphore::new(0));
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig {
                max_initialized_pool_hubs: 1,
                eviction_gate: Some(eviction_gate.clone()),
                ..PoolPublicationConfig::default()
            },
        ));
        let first = registry
            .attach(request(pool(1), "first.router.generate", "llama"))
            .await
            .unwrap();
        let second = registry
            .attach(request(pool(2), "second.router.generate", "llama"))
            .await
            .unwrap();
        let third = registry
            .attach(request(pool(3), "third.router.generate", "llama"))
            .await
            .unwrap();
        let first_hub = hub_slot(&registry, first.pool_id);
        let second_hub = hub_slot(&registry, second.pool_id);

        let subscription = subscribe(&registry, first.handle.identity()).await.unwrap();
        drop(subscription);
        let second_subscribe = tokio::spawn({
            let registry = registry.clone();
            let producer = second.handle.identity();
            async move { subscribe(&registry, producer).await }
        });
        tokio::time::timeout(Duration::from_secs(1), async {
            while first_hub.phase() != "evicting" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("pressure request must claim the idle hub");

        assert!(registry.fence_generation(
            second.pool_id,
            second.layout_generation,
            "test requester retirement during idle eviction",
        ));
        let second_result = tokio::time::timeout(Duration::from_secs(1), second_subscribe)
            .await
            .expect("retired requester must stop waiting for reclaimed capacity")
            .expect("requester task must join");
        let second_error = second_result
            .err()
            .expect("retired requester must reject publication admission");
        assert!(matches!(second_error, PublicationHubError::Unavailable(_)));
        tokio::time::timeout(Duration::from_secs(1), async {
            while second_hub.phase() != "retired" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("requester slot must finish retirement");
        assert_eq!(first_hub.phase(), "evicting");

        eviction_gate.add_permits(1);
        tokio::time::timeout(Duration::from_secs(1), async {
            while first_hub.phase() != "vacant"
                || registry.publication_hub_permits.available_permits() != 1
            {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("detached cleanup must release the idle victim and its capacity");

        let third_subscription = subscribe(&registry, third.handle.identity()).await.unwrap();
        assert_eq!(registry.publication_hub_permits.available_permits(), 0);

        drop(third_subscription);
        registry.detach(first).await.unwrap();
        registry.detach(second).await.unwrap();
        registry.detach(third).await.unwrap();
        assert_eq!(registry.publication_hub_permits.available_permits(), 1);
    }

    #[tokio::test]
    async fn terminal_hub_failure_withdraws_catalog_and_load_visibility() {
        let encoding_permits = Arc::new(Semaphore::new(1));
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig {
                hub: PublicationHubConfig {
                    encoding_permits: encoding_permits.clone(),
                    ..PublicationHubConfig::default()
                },
                ..PoolPublicationConfig::default()
            },
        ));
        let attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let mut subscription = subscribe(&registry, attachment.handle.identity())
            .await
            .unwrap();
        assert_eq!(registry.catalog().pools().len(), 1);
        assert_eq!(registry.load_snapshots().len(), 1);

        encoding_permits.close();
        attachment
            .handle
            .admit_event(0, stored_event(1))
            .await
            .unwrap();
        attachment.handle.flush().await.unwrap();
        tokio::time::timeout(Duration::from_secs(5), async {
            loop {
                if registry.catalog().pools().is_empty()
                    && registry.load_snapshots().is_empty()
                    && attachment.pool_cancel.is_cancelled()
                {
                    break;
                }
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("terminal publication failure must withdraw the generation");
        assert!(subscription.recv().await.is_err());

        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn stale_generation_cannot_fence_or_initialize_replacement_hub() {
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig::default(),
        ));
        let first = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let stale_generation = first.layout_generation;
        let stale_producer = first.handle.identity();
        registry.detach(first).await.unwrap();

        let replacement = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let current_producer = replacement.handle.identity();
        let current_hub = hub_slot(&registry, replacement.pool_id);
        assert_ne!(stale_producer, current_producer);

        assert!(!registry.fence_generation(
            replacement.pool_id,
            stale_generation,
            "stale terminal callback",
        ));
        assert!(!replacement.pool_cancel.is_cancelled());
        assert_eq!(registry.catalog().pools().len(), 1);

        assert!(matches!(
            registry.validate_active_producer(stale_producer),
            Err(PublicationHubError::ProducerMismatch(pool_id))
                if pool_id == replacement.pool_id
        ));
        let current_generation_cancel =
            registry.validate_active_producer(current_producer).unwrap();
        assert!(!current_generation_cancel.is_cancelled());

        let error = subscribe(&registry, stale_producer)
            .await
            .err()
            .expect("stale producer subscription must fail");
        assert_eq!(
            error,
            PublicationHubError::ProducerMismatch(replacement.pool_id)
        );
        assert_eq!(current_hub.phase(), "vacant");

        let subscription = subscribe(&registry, current_producer).await.unwrap();
        assert_eq!(
            subscription.snapshot().unwrap().identity(),
            current_producer
        );
        assert_eq!(current_hub.phase(), "ready");
        drop(subscription);
        registry.detach(replacement).await.unwrap();
        assert!(current_generation_cancel.is_cancelled());
    }

    #[tokio::test]
    async fn publication_source_streams_snapshot_before_concurrent_delta() {
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig::default(),
        ));
        let attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let producer = attachment.handle.identity();
        let lifecycle = CancellationToken::new();
        let source = publication_source(registry.clone(), lifecycle.clone(), 2);
        let mut stream = source.subscribe_pool(producer).await.unwrap();

        // Exercise the handoff race: a delta published before snapshot consumption must still
        // follow it.
        attachment
            .handle
            .admit_event(0, stored_event(1))
            .await
            .unwrap();
        attachment.handle.flush().await.unwrap();

        let snapshot = tokio::time::timeout(Duration::from_secs(1), stream.next())
            .await
            .expect("source snapshot must arrive")
            .expect("source stream must remain open")
            .expect("source snapshot must be valid");
        assert_eq!(snapshot.kind(), PublicationFrameKind::SnapshotChunk);
        assert_eq!(snapshot.identity(), producer);
        assert_eq!(snapshot.base_sequence(), snapshot.sequence());

        let delta = tokio::time::timeout(Duration::from_secs(1), stream.next())
            .await
            .expect("source delta must arrive")
            .expect("source stream must remain open")
            .expect("source delta must be valid");
        assert_eq!(delta.kind(), PublicationFrameKind::Delta);
        assert_eq!(delta.identity(), producer);
        assert_eq!(delta.base_sequence(), snapshot.sequence());
        assert_eq!(delta.sequence(), snapshot.sequence() + 1);

        lifecycle.cancel();
        tokio::time::timeout(Duration::from_secs(1), source.wait_for_shutdown())
            .await
            .expect("publication lifecycle must wake");
        drop(stream);
        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn publication_source_limit_and_retirement_are_generation_safe() {
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig::default(),
        ));
        let first = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let stale_producer = first.handle.identity();
        let lifecycle = CancellationToken::new();
        let source = publication_source(registry.clone(), lifecycle.clone(), 1);
        let mut stale_stream = source.subscribe_pool(stale_producer).await.unwrap();

        let error = source
            .subscribe_pool(stale_producer)
            .await
            .err()
            .expect("global stream limit must reject a second stream");
        assert_eq!(error.kind(), PublicationErrorKind::ResourceExhausted);

        registry.detach(first).await.unwrap();
        let retired = tokio::time::timeout(Duration::from_secs(1), stale_stream.next())
            .await
            .expect("retired source stream must wake")
            .expect("retired source stream must report its terminal error")
            .expect_err("retired source stream must require a fresh snapshot");
        assert!(matches!(
            retired.kind(),
            PublicationErrorKind::Unavailable | PublicationErrorKind::ResourceExhausted
        ));
        drop(stale_stream);

        let replacement = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let current_producer = replacement.handle.identity();
        assert_ne!(stale_producer, current_producer);
        let stale_error = source
            .subscribe_pool(stale_producer)
            .await
            .err()
            .expect("stale producer must not attach to the replacement");
        assert_eq!(stale_error.kind(), PublicationErrorKind::ProducerMismatch);
        let current_stream = source.subscribe_pool(current_producer).await.unwrap();

        lifecycle.cancel();
        tokio::time::timeout(Duration::from_secs(1), source.wait_for_shutdown())
            .await
            .expect("publication lifecycle must wake");
        drop(current_stream);
        registry.detach(replacement).await.unwrap();
    }

    #[tokio::test]
    async fn retirement_cancels_in_flight_hub_initialization() {
        let initialization_gate = Arc::new(Semaphore::new(0));
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig {
                hub: PublicationHubConfig {
                    initialization_gate: Some(initialization_gate),
                    ..PublicationHubConfig::default()
                },
                ..PoolPublicationConfig::default()
            },
        ));
        let attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let producer = attachment.handle.identity();
        let hub = hub_slot(&registry, attachment.pool_id);
        let subscriber = tokio::spawn({
            let registry = registry.clone();
            async move { subscribe(&registry, producer).await }
        });

        tokio::time::timeout(Duration::from_secs(1), async {
            while hub.phase() != "initializing" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("hub initialization must start");
        tokio::time::timeout(Duration::from_secs(1), registry.detach(attachment))
            .await
            .expect("retirement must cancel hub initialization")
            .unwrap();

        assert!(subscriber.await.unwrap().is_err());
        assert_eq!(hub.phase(), "retired");
        assert_eq!(registry.pool_count().await, 0);
        assert!(registry.catalog().pools().is_empty());
    }

    #[tokio::test]
    async fn cancellation_after_actor_snapshot_retires_the_publication_lease() {
        let post_snapshot_gate = Arc::new(Semaphore::new(0));
        let registry = Arc::new(PoolRegistry::new_with_publication_config(
            relay_identity(),
            config(),
            PoolPublicationConfig {
                hub: PublicationHubConfig {
                    post_snapshot_gate: Some(post_snapshot_gate),
                    ..PublicationHubConfig::default()
                },
                ..PoolPublicationConfig::default()
            },
        ));
        let attachment = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        let producer = attachment.handle.identity();
        let hub = hub_slot(&registry, attachment.pool_id);
        let subscriber = tokio::spawn({
            let registry = registry.clone();
            async move { subscribe(&registry, producer).await }
        });

        tokio::time::timeout(Duration::from_secs(5), async {
            while hub.phase() != "initializing" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("hub initialization must start");

        let mut event_id = 1;
        tokio::time::timeout(Duration::from_secs(5), async {
            loop {
                attachment
                    .handle
                    .admit_event(0, stored_event(event_id))
                    .await
                    .unwrap();
                attachment.handle.flush().await.unwrap();
                if attachment.handle.state_stats().await.unwrap().1 > 0 {
                    break;
                }
                event_id += 1;
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("actor publication lease must become active");

        attachment.pool_cancel.cancel();
        assert!(subscriber.await.unwrap().is_err());
        tokio::time::timeout(Duration::from_secs(5), async {
            while hub.phase() != "retired" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("cancelled hub initialization must retire");

        let sequence_before = attachment.handle.state_stats().await.unwrap().1;
        attachment
            .handle
            .admit_event(0, stored_event(event_id + 1))
            .await
            .unwrap();
        attachment.handle.flush().await.unwrap();
        let sequence_after = attachment.handle.state_stats().await.unwrap().1;
        assert_eq!(sequence_after, sequence_before);

        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn local_only_pool_publishes_no_serving_facts() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let mut attach_request = request(pool(1), "fast.router.generate", "llama");
        attach_request.serving_facts = None;
        let attachment = registry.attach(attach_request).await.unwrap();

        assert_eq!(registry.catalog().pools().len(), 1);
        assert!(registry.load_snapshots().is_empty());
        assert!(
            !registry
                .replace_load_capacity(
                    attachment.pool_id,
                    attachment.layout_generation,
                    &HashMap::new(),
                )
                .unwrap()
        );
        assert!(!registry.observe_load(
            attachment.pool_id,
            attachment.layout_generation,
            ActiveLoad::default(),
        ));

        attachment.handle.state_stats().await.unwrap();
        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn invalid_load_capacity_does_not_block_pool_attach() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let mut attach_request = request(pool(1), "fast.router.generate", "llama");
        attach_request.serving_facts = Some(PoolServingFacts {
            runtime_configs: HashMap::from([(
                1,
                ModelRuntimeConfig {
                    data_parallel_size: 0,
                    total_kv_blocks: Some(100),
                    ..ModelRuntimeConfig::default()
                },
            )]),
        });

        let attachment = registry.attach(attach_request).await.unwrap();

        assert_eq!(registry.catalog().pools().len(), 1);
        let snapshots = registry.load_snapshots();
        assert_eq!(snapshots.len(), 1);
        assert_eq!(snapshots[0].kv_used_blocks, None);
        assert_eq!(snapshots[0].total_kv_blocks, None);
        assert_eq!(snapshots[0].kv_expected_ranks, 0);
        assert!(snapshots[0].has_degraded_coverage());

        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn scheduler_only_load_is_accepted_without_changing_snapshot() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let mut attach_request = request(pool(1), "fast.router.generate", "llama");
        attach_request.serving_facts = Some(PoolServingFacts {
            runtime_configs: HashMap::from([(
                1,
                ModelRuntimeConfig {
                    total_kv_blocks: Some(100),
                    ..ModelRuntimeConfig::default()
                },
            )]),
        });
        let attachment = registry.attach(attach_request).await.unwrap();
        let before = registry.load_snapshots();
        let load_watch = registry.watch_load();
        assert!(!load_watch.has_changed().unwrap());

        assert!(registry.observe_load(
            attachment.pool_id,
            attachment.layout_generation,
            ActiveLoad {
                worker_id: 1,
                dp_rank: 0,
                active_decode_blocks: Some(30),
                active_prefill_tokens: Some(512),
                ..ActiveLoad::default()
            },
        ));

        assert_eq!(registry.load_snapshots(), before);
        assert!(!load_watch.has_changed().unwrap());
        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn invalid_capacity_refresh_withdraws_authoritative_load_snapshot() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let mut attach_request = request(pool(1), "fast.router.generate", "llama");
        attach_request.serving_facts = Some(PoolServingFacts {
            runtime_configs: HashMap::from([(
                1,
                ModelRuntimeConfig {
                    total_kv_blocks: Some(100),
                    ..ModelRuntimeConfig::default()
                },
            )]),
        });
        let attachment = registry.attach(attach_request).await.unwrap();
        assert!(registry.observe_load(
            attachment.pool_id,
            attachment.layout_generation,
            ActiveLoad {
                worker_id: 1,
                dp_rank: 0,
                kv_used_blocks: Some(40),
                ..ActiveLoad::default()
            },
        ));
        assert!(!registry.load_snapshots()[0].has_degraded_coverage());
        let mut load_watch = registry.watch_load();

        let error = registry
            .replace_load_capacity(
                attachment.pool_id,
                attachment.layout_generation,
                &HashMap::from([(
                    1,
                    ModelRuntimeConfig {
                        data_parallel_size: 0,
                        total_kv_blocks: Some(100),
                        ..ModelRuntimeConfig::default()
                    },
                )]),
            )
            .unwrap_err();
        assert!(error.to_string().contains("zero data_parallel_size"));
        assert!(load_watch.has_changed().unwrap());

        let snapshots = load_watch.borrow_and_update().clone();
        assert_eq!(snapshots.len(), 1);
        assert_eq!(snapshots[0].kv_used_blocks, None);
        assert_eq!(snapshots[0].total_kv_blocks, None);
        assert_eq!(snapshots[0].kv_expected_ranks, 0);
        assert!(snapshots[0].has_degraded_coverage());
        assert_eq!(registry.catalog().pools().len(), 1);

        registry.detach(attachment).await.unwrap();
    }

    #[tokio::test]
    async fn load_is_generation_scoped_and_withdrawn_with_the_pool() {
        let registry = PoolRegistry::new(relay_identity(), config());
        let runtime_configs = HashMap::from([(
            1,
            ModelRuntimeConfig {
                total_kv_blocks: Some(100),
                max_num_batched_tokens: Some(2_048),
                ..ModelRuntimeConfig::default()
            },
        )]);
        let attachment = registry
            .attach(PoolAttachRequest {
                pool_id: pool(1),
                endpoint: EndpointId::from("fast.router.generate"),
                registrations: vec![registration("llama")],
                query_semantics: query_semantics(),
                roles: vec![WorkerRole::Aggregated],
                serving_facts: Some(PoolServingFacts { runtime_configs }),
            })
            .await
            .unwrap();
        let old_generation = attachment.layout_generation;
        let old_producer = attachment.handle.identity();

        let initial_load = registry.load_snapshots();
        assert_eq!(initial_load.len(), 1);
        assert_eq!(initial_load[0].kv_expected_ranks, 1);
        assert_eq!(initial_load[0].kv_observed_ranks, 0);
        assert_eq!(initial_load[0].kv_used_blocks, None);
        assert_eq!(initial_load[0].total_kv_blocks, Some(100));
        assert!(initial_load[0].has_degraded_coverage());

        let mut load = ActiveLoad {
            worker_id: 1,
            dp_rank: 0,
            ..ActiveLoad::default()
        };
        load.kv_used_blocks = Some(40);
        load.active_decode_blocks = Some(30);
        load.active_prefill_tokens = Some(512);
        assert!(registry.observe_load(pool(1), old_generation, load));
        let observed = registry.load_snapshots()[0];
        assert_eq!(observed.kv_used_blocks, Some(40));
        assert_eq!(observed.total_kv_blocks, Some(100));
        assert!(!observed.has_degraded_coverage());

        assert!(
            registry
                .withdraw(pool(1), old_generation, PoolRetirementMode::Graceful)
                .await
        );
        assert!(registry.load_snapshots().is_empty());
        assert!(!registry.observe_load(
            pool(1),
            old_generation,
            ActiveLoad {
                worker_id: 1,
                dp_rank: 0,
                kv_used_blocks: Some(99),
                ..ActiveLoad::default()
            },
        ));
        registry.detach(attachment).await.unwrap();

        let replacement = registry
            .attach(request(pool(1), "fast.router.generate", "llama"))
            .await
            .unwrap();
        assert_ne!(replacement.layout_generation, old_generation);
        assert_ne!(replacement.handle.identity(), old_producer);
        assert!(
            !registry
                .replace_load_capacity(
                    pool(1),
                    old_generation,
                    &HashMap::from([(
                        1,
                        ModelRuntimeConfig {
                            total_kv_blocks: Some(999),
                            ..ModelRuntimeConfig::default()
                        },
                    )]),
                )
                .unwrap()
        );
        assert_eq!(registry.load_snapshots()[0].kv_expected_ranks, 0);
        registry.detach(replacement).await.unwrap();
    }
}