dynamo-llm 1.3.0

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

use std::{
    borrow::Cow,
    collections::{HashMap, HashSet},
    str::FromStr,
};

use serde::{Deserialize, Serialize, de::DeserializeOwned};
use validator::{Validate, ValidationError};

use dynamo_kv_router::protocols::KvTransferEnforcement;

/// Re-export from parsers crate so that `ModelRuntimeConfig` can use it
/// directly without type duplication.
pub use dynamo_parsers::tool_calling::StructuralTagSchemaMode;

// Reserve a topology namespace so generated taints can be rebuilt without touching caller taints.
pub const TOPOLOGY_TAINT_PREFIX: &str = "dynamo.topology/";

/// Canonical worker-taint form for topology metadata.
///
/// A topology domain/value pair such as `zone=us-east-1a` becomes
/// `dynamo.topology/zone=us-east-1a`.
pub fn topology_taint(domain: &str, value: &str) -> String {
    format!("{TOPOLOGY_TAINT_PREFIX}{domain}={value}")
}

/// Master switch for structural tag guided decoding.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "snake_case")]
pub enum StructuralTagMode {
    #[default]
    Off,
    On,
}

/// Controls when structural tags are activated based on `tool_choice`.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "snake_case")]
pub enum StructuralTagScope {
    #[default]
    Auto,
    Always,
}

pub const ENV_TOKENIZER_BACKEND: &str = "DYN_TOKENIZER";

/// Tokenizer backend used by the Rust preprocessor for BPE tokenizer.json models.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum TokenizerBackend {
    Default,
    Fastokens,
}

impl TokenizerBackend {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Default => "default",
            Self::Fastokens => "fastokens",
        }
    }

    pub fn is_fastokens(self) -> bool {
        matches!(self, Self::Fastokens)
    }

    pub fn from_env_or_default() -> Self {
        match std::env::var(ENV_TOKENIZER_BACKEND) {
            Ok(v) if v == "fastokens" => Self::Fastokens,
            Ok(v) if v == "default" || v.is_empty() => Self::Default,
            Ok(v) => {
                tracing::warn!(
                    value = %v,
                    "Unrecognized DYN_TOKENIZER value, expected 'fastokens' or 'default'; falling back to default"
                );
                Self::Default
            }
            Err(_) => Self::Default,
        }
    }
}

impl FromStr for TokenizerBackend {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "default" => Ok(Self::Default),
            "fastokens" => Ok(Self::Fastokens),
            _ => Err(format!(
                "invalid tokenizer backend '{value}' (expected 'default' or 'fastokens')"
            )),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct DisaggregatedEndpoint {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub bootstrap_host: Option<String>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub bootstrap_port: Option<u16>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Validate)]
#[validate(schema(function = "validate_model_runtime_config"))]
/// Runtime-resolved metadata published by a worker after its engine starts.
///
/// NOTE: This type is intended for facts that can only be known authoritatively at
/// runtime, such as the effective engine context limit, capacity, data-parallel
/// placement, and resolved service endpoints. Some legacy fields do not yet follow
/// this ownership boundary; avoid adding declarative model metadata here.
pub struct ModelRuntimeConfig {
    /// Effective context limit enforced by the running engine.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub context_length: Option<u32>,

    pub total_kv_blocks: Option<u64>,

    pub max_num_seqs: Option<u64>,

    pub max_num_batched_tokens: Option<u64>,

    pub tool_call_parser: Option<String>,

    pub reasoning_parser: Option<String>,

    /// Frontend tokenizer backend override. When unset, direct Rust callers can still use
    /// `DYN_TOKENIZER`; when set, this explicit value wins over process environment.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tokenizer_backend: Option<TokenizerBackend>,

    /// Whether structural tag guided decoding is enabled for tool calls.
    #[serde(default)]
    pub structural_tag_mode: StructuralTagMode,

    /// Controls when structural tags are activated based on tool_choice.
    #[serde(default)]
    pub structural_tag_scope: StructuralTagScope,

    /// Controls whether tools get real or generic schemas in structural tags.
    #[serde(default)]
    pub structural_tag_schema: StructuralTagSchemaMode,

    /// When true, strip tool definitions from the chat template when tool_choice is "none".
    #[serde(default = "default_exclude_tools_when_tool_choice_none")]
    pub exclude_tools_when_tool_choice_none: bool,

    /// Starting rank of data parallel ranks for this worker (0 if DP not enabled)
    #[serde(default = "default_data_parallel_start_rank")]
    pub data_parallel_start_rank: u32,

    /// Total number of data parallel ranks for this worker (1 if DP not enabled)
    #[serde(default = "default_data_parallel_size")]
    pub data_parallel_size: u32,

    /// Enable worker-local KV indexer for tracking this worker's own KV cache state (default: true)
    #[serde(default = "default_local_indexer")]
    pub enable_local_indexer: bool,

    /// Mapping of engine-specific runtime configs
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub runtime_data: HashMap<String, serde_json::Value>,

    /// Bootstrap endpoint for disaggregated serving (prefill workers publish this)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub disaggregated_endpoint: Option<DisaggregatedEndpoint>,

    #[serde(default = "default_eagle")]
    pub enable_eagle: bool,

    #[serde(default, skip_serializing_if = "HashSet::is_empty")]
    pub taints: HashSet<String>,

    /// Stable routing identity, set via the `DYN_STABLE_ROUTING_ID` env var. Used as
    /// the rendezvous-hash key so cache assignments survive a new ephemeral
    /// `worker_id`. `None` if unset.
    ///
    /// Recommended k8s wire-up (downward API on a StatefulSet pod):
    ///
    /// ```yaml
    /// env:
    ///   - name: DYN_STABLE_ROUTING_ID
    ///     valueFrom:
    ///       fieldRef:
    ///         fieldPath: metadata.name
    /// ```
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub stable_routing_id: Option<String>,

    /// Topology domain labels for this worker (e.g. {"zone": "us-east-1a", "rack": "rack1"}).
    /// Workers publish these as metadata and as additive canonical taints with the
    /// `dynamo.topology/<domain>=<value>` format.
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    #[validate(custom(function = "validate_topology_domains"))]
    pub topology_domains: HashMap<String, String>,

    /// Topology domain used for KV-cache transfer routing (e.g. "zone").
    #[serde(default, skip_serializing_if = "Option::is_none")]
    #[validate(custom(function = "validate_kv_transfer_domain"))]
    pub kv_transfer_domain: Option<String>,

    /// KV transfer topology enforcement mode selected by DGD (`required` or `preferred`).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub kv_transfer_enforcement: Option<KvTransferEnforcement>,

    /// Preferred-taint weight used when `kv_transfer_enforcement` is `preferred`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    #[validate(range(min = 0.0, max = 1.0))]
    pub kv_transfer_preferred_weight: Option<f32>,

    /// Per-worker LoRA adapter slot capacity (e.g. vLLM `--max-loras`, SGLang
    /// `--max-loras-per-batch`), advertised on the BASE worker registration so the LoRA
    /// allocation controller can see idle-but-LoRA-capable workers before any adapter is
    /// loaded on them. `None` for non-LoRA workers. Adapter (`card.lora`) registrations carry
    /// the same value via `LoraInfo::max_gpu_lora_count`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_gpu_lora_count: Option<u32>,
}

const fn default_data_parallel_start_rank() -> u32 {
    0
}

const fn default_data_parallel_size() -> u32 {
    1
}

const fn default_local_indexer() -> bool {
    true
}

const fn default_exclude_tools_when_tool_choice_none() -> bool {
    true
}

const fn default_eagle() -> bool {
    false
}

impl Default for ModelRuntimeConfig {
    fn default() -> Self {
        Self {
            context_length: None,
            total_kv_blocks: None,
            max_num_seqs: None,
            max_num_batched_tokens: None,
            tool_call_parser: None,
            reasoning_parser: None,
            tokenizer_backend: None,
            structural_tag_mode: StructuralTagMode::Off,
            structural_tag_scope: StructuralTagScope::Auto,
            structural_tag_schema: StructuralTagSchemaMode::Auto,
            exclude_tools_when_tool_choice_none: default_exclude_tools_when_tool_choice_none(),
            data_parallel_start_rank: default_data_parallel_start_rank(),
            data_parallel_size: default_data_parallel_size(),
            enable_local_indexer: true,
            runtime_data: HashMap::new(),
            disaggregated_endpoint: None,
            enable_eagle: false,
            taints: HashSet::new(),
            stable_routing_id: None,
            topology_domains: HashMap::new(),
            kv_transfer_domain: None,
            kv_transfer_enforcement: None,
            kv_transfer_preferred_weight: None,
            max_gpu_lora_count: None,
        }
    }
}

impl dynamo_kv_router::WorkerConfigLike for ModelRuntimeConfig {
    fn data_parallel_start_rank(&self) -> u32 {
        self.data_parallel_start_rank
    }

    fn data_parallel_size(&self) -> u32 {
        self.data_parallel_size
    }

    fn max_num_batched_tokens(&self) -> Option<u64> {
        self.max_num_batched_tokens
    }

    fn total_kv_blocks(&self) -> Option<u64> {
        self.total_kv_blocks
    }

    fn taints(&self) -> &HashSet<String> {
        &self.taints
    }

    fn stable_routing_id(&self) -> Option<&str> {
        self.stable_routing_id.as_deref()
    }

    fn topology_domains(&self) -> Option<&HashMap<String, String>> {
        if self.topology_domains.is_empty() {
            None
        } else {
            Some(&self.topology_domains)
        }
    }

    fn kv_transfer_domain(&self) -> Option<&str> {
        self.kv_transfer_domain.as_deref()
    }

    fn kv_transfer_enforcement(&self) -> Option<KvTransferEnforcement> {
        self.kv_transfer_enforcement
    }

    fn kv_transfer_preferred_weight(&self) -> Option<f32> {
        self.kv_transfer_preferred_weight
    }
}

fn validation_error(code: &'static str, message: impl Into<Cow<'static, str>>) -> ValidationError {
    let mut error = ValidationError::new(code);
    error.message = Some(message.into());
    error
}

fn validate_taint_component(
    component: &str,
    code_prefix: &'static str,
    name: &'static str,
) -> Result<(), ValidationError> {
    if component.trim().is_empty() {
        return Err(validation_error(
            code_prefix,
            format!("{name} must be non-empty"),
        ));
    }
    if component.trim() != component {
        return Err(validation_error(
            code_prefix,
            format!("{name} must not contain leading or trailing whitespace"),
        ));
    }
    if component.contains('=') {
        return Err(validation_error(
            code_prefix,
            format!("{name} must not contain '='"),
        ));
    }

    Ok(())
}

fn validate_topology_domains(
    topology_domains: &HashMap<String, String>,
) -> Result<(), ValidationError> {
    for (domain, value) in topology_domains {
        validate_taint_component(domain, "invalid_topology_domain", "topology_domains key")?;
        validate_taint_component(value, "invalid_topology_value", "topology_domains value")?;
    }

    Ok(())
}

fn validate_kv_transfer_domain(domain: &str) -> Result<(), ValidationError> {
    validate_taint_component(domain, "invalid_kv_transfer_domain", "kv_transfer_domain")
}

fn validate_model_runtime_config(config: &ModelRuntimeConfig) -> Result<(), ValidationError> {
    if let Some(domain) = &config.kv_transfer_domain
        && !config.topology_domains.contains_key(domain)
    {
        return Err(validation_error(
            "missing_kv_transfer_domain",
            "kv_transfer_domain must reference a key in topology_domains",
        ));
    }

    if config.kv_transfer_enforcement.is_some() && config.kv_transfer_domain.is_none() {
        return Err(validation_error(
            "missing_kv_transfer_domain",
            "kv_transfer_enforcement requires kv_transfer_domain",
        ));
    }

    if matches!(
        config.kv_transfer_enforcement,
        Some(KvTransferEnforcement::Preferred)
    ) && config.kv_transfer_preferred_weight.is_none()
    {
        return Err(validation_error(
            "missing_kv_transfer_preferred_weight",
            "kv_transfer_preferred_weight is required when kv_transfer_enforcement is preferred",
        ));
    }

    if config.kv_transfer_preferred_weight.is_some()
        && !matches!(
            config.kv_transfer_enforcement,
            Some(KvTransferEnforcement::Preferred)
        )
    {
        return Err(validation_error(
            "invalid_kv_transfer_preferred_weight",
            "kv_transfer_preferred_weight can only be set when kv_transfer_enforcement is preferred",
        ));
    }

    Ok(())
}

impl ModelRuntimeConfig {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn validate_config(&self) -> Result<(), String> {
        self.validate().map_err(|error| error.to_string())
    }

    pub fn set_engine_specific<T: Serialize>(&mut self, key: &str, value: T) -> anyhow::Result<()> {
        self.runtime_data
            .insert(key.to_string(), serde_json::to_value(value)?);
        Ok(())
    }

    pub fn get_engine_specific<T: DeserializeOwned>(&self, key: &str) -> anyhow::Result<Option<T>> {
        if let Some(value) = self.runtime_data.get(key) {
            Ok(Some(serde_json::from_value(value.clone())?))
        } else {
            Ok(None)
        }
    }

    pub fn effective_tokenizer_backend(&self) -> TokenizerBackend {
        self.tokenizer_backend
            .unwrap_or_else(TokenizerBackend::from_env_or_default)
    }

    pub fn set_tokenizer_backend(
        &mut self,
        tokenizer_backend: Option<TokenizerBackend>,
    ) -> &mut Self {
        self.tokenizer_backend = tokenizer_backend;
        self
    }

    /// Rebuild canonical topology taints derived from `topology_domains`.
    ///
    /// Existing caller-provided taints outside the reserved topology prefix are preserved; generated
    /// topology taints are refreshed in the same set so `RoutingConstraints` can match them through
    /// the standard worker taints path.
    pub fn add_topology_taints(&mut self) -> &mut Self {
        self.taints
            .retain(|taint| !taint.starts_with(TOPOLOGY_TAINT_PREFIX));
        self.taints
            .extend(self.topology_domains.iter().filter_map(|(domain, value)| {
                let domain = domain.trim();
                let value = value.trim();
                if domain.is_empty() || value.is_empty() {
                    None
                } else {
                    Some(topology_taint(domain, value))
                }
            }));
        self
    }

    /// Populate `stable_routing_id` from the `DYN_STABLE_ROUTING_ID` environment variable.
    ///
    /// Sets the field only if it is currently unset; returns `&mut self` for chaining. If
    /// `DYN_STABLE_ROUTING_ID` is unset or empty/whitespace-only, the field is left
    /// as `None`.
    ///
    /// See the doc on [`ModelRuntimeConfig::stable_routing_id`] for the recommended k8s
    /// downward-API recipe.
    pub fn populate_stable_routing_id_from_env(&mut self) -> &mut Self {
        if self.stable_routing_id.is_some() {
            return self;
        }
        let candidate = std::env::var("DYN_STABLE_ROUTING_ID")
            .ok()
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty());
        if let Some(value) = candidate {
            tracing::info!(stable_routing_id = %value, "populated stable_routing_id from DYN_STABLE_ROUTING_ID");
            self.stable_routing_id = Some(value);
        }
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // Env-touching tests use `temp_env` (snapshot + restore around the closure) and
    // `#[serial_test::serial]` (serialize against every other env-touching test in the
    // binary, not just this module). A module-local mutex would be insufficient because
    // `std::env::{set_var, remove_var}` race against `getenv` in any other thread.

    #[test]
    #[serial_test::serial]
    fn populates_from_dyn_env() {
        temp_env::with_vars([("DYN_STABLE_ROUTING_ID", Some("worker-3"))], || {
            let mut cfg = ModelRuntimeConfig::default();
            cfg.populate_stable_routing_id_from_env();
            assert_eq!(cfg.stable_routing_id.as_deref(), Some("worker-3"));
        });
    }

    #[test]
    #[serial_test::serial]
    fn preserves_caller_supplied_value() {
        temp_env::with_vars([("DYN_STABLE_ROUTING_ID", Some("from-env"))], || {
            let mut cfg = ModelRuntimeConfig {
                stable_routing_id: Some("explicit".to_string()),
                ..Default::default()
            };
            cfg.populate_stable_routing_id_from_env();
            assert_eq!(cfg.stable_routing_id.as_deref(), Some("explicit"));
        });
    }

    #[test]
    #[serial_test::serial]
    fn no_meaningful_env_leaves_field_none() {
        // Whitespace-only is rejected…
        temp_env::with_vars([("DYN_STABLE_ROUTING_ID", Some("   "))], || {
            let mut cfg = ModelRuntimeConfig::default();
            cfg.populate_stable_routing_id_from_env();
            assert!(cfg.stable_routing_id.is_none());
        });
        // …as is having the var unset.
        temp_env::with_vars_unset(["DYN_STABLE_ROUTING_ID"], || {
            let mut cfg = ModelRuntimeConfig::default();
            cfg.populate_stable_routing_id_from_env();
            assert!(cfg.stable_routing_id.is_none());
        });
    }

    #[test]
    fn max_gpu_lora_count_roundtrips_and_is_omitted_when_none() {
        // Worker LoRA capacity must survive the MDC -> discovery -> watcher wire so the frontend
        // can seed set_worker_capacity for idle LoRA-capable workers.
        let cfg = ModelRuntimeConfig {
            max_gpu_lora_count: Some(8),
            ..Default::default()
        };
        let json = serde_json::to_string(&cfg).unwrap();
        assert!(json.contains("\"max_gpu_lora_count\":8"));
        let parsed: ModelRuntimeConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.max_gpu_lora_count, Some(8));

        // Omitted from the wire (and defaults to None on read) for non-LoRA workers, so older
        // payloads without the field stay backward-compatible.
        let none_json = serde_json::to_string(&ModelRuntimeConfig::default()).unwrap();
        assert!(!none_json.contains("max_gpu_lora_count"));
        let from_legacy: ModelRuntimeConfig = serde_json::from_str("{}").unwrap();
        assert_eq!(from_legacy.max_gpu_lora_count, None);
    }

    #[test]
    fn roundtrips_through_serde_json() {
        let cfg = ModelRuntimeConfig {
            stable_routing_id: Some("worker-7".to_string()),
            ..Default::default()
        };
        let json = serde_json::to_string(&cfg).unwrap();
        assert!(json.contains("\"stable_routing_id\":\"worker-7\""));
        let parsed: ModelRuntimeConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.stable_routing_id.as_deref(), Some("worker-7"));
    }

    #[test]
    fn serde_skips_when_none() {
        let cfg = ModelRuntimeConfig::default();
        let json = serde_json::to_string(&cfg).unwrap();
        assert!(!json.contains("stable_routing_id"));
        assert!(!json.contains("context_length"));
    }

    #[test]
    #[serial_test::serial]
    fn tokenizer_backend_env_fallback() {
        temp_env::with_vars([(ENV_TOKENIZER_BACKEND, Some("fastokens"))], || {
            let cfg = ModelRuntimeConfig::default();
            assert_eq!(
                cfg.effective_tokenizer_backend(),
                TokenizerBackend::Fastokens
            );
        });

        temp_env::with_vars([(ENV_TOKENIZER_BACKEND, Some("default"))], || {
            let cfg = ModelRuntimeConfig::default();
            assert_eq!(cfg.effective_tokenizer_backend(), TokenizerBackend::Default);
        });

        temp_env::with_vars_unset([ENV_TOKENIZER_BACKEND], || {
            let cfg = ModelRuntimeConfig::default();
            assert_eq!(cfg.effective_tokenizer_backend(), TokenizerBackend::Default);
        });
    }

    #[test]
    #[serial_test::serial]
    fn tokenizer_backend_explicit_config_wins_over_env() {
        temp_env::with_vars([(ENV_TOKENIZER_BACKEND, Some("fastokens"))], || {
            let cfg = ModelRuntimeConfig {
                tokenizer_backend: Some(TokenizerBackend::Default),
                ..Default::default()
            };
            assert_eq!(cfg.effective_tokenizer_backend(), TokenizerBackend::Default);
        });

        temp_env::with_vars([(ENV_TOKENIZER_BACKEND, Some("default"))], || {
            let cfg = ModelRuntimeConfig {
                tokenizer_backend: Some(TokenizerBackend::Fastokens),
                ..Default::default()
            };
            assert_eq!(
                cfg.effective_tokenizer_backend(),
                TokenizerBackend::Fastokens
            );
        });
    }

    #[test]
    fn tokenizer_backend_roundtrips_through_serde_json() {
        let cfg = ModelRuntimeConfig {
            tokenizer_backend: Some(TokenizerBackend::Fastokens),
            ..Default::default()
        };
        let json = serde_json::to_string(&cfg).unwrap();
        assert!(json.contains("\"tokenizer_backend\":\"fastokens\""));
        let parsed: ModelRuntimeConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.tokenizer_backend, Some(TokenizerBackend::Fastokens));
    }

    #[test]
    fn test_serde_empty_topology_domains_omitted() {
        let config = ModelRuntimeConfig::default();
        let serialized = serde_json::to_string(&config).unwrap();

        // Empty topology_domains should not appear in serialized output
        assert!(
            !serialized.contains("topology_domains"),
            "empty topology_domains should be skipped during serialization, got: {serialized}"
        );
    }

    #[test]
    fn test_serde_backward_compat_deserialize_without_topology_domains() {
        // Simulate a config serialized before topology_domains existed
        let json = r#"{
            "total_kv_blocks": 100,
            "max_num_seqs": 32,
            "max_num_batched_tokens": null,
            "tool_call_parser": null,
            "reasoning_parser": null
        }"#;

        let config: ModelRuntimeConfig = serde_json::from_str(json).unwrap();
        assert!(config.topology_domains.is_empty());
        assert!(config.kv_transfer_domain.is_none());
        assert!(config.kv_transfer_enforcement.is_none());
        assert!(config.kv_transfer_preferred_weight.is_none());
    }

    #[test]
    fn test_serde_round_trip_preserves_topology_transfer_fields_and_taints() {
        let mut config = ModelRuntimeConfig {
            taints: HashSet::from(["caller/taint=value".to_string()]),
            topology_domains: HashMap::from([
                ("zone".to_string(), "us-west-2b".to_string()),
                ("rack".to_string(), "rack1".to_string()),
            ]),
            kv_transfer_domain: Some("zone".to_string()),
            kv_transfer_enforcement: Some(KvTransferEnforcement::Preferred),
            kv_transfer_preferred_weight: Some(0.85),
            ..Default::default()
        };
        config.add_topology_taints();

        let serialized = serde_json::to_string(&config).unwrap();
        let deserialized: ModelRuntimeConfig = serde_json::from_str(&serialized).unwrap();

        assert_eq!(deserialized.topology_domains.len(), 2);
        assert_eq!(deserialized.topology_domains["zone"], "us-west-2b");
        assert_eq!(deserialized.topology_domains["rack"], "rack1");
        assert_eq!(deserialized.kv_transfer_domain.as_deref(), Some("zone"));
        assert_eq!(
            deserialized.kv_transfer_enforcement,
            Some(KvTransferEnforcement::Preferred)
        );
        assert_eq!(deserialized.kv_transfer_preferred_weight, Some(0.85));
        assert!(deserialized.taints.contains("caller/taint=value"));
        assert!(
            deserialized
                .taints
                .contains("dynamo.topology/zone=us-west-2b")
        );
        assert!(deserialized.taints.contains("dynamo.topology/rack=rack1"));
    }

    #[test]
    fn test_serde_rejects_invalid_kv_transfer_enforcement() {
        let json = r#"{"kv_transfer_enforcement":"fallback"}"#;
        assert!(serde_json::from_str::<ModelRuntimeConfig>(json).is_err());
    }

    #[test]
    fn test_validate_config_accepts_kv_transfer_configs() {
        for config in [
            ModelRuntimeConfig {
                topology_domains: HashMap::from([("zone".to_string(), "us-east-1a".to_string())]),
                kv_transfer_domain: Some("zone".to_string()),
                kv_transfer_enforcement: Some(KvTransferEnforcement::Required),
                ..Default::default()
            },
            ModelRuntimeConfig {
                topology_domains: HashMap::from([("zone".to_string(), "us-east-1a".to_string())]),
                kv_transfer_domain: Some("zone".to_string()),
                kv_transfer_enforcement: Some(KvTransferEnforcement::Preferred),
                kv_transfer_preferred_weight: Some(0.5),
                ..Default::default()
            },
        ] {
            config.validate_config().unwrap();
        }
    }

    #[test]
    fn test_validate_config_rejects_invalid_topology_components() {
        for config in [
            ModelRuntimeConfig {
                topology_domains: HashMap::from([("".to_string(), "us-east-1a".to_string())]),
                ..Default::default()
            },
            ModelRuntimeConfig {
                topology_domains: HashMap::from([(
                    "zone=primary".to_string(),
                    "us-east-1a".to_string(),
                )]),
                ..Default::default()
            },
        ] {
            assert!(config.validate_config().is_err());
        }
    }

    #[test]
    fn test_validate_config_rejects_transfer_domain_mismatch() {
        let config = ModelRuntimeConfig {
            topology_domains: HashMap::from([("zone".to_string(), "us-east-1a".to_string())]),
            kv_transfer_domain: Some("rack".to_string()),
            ..Default::default()
        };

        assert!(config.validate_config().is_err());
    }

    #[test]
    fn test_validate_config_rejects_invalid_kv_transfer_combinations() {
        for config in [
            ModelRuntimeConfig {
                kv_transfer_enforcement: Some(KvTransferEnforcement::Required),
                ..Default::default()
            },
            ModelRuntimeConfig {
                topology_domains: HashMap::from([("zone".to_string(), "us-east-1a".to_string())]),
                kv_transfer_domain: Some("zone".to_string()),
                kv_transfer_enforcement: Some(KvTransferEnforcement::Preferred),
                ..Default::default()
            },
            ModelRuntimeConfig {
                topology_domains: HashMap::from([("zone".to_string(), "us-east-1a".to_string())]),
                kv_transfer_domain: Some("zone".to_string()),
                kv_transfer_enforcement: Some(KvTransferEnforcement::Required),
                kv_transfer_preferred_weight: Some(0.5),
                ..Default::default()
            },
        ] {
            assert!(config.validate_config().is_err());
        }
    }
}