refactor: rewrite the badge service to only calculate states for active tab

This also fixes an issue where the `difference` function didn't work and caused
all tabs to update for every single state update.

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

@@ -1,4 +1,4 @@
-import { map, Observable } from "rxjs";
+import { concat, defer, distinctUntilChanged, map, Observable } from "rxjs";
 
 import { PlatformUtilsService } from "@bitwarden/common/platform/abstractions/platform-utils.service";
 import { Utils } from "@bitwarden/common/platform/misc/utils";
@@ -26,9 +26,20 @@ export class DefaultBadgeBrowserApi implements BadgeBrowserApi {
   private badgeAction = BrowserApi.getBrowserAction();
   private sidebarAction = BrowserApi.getSidebarAction(self);
 
-  activeTab$ = fromChromeEvent(chrome.tabs.onActivated).pipe(
-    map(([tabActiveInfo]) => tabActiveInfo),
-  );
+  activeTab$ = concat(
+    defer(async () => {
+      const currentTab = await BrowserApi.getTabFromCurrentWindow();
+      if (currentTab == null || currentTab.id === undefined) {
+        return undefined;
+      }
+
+      return {
+        tabId: currentTab.id,
+        windowId: currentTab.windowId,
+      } satisfies chrome.tabs.TabActiveInfo;
+    }),
+    fromChromeEvent(chrome.tabs.onActivated).pipe(map(([activeInfo]) => activeInfo)),
+  ).pipe(distinctUntilChanged((a, b) => a?.tabId === b?.tabId && a?.windowId === b?.windowId));
 
   constructor(private platformUtilsService: PlatformUtilsService) {}
 

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

@@ -213,6 +213,22 @@ describe("BadgeService", () => {
           3: undefined,
         });
       });
+
+      it("only updates the active tab when setting state", async () => {
+        const state: BadgeState = {
+          text: "text",
+          backgroundColor: "color",
+          icon: BadgeIcon.Locked,
+        };
+        badgeApi.setState.mockReset();
+
+        await badgeService.setState("state-1", BadgeStatePriority.Default, state, tabId);
+        await badgeService.setState("state-2", BadgeStatePriority.Default, state, 2);
+        await badgeService.setState("state-2", BadgeStatePriority.Default, state, 2);
+
+        await new Promise((resolve) => setTimeout(resolve, 0));
+        expect(badgeApi.setState).toHaveBeenCalledTimes(1);
+      });
     });
   });
 

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

@@ -1,13 +1,4 @@
-import {
-  combineLatest,
-  concatMap,
-  distinctUntilChanged,
-  filter,
-  map,
-  pairwise,
-  startWith,
-  Subscription,
-} from "rxjs";
+import { concatMap, firstValueFrom, Subscription, withLatestFrom } from "rxjs";
 
 import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
 import {
@@ -17,7 +8,6 @@ import {
   StateProvider,
 } from "@bitwarden/common/platform/state";
 
-import { difference } from "./array-utils";
 import { BadgeBrowserApi, RawBadgeState } from "./badge-browser-api";
 import { DefaultBadgeState } from "./consts";
 import { BadgeStatePriority } from "./priority";
@@ -30,18 +20,18 @@ interface StateSetting {
 }
 
 const BADGE_STATES = new KeyDefinition(BADGE_MEMORY, "badgeStates", {
-  deserializer: (value: Record<string, StateSetting>) => value ?? {},
+  deserializer: (value: Record<string, StateSetting>) => value ?? { states: {} },
 });
 
 export class BadgeService {
-  private states: GlobalState<Record<string, StateSetting>>;
+  private serviceState: GlobalState<Record<string, StateSetting>>;
 
   constructor(
     private stateProvider: StateProvider,
     private badgeApi: BadgeBrowserApi,
     private logService: LogService,
   ) {
-    this.states = this.stateProvider.getGlobal(BADGE_STATES);
+    this.serviceState = this.stateProvider.getGlobal(BADGE_STATES);
   }
 
   /**
@@ -49,47 +39,15 @@ export class BadgeService {
    * Without this the service will not be able to update the badge state.
    */
   startListening(): Subscription {
-    return combineLatest({
-      states: this.states.state$.pipe(
-        startWith({}),
-        distinctUntilChanged(),
-        map((states) => new Set(states ? Object.values(states) : [])),
-        pairwise(),
-        map(([previous, current]) => {
-          const [removed, added] = difference(previous, current);
-          return { all: current, removed, added };
-        }),
-        filter(({ removed, added }) => removed.size > 0 || added.size > 0),
-      ),
-      activeTab: this.badgeApi.activeTab$.pipe(startWith(undefined)),
-    })
+    // React to tab changes
+    return this.badgeApi.activeTab$
       .pipe(
-        concatMap(async ({ states, activeTab }) => {
-          const changed = [...states.removed, ...states.added];
-
-          // If the active tab wasn't changed, we don't need to update the badge.
-          if (!changed.some((s) => s.tabId === activeTab?.tabId || s.tabId === undefined)) {
-            return;
-          }
-
-          try {
-            const state = this.calculateState(states.all, activeTab?.tabId);
-            await this.badgeApi.setState(state, activeTab?.tabId);
-          } catch (error) {
-            // This usually happens when the user opens a popout because of how the browser treats it
-            // as a tab in the same window but then won't let you set the badge state for it.
-            this.logService.warning("Failed to set badge state", error);
-          }
+        withLatestFrom(this.serviceState.state$),
+        concatMap(async ([activeTab, serviceState]) => {
+          await this.updateBadge(activeTab, serviceState, activeTab?.tabId);
         }),
       )
-      .subscribe({
-        error: (err: unknown) => {
-          this.logService.error(
-            "Fatal error in badge service observable, badge will fail to update",
-            err,
-          );
-        },
-      });
+      .subscribe();
   }
 
   /**
@@ -108,7 +66,13 @@ export class BadgeService {
    * @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) {
-    await this.states.update((s) => ({ ...s, [name]: { priority, state, tabId } }));
+    const newServiceState = await this.serviceState.update((s) => ({
+      ...s,
+      [name]: { priority, state, tabId },
+    }));
+
+    const activeTab = await firstValueFrom(this.badgeApi.activeTab$);
+    await this.updateBadge(activeTab, newServiceState, tabId);
   }
 
   /**
@@ -120,11 +84,21 @@ export class BadgeService {
    * @param name The name of the state to clear.
    */
   async clearState(name: string) {
-    await this.states.update((s) => {
+    let clearedState: StateSetting | undefined;
+
+    const newServiceState = await this.serviceState.update((s) => {
+      clearedState = s?.[name];
+
       const newStates = { ...s };
       delete newStates[name];
       return newStates;
     });
+
+    if (clearedState === undefined) {
+      return;
+    }
+    const activeTab = await firstValueFrom(this.badgeApi.activeTab$);
+    await this.updateBadge(activeTab, newServiceState, clearedState.tabId);
   }
 
   private calculateState(states: Set<StateSetting>, tabId?: number): RawBadgeState {
@@ -159,6 +133,40 @@ export class BadgeService {
       ...mergedState,
     };
   }
+
+  /**
+   * Common function deduplicating the logic for updating the badge with the current state.
+   * 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.
+   */
+  private async updateBadge(
+    activeTab: chrome.tabs.TabActiveInfo | null | undefined,
+    serviceState: Record<string, StateSetting> | null | undefined,
+    tabId: number | undefined,
+  ) {
+    if (activeTab === undefined) {
+      return; // If there is no active tab, we cannot set the badge state.
+    }
+
+    if (tabId !== activeTab?.tabId && tabId !== undefined) {
+      return; // No need to update the badge if the state is not for the active tab.
+    }
+
+    const newBadgeState = this.calculateState(
+      new Set(Object.values(serviceState ?? {})),
+      activeTab?.tabId,
+    );
+
+    try {
+      await this.badgeApi.setState(newBadgeState, activeTab?.tabId);
+    } catch (error) {
+      this.logService.error("Failed to set badge state", error);
+    }
+  }
 }
 
 /**

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

@@ -17,7 +17,7 @@ export class MockBadgeBrowserApi implements BadgeBrowserApi {
     });
   }
 
-  setState(state: RawBadgeState, tabId?: number): Promise<void> {
+  setState = jest.fn().mockImplementation((state: RawBadgeState, tabId?: number): Promise<void> => {
     if (tabId !== undefined) {
       this.specificStates[tabId] = state;
     } else {
@@ -25,7 +25,7 @@ export class MockBadgeBrowserApi implements BadgeBrowserApi {
     }
 
     return Promise.resolve();
-  }
+  });
 
   getTabs(): Promise<number[]> {
     return Promise.resolve(this.tabs);