diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 42a28d6..b6da164 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -163,8 +163,7 @@
 
         LockedUserState.get(context).runOnUserUnlocked(() -> {
             CustomWidgetManager cwm = CustomWidgetManager.INSTANCE.get(mContext);
-            cwm.setWidgetRefreshCallback(mModel::refreshAndBindWidgetsAndShortcuts);
-            mOnTerminateCallback.add(() -> cwm.setWidgetRefreshCallback(null));
+            mOnTerminateCallback.add(cwm.addWidgetRefreshCallback(mModel::rebindCallbacks)::close);
 
             IconObserver observer = new IconObserver();
             SafeCloseable iconChangeTracker = mIconProvider.registerIconChangeListener(
diff --git a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
index 9dddc18..20cce8f 100644
--- a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
+++ b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
@@ -39,8 +39,8 @@
 import com.android.launcher3.dagger.LauncherBaseAppComponent;
 import com.android.launcher3.util.DaggerSingletonObject;
 import com.android.launcher3.util.DaggerSingletonTracker;
-import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.PluginManagerWrapper;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.systemui.plugins.CustomWidgetPlugin;
@@ -51,7 +51,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Consumer;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.stream.Stream;
 
 import javax.inject.Inject;
@@ -70,7 +70,7 @@
     private final Context mContext;
     private final HashMap<ComponentName, CustomWidgetPlugin> mPlugins;
     private final List<CustomAppWidgetProviderInfo> mCustomWidgets;
-    private Consumer<PackageUserKey> mWidgetRefreshCallback;
+    private final List<Runnable> mWidgetRefreshCallbacks = new CopyOnWriteArrayList<>();
     private final @NonNull AppWidgetManager mAppWidgetManager;
 
     @Inject
@@ -111,23 +111,20 @@
 
     @Override
     public void onPluginConnected(CustomWidgetPlugin plugin, Context context) {
-        List<AppWidgetProviderInfo> providers = mAppWidgetManager
-                .getInstalledProvidersForProfile(Process.myUserHandle());
-        if (providers.isEmpty()) return;
-        Parcel parcel = Parcel.obtain();
-        providers.get(0).writeToParcel(parcel, 0);
-        parcel.setDataPosition(0);
-        CustomAppWidgetProviderInfo info = newInfo(plugin, parcel);
-        parcel.recycle();
-        mPlugins.put(info.provider, plugin);
-        mCustomWidgets.add(info);
+        CustomAppWidgetProviderInfo info = getAndAddInfo(new ComponentName(
+                PLUGIN_PKG, CLS_CUSTOM_WIDGET_PREFIX + plugin.getClass().getName()));
+        if (info != null) {
+            plugin.updateWidgetInfo(info, mContext);
+            mPlugins.put(info.provider, plugin);
+            mWidgetRefreshCallbacks.forEach(MAIN_EXECUTOR::execute);
+        }
     }
 
     @Override
     public void onPluginDisconnected(CustomWidgetPlugin plugin) {
-        ComponentName cn = getWidgetProviderComponent(plugin);
-        mPlugins.remove(cn);
-        mCustomWidgets.removeIf(w -> w.getComponent().equals(cn));
+        // Leave the providerInfo as plugins can get disconnected/reconnected multiple times
+        mPlugins.values().remove(plugin);
+        mWidgetRefreshCallbacks.forEach(MAIN_EXECUTOR::execute);
     }
 
     @VisibleForTesting
@@ -138,9 +135,11 @@
 
     /**
      * Inject a callback function to refresh the widgets.
+     * @return a closeable to remove this callback
      */
-    public void setWidgetRefreshCallback(Consumer<PackageUserKey> cb) {
-        mWidgetRefreshCallback = cb;
+    public SafeCloseable addWidgetRefreshCallback(Runnable callback) {
+        mWidgetRefreshCallbacks.add(callback);
+        return () -> mWidgetRefreshCallbacks.remove(callback);
     }
 
     /**
@@ -149,8 +148,9 @@
     public void onViewCreated(LauncherAppWidgetHostView view) {
         CustomAppWidgetProviderInfo info = (CustomAppWidgetProviderInfo) view.getAppWidgetInfo();
         CustomWidgetPlugin plugin = mPlugins.get(info.provider);
-        if (plugin == null) return;
-        plugin.onViewCreated(view);
+        if (plugin != null) {
+            plugin.onViewCreated(view);
+        }
     }
 
     /**
@@ -166,14 +166,13 @@
      */
     @Nullable
     public LauncherAppWidgetProviderInfo getWidgetProvider(ComponentName cn) {
-        return mCustomWidgets.stream()
+        LauncherAppWidgetProviderInfo info = mCustomWidgets.stream()
                 .filter(w -> w.getComponent().equals(cn)).findAny().orElse(null);
-    }
-
-    private CustomAppWidgetProviderInfo newInfo(CustomWidgetPlugin plugin, Parcel parcel) {
-        CustomAppWidgetProviderInfo info = new CustomAppWidgetProviderInfo(parcel, false);
-        info.provider = getWidgetProviderComponent(plugin);
-        plugin.updateWidgetInfo(info, mContext);
+        if (info == null) {
+            // If the info is not present, add a placeholder info since the
+            // plugin might get loaded later
+            info = getAndAddInfo(cn);
+        }
         return info;
     }
 
@@ -184,8 +183,24 @@
         return CUSTOM_WIDGET_ID - mCustomWidgets.indexOf(getWidgetProvider(componentName));
     }
 
-    private ComponentName getWidgetProviderComponent(CustomWidgetPlugin plugin) {
-        return new ComponentName(
-                PLUGIN_PKG, CLS_CUSTOM_WIDGET_PREFIX + plugin.getClass().getName());
+    @Nullable
+    private CustomAppWidgetProviderInfo getAndAddInfo(ComponentName cn) {
+        for (CustomAppWidgetProviderInfo info : mCustomWidgets) {
+            if (info.provider.equals(cn)) return info;
+        }
+
+        List<AppWidgetProviderInfo> providers = mAppWidgetManager
+                .getInstalledProvidersForProfile(Process.myUserHandle());
+        if (providers.isEmpty()) return null;
+        Parcel parcel = Parcel.obtain();
+        providers.get(0).writeToParcel(parcel, 0);
+        parcel.setDataPosition(0);
+        CustomAppWidgetProviderInfo info = new CustomAppWidgetProviderInfo(parcel, false);
+        parcel.recycle();
+
+        info.provider = cn;
+        info.initialLayout = 0;
+        mCustomWidgets.add(info);
+        return info;
     }
 }