public entry fun unstake(staker: &signer, module_owner: &signer, rapper_token: Object<Token>) acquires StakeInfo {
    let token_id = object::object_address(&rapper_token);
    
    
    let (wk, ha, ss, cr, wins) = one_shot::read_stats(token_id);
    one_shot::transfer_record_only(token_id, @battle_addr, staker_addr);
    object::transfer(module_owner, rapper_token, staker_addr); 
}
streets::stake(&alice, &battle_addr, rapperA);
streets::stake(&charlie, &battle_addr, rapperB);
streets::unstake(&bob, &battle_addr, rapperB);
 struct StakeInfo has key, store {
-    start_time_seconds: u64,
-    owner: address,
+    start_time_seconds: u64,
+    owner: address,
+    token_id: address, // bind token_id
 }
 public entry fun stake(...) {
-    move_to(staker, StakeInfo { start_time_seconds: now, owner: staker_addr });
+    move_to(staker, StakeInfo { start_time_seconds: now, owner: staker_addr, token_id });
 }
 public entry fun unstake(...) acquires StakeInfo {
-    let token_id = object::object_address(&rapper_token);
-    let (wk, ha, ss, cr, wins) = one_shot::read_stats(token_id);
-    one_shot::transfer_record_only(token_id, @battle_addr, staker_addr);
-    object::transfer(module_owner, rapper_token, staker_addr);
+    let si = borrow_global<StakeInfo>(staker_addr);
+    assert!(si.token_id == object::object_address(&rapper_token), E_TOKEN_NOT_STAKED);
+    let (wk, ha, ss, cr, wins) = one_shot::read_stats(si.token_id);
+    one_shot::transfer_record_only(si.token_id, @battle_addr, staker_addr);
+    // Return the custodied token from protocol storage
+    let stored_obj = internal_pop_custodied_object(si.token_id);
+    object::transfer(module_owner, stored_obj, staker_addr);
 }