diff --git a/common/src/app/common/types/tokens_theme_list.cljc b/common/src/app/common/types/tokens_theme_list.cljc
index 64dab3c04..f2e6058e2 100644
--- a/common/src/app/common/types/tokens_theme_list.cljc
+++ b/common/src/app/common/types/tokens_theme_list.cljc
@@ -7,7 +7,6 @@
 (ns app.common.types.tokens-theme-list
   (:require
    [app.common.data :as d]
-   [app.common.exceptions :as ex]
    [app.common.time :as dt]))
 
 (defn- touch
@@ -28,11 +27,9 @@
       (update :token-sets-index assoc id token-set)))
 
 (defn update-token-set
-  [file-data token-id f & args]
-  #_(ex/raise :type :not-implemented)
-  file-data)
+  [file-data token-set-id f & args]
+  (d/update-in-when file-data [:token-sets-index token-set-id] #(-> (apply f % args) (touch))))
 
 (defn delete-token-set
   [file-data token-id]
-  #_(ex/raise :type :not-implemented)
   file-data)
diff --git a/frontend/src/app/main/data/tokens.cljs b/frontend/src/app/main/data/tokens.cljs
index a82d7480f..50ae222a4 100644
--- a/frontend/src/app/main/data/tokens.cljs
+++ b/frontend/src/app/main/data/tokens.cljs
@@ -131,6 +131,9 @@
               new-token-set {:id (uuid/next)
                              :name "Global"
                              :tokens [(:id token)]}
+              selected-token-set-id (if create-set?
+                                      (:id new-token-set)
+                                      (:id token-set))
               changes (cond
                         create-set? (-> token-changes
                                         (pcb/add-token-set new-token-set))
@@ -138,9 +141,9 @@
                                                         token-set
                                                         (update token-set :tokens conj (:id token)))]
                                 (-> token-changes
-                                    (pcb/update-token-set token-set updated-token-set))))]
+                                    (pcb/update-token-set updated-token-set token-set))))]
           (rx/of
-           (set-selected-token-set-id (:id new-token-set))
+           (set-selected-token-set-id selected-token-set-id)
            (dch/commit-changes changes)))))))
 
 (defn delete-token
diff --git a/frontend/test/token_tests/logic/token_actions_test.cljs b/frontend/test/token_tests/logic/token_actions_test.cljs
index 790aadda0..cd026217f 100644
--- a/frontend/test/token_tests/logic/token_actions_test.cljs
+++ b/frontend/test/token_tests/logic/token_actions_test.cljs
@@ -59,10 +59,30 @@
             (t/testing "selects created workspace set and adds token to it"
               (t/is (some? selected-token-set))
               (t/is (= 1 (count (:tokens selected-token-set))))
-              (t/is (= (list border-radius-token) (map #(dissoc % :id :modified-at) set-tokens))))
+              (t/is (= (list border-radius-token) (map #(dissoc % :id :modified-at) set-tokens)))))))))))
+
+(t/deftest test-create-multiple-tokens
+  (t/testing "uses selected tokens set when creating multiple tokens"
+    (t/async
+     done
+     (let [file (setup-file)
+           store (ths/setup-store file)
+           events [(wdt/update-create-token border-radius-token)
+                   (wdt/update-create-token reference-border-radius-token)]]
+       (tohs/run-store-async
+        store done events
+        (fn [new-state]
+          (let [selected-token-set (wdt/get-selected-token-set new-state)
+                file' (ths/get-file-from-store new-state)
+                set-tokens (wdt/get-token-set-tokens selected-token-set file')]
+            (t/testing "selects created workspace set and adds token to it"
+              (t/is (some? selected-token-set))
+              (t/is (= 2 (count (:tokens selected-token-set))))
+              (t/is (= (list border-radius-token reference-border-radius-token) (map #(dissoc % :id :modified-at) set-tokens))))
             (reset! a new-state))))))))
 
 (comment
+  (wdt/get-selected-token-set @a)
   @a
   (t/run-tests)
   nil)