diff --git a/docs/real-time-balance-updates-flow.md b/docs/real-time-balance-updates-flow.md new file mode 100644 index 00000000000..3884643ab74 --- /dev/null +++ b/docs/real-time-balance-updates-flow.md @@ -0,0 +1,253 @@ +# Real-Time Balance Updates and Status Management Flow + +This document describes the architecture and flow for real-time balance updates and WebSocket status management in MetaMask Core, specifically focusing on the `AccountActivityService:balanceUpdated` and `AccountActivityService:statusChanged` events. + +## Overview + +The system provides real-time balance updates and intelligent polling management through a multi-layered architecture that combines WebSocket streaming with fallback HTTP polling. The key components work together to ensure users receive timely balance updates while optimizing network usage and battery consumption. + +## Architecture Components + +### 1. BackendWebSocketService + +- **Purpose**: Low-level WebSocket connection management +- **Responsibilities**: + - Maintains WebSocket connection with automatic reconnection + - Handles subscription management + - Routes incoming messages to registered callbacks + - Publishes connection state changes + +### 2. AccountActivityService + +- **Purpose**: High-level account activity monitoring +- **Responsibilities**: + - Subscribes to selected account activity + - Processes transaction and balance updates + - Emits `balanceUpdated` and `statusChanged` events + - Manages chain status based on WebSocket connectivity and system notifications + +### 3. TokenBalancesController + +- **Purpose**: Token balance state management and intelligent polling +- **Responsibilities**: + - Maintains token balance state for all accounts + - Implements per-chain configurable polling intervals + - Responds to real-time balance updates from AccountActivityService + - Dynamically adjusts polling based on WebSocket availability + - Imports newly detected tokens via TokenDetectionController + +## Event Flow + +### Balance Update Flow + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ BALANCE UPDATE FLOW │ +└─────────────────────────────────────────────────────────────────────────┘ + +1. WebSocket receives account activity message + ↓ +2. BackendWebSocketService routes message to registered callback + ↓ +3. AccountActivityService processes AccountActivityMessage + { + address: "0x123...", + tx: { hash: "0x...", chain: "eip155:1", status: "completed", ... }, + updates: [ + { + asset: { fungible: true, type: "eip155:1/erc20:0x...", unit: "USDT" }, + postBalance: { amount: "1254.75" }, + transfers: [{ from: "0x...", to: "0x...", amount: "500.00" }] + } + ] + } + ↓ +4. AccountActivityService publishes separate events: + - AccountActivityService:transactionUpdated (transaction data) + - AccountActivityService:balanceUpdated (balance updates) + ↓ +5. TokenBalancesController receives balanceUpdated event + ↓ +6. TokenBalancesController processes balance updates: + a. Parses CAIP chain ID (e.g., "eip155:1" → "0x1") + b. Parses asset types: + - ERC20 tokens: "eip155:1/erc20:0x..." → token address + - Native tokens: "eip155:1/slip44:60" → zero address + c. Validates addresses and checksums them + d. Checks if tokens are tracked (imported or detected) + ↓ +7. For tracked tokens: + - Updates tokenBalances state immediately + - Updates AccountTrackerController for native balances + ↓ +8. For untracked ERC20 tokens: + - Queues tokens for import via TokenDetectionController + - Triggers fallback polling to fetch newly imported token balances + ↓ +9. On errors: + - Falls back to HTTP polling for the affected chain +``` + +### Status Change Flow + +The system manages chain status through two primary mechanisms: + +#### A. WebSocket Connection State Changes + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ WEBSOCKET CONNECTION STATUS FLOW │ +└─────────────────────────────────────────────────────────────────────────┘ + +1. BackendWebSocketService detects connection state change + (CONNECTING → CONNECTED | DISCONNECTED | ERROR) + ↓ +2. BackendWebSocketService publishes: + BackendWebSocketService:connectionStateChanged + ↓ +3. AccountActivityService receives connection state change + ↓ +4. AccountActivityService determines affected chains: + - Fetches list of supported chains from backend API + - Example: ["eip155:1", "eip155:137", "eip155:56"] + ↓ +5. AccountActivityService publishes status based on connection state: + + IF state === CONNECTED: + → Publishes: statusChanged { chainIds: [...], status: 'up' } + → Triggers resubscription to selected account + + IF state === DISCONNECTED || ERROR: + → Publishes: statusChanged { chainIds: [...], status: 'down' } + ↓ +6. TokenBalancesController receives statusChanged event + ↓ +7. TokenBalancesController applies debouncing (5 second window) + - Accumulates status changes to prevent excessive updates + - Latest status wins for each chain + ↓ +8. After debounce period, processes accumulated changes: + - Converts CAIP format to hex (e.g., "eip155:1" → "0x1") + - Calculates new polling intervals: + * status = 'down' → Uses default interval (30 seconds) + * status = 'up' → Uses extended interval (5 minutes) + ↓ +9. Adds jitter delay (0 to default interval) + - Prevents synchronized requests across instances + ↓ +10. Updates chain polling configurations + - Triggers immediate balance fetch + - Restarts polling with new intervals +``` + +#### B. System Notifications (Per-Chain Status) + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ SYSTEM NOTIFICATION STATUS FLOW │ +└─────────────────────────────────────────────────────────────────────────┘ + +1. WebSocket receives system notification message + { + type: 'system', + chainIds: ['eip155:1'], // Specific affected chains + status: 'down' // or 'up' + } + ↓ +2. BackendWebSocketService routes to AccountActivityService + ↓ +3. AccountActivityService validates notification: + - Ensures chainIds array is present and valid + - Ensures status is present + ↓ +4. AccountActivityService publishes delta update: + AccountActivityService:statusChanged + { + chainIds: ['eip155:1'], // Only affected chains + status: 'down' + } + ↓ +5. TokenBalancesController processes (same as WebSocket flow above) +``` + +#### Status Change Event Format + +```typescript +// Event published by AccountActivityService +AccountActivityService:statusChanged +Payload: { + chainIds: string[]; // Array of CAIP chain IDs (e.g., ["eip155:1", "eip155:137"]) + status: 'up' | 'down'; // Connection status +} +``` + +## Polling Strategy + +The TokenBalancesController implements intelligent polling that adapts based on WebSocket availability: + +### Polling Intervals + +| Scenario | Interval | Reason | +| ----------------------------------------- | -------------------- | --------------------------------------------------- | +| WebSocket Connected (`status: 'up'`) | 5 minutes | Real-time updates available, polling is backup only | +| WebSocket Disconnected (`status: 'down'`) | 30 seconds (default) | Primary update mechanism, needs faster polling | +| Per-chain custom configuration | Configurable | Allows fine-tuning per chain requirements | + +### Debouncing Strategy + +To prevent excessive HTTP calls during unstable connections: + +1. **Accumulation Window**: 5 seconds + + - All status changes within this window are accumulated + - Latest status wins for each chain + +2. **Jitter Addition**: Random delay (0 to default interval) + + - Prevents synchronized requests across multiple instances + - Reduces backend load spikes + +3. **Batch Processing**: After debounce + jitter + - All accumulated changes applied at once + - Single polling configuration update + - Immediate balance fetch triggered + +### Per-Chain Polling Configuration + +TokenBalancesController supports per-chain polling intervals: + +```typescript +// Configure custom intervals for specific chains +tokenBalancesController.updateChainPollingConfigs({ + '0x1': { interval: 30000 }, // Ethereum: 30 seconds (default) + '0x89': { interval: 15000 }, // Polygon: 15 seconds (faster) + '0xa4b1': { interval: 60000 }, // Arbitrum: 1 minute (slower) +}); +``` + +## Token Discovery Flow + +When balance updates include previously unknown tokens: + +``` +1. TokenBalancesController receives balance update for unknown token + ↓ +2. Checks if token is tracked (in allTokens or allIgnoredTokens) + ↓ +3. If NOT tracked: + a. Queues token for import + b. Calls TokenDetectionController:addDetectedTokensViaWs + c. Token is added to detected tokens list + ↓ +4. Triggers balance fetch for the chain + ↓ +5. New token balance is fetched and state is updated +``` + +## References + +- [`TokenBalancesController.ts`](../packages/assets-controllers/src/TokenBalancesController.ts) - Main controller implementation +- [`AccountActivityService.ts`](../packages/core-backend/src/AccountActivityService.ts) - Account activity monitoring +- [`BackendWebSocketService.ts`](../packages/core-backend/src/BackendWebSocketService.ts) - WebSocket connection management +- [`types.ts`](../packages/core-backend/src/types.ts) - Type definitions +- [Core Backend README](../packages/core-backend/README.md) - Package overview diff --git a/eslint-warning-thresholds.json b/eslint-warning-thresholds.json index eb99ca3ffa4..254e13963dd 100644 --- a/eslint-warning-thresholds.json +++ b/eslint-warning-thresholds.json @@ -54,16 +54,6 @@ "packages/assets-controllers/src/Standards/NftStandards/ERC721/ERC721Standard.ts": { "prettier/prettier": 1 }, - "packages/assets-controllers/src/TokenDetectionController.test.ts": { - "import-x/namespace": 11, - "jsdoc/tag-lines": 1 - }, - "packages/assets-controllers/src/TokenDetectionController.ts": { - "@typescript-eslint/prefer-readonly": 3, - "jsdoc/check-tag-names": 8, - "jsdoc/tag-lines": 6, - "no-unused-private-class-members": 2 - }, "packages/assets-controllers/src/TokenListController.test.ts": { "import-x/namespace": 7, "import-x/order": 3, diff --git a/packages/assets-controllers/CHANGELOG.md b/packages/assets-controllers/CHANGELOG.md index 3e2b3121898..16448c727cb 100644 --- a/packages/assets-controllers/CHANGELOG.md +++ b/packages/assets-controllers/CHANGELOG.md @@ -7,8 +7,28 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -### Changed +### Added +- Add real-time balance updates via WebSocket integration with `AccountActivityService` to `TokenBalancesController` ([#6784](https://github.com/MetaMask/core/pull/6784)) + - Add `@metamask/core-backend` as a dependency and peer dependency ([#6784](https://github.com/MetaMask/core/pull/6784)) + - Controller now subscribes to `AccountActivityService:balanceUpdated` events for instant balance updates + - Add support for real-time balance updates for both ERC20 tokens and native tokens + - Add `TokenDetectionController:addDetectedTokensViaWs` action handler for adding tokens detected via WebSocket + - Controller now subscribes to `AccountActivityService:statusChanged` events to dynamically adjust polling intervals + - When WebSocket service is "up", polling interval increases to backup interval (5 minutes) + - When WebSocket service is "down", polling interval restores to default interval (30 seconds) + - Status changes are debounced (5 seconds) and jittered to prevent thundering herd + - Add fallback to polling when balance updates contain errors or unsupported asset types + +### Changed + +- **BREAKING:** `TokenBalancesController` messenger must now allow `AccountActivityService:balanceUpdated` and `AccountActivityService:statusChanged` events ([#6784](https://github.com/MetaMask/core/pull/6784)) +- **BREAKING:** `TokenBalancesController` messenger must now allow `TokenDetectionController:addDetectedTokensViaWs` action ([#6784](https://github.com/MetaMask/core/pull/6784)) +- **BREAKING:** Change `TokenBalancesController` default polling interval to 30 seconds (was 180 seconds) ([#6784](https://github.com/MetaMask/core/pull/6784)) + - With real-time WebSocket updates, the default interval only applies when WebSocket is disconnected + - When WebSocket is connected, polling automatically adjusts to 5 minutes as a backup +- Update `TokenBalancesController` README documentation to mention real-time balance updates via WebSocket and intelligent polling management ([#6784](https://github.com/MetaMask/core/pull/6784)) +- `TokenDetectionController` code cleanup: remove unused private properties and ESLint disable comments ([#6784](https://github.com/MetaMask/core/pull/6784)) - **Performance Optimization:** Remove collection API calls from NFT detection process ([#6762](https://github.com/MetaMask/core/pull/6762)) - Reduce NFT detection API calls by 83% (from 6 calls to 1 call per 100 tokens) by eliminating collection endpoint requests - Remove unused collection metadata fields: `contractDeployedAt`, `creator`, and `topBid` diff --git a/packages/assets-controllers/README.md b/packages/assets-controllers/README.md index 7f7ed3f26af..e112258c4e8 100644 --- a/packages/assets-controllers/README.md +++ b/packages/assets-controllers/README.md @@ -21,7 +21,7 @@ This package features the following controllers: - [**CurrencyRateController**](src/CurrencyRateController.ts) keeps a periodically updated value of the exchange rate from the currently selected "native" currency to another (handling testnet tokens specially). - [**DeFiPositionsController**](src/DeFiPositionsController/DeFiPositionsController.ts.ts) keeps a periodically updated value of the DeFi positions for the owner EVM addresses. - [**RatesController**](src/RatesController/RatesController.ts) keeps a periodically updated value for the exchange rates for different cryptocurrencies. The difference between the `RatesController` and `CurrencyRateController` is that the second one is coupled to the `NetworksController` and is EVM specific, whilst the first one can handle different blockchain currencies like BTC and SOL. -- [**TokenBalancesController**](src/TokenBalancesController.ts) keeps a periodically updated set of balances for the current set of ERC-20 tokens. +- [**TokenBalancesController**](src/TokenBalancesController.ts) keeps a periodically updated set of balances for the current set of ERC-20 tokens. It supports real-time balance updates via WebSocket and intelligent polling management. See [Real-Time Balance Updates Flow](./REAL_TIME_BALANCE_UPDATES.md) for details. - [**TokenDetectionController**](src/TokenDetectionController.ts) keeps a periodically updated list of ERC-20 tokens assigned to the currently selected address. - [**TokenListController**](src/TokenListController.ts) uses the MetaSwap API to keep a periodically updated list of known ERC-20 tokens along with their metadata. - [**TokenRatesController**](src/TokenRatesController.ts) keeps a periodically updated list of exchange rates for known ERC-20 tokens relative to the currently selected native currency. diff --git a/packages/assets-controllers/jest.config.js b/packages/assets-controllers/jest.config.js index 8baa2d75778..c7e580801e6 100644 --- a/packages/assets-controllers/jest.config.js +++ b/packages/assets-controllers/jest.config.js @@ -24,7 +24,7 @@ module.exports = merge(baseConfig, { coverageThreshold: { global: { branches: 90.5, - functions: 99.22, + functions: 99, lines: 98, statements: 98, }, diff --git a/packages/assets-controllers/package.json b/packages/assets-controllers/package.json index 33021679ac8..6e80d4ce8e8 100644 --- a/packages/assets-controllers/package.json +++ b/packages/assets-controllers/package.json @@ -84,6 +84,7 @@ "@metamask/accounts-controller": "^33.1.1", "@metamask/approval-controller": "^7.2.0", "@metamask/auto-changelog": "^3.4.4", + "@metamask/core-backend": "^1.0.1", "@metamask/ethjs-provider-http": "^0.3.0", "@metamask/keyring-controller": "^23.1.1", "@metamask/keyring-internal-api": "^9.0.0", @@ -114,6 +115,7 @@ "@metamask/account-tree-controller": "^1.0.0", "@metamask/accounts-controller": "^33.0.0", "@metamask/approval-controller": "^7.0.0", + "@metamask/core-backend": "^1.0.0", "@metamask/keyring-controller": "^23.0.0", "@metamask/network-controller": "^24.0.0", "@metamask/permission-controller": "^11.0.0", diff --git a/packages/assets-controllers/src/TokenBalancesController.test.ts b/packages/assets-controllers/src/TokenBalancesController.test.ts index ff0b11e4116..b6e3f4f4f8c 100644 --- a/packages/assets-controllers/src/TokenBalancesController.test.ts +++ b/packages/assets-controllers/src/TokenBalancesController.test.ts @@ -13,11 +13,16 @@ import type { AllowedActions, AllowedEvents, ChainIdHex, + ChecksumAddress, TokenBalancesControllerActions, TokenBalancesControllerEvents, TokenBalancesControllerState, } from './TokenBalancesController'; -import { TokenBalancesController } from './TokenBalancesController'; +import { + TokenBalancesController, + caipChainIdToHex, + parseAssetType, +} from './TokenBalancesController'; import type { TokensControllerState } from './TokensController'; import { advanceTime, flushPromises } from '../../../tests/helpers'; import { createMockInternalAccount } from '../../accounts-controller/src/tests/mocks'; @@ -41,7 +46,7 @@ const mockedSafelyExecuteWithTimeout = safelyExecuteWithTimeout as jest.Mock; const setupController = ({ config, - tokens = { allTokens: {}, allDetectedTokens: {} }, + tokens = { allTokens: {}, allDetectedTokens: {}, allIgnoredTokens: {} }, listAccounts = [], }: { config?: Partial[0]>; @@ -60,6 +65,7 @@ const setupController = ({ 'NetworkController:getNetworkClientById', 'PreferencesController:getState', 'TokensController:getState', + 'TokenDetectionController:addDetectedTokensViaWs', 'AccountsController:getSelectedAccount', 'AccountsController:listAccounts', 'AccountTrackerController:getState', @@ -71,6 +77,8 @@ const setupController = ({ 'PreferencesController:stateChange', 'TokensController:stateChange', 'KeyringController:accountRemoved', + 'AccountActivityService:balanceUpdated', + 'AccountActivityService:statusChanged', ], }); @@ -173,6 +181,85 @@ const setupController = ({ }; }; +describe('Utility Functions', () => { + describe('caipChainIdToHex', () => { + it('should convert valid CAIP chain ID to hex', () => { + expect(caipChainIdToHex('eip155:1')).toBe('0x1'); + expect(caipChainIdToHex('eip155:137')).toBe('0x89'); + expect(caipChainIdToHex('eip155:42161')).toBe('0xa4b1'); + }); + + it('should return hex string unchanged if already in hex format', () => { + expect(caipChainIdToHex('0x1')).toBe('0x1'); + expect(caipChainIdToHex('0x89')).toBe('0x89'); + expect(caipChainIdToHex('0xa4b1')).toBe('0xa4b1'); + }); + + it('should throw error for invalid CAIP chain ID format', () => { + expect(() => caipChainIdToHex('invalid-chain-id')).toThrow( + 'caipChainIdToHex - Failed to provide CAIP-2 or Hex chainId', + ); + expect(() => caipChainIdToHex('eip155')).toThrow( + 'caipChainIdToHex - Failed to provide CAIP-2 or Hex chainId', + ); + expect(() => caipChainIdToHex('not-caip-format')).toThrow( + 'caipChainIdToHex - Failed to provide CAIP-2 or Hex chainId', + ); + }); + + it('should throw error for empty string', () => { + expect(() => caipChainIdToHex('')).toThrow( + 'caipChainIdToHex - Failed to provide CAIP-2 or Hex chainId', + ); + }); + }); + + describe('parseAssetType', () => { + it('should parse ERC20 token asset type correctly', () => { + const result = parseAssetType( + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', + ); + expect(result).toStrictEqual([ + '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', + false, + ]); + }); + + it('should parse native token asset type (slip44) correctly', () => { + const result = parseAssetType('eip155:1/slip44:60'); + expect(result).toStrictEqual([ + '0x0000000000000000000000000000000000000000', + true, + ]); + }); + + it('should return null for invalid CAIP asset type format', () => { + expect(parseAssetType('not-a-caip-format')).toBeNull(); + expect(parseAssetType('eip155:1')).toBeNull(); + expect(parseAssetType('invalid/format')).toBeNull(); + expect(parseAssetType('')).toBeNull(); + }); + + it('should return null for unsupported asset namespace', () => { + const result = parseAssetType('eip155:1/unknown:0x123'); + expect(result).toBeNull(); + }); + + it('should handle different chain references', () => { + expect( + parseAssetType( + 'eip155:137/erc20:0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174', + ), + ).toStrictEqual(['0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174', false]); + + expect(parseAssetType('eip155:137/slip44:60')).toStrictEqual([ + '0x0000000000000000000000000000000000000000', + true, + ]); + }); + }); +}); + describe('TokenBalancesController', () => { let clock: sinon.SinonFakeTimers; @@ -675,6 +762,172 @@ describe('TokenBalancesController', () => { }); }); + it('should only update balances for tokens in allTokens or allIgnoredTokens', async () => { + const chainId = '0x1'; + const accountAddress = '0x0000000000000000000000000000000000000000'; + const trackedToken = '0x0000000000000000000000000000000000000001'; + const ignoredToken = '0x0000000000000000000000000000000000000002'; + const untrackedToken = '0x0000000000000000000000000000000000000003'; + + const tokens = { + allDetectedTokens: {}, + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: trackedToken, symbol: 'TRACKED', decimals: 18 }, + ], + }, + }, + allIgnoredTokens: { + [chainId]: { + [accountAddress]: [ignoredToken], + }, + }, + }; + + const { controller } = setupController({ tokens }); + + // Mock balance fetcher to return balances for all three tokens + const trackedBalance = new BN(1000); + const ignoredBalance = new BN(2000); + + jest + .spyOn(multicall, 'getTokenBalancesForMultipleAddresses') + .mockResolvedValue({ + tokenBalances: { + [trackedToken]: { + [accountAddress]: trackedBalance, + }, + [ignoredToken]: { + [accountAddress]: ignoredBalance, + }, + [NATIVE_TOKEN_ADDRESS]: { + [accountAddress]: new BN(0), + }, + }, + stakedBalances: { + [accountAddress]: new BN(0), + }, + }); + + await controller.updateBalances({ chainIds: [chainId] }); + + // Verify tracked token balance was updated + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[trackedToken], + ).toBe(toHex(trackedBalance)); + + // Verify ignored token balance was updated (ignored tokens should still be tracked) + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ignoredToken], + ).toBe(toHex(ignoredBalance)); + + // Verify untracked token balance was NOT updated + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + untrackedToken + ], + ).toBeUndefined(); + + // Verify native token is always updated regardless of tracking + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + NATIVE_TOKEN_ADDRESS + ], + ).toBe('0x0'); + }); + + it('should always update native token balances regardless of tracking status', async () => { + const chainId = '0x1'; + const accountAddress = '0x0000000000000000000000000000000000000000'; + + // Setup with no tracked tokens + const tokens = { + allDetectedTokens: {}, + allTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller } = setupController({ tokens }); + + const nativeBalance = new BN('1000000000000000000'); // 1 ETH + + jest + .spyOn(multicall, 'getTokenBalancesForMultipleAddresses') + .mockResolvedValue({ + tokenBalances: { + [NATIVE_TOKEN_ADDRESS]: { + [accountAddress]: nativeBalance, + }, + }, + stakedBalances: { + [accountAddress]: new BN(0), + }, + }); + + await controller.updateBalances({ chainIds: [chainId] }); + + // Verify native token balance was updated even though no tokens are tracked + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + NATIVE_TOKEN_ADDRESS + ], + ).toBe(toHex(nativeBalance)); + }); + + it('should filter untracked tokens from balance updates', async () => { + const chainId = '0x1'; + const accountAddress = '0x0000000000000000000000000000000000000000'; + const trackedToken = '0x0000000000000000000000000000000000000001'; + const untrackedToken = '0x0000000000000000000000000000000000000002'; + + const tokens = { + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: trackedToken, symbol: 'TRACKED', decimals: 18 }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller } = setupController({ tokens }); + + const trackedBalance = new BN(1000); + + jest + .spyOn(multicall, 'getTokenBalancesForMultipleAddresses') + .mockResolvedValue({ + tokenBalances: { + [trackedToken]: { + [accountAddress]: trackedBalance, + }, + [NATIVE_TOKEN_ADDRESS]: { + [accountAddress]: new BN(0), + }, + }, + stakedBalances: { + [accountAddress]: new BN(0), + }, + }); + + await controller.updateBalances({ chainIds: [chainId] }); + + // Verify tracked token balance was updated + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[trackedToken], + ).toBe(toHex(trackedBalance)); + + // Verify untracked token balance was NOT updated + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + untrackedToken + ], + ).toBeUndefined(); + }); + it('does not update balances when multi-account balances is enabled and all returned values did not change', async () => { const chainId = '0x1'; const account1 = '0x0000000000000000000000000000000000000001'; @@ -2269,7 +2522,7 @@ describe('TokenBalancesController', () => { expect(tokenAddressKeys[0]).toBe(tokenAddressProperChecksum); }); - it('should handle mixed case addresses in both allTokens and allDetectedTokens', async () => { + it('should handle mixed case addresses in allTokens', async () => { const chainId = '0x1'; const accountAddress = '0x0000000000000000000000000000000000000000'; const tokenAddress1Mixed = '0x581c3C1A2A4EBDE2A0Df29B5cf4c116E42945947'; @@ -2286,16 +2539,12 @@ describe('TokenBalancesController', () => { [chainId]: { [accountAddress]: [ { address: tokenAddress1Mixed, symbol: 'TK1', decimals: 18 }, - ], - }, - }, - allDetectedTokens: { - [chainId]: { - [accountAddress]: [ { address: tokenAddress2Mixed, symbol: 'TK2', decimals: 18 }, ], }, }, + allDetectedTokens: {}, + allIgnoredTokens: {}, }; const { controller } = setupController({ @@ -3540,15 +3789,6 @@ describe('TokenBalancesController', () => { // Should have attempted polls despite errors expect(pollSpy).toHaveBeenCalledTimes(2); - // Should have logged errors (both immediate and interval polling use the same error format) - expect(consoleSpy).toHaveBeenCalledWith( - expect.stringContaining( - 'Polling failed for chains 0x1 with interval 100:', - ), - expect.any(Error), - ); - expect(consoleSpy).toHaveBeenCalledTimes(2); // Should have been called twice - controller.stopAllPolling(); consoleSpy.mockRestore(); }); @@ -3655,12 +3895,11 @@ describe('TokenBalancesController', () => { }); it('should properly destroy controller and cleanup resources', () => { - const { controller, messenger } = setupController(); + const { controller } = setupController(); // Start some polling to create timers controller.startPolling({ chainIds: ['0x1'] }); - const unregisterSpy = jest.spyOn(messenger, 'unregisterActionHandler'); const superDestroySpy = jest.spyOn( Object.getPrototypeOf(Object.getPrototypeOf(controller)), 'destroy', @@ -3669,18 +3908,9 @@ describe('TokenBalancesController', () => { // Destroy the controller controller.destroy(); - // Should unregister action handlers - expect(unregisterSpy).toHaveBeenCalledWith( - 'TokenBalancesController:updateChainPollingConfigs', - ); - expect(unregisterSpy).toHaveBeenCalledWith( - 'TokenBalancesController:getChainPollingConfig', - ); - // Should call parent destroy expect(superDestroySpy).toHaveBeenCalled(); - unregisterSpy.mockRestore(); superDestroySpy.mockRestore(); }); @@ -3866,83 +4096,6 @@ describe('TokenBalancesController', () => { executePollSpy.mockRestore(); }); - it('should handle polling errors with console.warn', async () => { - const chainId = '0x1'; - const consoleWarnSpy = jest - .spyOn(console, 'warn') - .mockImplementation(() => { - return undefined; // Suppress console output during tests - }); - - const { controller } = setupController({ - config: { interval: 100 }, - }); - - // Mock _executePoll to throw errors - this will trigger lines 340-343 error handling - jest - .spyOn(controller, '_executePoll') - .mockRejectedValue(new Error('Test polling error')); - - // Use fake timers - jest.useFakeTimers(); - - // Start polling - this triggers immediate polling and error handling - controller.startPolling({ chainIds: [chainId] }); - - // Allow immediate polling error to be caught (lines 340-343) - await flushPromises(); - - // Advance timers to trigger interval polling and error handling - jest.advanceTimersByTime(150); - await flushPromises(); - - // Verify that console.warn was called for polling errors (covers lines 340-343) - expect(consoleWarnSpy).toHaveBeenCalledWith( - expect.stringContaining('Polling failed for chains'), - expect.any(Error), - ); - - // Verify multiple calls were made for different polling attempts - expect(consoleWarnSpy.mock.calls.length).toBeGreaterThanOrEqual(1); - - jest.useRealTimers(); - consoleWarnSpy.mockRestore(); - }); - - it('should handle outer catch blocks for polling function errors', async () => { - const chainId = '0x1'; - const consoleWarnSpy = jest - .spyOn(console, 'warn') - .mockImplementation(() => { - return undefined; // Suppress console output during tests - }); - - const { controller } = setupController({ - config: { interval: 100 }, - }); - - // Use fake timers - jest.useFakeTimers(); - - // Test covers the theoretical error handling paths (lines 349, 364) - // These may be unreachable due to internal try/catch, but we test the functionality - - // Start polling - controller.startPolling({ chainIds: [chainId] }); - - // Allow polling to run - await flushPromises(); - jest.advanceTimersByTime(150); - await flushPromises(); - - // Test that polling is functional - expect(controller).toBeDefined(); - expect(controller.state.tokenBalances).toStrictEqual({}); - - jest.useRealTimers(); - consoleWarnSpy.mockRestore(); - }); - it('should clear existing timer when starting polling for same interval', () => { const chainId1 = '0x1'; const chainId2 = '0x89'; // Polygon @@ -4135,6 +4288,890 @@ describe('TokenBalancesController', () => { }); }); + describe('AccountActivityService integration', () => { + it('should handle real-time balance updates for ERC20 tokens', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const tokenAddress = '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48'; // USDC + const chainId = '0x1'; + const account = createMockInternalAccount({ address: accountAddress }); + + // Setup with tracked token so it processes the balance immediately (account addresses are lowercase in allTokens) + const lowercaseAddress = accountAddress.toLowerCase(); + const tokens = { + allTokens: { + [chainId]: { + [lowercaseAddress]: [ + { + address: tokenAddress, + symbol: 'USDC', + decimals: 6, + }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller, messenger } = setupController({ + listAccounts: [account], + tokens, + }); + + // Emit balance update event + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: `eip155:1/erc20:${tokenAddress}`, + unit: 'USDC', + fungible: true, + }, + postBalance: { + amount: '0xf4240', // 1000000 in hex (1 USDC with 6 decimals) + }, + transfers: [], + }, + ], + }); + + // Wait for async update + await flushPromises(); + + // Verify balance was updated (account addresses are lowercase in state) + const checksumTokenAddress = tokenAddress; + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[checksumTokenAddress], + ).toBe('0xf4240'); + }); + + it('should handle real-time balance updates for native tokens and update AccountTracker', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const chainId = '0x1'; + const account = createMockInternalAccount({ address: accountAddress }); + + const { controller, messenger } = setupController({ + listAccounts: [account], + }); + + // Spy on AccountTrackerController calls + const updateNativeBalancesSpy = jest.fn(); + jest.spyOn(messenger, 'call').mockImplementation((( + action: string, + ...args: unknown[] + ) => { + updateNativeBalancesSpy(action, ...args); + return undefined; + }) as never); + + // Emit balance update event for native token + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: 'eip155:1/slip44:60', + unit: 'ETH', + fungible: true, + }, + postBalance: { + amount: '0xde0b6b3a7640000', // 1 ETH in wei + }, + transfers: [], + }, + ], + }); + + // Wait for async update + await flushPromises(); + + // Verify native balance was updated in TokenBalancesController (account addresses are lowercase in state) + const lowercaseAddr = accountAddress.toLowerCase(); + expect( + controller.state.tokenBalances[lowercaseAddr as ChecksumAddress]?.[ + chainId + ]?.[NATIVE_TOKEN_ADDRESS], + ).toBe('0xde0b6b3a7640000'); + + // Verify AccountTrackerController was called + expect(updateNativeBalancesSpy).toHaveBeenCalledWith( + 'AccountTrackerController:updateNativeBalances', + [ + { + address: lowercaseAddr, + chainId, + balance: '0xde0b6b3a7640000', + }, + ], + ); + }); + + it('should handle balance update errors and trigger fallback polling', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const account = createMockInternalAccount({ address: accountAddress }); + + const { controller, messenger } = setupController({ + listAccounts: [account], + }); + + const updateBalancesSpy = jest.spyOn(controller, 'updateBalances'); + + // Emit balance update event with error + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: 'eip155:1/slip44:60', + unit: 'ETH', + fungible: true, + }, + postBalance: { + amount: '0', + error: 'Network error', + }, + transfers: [], + }, + ], + }); + + // Wait for async update + await flushPromises(); + + // Verify fallback polling was triggered + expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + }); + + it('should handle unsupported asset types and trigger fallback polling', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const account = createMockInternalAccount({ address: accountAddress }); + + const { controller, messenger } = setupController({ + listAccounts: [account], + }); + + const updateBalancesSpy = jest.spyOn(controller, 'updateBalances'); + + // Emit balance update event with unsupported asset type + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: 'eip155:1/unknown:0x123', + unit: 'UNKNOWN', + fungible: true, + }, + postBalance: { + amount: '1000', + }, + transfers: [], + }, + ], + }); + + // Wait for async update + await flushPromises(); + + // Verify fallback polling was triggered + expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + }); + + it('should handle status change to "up" and increase polling interval', async () => { + jest.useFakeTimers(); + + const { controller, messenger } = setupController(); + + const updateConfigSpy = jest.spyOn( + controller, + 'updateChainPollingConfigs', + ); + + // Emit status change to "up" + messenger.publish('AccountActivityService:statusChanged', { + chainIds: ['eip155:1', 'eip155:137'], + status: 'up', + }); + + // Wait for debounce (5 seconds) + jest.advanceTimersByTime(5000); + await flushPromises(); + + // Wait for jitter (up to default interval) + jest.advanceTimersByTime(30000); + await flushPromises(); + + // Verify polling config was updated to backup interval (5 minutes) + expect(updateConfigSpy).toHaveBeenCalledWith( + expect.objectContaining({ + '0x1': { interval: 300000 }, + '0x89': { interval: 300000 }, + }), + { immediateUpdate: true }, + ); + + jest.useRealTimers(); + }); + + it('should handle status change to "down" and restore default polling interval', async () => { + jest.useFakeTimers(); + + const { controller, messenger } = setupController(); + + const updateConfigSpy = jest.spyOn( + controller, + 'updateChainPollingConfigs', + ); + + // Emit status change to "down" + messenger.publish('AccountActivityService:statusChanged', { + chainIds: ['eip155:1'], + status: 'down', + }); + + // Wait for debounce (5 seconds) + jest.advanceTimersByTime(5000); + await flushPromises(); + + // Wait for jitter (up to default interval) + jest.advanceTimersByTime(30000); + await flushPromises(); + + // Verify polling config was updated to default interval (30 seconds) + expect(updateConfigSpy).toHaveBeenCalledWith( + expect.objectContaining({ + '0x1': { interval: 30000 }, + }), + { immediateUpdate: true }, + ); + + jest.useRealTimers(); + }); + + it('should debounce rapid status changes', async () => { + jest.useFakeTimers(); + + const { controller, messenger } = setupController(); + + const updateConfigSpy = jest.spyOn( + controller, + 'updateChainPollingConfigs', + ); + + // Emit multiple rapid status changes + messenger.publish('AccountActivityService:statusChanged', { + chainIds: ['eip155:1'], + status: 'down', + }); + + jest.advanceTimersByTime(1000); + + messenger.publish('AccountActivityService:statusChanged', { + chainIds: ['eip155:1'], + status: 'up', + }); + + jest.advanceTimersByTime(1000); + + messenger.publish('AccountActivityService:statusChanged', { + chainIds: ['eip155:1'], + status: 'down', + }); + + // Wait for debounce (5 seconds) + jest.advanceTimersByTime(5000); + await flushPromises(); + + // Wait for jitter + jest.advanceTimersByTime(30000); + await flushPromises(); + + // Verify config was updated only once with the latest status + expect(updateConfigSpy).toHaveBeenCalledTimes(1); + expect(updateConfigSpy).toHaveBeenCalledWith( + expect.objectContaining({ + '0x1': { interval: 30000 }, // Latest status was "down" + }), + { immediateUpdate: true }, + ); + + jest.useRealTimers(); + }); + + it('should handle multiple chains in a single balance update', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const token1 = '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48'; // USDC + const token2 = '0xdAC17F958D2ee523a2206206994597C13D831ec7'; // USDT + const chainId = '0x1'; + const account = createMockInternalAccount({ address: accountAddress }); + + // Setup with both tokens tracked (account addresses are lowercase in allTokens) + const lowercaseAddress = accountAddress.toLowerCase(); + const tokens = { + allTokens: { + [chainId]: { + [lowercaseAddress]: [ + { + address: token1, + symbol: 'USDC', + decimals: 6, + }, + { + address: token2, + symbol: 'USDT', + decimals: 6, + }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller, messenger } = setupController({ + listAccounts: [account], + tokens, + }); + + // Emit balance update event with multiple tokens + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: `eip155:1/erc20:${token1}`, + unit: 'USDC', + fungible: true, + }, + postBalance: { + amount: '0xf4240', // 1000000 in hex + }, + transfers: [], + }, + { + asset: { + type: `eip155:1/erc20:${token2}`, + unit: 'USDT', + fungible: true, + }, + postBalance: { + amount: '0x1e8480', // 2000000 in hex + }, + transfers: [], + }, + ], + }); + + // Wait for async update + await flushPromises(); + + // Verify both balances were updated (account addresses are lowercase in state) + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + '0x1' + ]?.[token1], + ).toBe('0xf4240'); + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + '0x1' + ]?.[token2], + ).toBe('0x1e8480'); + }); + + it('should handle invalid token addresses and trigger fallback polling', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const account = createMockInternalAccount({ address: accountAddress }); + + const { controller, messenger } = setupController({ + listAccounts: [account], + }); + + const updateBalancesSpy = jest.spyOn(controller, 'updateBalances'); + + // Emit balance update with invalid address format + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: 'eip155:1/erc20:invalid-address', // Not a valid hex address + unit: 'INVALID', + fungible: true, + }, + postBalance: { amount: '1000000' }, + transfers: [], + }, + ], + }); + + await flushPromises(); + + expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + }); + + it('should handle status changes with hex chain ID format', async () => { + jest.useFakeTimers(); + + const { controller, messenger } = setupController(); + const updateConfigSpy = jest.spyOn( + controller, + 'updateChainPollingConfigs', + ); + + // Send status change with CAIP format (as expected from AccountActivityService) + messenger.publish('AccountActivityService:statusChanged', { + chainIds: ['eip155:1'], + status: 'down', + }); + + // Wait for debounce and jitter + jest.advanceTimersByTime(5000 + 30000); + await flushPromises(); + + expect(updateConfigSpy).toHaveBeenCalledWith( + expect.objectContaining({ '0x1': expect.any(Object) }), + expect.any(Object), + ); + + jest.useRealTimers(); + }); + + it('should call addTokens for new untracked tokens received via balance updates', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const newTokenAddress = '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48'; // New token not in allTokens + const chainId = '0x1'; + const account = createMockInternalAccount({ address: accountAddress }); + + // Setup with empty tokens state (no tokens tracked) + const tokens = { + allTokens: {}, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller, messenger } = setupController({ + listAccounts: [account], + tokens, + }); + + // Register and spy on addDetectedTokensViaWs action + const addTokensSpy = jest.fn().mockResolvedValue(undefined); + messenger.registerActionHandler( + 'TokenDetectionController:addDetectedTokensViaWs', + addTokensSpy, + ); + + // Emit balance update for untracked token + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: `eip155:1/erc20:${newTokenAddress}`, + unit: 'USDC', + fungible: true, + }, + postBalance: { + amount: '0xf4240', // 1000000 in hex + }, + transfers: [], + }, + ], + }); + + // Wait for async processing + await flushPromises(); + + // Verify addDetectedTokensViaWs was called with the new token addresses and chainId + expect(addTokensSpy).toHaveBeenCalledWith({ + tokensSlice: [newTokenAddress], + chainId, + }); + + // Verify balance was updated from websocket (account addresses are lowercase in state) + const lowercaseAddr2 = accountAddress.toLowerCase(); + expect( + controller.state.tokenBalances[lowercaseAddr2 as ChecksumAddress]?.[ + chainId + ]?.[newTokenAddress], + ).toBe('0xf4240'); + }); + + it('should process tracked tokens from allTokens without calling addTokens', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const trackedTokenAddress = '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48'; + const chainId = '0x1'; + const account = createMockInternalAccount({ address: accountAddress }); + + // Setup with tracked token in allTokens (account addresses are lowercase in allTokens) + const lowercaseAddress = accountAddress.toLowerCase(); + const tokens = { + allTokens: { + [chainId]: { + [lowercaseAddress]: [ + { + address: trackedTokenAddress, + symbol: 'USDC', + decimals: 6, + }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller, messenger } = setupController({ + listAccounts: [account], + tokens, + }); + + // Register spy on addDetectedTokensViaWs - should NOT be called + const addTokensSpy = jest.fn().mockResolvedValue(undefined); + messenger.registerActionHandler( + 'TokenDetectionController:addDetectedTokensViaWs', + addTokensSpy, + ); + + // Emit balance update for tracked token + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: `eip155:1/erc20:${trackedTokenAddress}`, + unit: 'USDC', + fungible: true, + }, + postBalance: { + amount: '0xf4240', // 1000000 in hex + }, + transfers: [], + }, + ], + }); + + // Wait for async processing + await flushPromises(); + + // Verify addTokens was NOT called since token is already tracked + expect(addTokensSpy).not.toHaveBeenCalled(); + + // Verify balance was updated (account addresses are lowercase in state) + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[trackedTokenAddress], + ).toBe('0xf4240'); + }); + + it('should process ignored tokens from allIgnoredTokens without calling addTokens', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const ignoredTokenAddress = '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48'; + const chainId = '0x1'; + const account = createMockInternalAccount({ address: accountAddress }); + + // Setup with token in allIgnoredTokens (account addresses are lowercase) + const lowercaseAddress = accountAddress.toLowerCase(); + const tokens = { + allTokens: {}, + allDetectedTokens: {}, + allIgnoredTokens: { + [chainId]: { + [lowercaseAddress]: [ignoredTokenAddress], + }, + }, + }; + + const { controller, messenger } = setupController({ + listAccounts: [account], + tokens, + }); + + // Register spy on addDetectedTokensViaWs - should NOT be called + const addTokensSpy = jest.fn().mockResolvedValue(undefined); + messenger.registerActionHandler( + 'TokenDetectionController:addDetectedTokensViaWs', + addTokensSpy, + ); + + // Emit balance update for ignored token + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: `eip155:1/erc20:${ignoredTokenAddress}`, + unit: 'USDC', + fungible: true, + }, + postBalance: { + amount: '0xf4240', // 1000000 in hex + }, + transfers: [], + }, + ], + }); + + // Wait for async processing + await flushPromises(); + + // Verify addTokens was NOT called since token is ignored (tracked) + expect(addTokensSpy).not.toHaveBeenCalled(); + + // Verify balance was still updated (ignored tokens should still have balances tracked, account addresses are lowercase in state) + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[ignoredTokenAddress], + ).toBe('0xf4240'); + }); + + it('should handle native tokens without checking if they are tracked', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const chainId = '0x1'; + const account = createMockInternalAccount({ address: accountAddress }); + + // Setup with empty tokens state + const tokens = { + allTokens: {}, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller, messenger } = setupController({ + listAccounts: [account], + tokens, + }); + + // Register spy on addDetectedTokensViaWs - should NOT be called for native tokens + const addTokensSpy = jest.fn().mockResolvedValue(undefined); + messenger.registerActionHandler( + 'TokenDetectionController:addDetectedTokensViaWs', + addTokensSpy, + ); + + // Emit balance update for native token + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: 'eip155:1/slip44:60', + unit: 'ETH', + fungible: true, + }, + postBalance: { + amount: '0xde0b6b3a7640000', // 1 ETH in wei + }, + transfers: [], + }, + ], + }); + + // Wait for async processing + await flushPromises(); + + // Verify addTokens was NOT called for native token + expect(addTokensSpy).not.toHaveBeenCalled(); + + // Verify native balance was updated (account addresses are lowercase in state) + const lowercaseAddr3 = accountAddress.toLowerCase(); + expect( + controller.state.tokenBalances[lowercaseAddr3 as ChecksumAddress]?.[ + chainId + ]?.[NATIVE_TOKEN_ADDRESS], + ).toBe('0xde0b6b3a7640000'); + }); + + it('should handle addTokens errors and trigger fallback polling', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const newTokenAddress = '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48'; + const account = createMockInternalAccount({ address: accountAddress }); + + const tokens = { + allTokens: {}, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller, messenger } = setupController({ + listAccounts: [account], + tokens, + }); + + const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); + + // Register addDetectedTokensViaWs to throw an error + const addTokensSpy = jest + .fn() + .mockRejectedValue(new Error('Failed to add token')); + messenger.registerActionHandler( + 'TokenDetectionController:addDetectedTokensViaWs', + addTokensSpy, + ); + + // Spy on updateBalances + const updateBalancesSpy = jest + .spyOn(controller, 'updateBalances') + .mockResolvedValue(); + + // Emit balance update for untracked token + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: `eip155:1/erc20:${newTokenAddress}`, + unit: 'USDC', + fungible: true, + }, + postBalance: { + amount: '0xf4240', // 1000000 in hex + }, + transfers: [], + }, + ], + }); + + // Wait for async processing + await flushPromises(); + + // Verify error was logged + expect(consoleSpy).toHaveBeenCalledWith( + 'Error updating balances from AccountActivityService for chain eip155:1, account 0x1234567890123456789012345678901234567890:', + expect.any(Error), + ); + + // Verify fallback polling was triggered (once in addTokens error handler) + expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + + consoleSpy.mockRestore(); + }); + + it('should process multiple tokens - some tracked, some untracked', async () => { + const accountAddress = '0x1234567890123456789012345678901234567890'; + const trackedToken = '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48'; + const untrackedToken = '0xdAC17F958D2ee523a2206206994597C13D831ec7'; + const chainId = '0x1'; + const account = createMockInternalAccount({ address: accountAddress }); + + // Setup with tracked token (account addresses are lowercase in allTokens) + const lowercaseAddress = accountAddress.toLowerCase(); + const tokens = { + allTokens: { + [chainId]: { + [lowercaseAddress]: [ + { + address: trackedToken, + symbol: 'USDC', + decimals: 6, + }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }; + + const { controller, messenger } = setupController({ + listAccounts: [account], + tokens, + }); + + const addTokensSpy = jest.fn().mockResolvedValue(undefined); + messenger.registerActionHandler( + 'TokenDetectionController:addDetectedTokensViaWs', + addTokensSpy, + ); + + // Emit balance update with both tracked and untracked tokens + messenger.publish('AccountActivityService:balanceUpdated', { + address: accountAddress, + chain: 'eip155:1', + updates: [ + { + asset: { + type: `eip155:1/erc20:${trackedToken}`, + unit: 'USDC', + fungible: true, + }, + postBalance: { + amount: '0xf4240', // 1000000 in hex + }, + transfers: [], + }, + { + asset: { + type: `eip155:1/erc20:${untrackedToken}`, + unit: 'USDT', + fungible: true, + }, + postBalance: { + amount: '0x1e8480', // 2000000 in hex + }, + transfers: [], + }, + ], + }); + + // Wait for async processing + await flushPromises(); + + // Verify addTokens was called only for the untracked token with networkClientId + expect(addTokensSpy).toHaveBeenCalledWith({ + tokensSlice: [untrackedToken], + chainId, + }); + + // Verify both token balances were updated from websocket (account addresses are lowercase in state) + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[trackedToken], + ).toBe('0xf4240'); + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[untrackedToken], + ).toBe('0x1e8480'); + }); + + it('should cleanup debouncing timer on destroy', () => { + jest.useFakeTimers(); + + const { controller, messenger } = setupController(); + const clearTimeoutSpy = jest.spyOn(global, 'clearTimeout'); + + // Create a pending status change + messenger.publish('AccountActivityService:statusChanged', { + chainIds: ['eip155:1'], + status: 'down', + }); + + controller.destroy(); + + expect(clearTimeoutSpy).toHaveBeenCalled(); + + jest.useRealTimers(); + clearTimeoutSpy.mockRestore(); + }); + }); + describe('metadata', () => { it('includes expected state in debug snapshots', () => { const { controller } = setupController(); diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index 0ba5491c0d1..c91f88b4b3e 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -14,6 +14,11 @@ import { toChecksumHexAddress, toHex, } from '@metamask/controller-utils'; +import type { + BalanceUpdate, + AccountActivityServiceBalanceUpdatedEvent, + AccountActivityServiceStatusChangedEvent, +} from '@metamask/core-backend'; import type { KeyringControllerAccountRemovedEvent } from '@metamask/keyring-controller'; import type { NetworkControllerGetNetworkClientByIdAction, @@ -27,7 +32,13 @@ import type { PreferencesControllerStateChangeEvent, } from '@metamask/preferences-controller'; import type { Hex } from '@metamask/utils'; -import { isStrictHexString } from '@metamask/utils'; +import { + isCaipAssetType, + isCaipChainId, + isStrictHexString, + parseCaipAssetType, + parseCaipChainId, +} from '@metamask/utils'; import { produce } from 'immer'; import { isEqual } from 'lodash'; @@ -43,6 +54,7 @@ import { type ProcessedBalance, } from './multi-chain-accounts-service/api-balance-fetcher'; import { RpcBalanceFetcher } from './rpc-service/rpc-balance-fetcher'; +import type { TokenDetectionControllerAddDetectedTokensViaWsAction } from './TokenDetectionController'; import type { TokensControllerGetStateAction, TokensControllerState, @@ -53,7 +65,8 @@ export type ChainIdHex = Hex; export type ChecksumAddress = Hex; const CONTROLLER = 'TokenBalancesController' as const; -const DEFAULT_INTERVAL_MS = 180_000; // 3 minutes +const DEFAULT_INTERVAL_MS = 30_000; // 30 seconds +const DEFAULT_WEBSOCKET_ACTIVE_POLLING_INTERVAL_MS = 300_000; // 5 minutes const metadata = { tokenBalances: { @@ -110,6 +123,7 @@ export type AllowedActions = | NetworkControllerGetNetworkClientByIdAction | NetworkControllerGetStateAction | TokensControllerGetStateAction + | TokenDetectionControllerAddDetectedTokensViaWsAction | PreferencesControllerGetStateAction | AccountsControllerGetSelectedAccountAction | AccountsControllerListAccountsAction @@ -121,7 +135,9 @@ export type AllowedEvents = | TokensControllerStateChangeEvent | PreferencesControllerStateChangeEvent | NetworkControllerStateChangeEvent - | KeyringControllerAccountRemovedEvent; + | KeyringControllerAccountRemovedEvent + | AccountActivityServiceBalanceUpdatedEvent + | AccountActivityServiceStatusChangedEvent; export type TokenBalancesControllerMessenger = RestrictedMessenger< typeof CONTROLLER, @@ -157,6 +173,8 @@ export type TokenBalancesControllerOptions = { /** Custom logger. */ log?: (...args: unknown[]) => void; platform?: 'extension' | 'mobile'; + /** Polling interval when WebSocket is active and providing real-time updates */ + websocketActivePollingInterval?: number; }; // endregion @@ -169,6 +187,53 @@ const ZERO_ADDRESS = const checksum = (addr: string): ChecksumAddress => toChecksumHexAddress(addr) as ChecksumAddress; + +/** + * Convert CAIP chain ID or hex chain ID to hex chain ID + * Handles both CAIP-2 format (e.g., "eip155:1") and hex format (e.g., "0x1") + * + * @param chainId - CAIP chain ID (e.g., "eip155:1") or hex chain ID (e.g., "0x1") + * @returns Hex chain ID (e.g., "0x1") + * @throws {Error} If chainId is neither a valid CAIP-2 chain ID nor a hex string + */ +export const caipChainIdToHex = (chainId: string): ChainIdHex => { + if (isStrictHexString(chainId)) { + return chainId; + } + + if (isCaipChainId(chainId)) { + return toHex(parseCaipChainId(chainId).reference); + } + + throw new Error('caipChainIdToHex - Failed to provide CAIP-2 or Hex chainId'); +}; + +/** + * Extract token address from asset type + * Returns tuple of [tokenAddress, isNativeToken] or null if invalid + * + * @param assetType - Asset type string (e.g., 'eip155:1/erc20:0x...' or 'eip155:1/slip44:60') + * @returns Tuple of [tokenAddress, isNativeToken] or null if invalid + */ +export const parseAssetType = (assetType: string): [string, boolean] | null => { + if (!isCaipAssetType(assetType)) { + return null; + } + + const parsed = parseCaipAssetType(assetType); + + // ERC20 token (e.g., "eip155:1/erc20:0x...") + if (parsed.assetNamespace === 'erc20') { + return [parsed.assetReference, false]; + } + + // Native token (e.g., "eip155:1/slip44:60") + if (parsed.assetNamespace === 'slip44') { + return [ZERO_ADDRESS, true]; + } + + return null; +}; // endregion // ──────────────────────────────────────────────────────────────────────────── @@ -192,9 +257,14 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ #detectedTokens: TokensControllerState['allDetectedTokens'] = {}; + #allIgnoredTokens: TokensControllerState['allIgnoredTokens'] = {}; + /** Default polling interval for chains without specific configuration */ readonly #defaultInterval: number; + /** Polling interval when WebSocket is active and providing real-time updates */ + readonly #websocketActivePollingInterval: number; + /** Per-chain polling configuration */ readonly #chainPollingConfig: Record; @@ -207,9 +277,19 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ /** Store original chainIds from startPolling to preserve intent */ #requestedChainIds: ChainIdHex[] = []; + /** Debouncing for rapid status changes to prevent excessive HTTP calls */ + readonly #statusChangeDebouncer: { + timer: NodeJS.Timeout | null; + pendingChanges: Map; + } = { + timer: null, + pendingChanges: new Map(), + }; + constructor({ messenger, interval = DEFAULT_INTERVAL_MS, + websocketActivePollingInterval = DEFAULT_WEBSOCKET_ACTIVE_POLLING_INTERVAL_MS, chainPollingIntervals = {}, state = {}, queryMultipleAccounts = true, @@ -228,6 +308,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ this.#queryAllAccounts = queryMultipleAccounts; this.#accountsApiChainIds = accountsApiChainIds; this.#defaultInterval = interval; + this.#websocketActivePollingInterval = websocketActivePollingInterval; this.#chainPollingConfig = { ...chainPollingIntervals }; // Strategy order: API first, then RPC fallback @@ -244,11 +325,11 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ this.setIntervalLength(interval); // initial token state & subscriptions - const { allTokens, allDetectedTokens } = this.messagingSystem.call( - 'TokensController:getState', - ); + const { allTokens, allDetectedTokens, allIgnoredTokens } = + this.messagingSystem.call('TokensController:getState'); this.#allTokens = allTokens; this.#detectedTokens = allDetectedTokens; + this.#allIgnoredTokens = allIgnoredTokens; this.messagingSystem.subscribe( 'TokensController:stateChange', @@ -277,6 +358,18 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ `TokenBalancesController:getChainPollingConfig`, this.getChainPollingConfig.bind(this), ); + + // Subscribe to AccountActivityService balance updates for real-time updates + this.messagingSystem.subscribe( + 'AccountActivityService:balanceUpdated', + this.#onAccountActivityBalanceUpdate.bind(this), + ); + + // Subscribe to AccountActivityService status changes for dynamic polling management + this.messagingSystem.subscribe( + 'AccountActivityService:statusChanged', + this.#onAccountActivityStatusChanged.bind(this), + ); } #chainIdsWithTokens(): ChainIdHex[] { @@ -736,6 +829,42 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ this.update(() => ({ tokenBalances: {} })); } + /** + * Helper method to check if a token is tracked (exists in allTokens or allIgnoredTokens) + * + * @param tokenAddress - The token address to check + * @param account - The account address + * @param chainId - The chain ID + * @returns True if the token is tracked (imported or ignored) + */ + #isTokenTracked( + tokenAddress: string, + account: ChecksumAddress, + chainId: ChainIdHex, + ): boolean { + const normalizedAddress = tokenAddress.toLowerCase(); + + // Check if token exists in allTokens + if ( + this.#allTokens?.[chainId]?.[account.toLowerCase()]?.some( + (token) => token.address.toLowerCase() === normalizedAddress, + ) + ) { + return true; + } + + // Check if token exists in allIgnoredTokens + if ( + this.#allIgnoredTokens?.[chainId]?.[account.toLowerCase()]?.some( + (addr) => addr.toLowerCase() === normalizedAddress, + ) + ) { + return true; + } + + return false; + } + readonly #onTokensChanged = async (state: TokensControllerState) => { const changed: ChainIdHex[] = []; let hasChanges = false; @@ -824,6 +953,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ this.#allTokens = state.allTokens; this.#detectedTokens = state.allDetectedTokens; + this.#allIgnoredTokens = state.allIgnoredTokens; // Only update balances for chains that still have tokens (and only if we haven't already updated state) if (changed.length && !hasChanges) { @@ -880,6 +1010,243 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ }); }; + // ──────────────────────────────────────────────────────────────────────────── + // AccountActivityService integration helpers + + /** + * Prepare balance updates from AccountActivityService + * Processes all updates and returns categorized results + * Throws an error if any updates have validation/parsing issues + * + * @param updates - Array of balance updates from AccountActivityService + * @param account - Lowercase account address (for consistency with tokenBalances state format) + * @param chainId - Hex chain ID + * @returns Object containing arrays of token balances, new token addresses to add, and native balance updates + * @throws Error if any balance update has validation or parsing errors + */ + #prepareBalanceUpdates( + updates: BalanceUpdate[], + account: ChecksumAddress, + chainId: ChainIdHex, + ): { + tokenBalances: { tokenAddress: ChecksumAddress; balance: Hex }[]; + newTokens: string[]; + nativeBalanceUpdates: { address: string; chainId: Hex; balance: Hex }[]; + } { + const tokenBalances: { tokenAddress: ChecksumAddress; balance: Hex }[] = []; + const newTokens: string[] = []; + const nativeBalanceUpdates: { + address: string; + chainId: Hex; + balance: Hex; + }[] = []; + + for (const update of updates) { + const { asset, postBalance } = update; + + // Throw if balance update has an error + if (postBalance.error) { + throw new Error('Balance update has error'); + } + + // Parse token address from asset type + const parsed = parseAssetType(asset.type); + if (!parsed) { + throw new Error('Failed to parse asset type'); + } + + const [tokenAddress, isNativeToken] = parsed; + + // Validate token address + if ( + !isStrictHexString(tokenAddress) || + !isValidHexAddress(tokenAddress) + ) { + throw new Error('Invalid token address'); + } + + const checksumTokenAddress = checksum(tokenAddress); + const isTracked = this.#isTokenTracked( + checksumTokenAddress, + account, + chainId, + ); + + // postBalance.amount is in hex format (raw units) + const balanceHex = postBalance.amount as Hex; + + // Add token balance (tracked tokens, ignored tokens, and native tokens all get balance updates) + tokenBalances.push({ + tokenAddress: checksumTokenAddress, + balance: balanceHex, + }); + + // Add native balance update if this is a native token + if (isNativeToken) { + nativeBalanceUpdates.push({ + address: account, + chainId, + balance: balanceHex, + }); + } + + // Handle untracked ERC20 tokens - queue for import + if (!isNativeToken && !isTracked) { + newTokens.push(checksumTokenAddress); + } + } + + return { tokenBalances, newTokens, nativeBalanceUpdates }; + } + + // ──────────────────────────────────────────────────────────────────────────── + // AccountActivityService event handlers + + /** + * Handle real-time balance updates from AccountActivityService + * Processes balance updates and updates the token balance state + * If any balance update has an error, triggers fallback polling for the chain + * + * @param options0 - Balance update parameters + * @param options0.address - Account address + * @param options0.chain - CAIP chain identifier + * @param options0.updates - Array of balance updates for the account + */ + readonly #onAccountActivityBalanceUpdate = async ({ + address, + chain, + updates, + }: { + address: string; + chain: string; + updates: BalanceUpdate[]; + }) => { + const chainId = caipChainIdToHex(chain); + const account = checksum(address); + + try { + // Process all balance updates at once + const { tokenBalances, newTokens, nativeBalanceUpdates } = + this.#prepareBalanceUpdates(updates, account, chainId); + + // Update state once with all token balances + if (tokenBalances.length > 0) { + this.update((state) => { + // Initialize account and chain structure + state.tokenBalances[account] ??= {}; + state.tokenBalances[account][chainId] ??= {}; + + // Apply all token balance updates + for (const { tokenAddress, balance } of tokenBalances) { + state.tokenBalances[account][chainId][tokenAddress] = balance; + } + }); + } + + // Update native balances in AccountTrackerController + if (nativeBalanceUpdates.length > 0) { + this.messagingSystem.call( + 'AccountTrackerController:updateNativeBalances', + nativeBalanceUpdates, + ); + } + + // Import any new tokens that were discovered (balance already updated from websocket) + if (newTokens.length > 0) { + await this.messagingSystem.call( + 'TokenDetectionController:addDetectedTokensViaWs', + { + tokensSlice: newTokens, + chainId: chainId as Hex, + }, + ); + } + } catch (error) { + console.warn( + `Error updating balances from AccountActivityService for chain ${chain}, account ${address}:`, + error, + ); + console.warn('Balance update data:', JSON.stringify(updates, null, 2)); + + // On error, trigger fallback polling + await this.updateBalances({ chainIds: [chainId] }).catch(() => { + // Silently handle polling errors + }); + } + }; + + /** + * Handle status changes from AccountActivityService + * Uses aggressive debouncing to prevent excessive HTTP calls from rapid up/down changes + * + * @param options0 - Status change event data + * @param options0.chainIds - Array of chain identifiers + * @param options0.status - Connection status ('up' for connected, 'down' for disconnected) + */ + readonly #onAccountActivityStatusChanged = ({ + chainIds, + status, + }: { + chainIds: string[]; + status: 'up' | 'down'; + }) => { + // Update pending changes (latest status wins for each chain) + for (const chainId of chainIds) { + this.#statusChangeDebouncer.pendingChanges.set(chainId, status); + } + + // Clear existing timer to extend debounce window + if (this.#statusChangeDebouncer.timer) { + clearTimeout(this.#statusChangeDebouncer.timer); + } + + // Set new timer - only process changes after activity settles + this.#statusChangeDebouncer.timer = setTimeout(() => { + this.#processAccumulatedStatusChanges(); + }, 5000); // 5-second debounce window + }; + + /** + * Process all accumulated status changes in one batch to minimize HTTP calls + */ + #processAccumulatedStatusChanges(): void { + const changes = Array.from( + this.#statusChangeDebouncer.pendingChanges.entries(), + ); + this.#statusChangeDebouncer.pendingChanges.clear(); + this.#statusChangeDebouncer.timer = null; + + if (changes.length === 0) { + return; + } + + // Calculate final polling configurations + const chainConfigs: Record = {}; + + for (const [chainId, status] of changes) { + // Convert CAIP format (eip155:1) to hex format (0x1) + // chainId is always in CAIP format from AccountActivityService + const hexChainId = caipChainIdToHex(chainId); + + if (status === 'down') { + // Chain is down - use default polling since no real-time updates available + chainConfigs[hexChainId] = { interval: this.#defaultInterval }; + } else { + // Chain is up - use longer intervals since WebSocket provides real-time updates + chainConfigs[hexChainId] = { + interval: this.#websocketActivePollingInterval, + }; + } + } + + // Add jitter to prevent synchronized requests across instances + const jitterDelay = Math.random() * this.#defaultInterval; // 0 to default interval + + setTimeout(() => { + this.updateChainPollingConfigs(chainConfigs, { immediateUpdate: true }); + }, jitterDelay); + } + /** * Clean up all timers and resources when controller is destroyed */ @@ -888,6 +1255,12 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ this.#intervalPollingTimers.forEach((timer) => clearInterval(timer)); this.#intervalPollingTimers.clear(); + // Clean up debouncing timer + if (this.#statusChangeDebouncer.timer) { + clearTimeout(this.#statusChangeDebouncer.timer); + this.#statusChangeDebouncer.timer = null; + } + // Unregister action handlers this.messagingSystem.unregisterActionHandler( `TokenBalancesController:updateChainPollingConfigs`, diff --git a/packages/assets-controllers/src/TokenDetectionController.test.ts b/packages/assets-controllers/src/TokenDetectionController.test.ts index f112e984cd2..ad443e7323c 100644 --- a/packages/assets-controllers/src/TokenDetectionController.test.ts +++ b/packages/assets-controllers/src/TokenDetectionController.test.ts @@ -1,4 +1,3 @@ -import type { AddApprovalRequest } from '@metamask/approval-controller'; import { Messenger } from '@metamask/base-controller'; import { ChainId, @@ -27,8 +26,7 @@ import { import type { Hex } from '@metamask/utils'; import BN from 'bn.js'; import nock from 'nock'; -import * as sinon from 'sinon'; -import { useFakeTimers } from 'sinon'; +import sinon from 'sinon'; import { formatAggregatorNames } from './assetsUtil'; import * as MutliChainAccountsServiceModule from './multi-chain-accounts-service'; @@ -150,10 +148,7 @@ const mockNetworkConfigurations: Record = { }, }; -type MainMessenger = Messenger< - AllowedActions | AddApprovalRequest, - AllowedEvents ->; +type MainMessenger = Messenger; /** * Builds a messenger that `TokenDetectionController` can use to communicate with other controllers. @@ -744,7 +739,7 @@ describe('TokenDetectionController', () => { describe('AccountsController:selectedAccountChange', () => { let clock: sinon.SinonFakeTimers; beforeEach(() => { - clock = useFakeTimers(); + clock = sinon.useFakeTimers(); }); afterEach(() => { @@ -2592,10 +2587,10 @@ describe('TokenDetectionController', () => { properties: { tokens: [`${sampleTokenA.symbol} - ${sampleTokenA.address}`], // TODO: Either fix this lint violation or explain why it's necessary to ignore. - // eslint-disable-next-line @typescript-eslint/naming-convention + token_standard: 'ERC20', // TODO: Either fix this lint violation or explain why it's necessary to ignore. - // eslint-disable-next-line @typescript-eslint/naming-convention + asset_type: 'TOKEN', }, }); @@ -2735,6 +2730,7 @@ describe('TokenDetectionController', () => { /** * Test Utility - Arrange and Act `detectTokens()` with the Accounts API feature * RPC flow will return `sampleTokenA` and the Accounts API flow will use `sampleTokenB` + * * @param props - options to modify these tests * @param props.overrideMockTokensCache - change the tokens cache * @param props.mockMultiChainAPI - change the Accounts API responses @@ -3535,6 +3531,347 @@ describe('TokenDetectionController', () => { }); }); }); + + describe('addDetectedTokensViaWs', () => { + it('should add tokens detected from websocket with metadata from cache', async () => { + const mockTokenAddress = '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48'; + const chainId = '0x1'; + + await withController( + { + options: { + disabled: false, + }, + }, + async ({ + controller, + mockTokenListGetState, + callActionSpy, + triggerTokenListStateChange, + }) => { + const tokenListState = { + ...getDefaultTokenListState(), + tokensChainsCache: { + [chainId]: { + timestamp: 0, + data: { + [mockTokenAddress]: { + name: 'USD Coin', + symbol: 'USDC', + decimals: 6, + address: mockTokenAddress, + aggregators: [], + iconUrl: 'https://example.com/usdc.png', + occurrences: 11, + }, + }, + }, + }, + }; + + mockTokenListGetState(tokenListState); + triggerTokenListStateChange(tokenListState); + + await controller.addDetectedTokensViaWs({ + tokensSlice: [mockTokenAddress], + chainId: chainId as Hex, + }); + + expect(callActionSpy).toHaveBeenCalledWith( + 'TokensController:addTokens', + [ + { + address: mockTokenAddress, + decimals: 6, + symbol: 'USDC', + aggregators: [], + image: 'https://example.com/usdc.png', + isERC721: false, + name: 'USD Coin', + }, + ], + 'mainnet', + ); + }, + ); + }); + + it('should skip tokens not found in cache and log warning', async () => { + const mockTokenAddress = '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48'; + const chainId = '0x1'; + + const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); + + await withController( + { + options: { + disabled: false, + }, + }, + async ({ + controller, + mockTokenListGetState, + callActionSpy, + triggerTokenListStateChange, + }) => { + // Empty token cache - token not found + const tokenListState = { + ...getDefaultTokenListState(), + tokensChainsCache: { + [chainId]: { + timestamp: 0, + data: {}, + }, + }, + }; + + mockTokenListGetState(tokenListState); + triggerTokenListStateChange(tokenListState); + + await controller.addDetectedTokensViaWs({ + tokensSlice: [mockTokenAddress], + chainId: chainId as Hex, + }); + + // Should log warning about missing token metadata + expect(consoleSpy).toHaveBeenCalledWith( + expect.stringContaining('Token metadata not found in cache'), + ); + + // Should not call addTokens if no tokens have metadata + expect(callActionSpy).not.toHaveBeenCalledWith( + 'TokensController:addTokens', + expect.anything(), + expect.anything(), + ); + + consoleSpy.mockRestore(); + }, + ); + }); + + it('should add all tokens provided without filtering (filtering is caller responsibility)', async () => { + const mockTokenAddress = '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48'; + const secondTokenAddress = '0x1f573d6fb3f13d689ff844b4ce37794d79a7ff1c'; + const chainId = '0x1'; + const selectedAccount = createMockInternalAccount({ + address: '0x0000000000000000000000000000000000000001', + }); + + await withController( + { + options: { + disabled: false, + }, + mocks: { + getSelectedAccount: selectedAccount, + getAccount: selectedAccount, + }, + }, + async ({ + controller, + mockTokenListGetState, + callActionSpy, + triggerTokenListStateChange, + }) => { + // Set up token list with both tokens + const tokenListState = { + ...getDefaultTokenListState(), + tokensChainsCache: { + [chainId]: { + timestamp: 0, + data: { + [mockTokenAddress]: { + name: 'USD Coin', + symbol: 'USDC', + decimals: 6, + address: mockTokenAddress, + aggregators: [], + iconUrl: 'https://example.com/usdc.png', + occurrences: 11, + }, + [secondTokenAddress]: { + name: 'Bancor', + symbol: 'BNT', + decimals: 18, + address: secondTokenAddress, + aggregators: [], + iconUrl: 'https://example.com/bnt.png', + occurrences: 11, + }, + }, + }, + }, + }; + + mockTokenListGetState(tokenListState); + triggerTokenListStateChange(tokenListState); + + // Add both tokens via websocket + await controller.addDetectedTokensViaWs({ + tokensSlice: [mockTokenAddress, secondTokenAddress], + chainId: chainId as Hex, + }); + + // Should add both tokens (no filtering in addDetectedTokensViaWs) + expect(callActionSpy).toHaveBeenCalledWith( + 'TokensController:addTokens', + [ + { + address: mockTokenAddress, + decimals: 6, + symbol: 'USDC', + aggregators: [], + image: 'https://example.com/usdc.png', + isERC721: false, + name: 'USD Coin', + }, + { + address: secondTokenAddress, + decimals: 18, + symbol: 'BNT', + aggregators: [], + image: 'https://example.com/bnt.png', + isERC721: false, + name: 'Bancor', + }, + ], + 'mainnet', + ); + }, + ); + }); + + it('should track metrics when adding tokens from websocket', async () => { + const mockTokenAddress = '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48'; + const chainId = '0x1'; + const mockTrackMetricsEvent = jest.fn(); + + await withController( + { + options: { + disabled: false, + trackMetaMetricsEvent: mockTrackMetricsEvent, + }, + }, + async ({ + controller, + mockTokenListGetState, + callActionSpy, + triggerTokenListStateChange, + }) => { + const tokenListState = { + ...getDefaultTokenListState(), + tokensChainsCache: { + [chainId]: { + timestamp: 0, + data: { + [mockTokenAddress]: { + name: 'USD Coin', + symbol: 'USDC', + decimals: 6, + address: mockTokenAddress, + aggregators: [], + iconUrl: 'https://example.com/usdc.png', + occurrences: 11, + }, + }, + }, + }, + }; + + mockTokenListGetState(tokenListState); + triggerTokenListStateChange(tokenListState); + + await controller.addDetectedTokensViaWs({ + tokensSlice: [mockTokenAddress], + chainId: chainId as Hex, + }); + + // Should track metrics event + expect(mockTrackMetricsEvent).toHaveBeenCalledWith({ + event: 'Token Detected', + category: 'Wallet', + properties: { + tokens: [`USDC - ${mockTokenAddress}`], + token_standard: 'ERC20', + asset_type: 'TOKEN', + }, + }); + + expect(callActionSpy).toHaveBeenCalledWith( + 'TokensController:addTokens', + expect.anything(), + expect.anything(), + ); + }, + ); + }); + + it('should be callable directly as a public method on the controller instance', async () => { + const mockTokenAddress = '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48'; + const chainId = '0x1'; + + await withController( + { + options: { + disabled: false, + }, + }, + async ({ + controller, + mockTokenListGetState, + callActionSpy, + triggerTokenListStateChange, + }) => { + const tokenListState = { + ...getDefaultTokenListState(), + tokensChainsCache: { + [chainId]: { + timestamp: 0, + data: { + [mockTokenAddress]: { + name: 'USD Coin', + symbol: 'USDC', + decimals: 6, + address: mockTokenAddress, + aggregators: [], + iconUrl: 'https://example.com/usdc.png', + occurrences: 11, + }, + }, + }, + }, + }; + + mockTokenListGetState(tokenListState); + triggerTokenListStateChange(tokenListState); + + // Call the public method directly on the controller instance + await controller.addDetectedTokensViaWs({ + tokensSlice: [mockTokenAddress], + chainId: chainId as Hex, + }); + + expect(callActionSpy).toHaveBeenCalledWith( + 'TokensController:addTokens', + [ + { + address: mockTokenAddress, + decimals: 6, + symbol: 'USDC', + aggregators: [], + image: 'https://example.com/usdc.png', + isERC721: false, + name: 'USD Coin', + }, + ], + 'mainnet', + ); + }, + ); + }); + }); }); /** @@ -3551,6 +3888,7 @@ function getTokensPath(chainId: Hex) { type WithControllerCallback = ({ controller, + messenger, mockGetAccount, mockGetSelectedAccount, mockKeyringGetState, @@ -3570,6 +3908,7 @@ type WithControllerCallback = ({ triggerTransactionConfirmed, }: { controller: TokenDetectionController; + messenger: MainMessenger; mockGetAccount: (internalAccount: InternalAccount) => void; mockGetSelectedAccount: (address: string) => void; mockKeyringGetState: (state: KeyringControllerState) => void; @@ -3739,6 +4078,7 @@ async function withController( try { return await fn({ controller, + messenger, mockGetAccount: (internalAccount: InternalAccount) => { mockGetAccount.mockReturnValue(internalAccount); }, diff --git a/packages/assets-controllers/src/TokenDetectionController.ts b/packages/assets-controllers/src/TokenDetectionController.ts index c577af329a2..62a4e47a375 100644 --- a/packages/assets-controllers/src/TokenDetectionController.ts +++ b/packages/assets-controllers/src/TokenDetectionController.ts @@ -95,6 +95,7 @@ export const STATIC_MAINNET_TOKEN_LIST = Object.entries( /** * Function that takes a TokensChainsCache object and maps chainId with TokenListMap. + * * @param tokensChainsCache - TokensChainsCache input object * @returns returns the map of chainId with TokenListMap */ @@ -118,8 +119,14 @@ export type TokenDetectionControllerGetStateAction = ControllerGetStateAction< TokenDetectionState >; +export type TokenDetectionControllerAddDetectedTokensViaWsAction = { + type: `TokenDetectionController:addDetectedTokensViaWs`; + handler: TokenDetectionController['addDetectedTokensViaWs']; +}; + export type TokenDetectionControllerActions = - TokenDetectionControllerGetStateAction; + | TokenDetectionControllerGetStateAction + | TokenDetectionControllerAddDetectedTokensViaWsAction; export type AllowedActions = | AccountsControllerGetSelectedAccountAction @@ -166,13 +173,19 @@ type TokenDetectionPollingInput = { /** * Controller that passively polls on a set interval for Tokens auto detection - * @property intervalId - Polling interval used to fetch new token rates - * @property selectedAddress - Vault selected address - * @property networkClientId - The network client ID of the current selected network - * @property disabled - Boolean to track if network requests are blocked - * @property isUnlocked - Boolean to track if the keyring state is unlocked - * @property isDetectionEnabledFromPreferences - Boolean to track if detection is enabled from PreferencesController - * @property isDetectionEnabledForNetwork - Boolean to track if detected is enabled for current network + * + * intervalId - Polling interval used to fetch new token rates + * + * selectedAddress - Vault selected address + * + * networkClientId - The network client ID of the current selected network + * + * disabled - Boolean to track if network requests are blocked + * + * isUnlocked - Boolean to track if the keyring state is unlocked + * + * isDetectionEnabledFromPreferences - Boolean to track if detection is enabled from PreferencesController + * */ export class TokenDetectionController extends StaticIntervalPollingController()< typeof controllerName, @@ -183,8 +196,6 @@ export class TokenDetectionController extends StaticIntervalPollingController boolean; - #isDetectionEnabledForNetwork: boolean; - readonly #getBalancesInSingleCall: AssetsContractController['getBalancesInSingleCall']; readonly #trackMetaMetricsEvent: (options: { @@ -206,16 +215,12 @@ export class TokenDetectionController extends StaticIntervalPollingController void; - #accountsAPI = { + readonly #accountsAPI = { isAccountsAPIEnabled: true, supportedNetworksCache: null as number[] | null, platform: '' as 'extension' | 'mobile', @@ -298,11 +303,7 @@ export class TokenDetectionController extends StaticIntervalPollingController void; @@ -319,15 +320,16 @@ export class TokenDetectionController extends StaticIntervalPollingController { this.#isUnlocked = true; await this.#restartTokenDetection(); @@ -375,8 +373,6 @@ export class TokenDetectionController extends StaticIntervalPollingController { const isEqualValues = this.#compareTokensChainsCache( tokensChainsCache, @@ -390,8 +386,6 @@ export class TokenDetectionController extends StaticIntervalPollingController { const selectedAccount = this.#getSelectedAccount(); const isDetectionChangedFromPreferences = @@ -409,8 +403,6 @@ export class TokenDetectionController extends StaticIntervalPollingController { const { networkConfigurationsByChainId } = this.messagingSystem.call( 'NetworkController:getState', @@ -455,7 +447,8 @@ export class TokenDetectionController extends StaticIntervalPollingController { + const tokensWithBalance: Token[] = []; + const eventTokensDetails: string[] = []; + + for (const nonZeroTokenAddress of tokensSlice) { + // Check map of validated tokens + const tokenData = + this.#tokensChainsCache[chainId]?.data?.[ + nonZeroTokenAddress.toLowerCase() + ]; + + if (!tokenData) { + console.warn( + `Token metadata not found in cache for ${nonZeroTokenAddress} on chain ${chainId}`, + ); + continue; + } + + const { decimals, symbol, aggregators, iconUrl, name, address } = + tokenData; + + // Push to lists + eventTokensDetails.push(`${symbol} - ${address}`); + tokensWithBalance.push({ + address, + decimals, + symbol, + aggregators, + image: iconUrl, + isERC721: false, + name, + }); + } + + // Perform addition + if (tokensWithBalance.length) { + this.#trackMetaMetricsEvent({ + event: 'Token Detected', + category: 'Wallet', + properties: { + tokens: eventTokensDetails, + token_standard: ERC20, + asset_type: ASSET_TYPES.TOKEN, + }, + }); + + const networkClientId = this.messagingSystem.call( + 'NetworkController:findNetworkClientIdByChainId', + chainId, + ); + + await this.messagingSystem.call( + 'TokensController:addTokens', + tokensWithBalance, + networkClientId, + ); + } + } + #getSelectedAccount() { return this.messagingSystem.call('AccountsController:getSelectedAccount'); } diff --git a/packages/assets-controllers/src/index.ts b/packages/assets-controllers/src/index.ts index 041cae84907..c7f64f8b2bd 100644 --- a/packages/assets-controllers/src/index.ts +++ b/packages/assets-controllers/src/index.ts @@ -87,6 +87,7 @@ export type { TokenDetectionControllerMessenger, TokenDetectionControllerActions, TokenDetectionControllerGetStateAction, + TokenDetectionControllerAddDetectedTokensViaWsAction, TokenDetectionControllerEvents, TokenDetectionControllerStateChangeEvent, } from './TokenDetectionController'; diff --git a/packages/assets-controllers/tsconfig.build.json b/packages/assets-controllers/tsconfig.build.json index bca6a835d37..629b833e22a 100644 --- a/packages/assets-controllers/tsconfig.build.json +++ b/packages/assets-controllers/tsconfig.build.json @@ -9,6 +9,7 @@ { "path": "../account-tree-controller/tsconfig.build.json" }, { "path": "../accounts-controller/tsconfig.build.json" }, { "path": "../approval-controller/tsconfig.build.json" }, + { "path": "../core-backend/tsconfig.build.json" }, { "path": "../base-controller/tsconfig.build.json" }, { "path": "../controller-utils/tsconfig.build.json" }, { "path": "../keyring-controller/tsconfig.build.json" }, diff --git a/packages/assets-controllers/tsconfig.json b/packages/assets-controllers/tsconfig.json index 2b0acd993f8..ae60fdfc0d7 100644 --- a/packages/assets-controllers/tsconfig.json +++ b/packages/assets-controllers/tsconfig.json @@ -8,6 +8,7 @@ { "path": "../account-tree-controller" }, { "path": "../accounts-controller" }, { "path": "../approval-controller" }, + { "path": "../core-backend" }, { "path": "../base-controller" }, { "path": "../controller-utils" }, { "path": "../keyring-controller" }, diff --git a/yarn.lock b/yarn.lock index 9ee1deba03e..894667ed4d3 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2594,6 +2594,7 @@ __metadata: "@metamask/base-controller": "npm:^8.4.1" "@metamask/contract-metadata": "npm:^2.4.0" "@metamask/controller-utils": "npm:^11.14.1" + "@metamask/core-backend": "npm:^1.0.1" "@metamask/eth-query": "npm:^4.0.0" "@metamask/ethjs-provider-http": "npm:^0.3.0" "@metamask/keyring-api": "npm:^21.0.0" @@ -2642,6 +2643,7 @@ __metadata: "@metamask/account-tree-controller": ^1.0.0 "@metamask/accounts-controller": ^33.0.0 "@metamask/approval-controller": ^7.0.0 + "@metamask/core-backend": ^1.0.0 "@metamask/keyring-controller": ^23.0.0 "@metamask/network-controller": ^24.0.0 "@metamask/permission-controller": ^11.0.0 @@ -2918,7 +2920,7 @@ __metadata: languageName: unknown linkType: soft -"@metamask/core-backend@workspace:packages/core-backend": +"@metamask/core-backend@npm:^1.0.1, @metamask/core-backend@workspace:packages/core-backend": version: 0.0.0-use.local resolution: "@metamask/core-backend@workspace:packages/core-backend" dependencies: