[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

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({
+    // Recalculate badge for a specific tab when it becomes active
-      account: this.accountService.activeAccount$,
+    this.badgeService.activeTabsUpdated$
-      enableBadgeCounter: this.badgeSettingsService.enableBadgeCounter$,
-      replaced: this.tabReplaced$,
-      ciphers: cipherViews$,
-    })
       .pipe(
-        mergeMap(async ({ account, enableBadgeCounter, replaced }) => {
+        withLatestFrom(
+          this.accountService.activeAccount$,
+          this.badgeSettingsService.enableBadgeCounter$,
+        ),
+        mergeMap(async ([tabs, account, enableBadgeCounter]) => {
           if (!account || !enableBadgeCounter) {
             return;
           }
 
-          await this.clearTabState(replaced.removedTabId);
+          for (const tab of tabs) {
-          await this.setTabState(replaced.addedTab, account.id);
+            await this.setTabState(tab, account.id);
-        }),
-      )
-      .subscribe();
-
-    combineLatest({
-      account: this.accountService.activeAccount$,
-      enableBadgeCounter: this.badgeSettingsService.enableBadgeCounter$,
-      tab: this.tabUpdated$,
-      ciphers: cipherViews$,
-    })
-      .pipe(
-        mergeMap(async ({ account, enableBadgeCounter, tab }) => {
-          if (!account || !enableBadgeCounter) {
-            return;
           }
-
-          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() {
+  private async setTabState(tab: Tab, userId: UserId) {
-    BrowserApi.addListener(chrome.tabs.onReplaced, async (addedTabId, removedTabId) => {
+    if (!tab.tabId) {
-      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) {
       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,
     );
   }
 

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>;
+  activeTabsUpdated$: Observable<Tab[]>;
-  // activeTabs$: Observable<chrome.tabs.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(
+  activeTabsUpdated$ = concat(
-    defer(async () => {
+    defer(async () => await this.getActiveTabs()),
-      const currentTab = await BrowserApi.getTabFromCurrentWindow();
-      if (currentTab == null || currentTab.id === undefined) {
-        return undefined;
-      }
-
-      return { tabId: currentTab.id, windowId: currentTab.windowId };
-    }),
     merge(
-      this.onTabActivated$,
+      this.onTabActivated$.pipe(
+        switchMap(async (activeInfo) => {
+          const tab = await BrowserApi.getTab(activeInfo.tabId);
+
+          if (tab == undefined || tab.id == undefined || tab.url == undefined) {
+            return [];
+          }
+
+          return [tabFromChromeTab(tab)];
+        }),
+      ),
       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(
-  ).pipe(shareReplay({ bufferSize: 1, refCount: true }));
+        map(([details]) => {
+          const toReturn: Tab[] =
+            details.transitionType === "reload" ? [{ tabId: details.tabId, url: details.url }] : [];
+          return toReturn;
+        }),
+      ),
+      // 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[]> {
+  async getActiveTabs(): Promise<Tab[]> {
-    return BrowserApi.getActiveTabs();
+    const tabs = await BrowserApi.getActiveTabs();
+    return tabs.filter((tab) => tab.id != undefined && tab.url != undefined).map(tabFromChromeTab);
   }
 
   constructor(private platformUtilsService: PlatformUtilsService) {}

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();
       });
 

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),
+        filter(([activeTabs]) => activeTabs.length > 0),
-        concatMap(async ([activeTab, serviceState]) => {
+        concatMap(async ([activeTabs, serviceState]) => {
-          await this.updateBadge(serviceState, activeTab!.tabId);
+          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) {

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);
+  private _activeTabsUpdated$ = new BehaviorSubject<Tab[]>([]);
-  activeTab$ = this._activeTab$.asObservable();
+  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,
+            tabId,
-            windowId: 1,
+            url: `https://example.com/${tabId}`,
-            active: true,
+          }) satisfies Tab,
-          }) as chrome.tabs.Tab,
       ),
     );
   }
 
   setActiveTabs(tabs: number[]) {
     this.activeTabs = tabs;
-  }
+    this._activeTabsUpdated$.next(
-
+      tabs.map((tabId) => ({ tabId, url: `https://example.com/${tabId}` })),
-  setLastActivatedTab(tabId: number) {
+    );
-    this._activeTab$.next({
-      tabId,
-      windowId: 1,
-    });
   }
 
   setState = jest.fn().mockImplementation((state: RawBadgeState, tabId?: number): Promise<void> => {