dynamo-llm 1.5.1

// 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, RouterHintWorkerMetadata},
    router_hint::{
        ROUTER_HINT_RUNTIME_CAPABILITY_KEY, ROUTER_HINT_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY,
        ROUTER_HINT_WORKER_TYPE_RUNTIME_KEY,
    },
};
use dynamo_runtime::{config::is_truthy, protocols::EndpointId};

use crate::protocols::openai::chat_completions::tool_parser_v2::unified_family_names;
use dynamo_parsers::tool_calling::parsers::get_available_tool_parsers;

/// 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/";

/// Runtime-data key for an engine-published token-overflow contract.
pub const TOKEN_BUDGET_RUNTIME_KEY: &str = "token_budget";

/// Runtime-data key indicating that a backend expects tool structural tags to
/// exclude reasoning and manages grammar activation around reasoning itself.
///
/// Absence means `false` for compatibility with workers that expect the
/// frontend's structural tag to model an already-opened reasoning block.
pub const TOOL_CALL_STRUCTURAL_TAG_EXCLUDES_REASONING_RUNTIME_KEY: &str =
    "tool_call_structural_tag_excludes_reasoning";

/// Describes which request-token overflows the frontend may reject early.
///
/// The combined limit already accounts for engine-reserved tokens. A false
/// flag delegates that overflow dimension to the backend, which remains
/// responsible for any clamping, truncation, or rejection.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub struct TokenBudget {
    pub combined_limit: u32,
    #[serde(default)]
    pub reject_prompt_overflow: bool,
    #[serde(default)]
    pub reject_total_overflow: bool,
}

/// 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";
pub const ENV_TOKENIZER_FALLBACK: &str = "DYN_TOKENIZER_FALLBACK";

/// Worker-advertised support for Dynamo's vLLM-compatible
/// `POST /inference/v1/generate` adapter.
///
/// This is deliberately a runtime capability rather than an inference from
/// `ModelType::Chat` / `ModelType::Completions`: other backends expose those
/// surfaces without implementing vLLM's Generate contract.
pub const VLLM_INFERENCE_V1_GENERATE_CAPABILITY: &str = "vllm_inference_v1_generate";

/// Worker-reported Qwen3 video prompt-expansion contract used by vLLM.
///
/// Absence disables exact video routing so a newer frontend remains safe with
/// older workers that predate this runtime contract.
pub const VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY: &str =
    "vllm_qwen_video_processor_contract";

/// Worker-reported vLLM setting that makes multimodal cache identities depend
/// on the active LoRA adapter. Missing and explicit `false` are equivalent.
pub const VLLM_ENABLE_TOWER_CONNECTOR_LORA_RUNTIME_KEY: &str = "vllm_enable_tower_connector_lora";

/// Worker-advertised support for Dynamo's SGLang-compatible `POST /generate`
/// adapter.
///
/// Keep this separate from [`VLLM_INFERENCE_V1_GENERATE_CAPABILITY`] so a
/// mixed-backend frontend never forwards one engine's opaque request envelope
/// to the other engine.
pub const SGLANG_GENERATE_CAPABILITY: &str = "sglang_generate";

/// 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,
    Basetenkenizer,
}

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

    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 == "basetenkenizer" => Self::Basetenkenizer,
            Ok(v) if v == "default" || v.is_empty() => Self::Default,
            Ok(v) => {
                tracing::warn!(
                    value = %v,
                    "Unrecognized DYN_TOKENIZER value, expected 'default', 'fastokens', or 'basetenkenizer'; 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),
            "basetenkenizer" => Ok(Self::Basetenkenizer),
            _ => Err(format!(
                "invalid tokenizer backend '{value}' (expected 'default', 'fastokens', or 'basetenkenizer')"
            )),
        }
    }
}

#[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>,
}

/// Controls how historical `function.arguments` are serialized before being
/// passed to the MiniJinja chat template.
#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolCallArgumentsFormat {
    /// Preserve arguments as a raw JSON string (default, backward-compatible).
    #[default]
    JsonString,
    /// Parse arguments into a JSON object before rendering.  Required for
    /// templates that iterate over key-value pairs (e.g. GLM-5.2).
    JsonObject,
}

#[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>,

    /// Physical KV-cache capacity for each router-visible data-parallel rank.
    /// This is per rank, never the aggregate capacity of the worker process.
    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>,

    /// Controls how historical `function.arguments` are presented to the MiniJinja
    /// chat template.  `JsonString` (default) preserves the raw JSON string, which
    /// is backward-compatible with all models.  `JsonObject` normalizes the string
    /// to a parsed object before rendering; required for models whose template
    /// iterates over argument key-value pairs (e.g. GLM-5.2).
    /// Also set to `JsonObject` automatically when `tool_call_parser` is `"glm47"`.
    #[serde(default)]
    pub tool_call_arguments_format: ToolCallArgumentsFormat,

    /// 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>,

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

    /// 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,

    /// Whether the running engine is configured to publish KV cache events.
    ///
    /// `None` indicates a legacy worker that does not declare this capability.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub kv_event_publishing_enabled: Option<bool>,

    /// Immutable KV event source mode for this worker lifecycle.
    ///
    /// Accepted values are `framework_v1` and `state_agent_v2`. Missing means the
    /// legacy Worker-only source. Unknown explicit values must disable KV-aware
    /// routing rather than falling back within the same worker lifecycle.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub kv_event_source_mode: Option<String>,

    /// Endpoint whose event sources describe this worker's KV state.
    ///
    /// When unset, consumers use the worker's serving endpoint. This keeps existing
    /// deployments wire-compatible while allowing KV-state ownership and request serving
    /// to be discovered independently.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub kv_state_endpoint: Option<EndpointId>,

    /// 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
}

pub(crate) 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,
            tool_call_arguments_format: ToolCallArgumentsFormat::JsonString,
            tokenizer_backend: None,
            tokenizer_fallback_enabled: 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,
            kv_event_publishing_enabled: None,
            kv_event_source_mode: None,
            kv_state_endpoint: None,
            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 ModelRuntimeConfig {
    /// Check whether a runtime boolean is explicitly enabled.
    ///
    /// Rust callers commonly store booleans, while compatibility cards may
    /// carry string-encoded flags. Both representations use Dynamo's canonical
    /// truthy vocabulary.
    pub(crate) fn runtime_flag_enabled(&self, key: &str) -> bool {
        match self.runtime_data.get(key) {
            Some(serde_json::Value::Bool(true)) => true,
            Some(serde_json::Value::String(value)) => is_truthy(value),
            _ => false,
        }
    }

    /// Check whether a runtime capability is explicitly enabled.
    pub(crate) fn supports_runtime_capability(&self, capability: &str) -> bool {
        self.runtime_flag_enabled(capability)
    }

    fn router_hints_enabled(&self) -> bool {
        self.supports_runtime_capability(ROUTER_HINT_RUNTIME_CAPABILITY_KEY)
    }

    fn router_hint_endpoint_for_dp_rank(&self, dp_rank: u32) -> Option<&str> {
        let dp_rank = dp_rank.to_string();
        let endpoint = self
            .runtime_data
            .get(ROUTER_HINT_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY)?
            .as_object()?
            .get(&dp_rank)?
            .as_str()?;
        (!endpoint.is_empty()).then_some(endpoint)
    }
}

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 router_hint_metadata_for_dp_rank(
        &self,
        dp_rank: u32,
    ) -> Option<RouterHintWorkerMetadata<'_>> {
        if !self.router_hints_enabled() {
            return None;
        }

        let worker_type = self
            .runtime_data
            .get(ROUTER_HINT_WORKER_TYPE_RUNTIME_KEY)?
            .as_str()?;
        if worker_type.is_empty() {
            return None;
        }

        Some(RouterHintWorkerMetadata {
            worker_type,
            source_control_endpoint: self.router_hint_endpoint_for_dp_rank(dp_rank),
        })
    }

    fn native_offloading_capacity_tokens(&self) -> Option<u64> {
        self.runtime_data
            .get("native_offloading_capacity")?
            .get("total_tokens")?
            .as_u64()
    }

    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(parser) = config
        .tool_call_parser
        .as_deref()
        .filter(|parser| !parser.is_empty())
    {
        let mut supported = get_available_tool_parsers();
        // Unified parser names live outside the v1 registry; normalize the union
        // for a stable error message.
        supported.extend_from_slice(unified_family_names());
        supported.sort_unstable();
        supported.dedup();
        if !supported.contains(&parser) {
            return Err(validation_error(
                "unsupported_tool_call_parser",
                format!(
                    "tool_call_parser '{parser}' is not supported; available parsers: {}",
                    supported.join(", ")
                ),
            ));
        }
    }

    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 is_tokenizer_fallback_enabled(&self) -> anyhow::Result<bool> {
        let enabled = if let Some(enabled) = self.tokenizer_fallback_enabled {
            enabled
        } else {
            match std::env::var(ENV_TOKENIZER_FALLBACK) {
                Ok(value) => dynamo_runtime::config::parse_bool(&value)
                    .map_err(|error| anyhow::anyhow!("{ENV_TOKENIZER_FALLBACK}: {error}"))?,
                Err(std::env::VarError::NotPresent) => true,
                Err(std::env::VarError::NotUnicode(_)) => {
                    anyhow::bail!("{ENV_TOKENIZER_FALLBACK} must contain valid UTF-8")
                }
            }
        };

        if enabled && self.effective_tokenizer_backend() != TokenizerBackend::Default {
            static TOKENIZER_FALLBACK_DEPRECATION_WARNED: std::sync::Once = std::sync::Once::new();
            TOKENIZER_FALLBACK_DEPRECATION_WARNED.call_once(|| {
                tracing::warn!(
                    "Automatic tokenizer fallback is deprecated and will be disabled by default in a future release. Set tokenizer fallback to false (`--no-tokenizer-fallback` in the frontend) to adopt the future behavior now."
                );
            });
        }

        Ok(enabled)
    }

    /// Resolve the KV-state endpoint, preserving the serving endpoint as the compatibility
    /// default for workers that do not advertise an explicit mapping.
    pub fn effective_kv_state_endpoint(&self, serving_endpoint: &EndpointId) -> EndpointId {
        self.kv_state_endpoint
            .clone()
            .unwrap_or_else(|| serving_endpoint.clone())
    }

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

    pub fn set_tokenizer_fallback_enabled(&mut self, enabled: Option<bool>) -> &mut Self {
        self.tokenizer_fallback_enabled = enabled;
        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::*;

    use crate::protocols::openai::chat_completions::tool_parser_v2::V2_FAMILIES;

    // 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 kv_state_endpoint_roundtrips_and_defaults_to_serving_endpoint() {
        let serving_endpoint = EndpointId::from("ns.worker.generate");
        let kv_state_endpoint = EndpointId::from("ns.kv.events");
        let cfg = ModelRuntimeConfig {
            kv_state_endpoint: Some(kv_state_endpoint.clone()),
            ..Default::default()
        };

        let json = serde_json::to_string(&cfg).unwrap();
        let parsed: ModelRuntimeConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.kv_state_endpoint, Some(kv_state_endpoint.clone()));
        assert_eq!(
            parsed.effective_kv_state_endpoint(&serving_endpoint),
            kv_state_endpoint
        );

        let legacy: ModelRuntimeConfig = serde_json::from_str("{}").unwrap();
        assert!(legacy.kv_state_endpoint.is_none());
        assert_eq!(
            legacy.effective_kv_state_endpoint(&serving_endpoint),
            serving_endpoint
        );
        assert!(
            !serde_json::to_string(&legacy)
                .unwrap()
                .contains("kv_state_endpoint")
        );
    }

    #[test]
    fn kv_event_publishing_capability_roundtrips_and_preserves_legacy_unknown() {
        for enabled in [true, false] {
            let cfg = ModelRuntimeConfig {
                kv_event_publishing_enabled: Some(enabled),
                ..Default::default()
            };
            let json = serde_json::to_string(&cfg).unwrap();
            let parsed: ModelRuntimeConfig = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed.kv_event_publishing_enabled, Some(enabled));
        }

        let legacy: ModelRuntimeConfig = serde_json::from_str("{}").unwrap();
        assert_eq!(legacy.kv_event_publishing_enabled, None);
        assert!(
            !serde_json::to_string(&legacy)
                .unwrap()
                .contains("kv_event_publishing_enabled")
        );
    }

    #[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("basetenkenizer"))], || {
            let cfg = ModelRuntimeConfig::default();
            assert_eq!(
                cfg.effective_tokenizer_backend(),
                TokenizerBackend::Basetenkenizer
            );
        });

        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
            );
        });

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

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

    #[test]
    #[serial_test::serial]
    fn tokenizer_fallback_env_default_and_explicit_config_precedence() {
        temp_env::with_vars([(ENV_TOKENIZER_FALLBACK, Some("false"))], || {
            let config = ModelRuntimeConfig::default();
            assert!(!config.is_tokenizer_fallback_enabled().unwrap());

            let config = ModelRuntimeConfig {
                tokenizer_fallback_enabled: Some(true),
                ..Default::default()
            };
            assert!(config.is_tokenizer_fallback_enabled().unwrap());
        });

        temp_env::with_vars([(ENV_TOKENIZER_FALLBACK, Some("true"))], || {
            let config = ModelRuntimeConfig {
                tokenizer_fallback_enabled: Some(false),
                ..Default::default()
            };
            assert!(!config.is_tokenizer_fallback_enabled().unwrap());
        });

        temp_env::with_vars_unset([ENV_TOKENIZER_FALLBACK], || {
            let config = ModelRuntimeConfig::default();
            assert!(config.is_tokenizer_fallback_enabled().unwrap());
        });

        temp_env::with_vars([(ENV_TOKENIZER_FALLBACK, Some("flase"))], || {
            let config = ModelRuntimeConfig::default();
            let error = config.is_tokenizer_fallback_enabled().unwrap_err();
            assert!(error.to_string().contains(ENV_TOKENIZER_FALLBACK));
        });
    }

    #[test]
    fn tokenizer_fallback_roundtrips_through_serde_json() {
        for enabled in [true, false] {
            let config = ModelRuntimeConfig {
                tokenizer_fallback_enabled: Some(enabled),
                ..Default::default()
            };
            let json = serde_json::to_string(&config).unwrap();
            assert!(json.contains(&format!("\"tokenizer_fallback_enabled\":{enabled}")));
            let parsed: ModelRuntimeConfig = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed.tokenizer_fallback_enabled, Some(enabled));
        }

        let json = serde_json::to_string(&ModelRuntimeConfig::default()).unwrap();
        assert!(!json.contains("tokenizer_fallback_enabled"));
        let legacy: ModelRuntimeConfig = serde_json::from_str("{}").unwrap();
        assert_eq!(legacy.tokenizer_fallback_enabled, None);
    }

    #[test]
    fn tokenizer_backend_string_values_are_strict() {
        for backend in [
            TokenizerBackend::Default,
            TokenizerBackend::Fastokens,
            TokenizerBackend::Basetenkenizer,
        ] {
            assert_eq!(backend.as_str().parse(), Ok(backend));
        }

        let error = "baseten".parse::<TokenizerBackend>().unwrap_err();
        assert!(error.contains("basetenkenizer"));
    }

    #[test]
    fn native_offloading_capacity_is_backend_neutral() {
        use dynamo_kv_router::WorkerConfigLike;

        let mut config = ModelRuntimeConfig::default();
        config
            .set_engine_specific(
                "native_offloading_capacity",
                serde_json::json!({"total_tokens": 300}),
            )
            .unwrap();

        assert_eq!(config.native_offloading_capacity_tokens(), Some(300));
    }

    #[test]
    fn router_hint_support_requires_explicit_true() {
        use dynamo_kv_router::WorkerConfigLike;

        let mut config = ModelRuntimeConfig::default();
        assert!(config.router_hint_metadata_for_dp_rank(0).is_none());

        config
            .set_engine_specific(ROUTER_HINT_RUNTIME_CAPABILITY_KEY, true)
            .unwrap();
        assert!(config.router_hint_metadata_for_dp_rank(0).is_none());

        config
            .set_engine_specific(ROUTER_HINT_WORKER_TYPE_RUNTIME_KEY, "prefill")
            .unwrap();
        let info = config.router_hint_metadata_for_dp_rank(0).unwrap();
        assert_eq!(info.worker_type, "prefill");
        assert!(info.source_control_endpoint.is_none());

        config
            .set_engine_specific(
                ROUTER_HINT_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY,
                serde_json::json!({"0": "tcp://127.0.0.1:23280"}),
            )
            .unwrap();
        let info = config.router_hint_metadata_for_dp_rank(0).unwrap();
        assert_eq!(info.worker_type, "prefill");
        assert_eq!(info.source_control_endpoint, Some("tcp://127.0.0.1:23280"));
        assert_eq!(
            config
                .router_hint_metadata_for_dp_rank(1)
                .unwrap()
                .source_control_endpoint,
            None
        );

        config
            .set_engine_specific(ROUTER_HINT_RUNTIME_CAPABILITY_KEY, "true")
            .unwrap();
        assert_eq!(
            config
                .router_hint_metadata_for_dp_rank(0)
                .unwrap()
                .worker_type,
            "prefill"
        );

        config
            .set_engine_specific(ROUTER_HINT_RUNTIME_CAPABILITY_KEY, "false")
            .unwrap();
        assert!(config.router_hint_metadata_for_dp_rank(0).is_none());
    }

    #[test]
    fn runtime_capability_support_accepts_boolean_and_string_truthy_values() {
        const CAPABILITY: &str = "test_capability";

        let mut config = ModelRuntimeConfig::default();
        assert!(!config.supports_runtime_capability(CAPABILITY));

        for enabled in [serde_json::json!(true), serde_json::json!(" yes ")] {
            config.runtime_data.insert(CAPABILITY.to_string(), enabled);
            assert!(config.supports_runtime_capability(CAPABILITY));
        }

        for disabled in [
            serde_json::json!(false),
            serde_json::json!("false"),
            serde_json::json!(1),
        ] {
            config.runtime_data.insert(CAPABILITY.to_string(), disabled);
            assert!(!config.supports_runtime_capability(CAPABILITY));
        }
    }

    #[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,
            "tool_call_arguments_format": "json_string"
        }"#;

        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());
        }
    }

    #[test]
    fn test_validate_config_checks_tool_call_parser() {
        let validate = |parser: &str| {
            ModelRuntimeConfig {
                tool_call_parser: Some(parser.to_string()),
                ..Default::default()
            }
            .validate_config()
        };

        // Every v2 or unified parser must remain valid at registration.
        for &parser in V2_FAMILIES.iter().chain(unified_family_names()) {
            assert!(validate(parser).is_ok(), "{parser} must be supported");
        }
        assert!(validate("").is_ok());

        let error = validate("not_registered").unwrap_err();
        assert!(error.contains("not_registered"));
        assert!(error.contains("muse_glimmer"));
    }
}