[PM-25488] Badge stays after lock when using pin (#16436)

* wip

* feat: add dynamic states

* feat: re-implement badge service with dynamic state functions

* feat: completely remove old static states

* feat: debounce calls to badge api per tab

* feat: use group-by to avoid re-setting all tabs on 1 tab change

* feat: simplify autofill badge updater

* feat: add hanging function test

* chore: clean up badge service

* feat: simplify private updateBadge

* feat: remove unnecessary Set usage

* fix: tests that broke after setState rename

* chore: clean up badge api

apps/browser/src/auth/services/auth-status-badge-updater.service.ts

@@ -17,8 +17,8 @@ export class AuthStatusBadgeUpdaterService {
     private accountService: AccountService,
     private authService: AuthService,
   ) {
-    this.accountService.activeAccount$
+    this.badgeService.setState(StateName, (_tab) =>
-      .pipe(
+      this.accountService.activeAccount$.pipe(
         switchMap((account) =>
           account
             ? this.authService.authStatusFor$(account.id)
@@ -27,30 +27,36 @@ export class AuthStatusBadgeUpdaterService {
         mergeMap(async (authStatus) => {
           switch (authStatus) {
             case AuthenticationStatus.LoggedOut: {
-              await this.badgeService.setState(StateName, BadgeStatePriority.High, {
+              return {
+                priority: BadgeStatePriority.High,
+                state: {
                   icon: BadgeIcon.LoggedOut,
                   backgroundColor: Unset,
                   text: Unset,
-              });
+                },
-              break;
+              };
             }
             case AuthenticationStatus.Locked: {
-              await this.badgeService.setState(StateName, BadgeStatePriority.High, {
+              return {
+                priority: BadgeStatePriority.High,
+                state: {
                   icon: BadgeIcon.Locked,
                   backgroundColor: Unset,
                   text: Unset,
-              });
+                },
-              break;
+              };
             }
             case AuthenticationStatus.Unlocked: {
-              await this.badgeService.setState(StateName, BadgeStatePriority.Low, {
+              return {
+                priority: BadgeStatePriority.Low,
+                state: {
                   icon: BadgeIcon.Unlocked,
-              });
+                },
-              break;
+              };
             }
           }
         }),
-      )
+      ),
-      .subscribe();
+    );
   }
 }

apps/browser/src/autofill/services/autofill-badge-updater.service.ts

@@ -1,4 +1,4 @@
-import { combineLatest, distinctUntilChanged, mergeMap, of, switchMap, withLatestFrom } from "rxjs";
+import { combineLatest, delay, distinctUntilChanged, mergeMap, of, switchMap } from "rxjs";
 
 import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
 import { BadgeSettingsServiceAbstraction } from "@bitwarden/common/autofill/services/badge-settings.service";
@@ -10,7 +10,7 @@ import { Tab } from "../../platform/badge/badge-browser-api";
 import { BadgeService } from "../../platform/badge/badge.service";
 import { BadgeStatePriority } from "../../platform/badge/priority";
 
-const StateName = (tabId: number) => `autofill-badge-${tabId}`;
+const StateName = "autofill-badge-updater";
 
 export class AutofillBadgeUpdaterService {
   constructor(
@@ -26,56 +26,30 @@ export class AutofillBadgeUpdaterService {
       switchMap((account) => (account?.id ? this.cipherService.ciphers$(account?.id) : of([]))),
     );
 
-    // Recalculate badges for all active tabs when ciphers or active account changes
+    this.badgeService.setState(StateName, (tab) => {
-    combineLatest({
+      return combineLatest({
         account: this.accountService.activeAccount$,
         enableBadgeCounter:
           this.badgeSettingsService.enableBadgeCounter$.pipe(distinctUntilChanged()),
-      ciphers: ciphers$,
+        ciphers: ciphers$.pipe(delay(100)), // Delay to allow cipherService.getAllDecryptedForUrl to pick up changes
-    })
+      }).pipe(
-      .pipe(
         mergeMap(async ({ account, enableBadgeCounter }) => {
-          if (!account) {
-            return;
-          }
-
-          const tabs = await this.badgeService.getActiveTabs();
-
-          for (const tab of tabs) {
-            if (!tab.tabId) {
-              continue;
-            }
-            if (enableBadgeCounter) {
-              await this.setTabState(tab, account.id);
-            } else {
-              await this.clearTabState(tab.tabId);
-            }
-          }
-        }),
-      )
-      .subscribe();
-
-    // Recalculate badge for a specific tab when it becomes active
-    this.badgeService.activeTabsUpdated$
-      .pipe(
-        withLatestFrom(
-          this.accountService.activeAccount$,
-          this.badgeSettingsService.enableBadgeCounter$,
-        ),
-        mergeMap(async ([tabs, account, enableBadgeCounter]) => {
           if (!account || !enableBadgeCounter) {
-            return;
+            return undefined;
           }
 
-          for (const tab of tabs) {
+          return {
-            await this.setTabState(tab, account.id);
+            state: {
-          }
+              text: await this.calculateCountText(tab, account.id),
+            },
+            priority: BadgeStatePriority.Default,
+          };
         }),
-      )
+      );
-      .subscribe();
+    });
   }
 
-  private async setTabState(tab: Tab, userId: UserId) {
+  private async calculateCountText(tab: Tab, userId: UserId) {
     if (!tab.tabId) {
       this.logService.warning("Tab event received but tab id is undefined");
       return;
@@ -85,22 +59,9 @@ export class AutofillBadgeUpdaterService {
     const cipherCount = ciphers.length;
 
     if (cipherCount === 0) {
-      await this.clearTabState(tab.tabId);
+      return undefined;
-      return;
     }
 
-    const countText = cipherCount > 9 ? "9+" : cipherCount.toString();
+    return cipherCount > 9 ? "9+" : cipherCount.toString();
-    await this.badgeService.setState(
-      StateName(tab.tabId),
-      BadgeStatePriority.Default,
-      {
-        text: countText,
-      },
-      tab.tabId,
-    );
-  }
-
-  private async clearTabState(tabId: number) {
-    await this.badgeService.clearState(StateName(tabId));
   }
 }

apps/browser/src/background/main.background.ts

@@ -1435,7 +1435,6 @@ export default class MainBackground {
     );
 
     this.badgeService = new BadgeService(
-      this.stateProvider,
       new DefaultBadgeBrowserApi(this.platformUtilsService),
       this.logService,
     );
@@ -1925,7 +1924,6 @@ export default class MainBackground {
       this.badgeService,
       this.accountService,
       this.cipherService,
-      this.logService,
       this.taskService,
     );
 

apps/browser/src/platform/badge/badge-browser-api.ts

@@ -1,4 +1,16 @@
-import { concat, defer, filter, map, merge, Observable, shareReplay, switchMap } from "rxjs";
+import {
+  concat,
+  concatMap,
+  defer,
+  filter,
+  map,
+  merge,
+  Observable,
+  of,
+  pairwise,
+  shareReplay,
+  switchMap,
+} from "rxjs";
 
 import { PlatformUtilsService } from "@bitwarden/common/platform/abstractions/platform-utils.service";
 import { Utils } from "@bitwarden/common/platform/misc/utils";
@@ -28,13 +40,37 @@ function tabFromChromeTab(tab: chrome.tabs.Tab): Tab {
 }
 
 export interface BadgeBrowserApi {
-  activeTabsUpdated$: Observable<Tab[]>;
+  /**
+   * An observable that emits all currently active tabs whenever one or more active tabs change.
+   */
+  activeTabs$: Observable<Tab[]>;
 
+  /**
+   * An observable that emits tab events such as updates and activations.
+   */
+  tabEvents$: Observable<TabEvent>;
+
+  /**
+   * Set the badge state for a specific tab.
+   * If the tabId is undefined the state will be applied to the browser action in general.
+   */
   setState(state: RawBadgeState, tabId?: number): Promise<void>;
-  getTabs(): Promise<number[]>;
-  getActiveTabs(): Promise<Tab[]>;
 }
 
+export type TabEvent =
+  | {
+      type: "updated";
+      tab: Tab;
+    }
+  | {
+      type: "activated";
+      tab: Tab;
+    }
+  | {
+      type: "deactivated";
+      tabId: number;
+    };
+
 export class DefaultBadgeBrowserApi implements BadgeBrowserApi {
   private badgeAction = BrowserApi.getBrowserAction();
   private sidebarAction = BrowserApi.getSidebarAction(self);
@@ -44,18 +80,25 @@ export class DefaultBadgeBrowserApi implements BadgeBrowserApi {
     shareReplay({ bufferSize: 1, refCount: true }),
   );
 
-  activeTabsUpdated$ = concat(
+  private createdOrUpdatedTabEvents$ = concat(
-    defer(async () => await this.getActiveTabs()),
+    defer(async () => await this.getActiveTabs()).pipe(
+      switchMap((activeTabs) => {
+        const tabEvents: TabEvent[] = activeTabs.map((tab) => ({
+          type: "activated",
+          tab,
+        }));
+        return of(...tabEvents);
+      }),
+    ),
     merge(
       this.onTabActivated$.pipe(
-        switchMap(async (activeInfo) => {
+        switchMap(async (activeInfo) => await BrowserApi.getTab(activeInfo.tabId)),
-          const tab = await BrowserApi.getTab(activeInfo.tabId);
+        filter(
-
+          (tab): tab is chrome.tabs.Tab =>
-          if (tab == undefined || tab.id == undefined || tab.url == undefined) {
+            !(tab == undefined || tab.id == undefined || tab.url == undefined),
-            return [];
+        ),
-          }
+        switchMap(async (tab) => {
-
+          return { type: "activated", tab: tabFromChromeTab(tab) } satisfies TabEvent;
-          return [tabFromChromeTab(tab)];
         }),
       ),
       fromChromeEvent(chrome.tabs.onUpdated).pipe(
@@ -64,22 +107,58 @@ export class DefaultBadgeBrowserApi implements BadgeBrowserApi {
             // Only emit if the url was updated
             changeInfo.url != undefined,
         ),
-        map(([_tabId, _changeInfo, tab]) => [tabFromChromeTab(tab)]),
+        map(
+          ([_tabId, _changeInfo, tab]) =>
+            ({ type: "updated", tab: tabFromChromeTab(tab) }) satisfies TabEvent,
+        ),
       ),
       fromChromeEvent(chrome.webNavigation.onCommitted).pipe(
+        filter(([details]) => details.transitionType === "reload"),
         map(([details]) => {
-          const toReturn: Tab[] =
+          return {
-            details.transitionType === "reload" ? [{ tabId: details.tabId, url: details.url }] : [];
+            type: "updated",
-          return toReturn;
+            tab: { tabId: details.tabId, url: details.url },
+          } satisfies TabEvent;
         }),
       ),
       // NOTE: We're only sharing the active tab changes, not the full list of active tabs.
       // This is so that any new subscriber will get the latest active tabs immediately, but
       // doesn't re-subscribe to chrome events.
     ).pipe(shareReplay({ bufferSize: 1, refCount: true })),
-  ).pipe(filter((tabs) => tabs.length > 0));
+  );
 
-  async getActiveTabs(): Promise<Tab[]> {
+  tabEvents$ = merge(
+    this.createdOrUpdatedTabEvents$,
+    this.createdOrUpdatedTabEvents$.pipe(
+      concatMap(async () => {
+        return this.getActiveTabs();
+      }),
+      pairwise(),
+      map(([previousTabs, currentTabs]) => {
+        const previousTabIds = previousTabs.map((t) => t.tabId);
+        const currentTabIds = currentTabs.map((t) => t.tabId);
+
+        const deactivatedTabIds = previousTabIds.filter((id) => !currentTabIds.includes(id));
+
+        return deactivatedTabIds.map(
+          (tabId) =>
+            ({
+              type: "deactivated",
+              tabId,
+            }) satisfies TabEvent,
+        );
+      }),
+      switchMap((events) => of(...events)),
+    ),
+  );
+
+  activeTabs$ = this.tabEvents$.pipe(
+    concatMap(async () => {
+      return this.getActiveTabs();
+    }),
+  );
+
+  private async getActiveTabs(): Promise<Tab[]> {
     const tabs = await BrowserApi.getActiveTabs();
     return tabs.filter((tab) => tab.id != undefined && tab.url != undefined).map(tabFromChromeTab);
   }
@@ -96,10 +175,6 @@ export class DefaultBadgeBrowserApi implements BadgeBrowserApi {
     ]);
   }
 
-  async getTabs(): Promise<number[]> {
-    return (await BrowserApi.tabsQuery({})).map((tab) => tab.id).filter((tab) => tab !== undefined);
-  }
-
   private setIcon(icon: IconPaths, tabId?: number) {
     return Promise.all([this.setActionIcon(icon, tabId), this.setSidebarActionIcon(icon, tabId)]);
   }

apps/browser/src/platform/badge/badge.service.spec.ts

@@ -1,11 +1,10 @@
 import { mock, MockProxy } from "jest-mock-extended";
-import { Subscription } from "rxjs";
+import { EMPTY, Observable, of, Subscription } from "rxjs";
 
 import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
-import { FakeAccountService, FakeStateProvider } from "@bitwarden/common/spec";
 
 import { RawBadgeState } from "./badge-browser-api";
-import { BadgeService } from "./badge.service";
+import { BadgeService, BadgeStateFunction } from "./badge.service";
 import { DefaultBadgeState } from "./consts";
 import { BadgeIcon } from "./icon";
 import { BadgeStatePriority } from "./priority";
@@ -14,7 +13,6 @@ import { MockBadgeBrowserApi } from "./test/mock-badge-browser-api";
 
 describe("BadgeService", () => {
   let badgeApi: MockBadgeBrowserApi;
-  let stateProvider: FakeStateProvider;
   let logService!: MockProxy<LogService>;
   let badgeService!: BadgeService;
 
@@ -22,16 +20,16 @@ describe("BadgeService", () => {
 
   beforeEach(() => {
     badgeApi = new MockBadgeBrowserApi();
-    stateProvider = new FakeStateProvider(new FakeAccountService({}));
     logService = mock<LogService>();
 
-    badgeService = new BadgeService(stateProvider, badgeApi, logService);
+    badgeService = new BadgeService(badgeApi, logService, 0);
   });
 
   afterEach(() => {
     badgeServiceSubscription?.unsubscribe();
   });
 
+  describe("static state", () => {
     describe("calling without tabId", () => {
       const tabId = 1;
 
@@ -49,7 +47,10 @@ describe("BadgeService", () => {
             icon: BadgeIcon.Locked,
           };
 
-        await badgeService.setState("state-name", BadgeStatePriority.Default, state);
+          await badgeService.setState(
+            "state-name",
+            GeneralStateFunction(BadgeStatePriority.Default, state),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates[tabId]).toEqual(state);
@@ -59,20 +60,39 @@ describe("BadgeService", () => {
           // This is a bit of a weird thing to do, but I don't think it's something we need to prohibit
           const state: BadgeState = {};
 
-        await badgeService.setState("state-name", BadgeStatePriority.Default, state);
+          await badgeService.setState(
+            "state-name",
+            GeneralStateFunction(BadgeStatePriority.Default, state),
+          );
+
+          await new Promise((resolve) => setTimeout(resolve, 0));
+          expect(badgeApi.specificStates[tabId]).toEqual(DefaultBadgeState);
+        });
+
+        it("sets default values even if state function never emits", async () => {
+          badgeService.setState("state-name", (_tab) => EMPTY);
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates[tabId]).toEqual(DefaultBadgeState);
         });
 
         it("merges states when multiple same-priority states have been set", async () => {
-        await badgeService.setState("state-1", BadgeStatePriority.Default, { text: "text" });
+          await badgeService.setState(
-        await badgeService.setState("state-2", BadgeStatePriority.Default, {
+            "state-1",
+            GeneralStateFunction(BadgeStatePriority.Default, { text: "text" }),
+          );
+          await badgeService.setState(
+            "state-2",
+            GeneralStateFunction(BadgeStatePriority.Default, {
               backgroundColor: "#fff",
-        });
+            }),
-        await badgeService.setState("state-3", BadgeStatePriority.Default, {
+          );
+          await badgeService.setState(
+            "state-3",
+            GeneralStateFunction(BadgeStatePriority.Default, {
               icon: BadgeIcon.Locked,
-        });
+            }),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           const expectedState: RawBadgeState = {
@@ -84,17 +104,26 @@ describe("BadgeService", () => {
         });
 
         it("overrides previous lower-priority state when higher-priority state is set", async () => {
-        await badgeService.setState("state-1", BadgeStatePriority.Low, {
+          await badgeService.setState(
+            "state-1",
+            GeneralStateFunction(BadgeStatePriority.Low, {
               text: "text",
               backgroundColor: "#fff",
               icon: BadgeIcon.Locked,
-        });
+            }),
-        await badgeService.setState("state-2", BadgeStatePriority.Default, {
+          );
+          await badgeService.setState(
+            "state-2",
+            GeneralStateFunction(BadgeStatePriority.Default, {
               text: "override",
-        });
+            }),
-        await badgeService.setState("state-3", BadgeStatePriority.High, {
+          );
+          await badgeService.setState(
+            "state-3",
+            GeneralStateFunction(BadgeStatePriority.High, {
               backgroundColor: "#aaa",
-        });
+            }),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           const expectedState: RawBadgeState = {
@@ -106,14 +135,20 @@ describe("BadgeService", () => {
         });
 
         it("removes override when a previously high-priority state is cleared", async () => {
-        await badgeService.setState("state-1", BadgeStatePriority.Low, {
+          await badgeService.setState(
+            "state-1",
+            GeneralStateFunction(BadgeStatePriority.Low, {
               text: "text",
               backgroundColor: "#fff",
               icon: BadgeIcon.Locked,
-        });
+            }),
-        await badgeService.setState("state-2", BadgeStatePriority.Default, {
+          );
+          await badgeService.setState(
+            "state-2",
+            GeneralStateFunction(BadgeStatePriority.Default, {
               text: "override",
-        });
+            }),
+          );
           await badgeService.clearState("state-2");
 
           await new Promise((resolve) => setTimeout(resolve, 0));
@@ -126,17 +161,26 @@ describe("BadgeService", () => {
         });
 
         it("sets default values when all states have been cleared", async () => {
-        await badgeService.setState("state-1", BadgeStatePriority.Low, {
+          await badgeService.setState(
+            "state-1",
+            GeneralStateFunction(BadgeStatePriority.Low, {
               text: "text",
               backgroundColor: "#fff",
               icon: BadgeIcon.Locked,
-        });
+            }),
-        await badgeService.setState("state-2", BadgeStatePriority.Default, {
+          );
+          await badgeService.setState(
+            "state-2",
+            GeneralStateFunction(BadgeStatePriority.Default, {
               text: "override",
-        });
+            }),
-        await badgeService.setState("state-3", BadgeStatePriority.High, {
+          );
+          await badgeService.setState(
+            "state-3",
+            GeneralStateFunction(BadgeStatePriority.High, {
               backgroundColor: "#aaa",
-        });
+            }),
+          );
           await badgeService.clearState("state-1");
           await badgeService.clearState("state-2");
           await badgeService.clearState("state-3");
@@ -146,14 +190,20 @@ describe("BadgeService", () => {
         });
 
         it("sets default value high-priority state contains Unset", async () => {
-        await badgeService.setState("state-1", BadgeStatePriority.Low, {
+          await badgeService.setState(
+            "state-1",
+            GeneralStateFunction(BadgeStatePriority.Low, {
               text: "text",
               backgroundColor: "#fff",
               icon: BadgeIcon.Locked,
-        });
+            }),
-        await badgeService.setState("state-3", BadgeStatePriority.High, {
+          );
+          await badgeService.setState(
+            "state-3",
+            GeneralStateFunction(BadgeStatePriority.High, {
               icon: Unset,
-        });
+            }),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           const expectedState: RawBadgeState = {
@@ -165,17 +215,26 @@ describe("BadgeService", () => {
         });
 
         it("ignores medium-priority Unset when high-priority contains a value", async () => {
-        await badgeService.setState("state-1", BadgeStatePriority.Low, {
+          await badgeService.setState(
+            "state-1",
+            GeneralStateFunction(BadgeStatePriority.Low, {
               text: "text",
               backgroundColor: "#fff",
               icon: BadgeIcon.Locked,
-        });
+            }),
-        await badgeService.setState("state-3", BadgeStatePriority.Default, {
+          );
+          await badgeService.setState(
+            "state-3",
+            GeneralStateFunction(BadgeStatePriority.Default, {
               icon: Unset,
-        });
+            }),
-        await badgeService.setState("state-3", BadgeStatePriority.High, {
+          );
+          await badgeService.setState(
+            "state-3",
+            GeneralStateFunction(BadgeStatePriority.High, {
               icon: BadgeIcon.Unlocked,
-        });
+            }),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           const expectedState: RawBadgeState = {
@@ -204,7 +263,10 @@ describe("BadgeService", () => {
             icon: BadgeIcon.Locked,
           };
 
-        await badgeService.setState("state-name", BadgeStatePriority.Default, state);
+          await badgeService.setState(
+            "state-name",
+            GeneralStateFunction(BadgeStatePriority.Default, state),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates).toEqual({
@@ -220,11 +282,22 @@ describe("BadgeService", () => {
             backgroundColor: "color",
             icon: BadgeIcon.Locked,
           };
-        badgeApi.setState.mockReset();
+          await badgeService.setState(
+            "state-1",
+            TabSpecificStateFunction(BadgeStatePriority.Default, state, tabId),
+          );
+          await badgeService.setState(
+            "state-2",
+            TabSpecificStateFunction(BadgeStatePriority.Default, state, 2),
+          );
+          await badgeService.setState(
+            "state-2",
+            TabSpecificStateFunction(BadgeStatePriority.Default, state, 2),
+          );
+          await new Promise((resolve) => setTimeout(resolve, 0));
 
-        await badgeService.setState("state-1", BadgeStatePriority.Default, state, tabId);
+          badgeApi.setState.mockReset();
-        await badgeService.setState("state-2", BadgeStatePriority.Default, state, 2);
+          badgeApi.updateTab(tabId);
-        await badgeService.setState("state-2", BadgeStatePriority.Default, state, 2);
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.setState).toHaveBeenCalledTimes(1);
@@ -248,7 +321,10 @@ describe("BadgeService", () => {
             icon: BadgeIcon.Locked,
           };
 
-        await badgeService.setState("state-name", BadgeStatePriority.Default, state);
+          await badgeService.setState(
+            "state-name",
+            GeneralStateFunction(BadgeStatePriority.Default, state),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates).toEqual({
@@ -264,11 +340,24 @@ describe("BadgeService", () => {
             backgroundColor: "color",
             icon: BadgeIcon.Locked,
           };
-        badgeApi.setState.mockReset();
 
-        await badgeService.setState("state-1", BadgeStatePriority.Default, state, 1);
+          await badgeService.setState(
-        await badgeService.setState("state-2", BadgeStatePriority.Default, state, 2);
+            "state-1",
-        await badgeService.setState("state-3", BadgeStatePriority.Default, state, 3);
+            TabSpecificStateFunction(BadgeStatePriority.Default, state, 1),
+          );
+          await badgeService.setState(
+            "state-2",
+            TabSpecificStateFunction(BadgeStatePriority.Default, state, 2),
+          );
+          await badgeService.setState(
+            "state-3",
+            TabSpecificStateFunction(BadgeStatePriority.Default, state, 3),
+          );
+          await new Promise((resolve) => setTimeout(resolve, 0));
+
+          badgeApi.setState.mockReset();
+          badgeApi.updateTab(activeTabIds[0]);
+          badgeApi.updateTab(activeTabIds[1]);
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.setState).toHaveBeenCalledTimes(2);
@@ -276,7 +365,7 @@ describe("BadgeService", () => {
       });
     });
 
-  describe("calling with tabId", () => {
+    describe("setting tab-specific states", () => {
       describe("given a single tab is open", () => {
         const tabId = 1;
 
@@ -293,7 +382,10 @@ describe("BadgeService", () => {
             icon: BadgeIcon.Locked,
           };
 
-        await badgeService.setState("state-name", BadgeStatePriority.Default, state, tabId);
+          await badgeService.setState(
+            "state-name",
+            TabSpecificStateFunction(BadgeStatePriority.Default, state, tabId),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates[tabId]).toEqual(state);
@@ -303,25 +395,42 @@ describe("BadgeService", () => {
           // This is a bit of a weird thing to do, but I don't think it's something we need to prohibit
           const state: BadgeState = {};
 
-        await badgeService.setState("state-name", BadgeStatePriority.Default, state, tabId);
+          await badgeService.setState(
+            "state-name",
+            TabSpecificStateFunction(BadgeStatePriority.Default, state, tabId),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates[tabId]).toEqual(DefaultBadgeState);
         });
 
         it("merges tabId specific state with general states", async () => {
-        await badgeService.setState("general-state", BadgeStatePriority.Default, { text: "text" });
+          await badgeService.setState(
+            "general-state",
+            TabSpecificStateFunction(
+              BadgeStatePriority.Default,
+              {
+                text: "text",
+              },
+              tabId,
+            ),
+          );
           await badgeService.setState(
             "specific-state",
+            TabSpecificStateFunction(
               BadgeStatePriority.Default,
               {
                 backgroundColor: "#fff",
               },
               tabId,
+            ),
           );
-        await badgeService.setState("general-state-2", BadgeStatePriority.Default, {
+          await badgeService.setState(
+            "general-state-2",
+            GeneralStateFunction(BadgeStatePriority.Default, {
               icon: BadgeIcon.Locked,
-        });
+            }),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates[tabId]).toEqual({
@@ -332,22 +441,29 @@ describe("BadgeService", () => {
         });
 
         it("merges states when multiple same-priority states with the same tabId have been set", async () => {
-        await badgeService.setState("state-1", BadgeStatePriority.Default, { text: "text" }, tabId);
+          await badgeService.setState(
+            "state-1",
+            TabSpecificStateFunction(BadgeStatePriority.Default, { text: "text" }, tabId),
+          );
           await badgeService.setState(
             "state-2",
+            TabSpecificStateFunction(
               BadgeStatePriority.Default,
               {
                 backgroundColor: "#fff",
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-3",
+            TabSpecificStateFunction(
               BadgeStatePriority.Default,
               {
                 icon: BadgeIcon.Locked,
               },
               tabId,
+            ),
           );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
@@ -362,6 +478,7 @@ describe("BadgeService", () => {
         it("overrides previous lower-priority state when higher-priority state with the same tabId is set", async () => {
           await badgeService.setState(
             "state-1",
+            TabSpecificStateFunction(
               BadgeStatePriority.Low,
               {
                 text: "text",
@@ -369,22 +486,27 @@ describe("BadgeService", () => {
                 icon: BadgeIcon.Locked,
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-2",
+            TabSpecificStateFunction(
               BadgeStatePriority.Default,
               {
                 text: "override",
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-3",
+            TabSpecificStateFunction(
               BadgeStatePriority.High,
               {
                 backgroundColor: "#aaa",
               },
               tabId,
+            ),
           );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
@@ -399,6 +521,7 @@ describe("BadgeService", () => {
         it("overrides lower-priority tab-specific state when higher-priority general state is set", async () => {
           await badgeService.setState(
             "state-1",
+            TabSpecificStateFunction(
               BadgeStatePriority.Low,
               {
                 text: "text",
@@ -406,13 +529,20 @@ describe("BadgeService", () => {
                 icon: BadgeIcon.Locked,
               },
               tabId,
+            ),
           );
-        await badgeService.setState("state-2", BadgeStatePriority.Default, {
+          await badgeService.setState(
+            "state-2",
+            GeneralStateFunction(BadgeStatePriority.Default, {
               text: "override",
-        });
+            }),
-        await badgeService.setState("state-3", BadgeStatePriority.High, {
+          );
+          await badgeService.setState(
+            "state-3",
+            GeneralStateFunction(BadgeStatePriority.High, {
               backgroundColor: "#aaa",
-        });
+            }),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates[tabId]).toEqual({
@@ -425,6 +555,7 @@ describe("BadgeService", () => {
         it("removes override when a previously high-priority state with the same tabId is cleared", async () => {
           await badgeService.setState(
             "state-1",
+            TabSpecificStateFunction(
               BadgeStatePriority.Low,
               {
                 text: "text",
@@ -432,14 +563,17 @@ describe("BadgeService", () => {
                 icon: BadgeIcon.Locked,
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-2",
+            TabSpecificStateFunction(
               BadgeStatePriority.Default,
               {
                 text: "override",
               },
               tabId,
+            ),
           );
           await badgeService.clearState("state-2");
 
@@ -454,6 +588,7 @@ describe("BadgeService", () => {
         it("sets default state when all states with the same tabId have been cleared", async () => {
           await badgeService.setState(
             "state-1",
+            TabSpecificStateFunction(
               BadgeStatePriority.Low,
               {
                 text: "text",
@@ -461,22 +596,27 @@ describe("BadgeService", () => {
                 icon: BadgeIcon.Locked,
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-2",
+            TabSpecificStateFunction(
               BadgeStatePriority.Default,
               {
                 text: "override",
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-3",
+            TabSpecificStateFunction(
               BadgeStatePriority.High,
               {
                 backgroundColor: "#aaa",
               },
               tabId,
+            ),
           );
           await badgeService.clearState("state-1");
           await badgeService.clearState("state-2");
@@ -489,6 +629,7 @@ describe("BadgeService", () => {
         it("sets default value when high-priority state contains Unset", async () => {
           await badgeService.setState(
             "state-1",
+            TabSpecificStateFunction(
               BadgeStatePriority.Low,
               {
                 text: "text",
@@ -496,14 +637,17 @@ describe("BadgeService", () => {
                 icon: BadgeIcon.Locked,
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-3",
+            TabSpecificStateFunction(
               BadgeStatePriority.High,
               {
                 icon: Unset,
               },
               tabId,
+            ),
           );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
@@ -517,6 +661,7 @@ describe("BadgeService", () => {
         it("ignores medium-priority Unset when high-priority contains a value", async () => {
           await badgeService.setState(
             "state-1",
+            TabSpecificStateFunction(
               BadgeStatePriority.Low,
               {
                 text: "text",
@@ -524,22 +669,27 @@ describe("BadgeService", () => {
                 icon: BadgeIcon.Locked,
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-3",
+            TabSpecificStateFunction(
               BadgeStatePriority.Default,
               {
                 icon: Unset,
               },
               tabId,
+            ),
           );
           await badgeService.setState(
             "state-3",
+            TabSpecificStateFunction(
               BadgeStatePriority.High,
               {
                 icon: BadgeIcon.Unlocked,
               },
               tabId,
+            ),
           );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
@@ -572,12 +722,13 @@ describe("BadgeService", () => {
             icon: BadgeIcon.Locked,
           };
 
-        await badgeService.setState("general-state", BadgeStatePriority.Default, generalState);
+          await badgeService.setState(
+            "general-state",
+            GeneralStateFunction(BadgeStatePriority.Default, generalState),
+          );
           await badgeService.setState(
             "tab-state",
-          BadgeStatePriority.Default,
+            TabSpecificStateFunction(BadgeStatePriority.Default, specificState, tabIds[0]),
-          specificState,
-          tabIds[0],
           );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
@@ -606,7 +757,10 @@ describe("BadgeService", () => {
             icon: BadgeIcon.Unlocked,
           };
 
-        await badgeService.setState("general-state", BadgeStatePriority.Default, generalState);
+          await badgeService.setState(
+            "general-state",
+            GeneralStateFunction(BadgeStatePriority.Default, generalState),
+          );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
           expect(badgeApi.specificStates).toEqual({
@@ -627,12 +781,13 @@ describe("BadgeService", () => {
             icon: BadgeIcon.Locked,
           };
 
-        await badgeService.setState("general-state", BadgeStatePriority.Default, generalState);
+          await badgeService.setState(
+            "general-state",
+            GeneralStateFunction(BadgeStatePriority.Default, generalState),
+          );
           await badgeService.setState(
             "tab-state",
-          BadgeStatePriority.Default,
+            TabSpecificStateFunction(BadgeStatePriority.Default, specificState, tabIds[0]),
-          specificState,
-          tabIds[0],
           );
 
           await new Promise((resolve) => setTimeout(resolve, 0));
@@ -642,6 +797,57 @@ describe("BadgeService", () => {
             [tabIds[2]]: undefined,
           });
         });
+
+        it("unsubscribes from state function when tab is deactivated", async () => {
+          let subscriptions = 0;
+          badgeService.setState("state", (tab) => {
+            return new Observable(() => {
+              subscriptions++;
+              return () => {
+                subscriptions--;
+              };
+            });
+          });
+
+          await new Promise((resolve) => setTimeout(resolve, 0));
+          expect(subscriptions).toBe(activeTabIds.length);
+
+          badgeApi.deactivateTab(activeTabIds[0]);
+          await new Promise((resolve) => setTimeout(resolve, 0));
+          expect(subscriptions).toBe(activeTabIds.length - 1);
         });
       });
     });
+  });
+});
+
+/**
+ * Creates a dynamic state function that only provides a state for a specific tab.
+ */
+function TabSpecificStateFunction(
+  priority: BadgeStatePriority,
+  state: BadgeState,
+  tabId: number,
+): BadgeStateFunction {
+  return (tab) => {
+    if (tab.tabId === tabId) {
+      return of({
+        priority,
+        state,
+      });
+    }
+
+    return EMPTY;
+  };
+}
+
+/**
+ * Creates a dynamic state function that provides the same state for all tabs.
+ */
+function GeneralStateFunction(priority: BadgeStatePriority, state: BadgeState): BadgeStateFunction {
+  return (_tab) =>
+    of({
+      priority,
+      state,
+    });
+}

apps/browser/src/platform/badge/badge.service.ts

@@ -1,69 +1,100 @@
-import { concatMap, filter, Subscription, withLatestFrom } from "rxjs";
+import {
+  BehaviorSubject,
+  combineLatest,
+  combineLatestWith,
+  concatMap,
+  debounceTime,
+  filter,
+  groupBy,
+  map,
+  mergeMap,
+  Observable,
+  of,
+  startWith,
+  Subscription,
+  switchMap,
+  takeUntil,
+} from "rxjs";
 
 import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
-import {
-  BADGE_MEMORY,
-  GlobalState,
-  KeyDefinition,
-  StateProvider,
-} from "@bitwarden/common/platform/state";
 
 import { BadgeBrowserApi, RawBadgeState, Tab } from "./badge-browser-api";
 import { DefaultBadgeState } from "./consts";
 import { BadgeStatePriority } from "./priority";
 import { BadgeState, Unset } from "./state";
 
-interface StateSetting {
+const BADGE_UPDATE_DEBOUNCE_MS = 100;
+
+export interface BadgeStateSetting {
   priority: BadgeStatePriority;
   state: BadgeState;
-  tabId?: number;
 }
 
-const BADGE_STATES = new KeyDefinition(BADGE_MEMORY, "badgeStates", {
-  deserializer: (value: Record<string, StateSetting>) => value ?? {},
-  cleanupDelayMs: 0,
-});
-
-export class BadgeService {
-  private serviceState: GlobalState<Record<string, StateSetting>>;
-
 /**
-   * An observable that emits whenever one or multiple tabs are updated and might need its state updated.
+ * A function that returns the badge state for a specific tab.
-   * Use this to know exactly which tabs to calculate the badge state for.
+ * Return `undefined` to clear any previously set state for the tab.
-   * This is not the same as `onActivated` which only emits when the active tab changes.
  */
-  activeTabsUpdated$ = this.badgeApi.activeTabsUpdated$;
+export type BadgeStateFunction = (tab: Tab) => Observable<BadgeStateSetting | undefined>;
 
-  getActiveTabs(): Promise<Tab[]> {
+export class BadgeService {
-    return this.badgeApi.getActiveTabs();
+  private stateFunctions = new BehaviorSubject<Record<string, BadgeStateFunction>>({});
-  }
 
   constructor(
-    private stateProvider: StateProvider,
     private badgeApi: BadgeBrowserApi,
     private logService: LogService,
-  ) {
+    private debounceTimeMs: number = BADGE_UPDATE_DEBOUNCE_MS,
-    this.serviceState = this.stateProvider.getGlobal(BADGE_STATES);
+  ) {}
-  }
 
   /**
    * Start listening for badge state changes.
    * Without this the service will not be able to update the badge state.
    */
   startListening(): Subscription {
-    // React to tab changes
+    // Default state function that always returns an empty state with lowest priority.
-    return this.badgeApi.activeTabsUpdated$
+    // This will ensure that there is always at least one state to consider when calculating the final badge state,
+    // so that the badge is cleared/set to default when no other states are set.
+    const defaultTabStateFunction: BadgeStateFunction = (_tab) =>
+      of({
+        priority: BadgeStatePriority.Low,
+        state: {},
+      });
+
+    return this.badgeApi.tabEvents$
       .pipe(
-        withLatestFrom(this.serviceState.state$),
+        groupBy((event) => (event.type === "deactivated" ? event.tabId : event.tab.tabId), {
-        filter(([activeTabs]) => activeTabs.length > 0),
+          duration: (group$) =>
-        concatMap(async ([activeTabs, serviceState]) => {
+            // Allow clean up of group when deactivated event arrives for this tabId
-          await Promise.all(activeTabs.map((tab) => this.updateBadge(serviceState, tab.tabId)));
+            group$.pipe(filter((evt) => evt.type === "deactivated")),
+        }),
+        mergeMap((group$) =>
+          group$.pipe(
+            // ignore deactivation events, only handle updates/activations
+            filter((evt) => evt.type !== "deactivated"),
+            map((evt) => evt.tab),
+            combineLatestWith(this.stateFunctions),
+            switchMap(([tab, dynamicStateFunctions]) => {
+              const functions = [...Object.values(dynamicStateFunctions), defaultTabStateFunction];
+
+              return combineLatest(functions.map((f) => f(tab).pipe(startWith(undefined)))).pipe(
+                map((states) => ({
+                  tab,
+                  states: states.filter((s): s is BadgeStateSetting => s !== undefined),
+                })),
+                debounceTime(this.debounceTimeMs),
+              );
+            }),
+            takeUntil(group$.pipe(filter((evt) => evt.type === "deactivated"))),
+          ),
+        ),
+
+        concatMap(async (tabUpdate) => {
+          await this.updateBadge(tabUpdate.states, tabUpdate.tab.tabId);
         }),
       )
       .subscribe({
         error: (error: unknown) => {
           this.logService.error(
-            "Fatal error in badge service observable, badge will fail to update",
+            "BadgeService: Fatal error updating badge state. Badge will no longer be updated.",
             error,
           );
         },
@@ -71,68 +102,45 @@ export class BadgeService {
   }
 
   /**
-   * Inform badge service of a new state that the badge should reflect.
+   * Register a function that takes an observable of active tab updates and returns an observable of state settings.
+   * This can be used to create dynamic badge states that react to tab changes.
+   * The returned observable should emit a new state setting whenever the badge state should be updated.
    *
-   * This will merge the new state with any existing states:
+   * This will merge all states:
    * - If the new state has a higher priority, it will override any lower priority states.
    * - If the new state has a lower priority, it will be ignored.
    * - If the name of the state is already in use, it will be updated.
    * - If the state has a `tabId` set, it will only apply to that tab.
    *   - States with `tabId` can still be overridden by states without `tabId` if they have a higher priority.
-   *
-   * @param name The name of the state. This is used to identify the state and will be used to clear it later.
-   * @param priority The priority of the state (higher numbers are higher priority, but setting arbitrary numbers is not supported).
-   * @param state The state to set.
-   * @param tabId Limit this badge state to a specific tab. If this is not set, the state will be applied to all tabs.
    */
-  async setState(name: string, priority: BadgeStatePriority, state: BadgeState, tabId?: number) {
+  setState(name: string, stateFunction: BadgeStateFunction) {
-    const newServiceState = await this.serviceState.update((s) => ({
+    this.stateFunctions.next({
-      ...s,
+      ...this.stateFunctions.value,
-      [name]: { priority, state, tabId },
+      [name]: stateFunction,
-    }));
+    });
-    await this.updateBadge(newServiceState, tabId);
   }
 
   /**
-   * Clear the state with the given name.
+   * Clear a state function previously registered with `setState`.
    *
-   * This will remove the state from the badge service and clear it from the badge.
+   * This will:
-   * If the state is not found, nothing will happen.
+   * - Stop the function from being called on future tab changes
+   * - Unsubscribe from any existing observables created by the function.
+   * - Clear any badge state previously set by the function.
    *
-   * @param name The name of the state to clear.
+   * @param name The name of the state function to clear.
    */
-  async clearState(name: string) {
+  clearState(name: string) {
-    let clearedState: StateSetting | undefined;
+    const currentDynamicStateFunctions = this.stateFunctions.value;
-
+    const newDynamicStateFunctions = { ...currentDynamicStateFunctions };
-    const newServiceState = await this.serviceState.update((s) => {
+    delete newDynamicStateFunctions[name];
-      clearedState = s?.[name];
+    this.stateFunctions.next(newDynamicStateFunctions);
-
-      const newStates = { ...s };
-      delete newStates[name];
-      return newStates;
-    });
-
-    if (clearedState === undefined) {
-      return;
-    }
-    await this.updateBadge(newServiceState, clearedState.tabId);
   }
 
-  private calculateState(states: Set<StateSetting>, tabId?: number): RawBadgeState {
+  private calculateState(states: BadgeStateSetting[]): RawBadgeState {
-    const sortedStates = [...states].sort((a, b) => a.priority - b.priority);
+    const sortedStates = states.sort((a, b) => a.priority - b.priority);
 
-    let filteredStates = sortedStates;
+    const mergedState = sortedStates
-    if (tabId !== undefined) {
-      // Filter out states that are not applicable to the current tab.
-      // If a state has no tabId, it is considered applicable to all tabs.
-      // If a state has a tabId, it is only applicable to that tab.
-      filteredStates = sortedStates.filter((s) => s.tabId === tabId || s.tabId === undefined);
-    } else {
-      // If no tabId is provided, we only want states that are not tab-specific.
-      filteredStates = sortedStates.filter((s) => s.tabId === undefined);
-    }
-
-    const mergedState = filteredStates
       .map((s) => s.state)
       .reduce<Partial<RawBadgeState>>((acc: Partial<RawBadgeState>, state: BadgeState) => {
         const newState = { ...acc };
@@ -156,45 +164,18 @@ export class BadgeService {
    * This will only update the badge if the active tab is the same as the tabId of the latest change.
    * If the active tab is not set, it will not update the badge.
    *
-   * @param activeTab The currently active tab.
    * @param serviceState The current state of the badge service. If this is null or undefined, an empty set will be assumed.
    * @param tabId Tab id for which the the latest state change applied to. Set this to activeTab.tabId to force an update.
+   * @param activeTabs The currently active tabs. If not provided, it will be fetched from the badge API.
    */
-  private async updateBadge(
+  private async updateBadge(serviceState: BadgeStateSetting[], tabId: number) {
-    serviceState: Record<string, StateSetting> | null | undefined,
+    const newBadgeState = this.calculateState(serviceState);
-    tabId: number | undefined,
-  ) {
-    const activeTabs = await this.badgeApi.getActiveTabs();
-    if (tabId !== undefined && !activeTabs.some((tab) => tab.tabId === tabId)) {
-      return; // No need to update the badge if the state is not for the active tab.
-    }
-
-    const tabIdsToUpdate = tabId ? [tabId] : activeTabs.map((tab) => tab.tabId);
-
-    for (const tabId of tabIdsToUpdate) {
-      if (tabId === undefined) {
-        continue; // Skip if tab id is undefined.
-      }
-
-      const newBadgeState = this.calculateState(new Set(Object.values(serviceState ?? {})), tabId);
     try {
       await this.badgeApi.setState(newBadgeState, tabId);
     } catch (error) {
       this.logService.error("Failed to set badge state", error);
     }
   }
-
-    if (tabId === undefined) {
-      // If no tabId was provided we should also update the general badge state
-      const newBadgeState = this.calculateState(new Set(Object.values(serviceState ?? {})));
-
-      try {
-        await this.badgeApi.setState(newBadgeState, tabId);
-      } catch (error) {
-        this.logService.error("Failed to set general badge state", error);
-      }
-    }
-  }
 }
 
 /**

apps/browser/src/platform/badge/scope.ts

@@ -0,0 +1,23 @@
+export const BadgeStateScope = {
+  /**
+   * The state is global and applies to all users.
+   */
+  Global: { type: "global" } satisfies BadgeStateScope,
+  /**
+   * The state is for a specific user and only applies to that user when they are unlocked.
+   */
+  UserUnlocked: (userId: string) =>
+    ({
+      type: "user_unlocked",
+      userId,
+    }) satisfies BadgeStateScope,
+} as const;
+
+export type BadgeStateScope =
+  | {
+      type: "global";
+    }
+  | {
+      type: "user_unlocked";
+      userId: string;
+    };

apps/browser/src/platform/badge/state.ts

@@ -1,6 +1,8 @@
 import { BadgeIcon } from "./icon";
 
-export const Unset = Symbol("Unset badge state");
+const UnsetValue = Symbol("Unset badge state");
+
+export const Unset = UnsetValue as typeof UnsetValue;
 export type Unset = typeof Unset;
 
 export type BadgeState = {

apps/browser/src/platform/badge/test/mock-badge-browser-api.ts

@@ -1,33 +1,50 @@
-import { BehaviorSubject } from "rxjs";
+import { BehaviorSubject, concat, defer, of, Subject, switchMap } from "rxjs";
 
-import { BadgeBrowserApi, RawBadgeState, Tab } from "../badge-browser-api";
+import { BadgeBrowserApi, RawBadgeState, Tab, TabEvent } from "../badge-browser-api";
 
 export class MockBadgeBrowserApi implements BadgeBrowserApi {
-  private _activeTabsUpdated$ = new BehaviorSubject<Tab[]>([]);
+  private _activeTabs$ = new BehaviorSubject<Tab[]>([]);
-  activeTabsUpdated$ = this._activeTabsUpdated$.asObservable();
+  private _tabEvents$ = new Subject<TabEvent>();
+  activeTabs$ = this._activeTabs$.asObservable();
 
   specificStates: Record<number, RawBadgeState> = {};
   generalState?: RawBadgeState;
   tabs: number[] = [];
-  activeTabs: number[] = [];
 
-  getActiveTabs(): Promise<Tab[]> {
+  tabEvents$ = concat(
-    return Promise.resolve(
+    defer(() => [this.activeTabs]).pipe(
-      this.activeTabs.map(
+      switchMap((activeTabs) => {
-        (tabId) =>
+        const tabEvents: TabEvent[] = activeTabs.map((tab) => ({
-          ({
+          type: "activated",
-            tabId,
+          tab,
-            url: `https://example.com/${tabId}`,
+        }));
-          }) satisfies Tab,
+        return of(...tabEvents);
+      }),
     ),
+    this._tabEvents$.asObservable(),
   );
+
+  get activeTabs() {
+    return this._activeTabs$.value;
   }
 
   setActiveTabs(tabs: number[]) {
-    this.activeTabs = tabs;
+    this._activeTabs$.next(tabs.map((tabId) => ({ tabId, url: `https://example.com/${tabId}` })));
-    this._activeTabsUpdated$.next(
+
-      tabs.map((tabId) => ({ tabId, url: `https://example.com/${tabId}` })),
+    tabs.forEach((tabId) => {
-    );
+      this._tabEvents$.next({
+        type: "activated",
+        tab: { tabId, url: `https://example.com/${tabId}` },
+      });
+    });
+  }
+
+  updateTab(tabId: number) {
+    this._tabEvents$.next({ type: "updated", tab: { tabId, url: `https://example.com/${tabId}` } });
+  }
+
+  deactivateTab(tabId: number) {
+    this._tabEvents$.next({ type: "deactivated", tabId });
   }
 
   setState = jest.fn().mockImplementation((state: RawBadgeState, tabId?: number): Promise<void> => {
@@ -39,8 +56,4 @@ export class MockBadgeBrowserApi implements BadgeBrowserApi {
 
     return Promise.resolve();
   });
-
-  getTabs(): Promise<number[]> {
-    return Promise.resolve(this.tabs);
-  }
 }

apps/browser/src/vault/services/at-risk-cipher-badge-updater.service.spec.ts

@@ -1,12 +1,11 @@
-import { BehaviorSubject } from "rxjs";
+import { BehaviorSubject, firstValueFrom } from "rxjs";
 
 import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
 import { CipherService } from "@bitwarden/common/vault/abstractions/cipher.service";
-import { SecurityTask, TaskService } from "@bitwarden/common/vault/tasks";
+import { SecurityTask, SecurityTaskType, TaskService } from "@bitwarden/common/vault/tasks";
-import { LogService } from "@bitwarden/logging";
-import { UserId } from "@bitwarden/user-core";
 
-import { BadgeService } from "../../platform/badge/badge.service";
+import { Tab } from "../../platform/badge/badge-browser-api";
+import { BadgeService, BadgeStateFunction } from "../../platform/badge/badge.service";
 import { BadgeIcon } from "../../platform/badge/icon";
 import { BadgeStatePriority } from "../../platform/badge/priority";
 import { Unset } from "../../platform/badge/state";
@@ -18,34 +17,32 @@ describe("AtRiskCipherBadgeUpdaterService", () => {
   let service: AtRiskCipherBadgeUpdaterService;
 
   let setState: jest.Mock;
-  let clearState: jest.Mock;
-  let warning: jest.Mock;
   let getAllDecryptedForUrl: jest.Mock;
   let getTab: jest.Mock;
   let addListener: jest.Mock;
 
-  const activeAccount$ = new BehaviorSubject({ id: "test-account-id" });
+  let activeAccount$: BehaviorSubject<{ id: string }>;
-  const cipherViews$ = new BehaviorSubject([]);
+  let cipherViews$: BehaviorSubject<Array<{ id: string; isDeleted?: boolean }>>;
-  const pendingTasks$ = new BehaviorSubject<SecurityTask[]>([]);
+  let pendingTasks$: BehaviorSubject<SecurityTask[]>;
-  const userId = "test-user-id" as UserId;
 
   beforeEach(async () => {
     setState = jest.fn().mockResolvedValue(undefined);
-    clearState = jest.fn().mockResolvedValue(undefined);
-    warning = jest.fn();
     getAllDecryptedForUrl = jest.fn().mockResolvedValue([]);
     getTab = jest.fn();
     addListener = jest.fn();
 
+    activeAccount$ = new BehaviorSubject({ id: "test-account-id" });
+    cipherViews$ = new BehaviorSubject<Array<{ id: string; isDeleted?: boolean }>>([]);
+    pendingTasks$ = new BehaviorSubject<SecurityTask[]>([]);
+
     jest.spyOn(BrowserApi, "addListener").mockImplementation(addListener);
     jest.spyOn(BrowserApi, "getTab").mockImplementation(getTab);
 
     service = new AtRiskCipherBadgeUpdaterService(
-      { setState, clearState } as unknown as BadgeService,
+      { setState } as unknown as BadgeService,
       { activeAccount$ } as unknown as AccountService,
-      { cipherViews$, getAllDecryptedForUrl } as unknown as CipherService,
+      { cipherViews$: () => cipherViews$, getAllDecryptedForUrl } as unknown as CipherService,
-      { warning } as unknown as LogService,
+      { pendingTasks$: () => pendingTasks$ } as unknown as TaskService,
-      { pendingTasks$ } as unknown as TaskService,
     );
 
     await service.init();
@@ -55,30 +52,41 @@ describe("AtRiskCipherBadgeUpdaterService", () => {
     jest.restoreAllMocks();
   });
 
+  it("registers dynamic state function on init", () => {
+    expect(setState).toHaveBeenCalledWith("at-risk-cipher-badge", expect.any(Function));
+  });
+
   it("clears the tab state when there are no ciphers and no pending tasks", async () => {
-    const tab = { id: 1 } as chrome.tabs.Tab;
+    const tab: Tab = { tabId: 1, url: "https://bitwarden.com" };
+    const stateFunction = setState.mock.calls[0][1];
 
-    await service["setTabState"](tab, userId, []);
+    const state = await firstValueFrom(stateFunction(tab));
 
-    expect(clearState).toHaveBeenCalledWith("at-risk-cipher-badge-1");
+    expect(state).toBeUndefined();
   });
 
   it("sets state when there are pending tasks for the tab", async () => {
-    const tab = { id: 3, url: "https://bitwarden.com" } as chrome.tabs.Tab;
+    const tab: Tab = { tabId: 3, url: "https://bitwarden.com" };
-    const pendingTasks: SecurityTask[] = [{ id: "task1", cipherId: "cipher1" } as SecurityTask];
+    const stateFunction: BadgeStateFunction = setState.mock.calls[0][1];
+    const pendingTasks: SecurityTask[] = [
+      {
+        id: "task1",
+        cipherId: "cipher1",
+        type: SecurityTaskType.UpdateAtRiskCredential,
+      } as SecurityTask,
+    ];
+    pendingTasks$.next(pendingTasks);
     getAllDecryptedForUrl.mockResolvedValueOnce([{ id: "cipher1" }]);
 
-    await service["setTabState"](tab, userId, pendingTasks);
+    const state = await firstValueFrom(stateFunction(tab));
 
-    expect(setState).toHaveBeenCalledWith(
+    expect(state).toEqual({
-      "at-risk-cipher-badge-3",
+      priority: BadgeStatePriority.High,
-      BadgeStatePriority.High,
+      state: {
-      {
         icon: BadgeIcon.Berry,
         text: Unset,
         backgroundColor: Unset,
       },
-      3,
+    });
-    );
   });
 });

apps/browser/src/vault/services/at-risk-cipher-badge-updater.service.ts

@@ -1,26 +1,18 @@
-import { combineLatest, map, mergeMap, of, Subject, switchMap } from "rxjs";
+import { combineLatest, concatMap, map, of, switchMap } from "rxjs";
 
 import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
-import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
-import { UserId } from "@bitwarden/common/types/guid";
 import { CipherService } from "@bitwarden/common/vault/abstractions/cipher.service";
-import { SecurityTask, SecurityTaskType, TaskService } from "@bitwarden/common/vault/tasks";
+import { SecurityTaskType, TaskService } from "@bitwarden/common/vault/tasks";
 import { filterOutNullish } from "@bitwarden/common/vault/utils/observable-utilities";
 
 import { BadgeService } from "../../platform/badge/badge.service";
 import { BadgeIcon } from "../../platform/badge/icon";
 import { BadgeStatePriority } from "../../platform/badge/priority";
 import { Unset } from "../../platform/badge/state";
-import { BrowserApi } from "../../platform/browser/browser-api";
 
-const StateName = (tabId: number) => `at-risk-cipher-badge-${tabId}`;
+const StateName = "at-risk-cipher-badge";
 
 export class AtRiskCipherBadgeUpdaterService {
-  private tabReplaced$ = new Subject<{ addedTab: chrome.tabs.Tab; removedTabId: number }>();
-  private tabUpdated$ = new Subject<chrome.tabs.Tab>();
-  private tabRemoved$ = new Subject<number>();
-  private tabActivated$ = new Subject<chrome.tabs.Tab>();
-
   private activeUserData$ = this.accountService.activeAccount$.pipe(
     filterOutNullish(),
     switchMap((user) =>
@@ -40,96 +32,13 @@ export class AtRiskCipherBadgeUpdaterService {
     private badgeService: BadgeService,
     private accountService: AccountService,
     private cipherService: CipherService,
-    private logService: LogService,
     private taskService: TaskService,
-  ) {
+  ) {}
-    combineLatest({
-      replaced: this.tabReplaced$,
-      activeUserData: this.activeUserData$,
-    })
-      .pipe(
-        mergeMap(async ({ replaced, activeUserData: [userId, pendingTasks] }) => {
-          await this.clearTabState(replaced.removedTabId);
-          await this.setTabState(replaced.addedTab, userId, pendingTasks);
-        }),
-      )
-      .subscribe(() => {});
-
-    combineLatest({
-      tab: this.tabActivated$,
-      activeUserData: this.activeUserData$,
-    })
-      .pipe(
-        mergeMap(async ({ tab, activeUserData: [userId, pendingTasks] }) => {
-          await this.setTabState(tab, userId, pendingTasks);
-        }),
-      )
-      .subscribe();
-
-    combineLatest({
-      tab: this.tabUpdated$,
-      activeUserData: this.activeUserData$,
-    })
-      .pipe(
-        mergeMap(async ({ tab, activeUserData: [userId, pendingTasks] }) => {
-          await this.setTabState(tab, userId, pendingTasks);
-        }),
-      )
-      .subscribe();
-
-    this.tabRemoved$
-      .pipe(
-        mergeMap(async (tabId) => {
-          await this.clearTabState(tabId);
-        }),
-      )
-      .subscribe();
-  }
 
   init() {
-    BrowserApi.addListener(chrome.tabs.onReplaced, async (addedTabId, removedTabId) => {
+    this.badgeService.setState(StateName, (tab) => {
-      const newTab = await BrowserApi.getTab(addedTabId);
+      return this.activeUserData$.pipe(
-      if (!newTab) {
+        concatMap(async ([userId, pendingTasks]) => {
-        this.logService.warning(
-          `Tab replaced event received but new tab not found (id: ${addedTabId})`,
-        );
-        return;
-      }
-
-      this.tabReplaced$.next({
-        removedTabId,
-        addedTab: newTab,
-      });
-    });
-
-    BrowserApi.addListener(chrome.tabs.onUpdated, (_, changeInfo, tab) => {
-      if (changeInfo.url) {
-        this.tabUpdated$.next(tab);
-      }
-    });
-
-    BrowserApi.addListener(chrome.tabs.onActivated, async (activeInfo) => {
-      const tab = await BrowserApi.getTab(activeInfo.tabId);
-      if (!tab) {
-        this.logService.warning(
-          `Tab activated event received but tab not found (id: ${activeInfo.tabId})`,
-        );
-        return;
-      }
-
-      this.tabActivated$.next(tab);
-    });
-
-    BrowserApi.addListener(chrome.tabs.onRemoved, (tabId, _) => this.tabRemoved$.next(tabId));
-  }
-
-  /** Sets the pending task state for the tab */
-  private async setTabState(tab: chrome.tabs.Tab, userId: UserId, pendingTasks: SecurityTask[]) {
-    if (!tab.id) {
-      this.logService.warning("Tab event received but tab id is undefined");
-      return;
-    }
-
           const ciphers = tab.url
             ? await this.cipherService.getAllDecryptedForUrl(tab.url, userId, [], undefined, true)
             : [];
@@ -139,25 +48,20 @@ export class AtRiskCipherBadgeUpdaterService {
           );
 
           if (!hasPendingTasksForTab) {
-      await this.clearTabState(tab.id);
+            return undefined;
-      return;
           }
 
-    await this.badgeService.setState(
+          return {
-      StateName(tab.id),
+            priority: BadgeStatePriority.High,
-      BadgeStatePriority.High,
+            state: {
-      {
               icon: BadgeIcon.Berry,
               // Unset text and background color to use default badge appearance
               text: Unset,
               backgroundColor: Unset,
             },
-      tab.id,
+          };
+        }),
       );
-  }
+    });
-
-  /** Clears the pending task state from a tab */
-  private async clearTabState(tabId: number) {
-    await this.badgeService.clearState(StateName(tabId));
   }
 }

libs/state/src/core/state-definitions.ts

@@ -111,9 +111,6 @@ export const NEW_WEB_LAYOUT_BANNER_DISK = new StateDefinition("newWebLayoutBanne
 export const APPLICATION_ID_DISK = new StateDefinition("applicationId", "disk", {
   web: "disk-local",
 });
-export const BADGE_MEMORY = new StateDefinition("badge", "memory", {
-  browser: "memory-large-object",
-});
 export const BIOMETRIC_SETTINGS_DISK = new StateDefinition("biometricSettings", "disk");
 export const CLEAR_EVENT_DISK = new StateDefinition("clearEvent", "disk");
 export const CONFIG_DISK = new StateDefinition("config", "disk", {