feat: add wasm restart metadata handling

echobt · echobt · commit 90beef13d4cd · 2026-02-14T22:31:58.000Z
Add restart metadata to WASM configs and registry entries.
Emit restart lifecycle events and track config versions.
Bump state versioning and schema registry defaults.
diff --git a/crates/challenge-registry/src/lifecycle.rs b/crates/challenge-registry/src/lifecycle.rs
@@ -46,6 +46,14 @@ pub enum LifecycleEvent {
         old_version: ChallengeVersion,
         new_version: ChallengeVersion,
     },
+    /// Challenge restart configuration changed
+    Restarted {
+        challenge_id: ChallengeId,
+        previous_restart_id: Option<String>,
+        new_restart_id: Option<String>,
+        previous_config_version: u64,
+        new_config_version: u64,
+    },
 }
 
 /// Manages challenge lifecycle transitions
@@ -112,7 +120,25 @@ impl ChallengeLifecycle {
         self.auto_restart
     }
 
-    /// Get max restart attempts
+    /// Check if restart configuration should trigger a restart
+    pub fn restart_required(
+        &self,
+        previous_restart_id: Option<&str>,
+        new_restart_id: Option<&str>,
+        previous_config_version: u64,
+        new_config_version: u64,
+    ) -> bool {
+        if previous_config_version != new_config_version {
+            return true;
+        }
+
+        match (previous_restart_id, new_restart_id) {
+            (Some(prev), Some(next)) => prev != next,
+            (None, Some(_)) => true,
+            (Some(_), None) => true,
+            (None, None) => false,
+        }
+    }
     pub fn max_restart_attempts(&self) -> u32 {
         self.max_restart_attempts
     }
@@ -307,6 +333,43 @@ mod tests {
             }
             _ => panic!("Expected VersionChanged event"),
         }
+
+        // Test Restarted event
+        let restarted_event = LifecycleEvent::Restarted {
+            challenge_id,
+            previous_restart_id: Some("old".to_string()),
+            new_restart_id: Some("new".to_string()),
+            previous_config_version: 1,
+            new_config_version: 2,
+        };
+        match restarted_event {
+            LifecycleEvent::Restarted {
+                challenge_id: id,
+                previous_restart_id,
+                new_restart_id,
+                previous_config_version,
+                new_config_version,
+            } => {
+                assert_eq!(id, challenge_id);
+                assert_eq!(previous_restart_id, Some("old".to_string()));
+                assert_eq!(new_restart_id, Some("new".to_string()));
+                assert_eq!(previous_config_version, 1);
+                assert_eq!(new_config_version, 2);
+            }
+            _ => panic!("Expected Restarted event"),
+        }
+    }
+
+    #[test]
+    fn test_restart_required() {
+        let lifecycle = ChallengeLifecycle::new();
+
+        assert!(lifecycle.restart_required(Some("a"), Some("b"), 0, 0));
+        assert!(lifecycle.restart_required(None, Some("b"), 0, 0));
+        assert!(lifecycle.restart_required(Some("a"), None, 0, 0));
+        assert!(!lifecycle.restart_required(None, None, 0, 0));
+        assert!(lifecycle.restart_required(Some("a"), Some("a"), 1, 2));
+        assert!(!lifecycle.restart_required(Some("a"), Some("a"), 2, 2));
     }
 
     #[test]
diff --git a/crates/challenge-registry/src/registry.rs b/crates/challenge-registry/src/registry.rs
@@ -25,6 +25,12 @@ pub struct WasmModuleMetadata {
     /// Network policy for WASM execution
     #[serde(default)]
     pub network_policy: NetworkPolicy,
+    /// Restartable configuration identifier
+    #[serde(default)]
+    pub restart_id: Option<String>,
+    /// Configuration version for hot-restarts
+    #[serde(default)]
+    pub config_version: u64,
 }
 
 impl WasmModuleMetadata {
@@ -39,6 +45,8 @@ impl WasmModuleMetadata {
             module_location,
             entrypoint,
             network_policy,
+            restart_id: None,
+            config_version: 0,
         }
     }
 }
@@ -59,6 +67,12 @@ pub struct ChallengeEntry {
     /// WASM module metadata
     #[serde(default)]
     pub wasm_module: Option<WasmModuleMetadata>,
+    /// Restartable configuration identifier
+    #[serde(default)]
+    pub restart_id: Option<String>,
+    /// Configuration version for hot-restarts
+    #[serde(default)]
+    pub config_version: u64,
     /// Current lifecycle state
     pub lifecycle_state: LifecycleState,
     /// Health status
@@ -81,6 +95,8 @@ impl ChallengeEntry {
             docker_image,
             endpoint: None,
             wasm_module: None,
+            restart_id: None,
+            config_version: 0,
             lifecycle_state: LifecycleState::Registered,
             health_status: HealthStatus::Unknown,
             registered_at: now,
@@ -297,6 +313,58 @@ impl ChallengeRegistry {
         Ok(old_version)
     }
 
+    /// Update restart configuration metadata
+    pub fn update_restart_config(
+        &self,
+        id: &ChallengeId,
+        restart_id: Option<String>,
+        config_version: u64,
+    ) -> RegistryResult<(Option<String>, u64)> {
+        let mut challenges = self.challenges.write();
+        let registered = challenges
+            .get_mut(id)
+            .ok_or_else(|| RegistryError::ChallengeNotFound(id.to_string()))?;
+
+        let previous_restart_id = registered.entry.restart_id.clone();
+        let previous_config_version = registered.entry.config_version;
+
+        let restart_required = self.lifecycle.restart_required(
+            previous_restart_id.as_deref(),
+            restart_id.as_deref(),
+            previous_config_version,
+            config_version,
+        );
+
+        registered.entry.restart_id = restart_id.clone();
+        registered.entry.config_version = config_version;
+        registered.entry.updated_at = chrono::Utc::now().timestamp_millis();
+
+        if let Some(wasm_module) = registered.entry.wasm_module.as_mut() {
+            wasm_module.restart_id = restart_id.clone();
+            wasm_module.config_version = config_version;
+        }
+
+        if restart_required {
+            info!(
+                challenge_id = %id,
+                previous_restart_id = ?previous_restart_id,
+                new_restart_id = ?restart_id,
+                previous_config_version = previous_config_version,
+                new_config_version = config_version,
+                "Challenge restart configuration updated"
+            );
+            self.emit_event(LifecycleEvent::Restarted {
+                challenge_id: *id,
+                previous_restart_id: previous_restart_id.clone(),
+                new_restart_id: restart_id,
+                previous_config_version,
+                new_config_version: config_version,
+            });
+        }
+
+        Ok((previous_restart_id, previous_config_version))
+    }
+
     /// Get state store for a challenge
     pub fn state_store(&self, id: &ChallengeId) -> Option<Arc<StateStore>> {
         self.challenges
@@ -345,7 +413,6 @@ impl Default for ChallengeRegistry {
 #[cfg(test)]
 mod tests {
     use super::*;
-
     #[test]
     fn test_register_challenge() {
         let registry = ChallengeRegistry::new();
@@ -446,11 +513,40 @@ mod tests {
         assert_eq!(challenge.entry.version, ChallengeVersion::new(1, 1, 0));
     }
 
+    #[test]
+    fn test_update_restart_config() {
+        let registry = ChallengeRegistry::new();
+        let entry = ChallengeEntry::new(
+            "test".to_string(),
+            ChallengeVersion::new(1, 0, 0),
+            "test:latest".to_string(),
+        )
+        .with_wasm_module(WasmModuleMetadata::new(
+            "hash".to_string(),
+            "module.wasm".to_string(),
+            "evaluate".to_string(),
+            NetworkPolicy::default(),
+        ));
+
+        let id = registry.register(entry).unwrap();
+        let previous = registry
+            .update_restart_config(&id, Some("restart-1".to_string()), 1)
+            .unwrap();
+
+        assert_eq!(previous, (None, 0));
+
+        let challenge = registry.get(&id).unwrap();
+        assert_eq!(challenge.entry.restart_id, Some("restart-1".to_string()));
+        assert_eq!(challenge.entry.config_version, 1);
+        let wasm_module = challenge.entry.wasm_module.unwrap();
+        assert_eq!(wasm_module.restart_id, Some("restart-1".to_string()));
+        assert_eq!(wasm_module.config_version, 1);
+    }
+
     #[test]
     fn test_list_active() {
         let registry = ChallengeRegistry::new();
 
-        // Register two challenges
         let entry1 = ChallengeEntry::new(
             "active".to_string(),
             ChallengeVersion::new(1, 0, 0),
@@ -465,7 +561,6 @@ mod tests {
         let id1 = registry.register(entry1).unwrap();
         registry.register(entry2).unwrap();
 
-        // Make first one active
         registry
             .update_state(&id1, LifecycleState::Running)
             .unwrap();
diff --git a/crates/core/src/challenge.rs b/crates/core/src/challenge.rs
@@ -86,6 +86,7 @@ impl Challenge {
 
 /// Challenge configuration
 #[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(default)]
 pub struct ChallengeConfig {
     /// Mechanism ID on Bittensor (1, 2, 3... - 0 is reserved for default)
     /// Each challenge has its own mechanism for weight setting
@@ -168,13 +169,17 @@ impl WasmModuleMetadata {
 
 /// WASM execution configuration
 #[derive(Clone, Debug, Serialize, Deserialize, Default)]
+#[serde(default)]
 pub struct WasmConfig {
     /// Network policy for WASM host functions
     #[serde(default)]
     pub network_policy: NetworkPolicy,
     /// Restartable configuration identifier
     #[serde(default)]
     pub restart_id: String,
+    /// Configuration version for hot-restarts
+    #[serde(default)]
+    pub config_version: u64,
 }
 
 /// WASM-only challenge configuration stored in chain state
diff --git a/crates/core/src/schema_guard.rs b/crates/core/src/schema_guard.rs
@@ -111,6 +111,21 @@ pub fn expected_schema_hashes() -> BTreeMap<u32, SchemaRegistry> {
         description: "Added wasm_challenge_configs to ChainState",
     });
 
+    // Version 5: Added WASM restart metadata
+    registry.insert(5, SchemaRegistry {
+        version: 5,
+        validator_info_hash: const_hash("ValidatorInfo:hotkey:Hotkey,stake:Stake,is_active:bool,last_seen:DateTime,peer_id:Option<String>,x25519_pubkey:Option<String>"),
+        chain_state_hash: const_hash("ChainState:block_height:u64,epoch:u64,config:NetworkConfig,sudo_key:Hotkey,validators:HashMap<Hotkey,ValidatorInfo>,challenges:HashMap<ChallengeId,Challenge>,challenge_configs:HashMap,wasm_challenge_configs:HashMap,mechanism_configs:HashMap,challenge_weights:HashMap,required_version:Option,pending_jobs:Vec,state_hash:[u8;32],last_updated:DateTime,registered_hotkeys:HashSet<Hotkey>"),
+        description: "Added WASM restart metadata",
+    });
+    // Version 4: Added wasm_challenge_configs to ChainState
+    registry.insert(4, SchemaRegistry {
+        version: 4,
+        validator_info_hash: const_hash("ValidatorInfo:hotkey:Hotkey,stake:Stake,is_active:bool,last_seen:DateTime,peer_id:Option<String>,x25519_pubkey:Option<String>"),
+        chain_state_hash: const_hash("ChainState:block_height:u64,epoch:u64,config:NetworkConfig,sudo_key:Hotkey,validators:HashMap<Hotkey,ValidatorInfo>,challenges:HashMap<ChallengeId,Challenge>,challenge_configs:HashMap,wasm_challenge_configs:HashMap,mechanism_configs:HashMap,challenge_weights:HashMap,required_version:Option,pending_jobs:Vec,state_hash:[u8;32],last_updated:DateTime,registered_hotkeys:HashSet<Hotkey>"),
+        description: "Added wasm_challenge_configs to ChainState",
+    });
+
     registry
 }
 
diff --git a/crates/core/src/state_versioning.rs b/crates/core/src/state_versioning.rs
@@ -32,7 +32,8 @@ use tracing::{info, warn};
 /// V2: Added registered_hotkeys
 /// V3: Added x25519_pubkey to ValidatorInfo
 /// V4: Added wasm_challenge_configs
-pub const CURRENT_STATE_VERSION: u32 = 4;
+/// V5: Added WASM restart metadata
+pub const CURRENT_STATE_VERSION: u32 = 5;
 
 /// Minimum supported version for migration
 pub const MIN_SUPPORTED_VERSION: u32 = 1;
@@ -223,42 +224,54 @@ impl ChainStateV2 {
 fn migrate_state(version: u32, data: &[u8]) -> Result<crate::ChainState> {
     match version {
         1 => {
-            // V1 -> V4: Add registered_hotkeys, x25519_pubkey, wasm_challenge_configs
+            // V1 -> V5: Add registered_hotkeys, x25519_pubkey, wasm_challenge_configs
             let v1: ChainStateV1 = bincode::deserialize(data).map_err(|e| {
                 crate::MiniChainError::Serialization(format!("V1 migration failed: {}", e))
             })?;
             info!(
-                "Migrated state V1->V4: block_height={}, validators={}",
+                "Migrated state V1->V5: block_height={}, validators={}",
                 v1.block_height,
                 v1.validators.len()
             );
             Ok(v1.migrate())
         }
         2 => {
-            // V2 -> V4: Add x25519_pubkey to ValidatorInfo and wasm_challenge_configs
+            // V2 -> V5: Add x25519_pubkey to ValidatorInfo and wasm_challenge_configs
             let v2: ChainStateV2 = bincode::deserialize(data).map_err(|e| {
                 crate::MiniChainError::Serialization(format!("V2 migration failed: {}", e))
             })?;
             info!(
-                "Migrated state V2->V4: block_height={}, validators={}",
+                "Migrated state V2->V5: block_height={}, validators={}",
                 v2.block_height,
                 v2.validators.len()
             );
             Ok(v2.migrate())
         }
         3 => {
-            // V3 -> V4: Add wasm_challenge_configs
+            // V3 -> V5: Add wasm_challenge_configs
             let mut v3: crate::ChainState = bincode::deserialize(data).map_err(|e| {
                 crate::MiniChainError::Serialization(format!("V3 migration failed: {}", e))
             })?;
             v3.wasm_challenge_configs = HashMap::new();
             info!(
-                "Migrated state V3->V4: block_height={}, validators={}",
+                "Migrated state V3->V5: block_height={}, validators={}",
                 v3.block_height,
                 v3.validators.len()
             );
             Ok(v3)
         }
+        4 => {
+            // V4 -> V5: Added WASM restart metadata (handled by serde defaults)
+            let v4: crate::ChainState = bincode::deserialize(data).map_err(|e| {
+                crate::MiniChainError::Serialization(format!("V4 migration failed: {}", e))
+            })?;
+            info!(
+                "Migrated state V4->V5: block_height={}, validators={}",
+                v4.block_height,
+                v4.validators.len()
+            );
+            Ok(v4)
+        }
         _ => Err(crate::MiniChainError::Serialization(format!(
             "Unknown state version: {}",
             version
@@ -456,7 +469,7 @@ mod tests {
     #[test]
     fn test_version_constants() {
         const _: () = assert!(CURRENT_STATE_VERSION >= MIN_SUPPORTED_VERSION);
-        assert_eq!(CURRENT_STATE_VERSION, 4);
+        assert_eq!(CURRENT_STATE_VERSION, 5);
     }
 
     #[test]
