[PM-24615][PM-24999] Adjust autofill badge updater to only calculate the active tab (#16163)

* wip * feat: refactor how we react to tab changes * feat: always begin me emitting all active tabs * feat: only calculate autofill for active tabs * fix: bug not properly listening to reloads * wip * fix: clean up * fix: clean up
2025-12-11 13:53:34 +00:00 · 2025-09-05 09:18:10 +02:00
parent da70fa541c
commit 1d9ebb028e
5 changed files with 101 additions and 136 deletions
--- a/apps/browser/src/autofill/services/autofill-badge-updater.service.ts
+++ b/apps/browser/src/autofill/services/autofill-badge-updater.service.ts
@@ -1,4 +1,4 @@
-import { combineLatest, distinctUntilChanged, mergeMap, of, Subject, switchMap } from "rxjs";
+import { combineLatest, distinctUntilChanged, mergeMap, of, switchMap, withLatestFrom } from "rxjs";

 import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
 import { BadgeSettingsServiceAbstraction } from "@bitwarden/common/autofill/services/badge-settings.service";
@@ -6,134 +6,77 @@ 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 { Tab } from "../../platform/badge/badge-browser-api";
 import { BadgeService } from "../../platform/badge/badge.service";
 import { BadgeStatePriority } from "../../platform/badge/priority";
-import { BrowserApi } from "../../platform/browser/browser-api";

 const StateName = (tabId: number) => `autofill-badge-${tabId}`;

 export class AutofillBadgeUpdaterService {
-  private tabReplaced$ = new Subject<{ addedTab: chrome.tabs.Tab; removedTabId: number }>();
-  private tabUpdated$ = new Subject<chrome.tabs.Tab>();
-  private tabRemoved$ = new Subject<number>();
-
  constructor(
    private badgeService: BadgeService,
    private accountService: AccountService,
    private cipherService: CipherService,
    private badgeSettingsService: BadgeSettingsServiceAbstraction,
    private logService: LogService,
-  ) {
-    const cipherViews$ = this.accountService.activeAccount$.pipe(
-      switchMap((account) => (account?.id ? this.cipherService.cipherViews$(account?.id) : of([]))),
+  ) {}
+
+  init() {
+    const ciphers$ = this.accountService.activeAccount$.pipe(
+      switchMap((account) => (account?.id ? this.cipherService.ciphers$(account?.id) : of([]))),
    );

+    // Recalculate badges for all active tabs when ciphers or active account changes
    combineLatest({
      account: this.accountService.activeAccount$,
      enableBadgeCounter:
        this.badgeSettingsService.enableBadgeCounter$.pipe(distinctUntilChanged()),
-      ciphers: cipherViews$,
+      ciphers: ciphers$,
    })
      .pipe(
-        mergeMap(async ({ account, enableBadgeCounter, ciphers }) => {
+        mergeMap(async ({ account, enableBadgeCounter }) => {
          if (!account) {
            return;
          }

-          const tabs = await BrowserApi.tabsQuery({});
+          const tabs = await this.badgeService.getActiveTabs();
+
          for (const tab of tabs) {
-            if (!tab.id) {
+            if (!tab.tabId) {
              continue;
            }
-
            if (enableBadgeCounter) {
              await this.setTabState(tab, account.id);
            } else {
-              await this.clearTabState(tab.id);
+              await this.clearTabState(tab.tabId);
            }
          }
        }),
      )
      .subscribe();

-    combineLatest({
-      account: this.accountService.activeAccount$,
-      enableBadgeCounter: this.badgeSettingsService.enableBadgeCounter$,
-      replaced: this.tabReplaced$,
-      ciphers: cipherViews$,
-    })
+    // Recalculate badge for a specific tab when it becomes active
+    this.badgeService.activeTabsUpdated$
      .pipe(
-        mergeMap(async ({ account, enableBadgeCounter, replaced }) => {
-          if (!account || !enableBadgeCounter) {
-            return;
-          }
-
-          await this.clearTabState(replaced.removedTabId);
-          await this.setTabState(replaced.addedTab, account.id);
-        }),
-      )
-      .subscribe();
-
-    combineLatest({
-      account: this.accountService.activeAccount$,
-      enableBadgeCounter: this.badgeSettingsService.enableBadgeCounter$,
-      tab: this.tabUpdated$,
-      ciphers: cipherViews$,
-    })
-      .pipe(
-        mergeMap(async ({ account, enableBadgeCounter, tab }) => {
+        withLatestFrom(
+          this.accountService.activeAccount$,
+          this.badgeSettingsService.enableBadgeCounter$,
+        ),
+        mergeMap(async ([tabs, account, enableBadgeCounter]) => {
          if (!account || !enableBadgeCounter) {
            return;
          }

+          for (const tab of tabs) {
            await this.setTabState(tab, account.id);
-        }),
-      )
-      .subscribe();
-
-    combineLatest({
-      account: this.accountService.activeAccount$,
-      enableBadgeCounter: this.badgeSettingsService.enableBadgeCounter$,
-      tabId: this.tabRemoved$,
-      ciphers: cipherViews$,
-    })
-      .pipe(
-        mergeMap(async ({ account, enableBadgeCounter, tabId }) => {
-          if (!account || !enableBadgeCounter) {
-            return;
          }
-
-          await this.clearTabState(tabId);
        }),
      )
      .subscribe();
  }

-  init() {
-    BrowserApi.addListener(chrome.tabs.onReplaced, async (addedTabId, removedTabId) => {
-      const newTab = await BrowserApi.getTab(addedTabId);
-      if (!newTab) {
-        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.onRemoved, (tabId, _) => this.tabRemoved$.next(tabId));
-  }
-
-  private async setTabState(tab: chrome.tabs.Tab, userId: UserId) {
-    if (!tab.id) {
+  private async setTabState(tab: Tab, userId: UserId) {
+    if (!tab.tabId) {
      this.logService.warning("Tab event received but tab id is undefined");
      return;
    }
@@ -142,18 +85,18 @@ export class AutofillBadgeUpdaterService {
    const cipherCount = ciphers.length;

    if (cipherCount === 0) {
-      await this.clearTabState(tab.id);
+      await this.clearTabState(tab.tabId);
      return;
    }

    const countText = cipherCount > 9 ? "9+" : cipherCount.toString();
    await this.badgeService.setState(
-      StateName(tab.id),
+      StateName(tab.tabId),
      BadgeStatePriority.Default,
      {
        text: countText,
      },
-      tab.id,
+      tab.tabId,
    );
  }

--- a/apps/browser/src/platform/badge/badge-browser-api.ts
+++ b/apps/browser/src/platform/badge/badge-browser-api.ts
@@ -15,13 +15,24 @@ export interface RawBadgeState {
  icon: BadgeIcon;
 }

+export interface Tab {
+  tabId: number;
+  url: string;
+}
+
+function tabFromChromeTab(tab: chrome.tabs.Tab): Tab {
+  return {
+    tabId: tab.id!,
+    url: tab.url!,
+  };
+}
+
 export interface BadgeBrowserApi {
-  activeTab$: Observable<chrome.tabs.TabActiveInfo | undefined>;
-  // activeTabs$: Observable<chrome.tabs.Tab[]>;
+  activeTabsUpdated$: Observable<Tab[]>;

  setState(state: RawBadgeState, tabId?: number): Promise<void>;
  getTabs(): Promise<number[]>;
-  getActiveTabs(): Promise<chrome.tabs.Tab[]>;
+  getActiveTabs(): Promise<Tab[]>;
 }

 export class DefaultBadgeBrowserApi implements BadgeBrowserApi {
@@ -29,34 +40,48 @@ export class DefaultBadgeBrowserApi implements BadgeBrowserApi {
  private sidebarAction = BrowserApi.getSidebarAction(self);

  private onTabActivated$ = fromChromeEvent(chrome.tabs.onActivated).pipe(
-    switchMap(async ([activeInfo]) => activeInfo),
+    map(([activeInfo]) => activeInfo),
    shareReplay({ bufferSize: 1, refCount: true }),
  );

-  activeTab$ = concat(
-    defer(async () => {
-      const currentTab = await BrowserApi.getTabFromCurrentWindow();
-      if (currentTab == null || currentTab.id === undefined) {
-        return undefined;
+  activeTabsUpdated$ = concat(
+    defer(async () => await this.getActiveTabs()),
+    merge(
+      this.onTabActivated$.pipe(
+        switchMap(async (activeInfo) => {
+          const tab = await BrowserApi.getTab(activeInfo.tabId);
+
+          if (tab == undefined || tab.id == undefined || tab.url == undefined) {
+            return [];
          }

-      return { tabId: currentTab.id, windowId: currentTab.windowId };
+          return [tabFromChromeTab(tab)];
        }),
-    merge(
-      this.onTabActivated$,
+      ),
      fromChromeEvent(chrome.tabs.onUpdated).pipe(
        filter(
          ([_, changeInfo]) =>
            // Only emit if the url was updated
            changeInfo.url != undefined,
        ),
-        map(([tabId, _changeInfo, tab]) => ({ tabId, windowId: tab.windowId })),
+        map(([_tabId, _changeInfo, tab]) => [tabFromChromeTab(tab)]),
      ),
+      fromChromeEvent(chrome.webNavigation.onCommitted).pipe(
+        map(([details]) => {
+          const toReturn: Tab[] =
+            details.transitionType === "reload" ? [{ tabId: details.tabId, url: details.url }] : [];
+          return toReturn;
+        }),
      ),
-  ).pipe(shareReplay({ bufferSize: 1, refCount: true }));
+      // 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));

-  getActiveTabs(): Promise<chrome.tabs.Tab[]> {
-    return BrowserApi.getActiveTabs();
+  async getActiveTabs(): Promise<Tab[]> {
+    const tabs = await BrowserApi.getActiveTabs();
+    return tabs.filter((tab) => tab.id != undefined && tab.url != undefined).map(tabFromChromeTab);
  }

  constructor(private platformUtilsService: PlatformUtilsService) {}
--- a/apps/browser/src/platform/badge/badge.service.spec.ts
+++ b/apps/browser/src/platform/badge/badge.service.spec.ts
@@ -39,7 +39,6 @@ describe("BadgeService", () => {
      beforeEach(() => {
        badgeApi.tabs = [tabId];
        badgeApi.setActiveTabs([tabId]);
-        badgeApi.setLastActivatedTab(tabId);
        badgeServiceSubscription = badgeService.startListening();
      });

@@ -195,7 +194,6 @@ describe("BadgeService", () => {
      beforeEach(() => {
        badgeApi.tabs = tabIds;
        badgeApi.setActiveTabs([tabId]);
-        badgeApi.setLastActivatedTab(tabId);
        badgeServiceSubscription = badgeService.startListening();
      });

@@ -240,7 +238,6 @@ describe("BadgeService", () => {
      beforeEach(() => {
        badgeApi.tabs = tabIds;
        badgeApi.setActiveTabs(activeTabIds);
-        badgeApi.setLastActivatedTab(1);
        badgeServiceSubscription = badgeService.startListening();
      });

@@ -286,7 +283,6 @@ describe("BadgeService", () => {
      beforeEach(() => {
        badgeApi.tabs = [tabId];
        badgeApi.setActiveTabs([tabId]);
-        badgeApi.setLastActivatedTab(tabId);
        badgeServiceSubscription = badgeService.startListening();
      });

@@ -562,7 +558,6 @@ describe("BadgeService", () => {
      beforeEach(() => {
        badgeApi.tabs = tabIds;
        badgeApi.setActiveTabs([tabId]);
-        badgeApi.setLastActivatedTab(tabId);
        badgeServiceSubscription = badgeService.startListening();
      });

@@ -595,14 +590,12 @@ describe("BadgeService", () => {
    });

    describe("given multiple tabs are open and multiple are active", () => {
-      const tabId = 1;
      const activeTabIds = [1, 2];
      const tabIds = [1, 2, 3];

      beforeEach(() => {
        badgeApi.tabs = tabIds;
        badgeApi.setActiveTabs(activeTabIds);
-        badgeApi.setLastActivatedTab(tabId);
        badgeServiceSubscription = badgeService.startListening();
      });

--- a/apps/browser/src/platform/badge/badge.service.ts
+++ b/apps/browser/src/platform/badge/badge.service.ts
@@ -8,7 +8,7 @@ import {
  StateProvider,
 } from "@bitwarden/common/platform/state";

-import { BadgeBrowserApi, RawBadgeState } from "./badge-browser-api";
+import { BadgeBrowserApi, RawBadgeState, Tab } from "./badge-browser-api";
 import { DefaultBadgeState } from "./consts";
 import { BadgeStatePriority } from "./priority";
 import { BadgeState, Unset } from "./state";
@@ -21,11 +21,23 @@ interface StateSetting {

 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.
+   * Use this to know exactly which tabs to calculate the badge state for.
+   * This is not the same as `onActivated` which only emits when the active tab changes.
+   */
+  activeTabsUpdated$ = this.badgeApi.activeTabsUpdated$;
+
+  getActiveTabs(): Promise<Tab[]> {
+    return this.badgeApi.getActiveTabs();
+  }
+
  constructor(
    private stateProvider: StateProvider,
    private badgeApi: BadgeBrowserApi,
@@ -40,12 +52,12 @@ export class BadgeService {
   */
  startListening(): Subscription {
    // React to tab changes
-    return this.badgeApi.activeTab$
+    return this.badgeApi.activeTabsUpdated$
      .pipe(
        withLatestFrom(this.serviceState.state$),
-        filter(([activeTab]) => activeTab != undefined),
-        concatMap(async ([activeTab, serviceState]) => {
-          await this.updateBadge(serviceState, activeTab!.tabId);
+        filter(([activeTabs]) => activeTabs.length > 0),
+        concatMap(async ([activeTabs, serviceState]) => {
+          await Promise.all(activeTabs.map((tab) => this.updateBadge(serviceState, tab.tabId)));
        }),
      )
      .subscribe({
@@ -78,7 +90,6 @@ export class BadgeService {
      ...s,
      [name]: { priority, state, tabId },
    }));
-
    await this.updateBadge(newServiceState, tabId);
  }

@@ -104,7 +115,6 @@ export class BadgeService {
    if (clearedState === undefined) {
      return;
    }
-    // const activeTabs = await firstValueFrom(this.badgeApi.activeTabs$);
    await this.updateBadge(newServiceState, clearedState.tabId);
  }

@@ -151,16 +161,15 @@ export class BadgeService {
   * @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(
-    // activeTabs: chrome.tabs.Tab[],
    serviceState: Record<string, StateSetting> | null | undefined,
    tabId: number | undefined,
  ) {
    const activeTabs = await this.badgeApi.getActiveTabs();
-    if (tabId !== undefined && !activeTabs.some((tab) => tab.id === tabId)) {
+    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.id);
+    const tabIdsToUpdate = tabId ? [tabId] : activeTabs.map((tab) => tab.tabId);

    for (const tabId of tabIdsToUpdate) {
      if (tabId === undefined) {
--- a/apps/browser/src/platform/badge/test/mock-badge-browser-api.ts
+++ b/apps/browser/src/platform/badge/test/mock-badge-browser-api.ts
@@ -1,38 +1,33 @@
 import { BehaviorSubject } from "rxjs";

-import { BadgeBrowserApi, RawBadgeState } from "../badge-browser-api";
+import { BadgeBrowserApi, RawBadgeState, Tab } from "../badge-browser-api";

 export class MockBadgeBrowserApi implements BadgeBrowserApi {
-  private _activeTab$ = new BehaviorSubject<chrome.tabs.TabActiveInfo | undefined>(undefined);
-  activeTab$ = this._activeTab$.asObservable();
+  private _activeTabsUpdated$ = new BehaviorSubject<Tab[]>([]);
+  activeTabsUpdated$ = this._activeTabsUpdated$.asObservable();

  specificStates: Record<number, RawBadgeState> = {};
  generalState?: RawBadgeState;
  tabs: number[] = [];
  activeTabs: number[] = [];

-  getActiveTabs(): Promise<chrome.tabs.Tab[]> {
+  getActiveTabs(): Promise<Tab[]> {
    return Promise.resolve(
      this.activeTabs.map(
        (tabId) =>
          ({
-            id: tabId,
-            windowId: 1,
-            active: true,
-          }) as chrome.tabs.Tab,
+            tabId,
+            url: `https://example.com/${tabId}`,
+          }) satisfies Tab,
      ),
    );
  }

  setActiveTabs(tabs: number[]) {
    this.activeTabs = tabs;
-  }
-
-  setLastActivatedTab(tabId: number) {
-    this._activeTab$.next({
-      tabId,
-      windowId: 1,
-    });
+    this._activeTabsUpdated$.next(
+      tabs.map((tabId) => ({ tabId, url: `https://example.com/${tabId}` })),
+    );
  }

  setState = jest.fn().mockImplementation((state: RawBadgeState, tabId?: number): Promise<void> => {