diff --git a/packages/ui-store/src/ExtensionSession.ts b/packages/ui-store/src/ExtensionSession.ts
index 70246e7..6a92745 100644
--- a/packages/ui-store/src/ExtensionSession.ts
+++ b/packages/ui-store/src/ExtensionSession.ts
@@ -33,12 +33,10 @@ export class ExtensionSession extends BaseStore {
   }
 
   public remove(key: string, fn?: () => void) {
-    session.remove(key);
-    fn && fn();
+    session.remove(key, fn);
   }
 
   public set(key: string, value: unknown, fn?: () => void) {
-    session.set({ [key]: value });
-    fn && fn();
+    session.set({ [key]: value }, fn);
   }
 }
diff --git a/packages/ui-store/src/ExtensionStore.ts b/packages/ui-store/src/ExtensionStore.ts
index d5d7bc9..fbbc949 100644
--- a/packages/ui-store/src/ExtensionStore.ts
+++ b/packages/ui-store/src/ExtensionStore.ts
@@ -33,12 +33,10 @@ export class ExtensionStore extends BaseStore {
   }
 
   public remove(key: string, fn?: () => void) {
-    storage.remove(key);
-    fn && fn();
+    storage.remove(key, fn);
   }
 
   public set(key: string, value: unknown, fn?: () => void) {
-    storage.set({ [key]: value });
-    fn && fn();
+    storage.set({ [key]: value }, fn);
   }
 }