diff --git a/frontend/src/app/main/ui/workspace/tokens/modals/themes.cljs b/frontend/src/app/main/ui/workspace/tokens/modals/themes.cljs
index 581b87643d..eac37fe8ec 100644
--- a/frontend/src/app/main/ui/workspace/tokens/modals/themes.cljs
+++ b/frontend/src/app/main/ui/workspace/tokens/modals/themes.cljs
@@ -96,9 +96,24 @@
 
 (mf/defc edit-theme
   [{:keys [state]}]
-  (let [{:keys [theme]} @state]
+  (let [{:keys [theme]} @state
+        token-sets (mf/deref refs/workspace-token-sets)
+        selected-token-set-id (mf/deref refs/workspace-selected-token-set-id)
+        token-set-selected? (mf/use-callback
+                             (mf/deps selected-token-set-id)
+                             (fn [id]
+                               (= id selected-token-set-id)))
+        active-token-set-ids (mf/deref refs/workspace-active-set-ids)
+        token-set-active? (mf/use-callback
+                           (mf/deps active-token-set-ids)
+                           (fn [id]
+                             (get active-token-set-ids id)))]
     [:div {:class (stl/css :sets-list-wrapper)}
-     [:& wts/sets-list]]))
+     [:& wts/controlled-sets-list
+      {:token-sets token-sets
+       :selected-token-set-id selected-token-set-id
+       :token-set-selected? token-set-selected?
+       :token-set-active? token-set-active?}]]))
 
 (mf/defc themes
   [{:keys [] :as _args}]
diff --git a/frontend/src/app/main/ui/workspace/tokens/sets.cljs b/frontend/src/app/main/ui/workspace/tokens/sets.cljs
index 852844bf51..fb1780469c 100644
--- a/frontend/src/app/main/ui/workspace/tokens/sets.cljs
+++ b/frontend/src/app/main/ui/workspace/tokens/sets.cljs
@@ -64,6 +64,15 @@
                                  :set-id child-id
                                  :selected-set-id selected-token-set-id)])])]))
 
+(mf/defc controlled-sets-list
+  [{:keys [token-sets] :as props}]
+  [:ul {:class (stl/css :sets-list)}
+   (for [[id token-set] token-sets]
+     [:& sets-tree (-> (assoc props
+                              :key id
+                              :token-set token-set)
+                       (dissoc :token-sets))])])
+
 (mf/defc sets-list
   [{:keys []}]
   (let [token-sets (mf/deref refs/workspace-token-sets)
@@ -77,11 +86,8 @@
                            (mf/deps active-token-set-ids)
                            (fn [id]
                              (get active-token-set-ids id)))]
-    [:ul {:class (stl/css :sets-list)}
-     (for [[id token-set] token-sets]
-       [:& sets-tree
-        {:key id
-         :token-set token-set
-         :selected-token-set-id selected-token-set-id
-         :token-set-selected? token-set-selected?
-         :token-set-active? token-set-active?}])]))
+    [:& controlled-sets-list
+     {:token-sets token-sets
+      :selected-token-set-id selected-token-set-id
+      :token-set-selected? token-set-selected?
+      :token-set-active? token-set-active?}]))