[model-gateway] Refine TokenizerRegisty.load() to handle duplication (#17230)

sgl-model-gateway/src/core/steps/mod.rs

@@ -22,7 +22,7 @@ pub use mcp_registration::{
 };
 pub use tokenizer_registration::{
     create_tokenizer_registration_workflow, create_tokenizer_workflow_data, LoadTokenizerStep,
-    TokenizerConfigRequest, TokenizerRemovalRequest, ValidateTokenizerConfigStep,
+    TokenizerConfigRequest, TokenizerRemovalRequest,
 };
 pub use wasm_module_registration::{
     create_wasm_module_registration_workflow, create_wasm_registration_workflow_data,

sgl-model-gateway/src/core/steps/tokenizer_registration.rs

@@ -11,7 +11,7 @@ use std::{sync::Arc, time::Duration};
 
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
-use tracing::{debug, error, info};
+use tracing::{error, info};
 
 use super::workflow_data::TokenizerWorkflowData;
 use crate::{
@@ -20,6 +20,7 @@ use crate::{
     tokenizer::{
         cache::{CacheConfig, CachedTokenizer},
         factory,
+        registry::LoadOutcome,
         traits::Tokenizer,
     },
     workflow::{
@@ -42,6 +43,11 @@ pub struct TokenizerConfigRequest {
     /// Optional cache configuration. If provided, wraps tokenizer with CachedTokenizer.
     #[serde(default)]
     pub cache_config: Option<TokenizerCacheConfig>,
+    /// If true, the workflow fails when a tokenizer with the same name already exists.
+    /// If false (default), the workflow succeeds and returns the existing tokenizer's ID.
+    /// API callers should set this to true.
+    #[serde(default)]
+    pub fail_on_duplicate: bool,
 }
 
 /// Configuration for removing a tokenizer
@@ -55,61 +61,13 @@ pub struct TokenizerRemovalRequest {
 // Workflow Steps
 // ============================================================================
 
-/// Step 1: Validate the tokenizer configuration
+/// Load the tokenizer from source (local path or HuggingFace)
-pub struct ValidateTokenizerConfigStep;
+///
-
+/// This step handles:
-#[async_trait]
+/// - Input validation (via registry.load())
-impl StepExecutor<TokenizerWorkflowData> for ValidateTokenizerConfigStep {
+/// - Deduplication (returns success if already exists)
-    async fn execute(
+/// - Loading from local path or HuggingFace
-        &self,
+/// - Optional caching layer wrapping
-        context: &mut WorkflowContext<TokenizerWorkflowData>,
-    ) -> WorkflowResult<StepResult> {
-        let config = &context.data.config;
-        let app_context = context
-            .data
-            .app_context
-            .as_ref()
-            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
-
-        debug!(
-            "Validating tokenizer config: name={}, source={}",
-            config.name, config.source
-        );
-
-        // Validate name is not empty
-        if config.name.is_empty() {
-            return Err(WorkflowError::StepFailed {
-                step_id: StepId::new("validate_config"),
-                message: "Tokenizer name cannot be empty".to_string(),
-            });
-        }
-
-        // Validate source is not empty
-        if config.source.is_empty() {
-            return Err(WorkflowError::StepFailed {
-                step_id: StepId::new("validate_config"),
-                message: "Tokenizer source cannot be empty".to_string(),
-            });
-        }
-
-        // Check if tokenizer already exists
-        if app_context.tokenizer_registry.contains(&config.name) {
-            return Err(WorkflowError::StepFailed {
-                step_id: StepId::new("validate_config"),
-                message: format!("Tokenizer '{}' already exists", config.name),
-            });
-        }
-
-        debug!("Tokenizer config validated successfully");
-        Ok(StepResult::Success)
-    }
-
-    fn is_retryable(&self, _error: &WorkflowError) -> bool {
-        false // Validation errors are not retryable
-    }
-}
-
-/// Step 2: Load the tokenizer from source (local path or HuggingFace)
 pub struct LoadTokenizerStep;
 
 #[async_trait]
@@ -145,7 +103,8 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
         let chat_template = config.chat_template_path.clone();
         let cache_config = config.cache_config.clone();
 
-        // Load the tokenizer using the registry's load method (handles deduplication)
+        // Load the tokenizer using the registry's load method
+        // This handles: validation, deduplication, and loading
         let result = app_context
             .tokenizer_registry
             .load(&id, &name, &source, || {
@@ -181,17 +140,38 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
             .await;
 
         match result {
-            Ok(loaded_id) => {
+            Ok(outcome) => {
+                let loaded_id = outcome.id();
+
                 // Get vocab size for logging
                 let vocab_size = app_context
                     .tokenizer_registry
-                    .get_by_id(&loaded_id)
+                    .get_by_id(loaded_id)
                     .map(|e| e.tokenizer.vocab_size());
 
-                info!(
+                match &outcome {
-                    "Successfully loaded tokenizer '{}' (id: {}) with vocab_size: {:?}",
+                    LoadOutcome::Loaded { id } => {
-                    name, loaded_id, vocab_size
+                        info!(
-                );
+                            "Successfully loaded tokenizer '{}' (id: {}) with vocab_size: {:?}",
+                            name, id, vocab_size
+                        );
+                    }
+                    LoadOutcome::AlreadyExists { id } => {
+                        if config.fail_on_duplicate {
+                            return Err(WorkflowError::StepFailed {
+                                step_id: StepId::new("load_tokenizer"),
+                                message: format!(
+                                    "Tokenizer '{}' already exists (id: {})",
+                                    name, id
+                                ),
+                            });
+                        }
+                        info!(
+                            "Tokenizer '{}' already exists (id: {}), skipping load",
+                            name, id
+                        );
+                    }
+                }
 
                 // Store vocab size in typed data
                 if let Some(size) = vocab_size {
@@ -204,7 +184,7 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
                 error!("Failed to load tokenizer '{}': {}", name, e);
                 Err(WorkflowError::StepFailed {
                     step_id: StepId::new("load_tokenizer"),
-                    message: e,
+                    message: e.to_string(),
                 })
             }
         }
@@ -221,38 +201,29 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
 
 /// Create the tokenizer registration workflow
 ///
-/// This workflow:
+/// This workflow loads and registers a tokenizer. The single LoadTokenizerStep handles:
-/// - Validates the tokenizer configuration
+/// - Input validation (empty name/source)
-/// - Loads the tokenizer from local path or HuggingFace
+/// - Deduplication (returns success if already exists)
+/// - Loading from local path or HuggingFace
+/// - Optional caching layer wrapping
 ///
-/// Workflow configuration:
+/// Configuration:
-/// - ValidateConfig: No retry, 5s timeout (fast validation)
+/// - 3 retries with 2s backoff (for network issues)
-/// - LoadTokenizer: 3 retries, 5min timeout (may need to download from HuggingFace)
+/// - 5 minute timeout (HuggingFace downloads can be slow)
 pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition<TokenizerWorkflowData> {
-    WorkflowDefinition::new("tokenizer_registration", "Tokenizer Registration")
+    WorkflowDefinition::new("tokenizer_registration", "Tokenizer Registration").add_step(
-        .add_step(
+        StepDefinition::new(
-            StepDefinition::new(
+            "load_tokenizer",
-                "validate_config",
+            "Load Tokenizer",
-                "Validate Configuration",
+            Arc::new(LoadTokenizerStep),
-                Arc::new(ValidateTokenizerConfigStep),
-            )
-            .with_timeout(Duration::from_secs(5))
-            .with_failure_action(FailureAction::FailWorkflow),
-        )
-        .add_step(
-            StepDefinition::new(
-                "load_tokenizer",
-                "Load Tokenizer",
-                Arc::new(LoadTokenizerStep),
-            )
-            .with_retry(RetryPolicy {
-                max_attempts: 3,
-                backoff: BackoffStrategy::Fixed(Duration::from_secs(2)),
-            })
-            .with_timeout(Duration::from_secs(300)) // 5 min for HuggingFace downloads
-            .with_failure_action(FailureAction::FailWorkflow)
-            .depends_on(&["validate_config"]),
         )
+        .with_retry(RetryPolicy {
+            max_attempts: 3,
+            backoff: BackoffStrategy::Fixed(Duration::from_secs(2)),
+        })
+        .with_timeout(Duration::from_secs(300)) // 5 min for HuggingFace downloads
+        .with_failure_action(FailureAction::FailWorkflow),
+    )
 }
 
 /// Helper to create initial workflow data for tokenizer registration
@@ -279,6 +250,7 @@ mod tests {
             source: "meta-llama/Llama-2-7b-hf".to_string(),
             chat_template_path: None,
             cache_config: None,
+            fail_on_duplicate: false,
         };
 
         let json = serde_json::to_string(&config).unwrap();
@@ -289,6 +261,35 @@ mod tests {
         assert_eq!(parsed.source, "meta-llama/Llama-2-7b-hf");
         assert!(parsed.chat_template_path.is_none());
         assert!(parsed.cache_config.is_none());
+        assert!(!parsed.fail_on_duplicate);
+    }
+
+    #[test]
+    fn test_tokenizer_config_request_fail_on_duplicate_defaults_to_false() {
+        // Test that fail_on_duplicate defaults to false when not specified in JSON
+        let json = r#"{
+            "id": "test-uuid",
+            "name": "test-model",
+            "source": "/path/to/tokenizer"
+        }"#;
+        let parsed: TokenizerConfigRequest = serde_json::from_str(json).unwrap();
+        assert!(!parsed.fail_on_duplicate);
+    }
+
+    #[test]
+    fn test_tokenizer_config_request_fail_on_duplicate_true() {
+        let config = TokenizerConfigRequest {
+            id: "test-uuid-1234".to_string(),
+            name: "test-model".to_string(),
+            source: "meta-llama/Llama-2-7b-hf".to_string(),
+            chat_template_path: None,
+            cache_config: None,
+            fail_on_duplicate: true,
+        };
+
+        let json = serde_json::to_string(&config).unwrap();
+        let parsed: TokenizerConfigRequest = serde_json::from_str(&json).unwrap();
+        assert!(parsed.fail_on_duplicate);
     }
 
     #[test]
@@ -304,6 +305,7 @@ mod tests {
                 enable_l1: false,
                 l1_max_memory: 0,
             }),
+            fail_on_duplicate: false,
         };
 
         let json = serde_json::to_string(&config).unwrap();

sgl-model-gateway/src/core/steps/worker/local/mod.rs

@@ -4,9 +4,9 @@ mod discover_dp;
 mod discover_metadata;
 mod find_worker_to_update;
 mod find_workers_to_remove;
-mod register_tokenizer;
 mod remove_from_policy_registry;
 mod remove_from_worker_registry;
+mod submit_tokenizer_job;
 mod update_policies_for_worker;
 mod update_remaining_policies;
 mod update_worker_properties;
@@ -27,9 +27,9 @@ pub use discover_dp::{get_dp_info, DiscoverDPInfoStep, DpInfo};
 pub use discover_metadata::DiscoverMetadataStep;
 pub use find_worker_to_update::FindWorkerToUpdateStep;
 pub use find_workers_to_remove::{FindWorkersToRemoveStep, WorkerRemovalRequest};
-pub use register_tokenizer::SubmitTokenizerJobStep;
 pub use remove_from_policy_registry::RemoveFromPolicyRegistryStep;
 pub use remove_from_worker_registry::RemoveFromWorkerRegistryStep;
+pub use submit_tokenizer_job::SubmitTokenizerJobStep;
 pub use update_policies_for_worker::UpdatePoliciesForWorkerStep;
 pub use update_remaining_policies::UpdateRemainingPoliciesStep;
 pub use update_worker_properties::UpdateWorkerPropertiesStep;

sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs

@@ -1,8 +1,8 @@
 //! Tokenizer registration step for local workers.
 //!
 //! This step submits a Job::AddTokenizer to the job queue, which triggers the
-//! tokenizer_registration workflow. This ensures all tokenizer registrations
+//! tokenizer_registration workflow. The workflow handles validation, deduplication,
-//! go through the same workflow with consistent behavior (validation, caching).
+//! and caching - this step just submits the job.
 
 use async_trait::async_trait;
 use tracing::{debug, info, warn};
@@ -93,14 +93,9 @@ impl StepExecutor<LocalWorkerWorkflowData> for SubmitTokenizerJobStep {
                 continue;
             };
 
-            // Check if tokenizer already exists for this model
+            // Note: We don't check if tokenizer already exists here.
-            if app_context.tokenizer_registry.contains(&model_id) {
+            // The registry.load() handles deduplication gracefully (returns AlreadyExists).
-                debug!(
+            // This simplifies the code and ensures consistent behavior.
-                    "Tokenizer already registered for model {}, skipping",
-                    model_id
-                );
-                continue;
-            }
 
             info!(
                 "Submitting tokenizer registration job for model {} from {}",
@@ -114,6 +109,7 @@ impl StepExecutor<LocalWorkerWorkflowData> for SubmitTokenizerJobStep {
                 source: tokenizer_path,
                 chat_template_path: chat_template.clone(),
                 cache_config: cache_config.clone(),
+                fail_on_duplicate: false,
             };
 
             // Submit job (fire-and-forget, don't wait for completion)

sgl-model-gateway/src/routers/tokenize/handlers.rs

@@ -235,6 +235,7 @@ pub async fn add_tokenizer(context: &Arc<AppContext>, request: AddTokenizerReque
         source: request.source.clone(),
         chat_template_path: request.chat_template_path.clone(),
         cache_config: None,
+        fail_on_duplicate: true,
     };
 
     let job = Job::AddTokenizer {

sgl-model-gateway/src/server.rs

@@ -852,6 +852,7 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
             source: tokenizer_source.clone(),
             chat_template_path: config.router_config.chat_template.clone(),
             cache_config: config.router_config.tokenizer_cache.to_option(),
+            fail_on_duplicate: false,
         };
 
         let job = Job::AddTokenizer {

sgl-model-gateway/src/tokenizer/mod.rs

@@ -27,7 +27,7 @@ mod tests;
 use factory::{create_tokenizer_from_file, create_tokenizer_with_chat_template};
 // Re-export types used outside this module
 pub use huggingface::HuggingFaceTokenizer;
-pub use registry::TokenizerRegistry;
+pub use registry::{LoadError, LoadOutcome, TokenizerRegistry};
 pub use stop::StopSequenceDecoder;
 pub use stream::DecodeStream;
 pub use traits::{Decoder, Encoder, Encoding, SpecialTokens};

sgl-model-gateway/src/tokenizer/registry.rs

@@ -19,12 +19,51 @@
 use std::sync::Arc;
 
 use dashmap::DashMap;
+use thiserror::Error;
 use tokio::sync::Mutex;
 use tracing::{debug, info};
 use uuid::Uuid;
 
 use super::traits::Tokenizer;
 
+/// Outcome of a tokenizer load operation
+#[derive(Debug, Clone)]
+pub enum LoadOutcome {
+    /// Tokenizer was newly loaded and registered
+    Loaded { id: String },
+    /// Tokenizer already existed, returning existing ID
+    AlreadyExists { id: String },
+}
+
+impl LoadOutcome {
+    /// Get the ID regardless of outcome
+    pub fn id(&self) -> &str {
+        match self {
+            LoadOutcome::Loaded { id } => id,
+            LoadOutcome::AlreadyExists { id } => id,
+        }
+    }
+
+    /// Returns true if the tokenizer was newly loaded
+    pub fn is_newly_loaded(&self) -> bool {
+        matches!(self, LoadOutcome::Loaded { .. })
+    }
+}
+
+/// Error type for tokenizer loading operations
+#[derive(Debug, Error)]
+pub enum LoadError {
+    /// Name cannot be empty
+    #[error("tokenizer name cannot be empty")]
+    EmptyName,
+    /// Source cannot be empty
+    #[error("tokenizer source cannot be empty")]
+    EmptySource,
+    /// Loading failed
+    #[error("{0}")]
+    LoadFailed(String),
+}
+
 /// Metadata and tokenizer instance for a registered tokenizer
 #[derive(Clone)]
 pub struct TokenizerEntry {
@@ -53,6 +92,19 @@ pub struct TokenizerRegistry {
     loading_locks: DashMap<String, Arc<Mutex<()>>>,
 }
 
+/// RAII guard that removes the loading lock entry on drop.
+/// Ensures cleanup happens on normal completion, early return, or panic.
+struct LoadingLockGuard<'a> {
+    locks: &'a DashMap<String, Arc<Mutex<()>>>,
+    key: String,
+}
+
+impl Drop for LoadingLockGuard<'_> {
+    fn drop(&mut self) {
+        self.locks.remove(&self.key);
+    }
+}
+
 impl TokenizerRegistry {
     /// Create a new empty registry
     pub fn new() -> Self {
@@ -68,36 +120,46 @@ impl TokenizerRegistry {
         Uuid::new_v4().to_string()
     }
 
-    /// Load and register a tokenizer with a pre-generated ID
+    /// Load and register a tokenizer
     ///
-    /// If the tokenizer is already loaded (by name), returns the existing ID.
+    /// Validates inputs, handles deduplication, and loads the tokenizer if needed.
-    /// Otherwise, uses the provided loader function to load it.
     /// Per-key locking ensures only one load happens per name, preventing race conditions.
     ///
     /// # Arguments
-    /// * `id` - Pre-generated UUID for this tokenizer
+    /// * `id` - Pre-generated UUID (use `generate_id()` to create one)
-    /// * `name` - User-provided name
+    /// * `name` - User-provided name (used for deduplication, must not be empty)
-    /// * `source` - Source path or HuggingFace model ID
+    /// * `source` - Source path or HuggingFace model ID (must not be empty)
     /// * `loader` - Async function that loads the tokenizer
     ///
     /// # Returns
-    /// * `Ok(id)` - Successfully loaded and registered (returns the ID)
+    /// * `Ok(LoadOutcome::Loaded { id })` - Tokenizer was newly loaded
-    /// * `Err(message)` - Error message if loading fails
+    /// * `Ok(LoadOutcome::AlreadyExists { id })` - Tokenizer already existed
+    /// * `Err(LoadError)` - Validation failed or loading failed
     pub async fn load<F, Fut>(
         &self,
         id: &str,
         name: &str,
         source: &str,
         loader: F,
-    ) -> Result<String, String>
+    ) -> Result<LoadOutcome, LoadError>
     where
         F: FnOnce() -> Fut,
         Fut: std::future::Future<Output = Result<Arc<dyn Tokenizer>, String>>,
     {
+        // Validate inputs
+        if name.is_empty() {
+            return Err(LoadError::EmptyName);
+        }
+        if source.is_empty() {
+            return Err(LoadError::EmptySource);
+        }
+
         // Fast path: already loaded by name
         if let Some(existing_id) = self.name_to_id.get(name) {
             debug!("Tokenizer already registered for name: {}", name);
-            return Ok(existing_id.clone());
+            return Ok(LoadOutcome::AlreadyExists {
+                id: existing_id.clone(),
+            });
         }
 
         debug!("Tokenizer cache miss for name: {}", name);
@@ -109,24 +171,27 @@ impl TokenizerRegistry {
             .or_insert_with(|| Arc::new(Mutex::new(())))
             .clone();
 
-        let _guard = lock.lock().await;
+        let _mutex_guard = lock.lock().await;
+        let _lock_cleanup = LoadingLockGuard {
+            locks: &self.loading_locks,
+            key: name.to_string(),
+        };
 
         // Double-check after acquiring lock (another thread may have loaded it)
         if let Some(existing_id) = self.name_to_id.get(name) {
             debug!("Tokenizer loaded by another thread for name: {}", name);
-            return Ok(existing_id.clone());
+            return Ok(LoadOutcome::AlreadyExists {
+                id: existing_id.clone(),
+            });
         }
 
         // Load tokenizer
         info!("Loading tokenizer '{}' from source: {}", name, source);
         let result = loader().await;
 
-        // Always clean up the lock, whether loading succeeded or failed
+        let tokenizer = result.map_err(LoadError::LoadFailed)?;
-        self.loading_locks.remove(name);
 
-        let tokenizer = result?;
+        // Create entry with provided ID
-
-        // Create entry
         let entry = TokenizerEntry {
             id: id.to_string(),
             name: name.to_string(),
@@ -143,18 +208,21 @@ impl TokenizerRegistry {
             name, id
         );
 
-        Ok(id.to_string())
+        Ok(LoadOutcome::Loaded { id: id.to_string() })
     }
 
-    /// Register a pre-loaded tokenizer with a pre-generated ID
+    /// Register a preloaded tokenizer with a pre-generated ID
     ///
     /// Atomically inserts a tokenizer into the registry only if no tokenizer
     /// with the same name exists. Returns the ID if successful.
     ///
+    /// This is primarily used for testing. Production code should use `load()`.
+    ///
     /// # Returns
     /// * `Some(id)` - If the tokenizer was successfully registered
     /// * `None` - If a tokenizer with this name already existed
-    pub fn register(
+    #[cfg(test)]
+    pub(crate) fn register(
         &self,
         id: &str,
         name: &str,
@@ -292,13 +360,17 @@ mod tests {
 
         // Load and register a tokenizer
         let id = TokenizerRegistry::generate_id();
-        registry
+        let outcome = registry
             .load(&id, "model1", "path/to/model", || async {
                 Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
             })
             .await
             .unwrap();
 
+        // Verify LoadOutcome::Loaded
+        assert!(outcome.is_newly_loaded());
+        assert_eq!(outcome.id(), id);
+
         // Verify it's loaded
         assert!(!registry.is_empty());
         assert_eq!(registry.len(), 1);
@@ -317,6 +389,65 @@ mod tests {
         assert!(registry.is_empty());
     }
 
+    #[tokio::test]
+    async fn test_load_returns_already_exists() {
+        let registry = TokenizerRegistry::new();
+        let id1 = TokenizerRegistry::generate_id();
+        let id2 = TokenizerRegistry::generate_id();
+
+        // First load should return Loaded
+        let outcome1 = registry
+            .load(&id1, "model1", "source1", || async {
+                Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
+            })
+            .await
+            .unwrap();
+        assert!(outcome1.is_newly_loaded());
+        assert_eq!(outcome1.id(), id1);
+
+        // Second load with same name should return AlreadyExists with ORIGINAL id
+        let outcome2 = registry
+            .load(&id2, "model1", "source2", || async {
+                panic!("Loader should not be called for duplicate name");
+            })
+            .await
+            .unwrap();
+        assert!(!outcome2.is_newly_loaded());
+        assert_eq!(outcome2.id(), id1); // Returns the original ID, not id2
+
+        // Registry still has only one entry
+        assert_eq!(registry.len(), 1);
+
+        // Original source is preserved
+        let entry = registry.get_by_name("model1").unwrap();
+        assert_eq!(entry.source, "source1");
+    }
+
+    #[tokio::test]
+    async fn test_load_validation() {
+        let registry = TokenizerRegistry::new();
+        let id = TokenizerRegistry::generate_id();
+
+        // Empty name should fail
+        let result = registry
+            .load(&id, "", "source", || async {
+                panic!("Loader should not be called for invalid input");
+            })
+            .await;
+        assert!(matches!(result, Err(LoadError::EmptyName)));
+
+        // Empty source should fail
+        let result = registry
+            .load(&id, "model", "", || async {
+                panic!("Loader should not be called for invalid input");
+            })
+            .await;
+        assert!(matches!(result, Err(LoadError::EmptySource)));
+
+        // Registry should be empty (nothing was loaded)
+        assert!(registry.is_empty());
+    }
+
     #[tokio::test]
     async fn test_load_prevents_duplicate_loading() {
         let registry = Arc::new(TokenizerRegistry::new());
@@ -459,4 +590,88 @@ mod tests {
         assert!(result3.is_some());
         assert_eq!(registry.len(), 2);
     }
+
+    #[tokio::test]
+    async fn test_loading_lock_cleanup_on_panic() {
+        let registry = Arc::new(TokenizerRegistry::new());
+
+        // Spawn a task that will panic during loading
+        let registry_clone = registry.clone();
+        let handle = tokio::spawn(async move {
+            registry_clone
+                .load(
+                    &TokenizerRegistry::generate_id(),
+                    "panic-model",
+                    "source",
+                    || async {
+                        panic!("Simulated panic during tokenizer loading");
+                    },
+                )
+                .await
+        });
+
+        // Wait for the task - it should panic
+        let result = handle.await;
+        assert!(result.is_err(), "Task should have panicked");
+
+        // The RAII guard should have cleaned up the loading lock.
+        // Verify by attempting another load with the same name - it should work,
+        // not deadlock or fail due to stale lock.
+        let id = TokenizerRegistry::generate_id();
+        let outcome = registry
+            .load(&id, "panic-model", "source", || async {
+                Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
+            })
+            .await;
+
+        // Should succeed - the lock was properly cleaned up
+        assert!(outcome.is_ok(), "Load should succeed after panic cleanup");
+        assert!(outcome.unwrap().is_newly_loaded());
+        assert_eq!(registry.len(), 1);
+        assert!(registry.contains("panic-model"));
+    }
+
+    #[tokio::test]
+    async fn test_loading_lock_cleanup_on_early_return() {
+        let registry = Arc::new(TokenizerRegistry::new());
+
+        // Load a tokenizer
+        let id1 = TokenizerRegistry::generate_id();
+        registry
+            .load(&id1, "model1", "source1", || async {
+                Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
+            })
+            .await
+            .unwrap();
+
+        // Now simulate concurrent load attempts where one thread wins
+        // and another thread sees "already exists" in the double-check.
+        // The RAII guard should clean up the lock on early return.
+
+        // First, verify loading_locks is empty after successful load
+        // by checking that we can load a different model without issues
+        let id2 = TokenizerRegistry::generate_id();
+        let outcome = registry
+            .load(&id2, "model2", "source2", || async {
+                Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
+            })
+            .await
+            .unwrap();
+
+        assert!(outcome.is_newly_loaded());
+        assert_eq!(registry.len(), 2);
+
+        // Try to load model1 again - should return AlreadyExists
+        // and the lock should be cleaned up (not leak)
+        let id3 = TokenizerRegistry::generate_id();
+        let outcome = registry
+            .load(&id3, "model1", "source1", || async {
+                panic!("Loader should not be called for existing model");
+            })
+            .await
+            .unwrap();
+
+        assert!(!outcome.is_newly_loaded());
+        assert_eq!(outcome.id(), id1); // Returns original ID
+    }
 }