fix(background): update icon change logic; using TabState for sessions

src/background/services/monetization.ts

@@ -24,10 +24,6 @@ import { ALLOWED_PROTOCOLS } from '@/shared/defines'
 import type { PopupStore, Storage } from '@/shared/types'
 
 export class MonetizationService {
-  private sessions: {
-    [tabId: number]: Map<string, PaymentSession>
-  }
-
   constructor(
     private logger: Logger,
     private t: Translation,
@@ -37,7 +33,6 @@ export class MonetizationService {
     private events: EventsService,
     private tabState: TabState
   ) {
-    this.sessions = {}
     this.registerEventListeners()
   }
 
@@ -66,12 +61,8 @@ export class MonetizationService {
       return
     }
     const { tabId, frameId, url } = getSender(sender)
+    const sessions = this.tabState.getSessions(tabId)
 
-    if (this.sessions[tabId] == null) {
-      this.sessions[tabId] = new Map()
-    }
-
-    const sessions = this.sessions[tabId]
     const sessionsCount = sessions.size + payload.length
     const rate = computeRate(rateOfPay, sessionsCount)
 
@@ -106,9 +97,8 @@ export class MonetizationService {
   }
 
   async stopPaymentSessionsByTabId(tabId: number) {
-    const sessions = this.sessions[tabId]
-
-    if (!sessions?.size) {
+    const sessions = this.tabState.getSessions(tabId)
+    if (!sessions.size) {
       this.logger.debug(`No active sessions found for tab ${tabId}.`)
       return
     }
@@ -123,9 +113,9 @@ export class MonetizationService {
     sender: Runtime.MessageSender
   ) {
     const tabId = getTabId(sender)
-    const sessions = this.sessions[tabId]
+    const sessions = this.tabState.getSessions(tabId)
 
-    if (!sessions) {
+    if (!sessions.size) {
       this.logger.debug(`No active sessions found for tab ${tabId}.`)
       return
     }
@@ -155,9 +145,9 @@ export class MonetizationService {
     sender: Runtime.MessageSender
   ) {
     const tabId = getTabId(sender)
-    const sessions = this.sessions[tabId]
+    const sessions = this.tabState.getSessions(tabId)
 
-    if (!sessions?.size) {
+    if (!sessions.size) {
       this.logger.debug(`No active sessions found for tab ${tabId}.`)
       return
     }
@@ -177,8 +167,8 @@ export class MonetizationService {
   }
 
   async resumePaymentSessionsByTabId(tabId: number) {
-    const sessions = this.sessions[tabId]
-    if (!sessions?.size) {
+    const sessions = this.tabState.getSessions(tabId)
+    if (!sessions.size) {
       this.logger.debug(`No active sessions found for tab ${tabId}.`)
       return
     }
@@ -209,9 +199,9 @@ export class MonetizationService {
 
   clearTabSessions(tabId: number) {
     this.logger.debug(`Attempting to clear sessions for tab ${tabId}.`)
-    const sessions = this.sessions[tabId]
+    const sessions = this.tabState.getSessions(tabId)
 
-    if (!sessions) {
+    if (!sessions.size) {
       this.logger.debug(`No active sessions found for tab ${tabId}.`)
       return
     }
@@ -220,7 +210,6 @@ export class MonetizationService {
       session.stop()
     }
 
-    delete this.sessions[tabId]
     this.tabState.clearByTabId(tabId)
 
     this.logger.debug(`Cleared ${sessions.size} sessions for tab ${tabId}.`)
@@ -231,10 +220,8 @@ export class MonetizationService {
     if (!tab || !tab.id) {
       throw new Error('Could not find active tab.')
     }
-
-    const sessions = this.sessions[tab.id]
-
-    if (!sessions?.size) {
+    const sessions = this.tabState.getSessions(tab.id)
+    if (!sessions.size) {
       throw new Error('This website is not monetized.')
     }
 
@@ -284,13 +271,9 @@ export class MonetizationService {
   private onRateOfPayUpdate() {
     this.events.on('storage.rate_of_pay_update', ({ rate }) => {
       this.logger.debug("Received event='storage.rate_of_pay_update'")
-      Object.keys(this.sessions).forEach((tabId) => {
-        const tabSessions = this.sessions[tabId as unknown as number]
-        this.logger.debug(`Re-evaluating sessions amount for tab=${tabId}`)
-        for (const session of tabSessions.values()) {
-          session.adjustSessionAmount(rate)
-        }
-      })
+      for (const session of this.tabState.getAllSessions()) {
+        session.adjustSessionAmount(rate)
+      }
     })
   }
 
@@ -313,10 +296,8 @@ export class MonetizationService {
   }
 
   private stopAllSessions() {
-    for (const sessions of Object.values(this.sessions)) {
-      for (const session of sessions.values()) {
-        session.stop()
-      }
+    for (const session of this.tabState.getAllSessions()) {
+      session.stop()
     }
     this.logger.debug(`All payment sessions stopped.`)
   }
@@ -351,8 +332,7 @@ export class MonetizationService {
         // noop
       }
     }
-
-    const isSiteMonetized = tab?.id ? this.sessions[tab.id]?.size > 0 : false
+    const isSiteMonetized = this.tabState.getSessions(tab.id!).size > 0
 
     return {
       ...dataFromStorage,

src/background/services/sendToPopup.ts

@@ -8,6 +8,7 @@ import {
 export class SendToPopup {
   private isConnected = false
   private port: Runtime.Port
+  private queue = new Map<keyof BackgroundToPopupMessagesMap, any>()
 
   constructor(private browser: Browser) {}
 
@@ -19,6 +20,10 @@ export class SendToPopup {
       }
       this.port = port
       this.isConnected = true
+      for (const [type, data] of this.queue) {
+        this.send(type, data)
+        this.queue.delete(type)
+      }
       port.onDisconnect.addListener(() => {
         this.isConnected = false
       })
@@ -34,6 +39,7 @@ export class SendToPopup {
     data: BackgroundToPopupMessagesMap[T]
   ) {
     if (!this.isConnected) {
+      this.queue.set(type, data)
       return
     }
     const message = { type, data } as BackgroundToPopupMessage

src/background/services/tabEvents.ts

@@ -1,10 +1,15 @@
 import browser from 'webextension-polyfill'
 import type { Browser, Runtime, Tabs } from 'webextension-polyfill'
-import { MonetizationService } from './monetization'
-import { StorageService } from './storage'
 import { IsTabMonetizedPayload } from '@/shared/messages'
 import { getTabId } from '../utils'
 import { isOkState, type Translation } from '@/shared/helpers'
+import type {
+  MonetizationService,
+  SendToPopup,
+  StorageService,
+  TabState
+} from '.'
+import type { Storage, TabId } from '@/shared/types'
 
 const runtime = browser.runtime
 const ICONS = {
@@ -50,15 +55,24 @@ const ICONS = {
   }
 }
 
+type CallbackTabOnActivated = Parameters<
+  Browser['tabs']['onActivated']['addListener']
+>[0]
+type CallbackTabOnCreated = Parameters<
+  Browser['tabs']['onCreated']['addListener']
+>[0]
+
 export class TabEvents {
   constructor(
     private monetizationService: MonetizationService,
     private storage: StorageService,
+    private tabState: TabState,
+    private sendToPopup: SendToPopup,
     private t: Translation,
     private browser: Browser
   ) {}
   clearTabSessions = (
-    tabId: number,
+    tabId: TabId,
     changeInfo: Tabs.OnUpdatedChangeInfoType | Tabs.OnRemovedRemoveInfoType
   ) => {
     if (
@@ -69,34 +83,58 @@ export class TabEvents {
     }
   }
 
-  private changeIcon = async () => {
-    const { enabled } = await this.storage.get(['enabled'])
-    const iconData = enabled ? ICONS.default : ICONS.default_gray
-    await this.browser.action.setIcon({ path: iconData })
+  private updateVisualIndicators = async (
+    tabId?: TabId,
+    isTabMonetized?: boolean
+  ) => {
+    const { enabled, state } = await this.storage.get(['enabled', 'state'])
+
+    const { path, title, isMonetized } = this.getIconAndTooltip({
+      enabled,
+      state,
+      tabId,
+      isTabMonetized
+    })
+
+    this.sendToPopup.send('SET_IS_MONETIZED', isMonetized)
+    await this.browser.action.setIcon({ path, tabId })
+    await this.browser.action.setTitle({ title, tabId })
   }
 
-  onActivatedTab = async () => {
-    await this.changeIcon()
+  onActivatedTab: CallbackTabOnActivated = async (info) => {
+    await this.updateVisualIndicators(info.tabId)
   }
 
-  onCreatedTab = async () => {
-    await this.changeIcon()
+  onCreatedTab: CallbackTabOnCreated = async (tab) => {
+    await this.updateVisualIndicators(tab.id)
   }
 
   onUpdatedTab = async (
     payload?: IsTabMonetizedPayload | null,
     sender?: Runtime.MessageSender
   ) => {
-    const { enabled, state } = await this.storage.get(['enabled', 'state'])
+    const tabId = sender && getTabId(sender)
+    await this.updateVisualIndicators(tabId, payload?.value)
+  }
 
+  private getIconAndTooltip({
+    tabId,
+    enabled,
+    state,
+    isTabMonetized = tabId ? this.tabState.getSessions(tabId).size > 0 : false
+  }: {
+    enabled: Storage['enabled']
+    state: Storage['state']
+    tabId?: TabId
+    isTabMonetized?: boolean
+  }) {
     let title = this.t('appName')
     let iconData = ICONS.default
     if (!isOkState(state)) {
       iconData = enabled ? ICONS.enabled_warn : ICONS.disabled_warn
       const tabStateText = this.t('icon_state_actionRequired')
       title = `${title} - ${tabStateText}`
-    } else if (payload) {
-      const { value: isTabMonetized } = payload
+    } else {
       if (enabled) {
         iconData = isTabMonetized
           ? ICONS.enabled_hasLinks
@@ -111,9 +149,11 @@ export class TabEvents {
         : this.t('icon_state_monetizationInactive')
       title = `${title} - ${tabStateText}`
     }
-    const tabId = sender && getTabId(sender)
 
-    await this.browser.action.setIcon({ path: iconData, tabId })
-    await this.browser.action.setTitle({ title, tabId })
+    return {
+      path: iconData,
+      isMonetized: isTabMonetized,
+      title
+    }
   }
 }

src/background/services/tabState.ts

@@ -1,5 +1,6 @@
 import type { MonetizationEventDetails } from '@/shared/messages'
-import type { Tabs } from 'webextension-polyfill'
+import type { TabId } from '@/shared/types'
+import type { PaymentSession } from './paymentSession'
 
 type State = {
   monetizationEvent: MonetizationEventDetails
@@ -13,10 +14,11 @@ interface SaveOverpayingDetails {
   intervalInMs: number
 }
 
-type TabId = NonNullable<Tabs.Tab['id']>
+type SessionId = string
 
 export class TabState {
   private state = new Map<TabId, Map<string, State>>()
+  private sessions = new Map<TabId, Map<SessionId, PaymentSession>>()
 
   constructor() {}
 
@@ -73,7 +75,19 @@ export class TabState {
     }
   }
 
+  getSessions(tabId: TabId) {
+    if (!this.sessions.has(tabId)) {
+      this.sessions.set(tabId, new Map())
+    }
+    return this.sessions.get(tabId)!
+  }
+
+  getAllSessions() {
+    return [...this.sessions.values()].flatMap((s) => [...s.values()])
+  }
+
   clearByTabId(tabId: TabId) {
     this.state.delete(tabId)
+    this.sessions.delete(tabId)
   }
 }