diff --git a/oxlint-suppressions.json b/oxlint-suppressions.json index 033bcfb9753..c75cc4c8941 100644 --- a/oxlint-suppressions.json +++ b/oxlint-suppressions.json @@ -407,17 +407,17 @@ "count": 3 }, "typescript/no-unsafe-assignment": { - "count": 17 + "count": 22 }, "typescript/no-unsafe-call": { - "count": 26 + "count": 4 }, "typescript/no-unsafe-member-access": { - "count": 11 + "count": 19 } }, "packages/assets-controller/src/AssetsController.ts": { - "no-unused-vars": { + "typescript/no-unnecessary-type-assertion": { "count": 1 }, "typescript/no-unsafe-assignment": { @@ -472,10 +472,10 @@ "count": 2 }, "typescript/no-unsafe-assignment": { - "count": 6 + "count": 13 }, "typescript/no-unsafe-member-access": { - "count": 8 + "count": 14 }, "typescript/no-unsafe-return": { "count": 2 @@ -493,30 +493,25 @@ "id-length": { "count": 4 }, - "no-unused-vars": { - "count": 1 - }, "typescript/no-unsafe-assignment": { - "count": 14 + "count": 12 }, "typescript/no-unsafe-member-access": { "count": 15 } }, - "packages/assets-controller/src/data-sources/RpcDataSource.ts": { - "typescript/no-unsafe-assignment": { - "count": 3 - } - }, "packages/assets-controller/src/data-sources/SnapDataSource.test.ts": { "no-useless-return": { "count": 1 }, "typescript/no-unsafe-assignment": { - "count": 4 + "count": 6 }, "typescript/no-unsafe-member-access": { - "count": 3 + "count": 1 + }, + "typescript/no-unsafe-return": { + "count": 1 } }, "packages/assets-controller/src/data-sources/SnapDataSource.ts": { @@ -615,26 +610,6 @@ "count": 2 } }, - "packages/assets-controller/src/middlewares/RpcFallbackMiddleware.test.ts": { - "no-unused-vars": { - "count": 1 - }, - "typescript/no-unsafe-argument": { - "count": 2 - }, - "typescript/no-unsafe-assignment": { - "count": 11 - }, - "typescript/no-unsafe-call": { - "count": 1 - }, - "typescript/no-unsafe-member-access": { - "count": 22 - }, - "typescript/no-unsafe-return": { - "count": 24 - } - }, "packages/assets-controller/src/migrations/healAssetsInfoMetadata.test.ts": { "typescript/no-unsafe-assignment": { "count": 2 diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 625b93a9305..5d8323aa3f5 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -7,10 +7,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add `getZeroNativeAssetBalance` and `getZeroTokenAssetBalance` so Stellar native zeros include spendable/reserve metadata and Stellar token zeros include empty trustline metadata. `getZeroAssetBalance` picks between them with `isNativeAssetId` + ### Changed +- **BREAKING:** Split asset fetching into two paths behind `assetsAccountsApiV6` ([#9651](https://github.com/MetaMask/core/pull/9651)) + - **Architecture:** The v5 path keeps production behavior — the API decides the returned set, and results are merged. The v6 path has the client declare the visible set (`includeAssetIds` / `excludeAssetIds`) and write an authoritative `full` snapshot for covered chains. Visibility is computed from controller state and shared by Accounts API, Snap, RPC, and RPC fallback. The flag is read only in `AssetsController` and injected as `isBalanceV6Enabled`. Hide/unhide re-evaluates live subscriptions so the next poll uses the new set. + - **Why this is breaking:** + - `'update'` is removed from `AssetsUpdateMode`; use `'full'` or `'merge'` + - `getAssets` no longer accepts `updateMode`; the data source sets it on the response + - `getAssetsState` is removed from pipeline `Context`, `SubscriptionRequest`, and `PriceDataSource.fetch`; inject it on data-source and middleware constructors instead + - `getAssetsState` and `getAssetVisibility` are required on Accounts API, Snap, and RPC data sources + - `isBalanceV6Enabled` is required on `SnapDataSource` + - `unhideAsset` is now async and force-fetches the asset's chain, matching `addCustomAsset` - Bump `@metamask/transaction-controller` from `^72.0.0` to `^72.0.1` ([#10462](https://github.com/MetaMask/core/pull/10462)) +### Fixed + +- Treat `assetsAccountsApiV6` as enabled when it is `true`, not a nested `{ value }` object ([#9651](https://github.com/MetaMask/core/pull/9651)) + ## [16.1.2] ### Changed diff --git a/packages/assets-controller/docs/decisions/2026-09-23-accounts-api-v6-integration.md b/packages/assets-controller/docs/decisions/2026-09-23-accounts-api-v6-integration.md new file mode 100644 index 00000000000..43823477bbf --- /dev/null +++ b/packages/assets-controller/docs/decisions/2026-09-23-accounts-api-v6-integration.md @@ -0,0 +1,462 @@ +# Accounts API v6 integration + +- Date: 2026-09-23 +- Status: Accepted +- Package: `@metamask/assets-controller` + +How AssetsController manages the visible asset set across Accounts API v5 +and v6 while both paths ship behind `assetsAccountsApiV6`. + +## What this change is about + +The AssetsController collects token balances, metadata, and prices for every +account and chain, and stores the result for the UI to render. The hard part is +not fetching a balance — it is agreeing on **which tokens the user should see**, +and keeping that set intact as several sources write into the same state. + +The two Accounts API versions answer that question differently: + +- **v5**: the API decides what to return. The client cannot ask for a specific + token, so anything the API does not index (a pin, mUSD at zero) is patched in + afterwards — by RPC re-reads, a supplemental custom-asset poll, and custom + asset graduation. Responses are overlaid onto state (`merge`), so nothing is + ever removed and a stale token can linger. +- **v6**: the client declares the visible set up front as `includeAssetIds` / + `excludeAssetIds`, and the answer is authoritative for the chains it covers + (`full`). That removes the patch-ups, but it means every visible asset must + appear in the snapshot — otherwise the replace would drop it. + +Both paths ship together and stay fully separate. Which one runs is decided at +runtime by the `assetsAccountsApiV6` remote feature flag: + +- flag off, missing, or unreadable: the legacy **v5** path (today's production + behavior) +- `assetsAccountsApiV6: true`: the new **v6** path + +The flag is read in exactly one place, `AssetsController.#isBalanceV6Enabled()`, +and injected into `AccountsApiDataSource`, `SnapDataSource`, `RpcDataSource` +(and through it `BalanceFetcher`), and `RpcFallbackMiddleware`. + +## Vocabulary used below + +| Term | Meaning | +| ----------------------------- | --------------------------------------------------------------------------------------------------------------------------------- | +| Pinned asset (custom asset) | A token the user added manually (Manage Token flow). We always want to fetch it, even when the API does not return it on its own. | +| Hidden asset | A token the user chose to hide, so it should not be fetched or shown. | +| Default tracked asset | A token the client always tracks on a chain even at zero (for example mUSD). | +| Visible asset | Native, pin, or default tracked asset that is not hidden. Staking vault IDs are not visible; a dedicated source owns them. | +| Data source | Where balances come from: the Accounts API, RPC nodes, a Snap, or staking. | +| Middleware | A step that enriches or repairs data after it is fetched (token detection, prices, metadata, RPC fallback). | +| Fast lane | The first, quick round of fetching. Its result is written to state right away so the UI can render. | +| Background lane | The slower sources that run afterwards; their results are merged in. | +| Update mode `merge` vs `full` | `merge` keeps assets already in state that the response did not mention. `full` replaces the whole slice the source covers. | +| Failed chain | A chain the source could not answer in full. Its balances are dropped from the response and existing state is left unchanged. | +| Basic functionality | The user setting that turns off network calls to MetaMask services. | + +## Why there are two paths + +Rather than scattering flag checks through one shared flow, the controller now +has two named flows: + +- the **v5 path** preserves the old behavior exactly, so turning the flag off + matches production +- the **v6 path** isolates the new behavior, so it can be evaluated on its own + and the old path can later be deleted in one clean sweep + +The split shows up in three places: a forced refresh, setting up live updates, +and handling an update once it arrives. + +```mermaid +flowchart LR + A["Refresh · Subscribe · Incoming update"] --> F{assetsAccountsApiV6} + F -->|false| V5["v5 path"] + F -->|true| V6["v6 path"] +``` + +## Shared visibility (v6) + +v6 sources do not take pins or hides from the request. They all call +`getAssetVisibility(accountIds, chainIds)` (`utils/assetVisibility.ts`), which +reads controller state and returns: + +- `visibleAssetIds`: natives (except chains with no native, e.g. Tempo), pins + for those accounts, and default tracked assets, minus hidden and staking IDs +- `hiddenAssetIds`: hidden preferences on those chains + +Accounts API sends those lists as `includeAssetIds` / `excludeAssetIds`. RPC +and Snap fetch the visible set and skip hidden IDs. Snap also asks the keyring +for those visible IDs (not only `listAccountAssets` holdings). Because a `full` +snapshot must mention every visible asset, Snap then fills any ID still +missing from a non-empty snap response — last-known amount if present, +otherwise `0`. An empty `getAccountBalances` result is treated as a failed +fetch, not a zero snapshot. RPC requests the visible set plus already-tracked +ERC-20s (detected tokens already in `assetsBalance`), always stamps `full`, +and carries the previous amount for any asset whose `balanceOf` or decimals +lookup failed so one unreadable token cannot wipe the rest of the chain. + +`hideAsset` / `removeCustomAsset` re-run `#subscribeAssets` so the next poll +sees the new lists. `addCustomAsset` and `unhideAsset` also force-fetch that +token's chain (every pin from state, not only the new token) so a `full` +snapshot restores the balance immediately and cannot wipe the other pins. + +## 1. Forced refresh + +Triggered by `getAssets(..., { forceUpdate: true })`. Both paths have the same +shape: build a request, run the fast lane, write state, then run the background +lane. + +```mermaid +flowchart TB + subgraph v5fast ["v5 fast lane"] + direction TB + V5P["Accounts API v5 + Staked"] --> V5G[CustomAssetGraduation] --> V5F[RpcFallback] --> V5D[Detection] --> V5T["Token + Price"] --> V5S["State: merge + replaceCoveredChainBalances"] + end + subgraph v6fast ["v6 fast lane"] + direction TB + V6P["Accounts API v6 + Staked
includeAssetIds / excludeAssetIds"] --> V6F["RpcFallback
errored chains"] --> V6D[Detection] --> V6T["Token + Price"] --> V6S["State: full for covered chains"] + end + subgraph bgv5 ["v5 background lane"] + direction TB + B5["Snap + RPC
updateMode merge"] --> B5D[Detection] --> B5T["Token + Price"] --> B5S["State: merge"] + end + subgraph bgv6 ["v6 background lane"] + direction TB + B6["Snap + RPC
updateMode full"] --> B6D[Detection] --> B6T["Token + Price"] --> B6S["State: full for covered chains"] + end + v5fast --> bgv5 + v6fast --> bgv6 +``` + +The requests differ as well: + +- **v5** puts every pinned asset of the requested accounts on + `request.customAssets`, without scoping them to chains, and does not mention + hidden assets. RPC fallback may add stale tracked assets to that list. Fast + fetch then writes state as `merge` plus `replaceCoveredChainBalances`, which + restores custom assets omitted by the API. +- **v6** does not copy pins or hides onto the request. The Accounts API reads + visibility from state. Detected ERC-20 balances come back with the chain, so + `updateMode: 'full'` is a complete snapshot for successful chains. A chain + that is unprocessed, or that did not resolve every `includeAssetId`, is put + in `errors` and contributes no balances. RPC fallback retries those **chains** + in full from state. RPC always stamps `full`: resolved tokens get a fresh + amount, unread tokens keep the amount already in state, and a chain that + resolved nothing is listed in `errors` instead of writing an empty slice. + Snap `#fetchV6` requests visible assets with listed holdings, stamps `full`, + and fills omitted visible assets from current state (`0` if none). An empty + `getAccountBalances` map (or a thrown fetch) contributes nothing, so last- + known amounts are kept. Background Snap and + RPC both stamp `full`. `mergeDataResponses` promotes `full` if any source did. + +When basic functionality is off, both fast lanes shrink to +`Staked -> Detection`. The background lane is RPC only in both paths +(`merge` on v5; `full` on v6). + +## 2. Setting up live updates (subscribe) + +Both paths first hand each chain to the data source with the highest priority +for it (Accounts API, then Snap, then RPC, plus Account Activity and staking +on the chains they support). + +```mermaid +flowchart TB + subgraph shared ["Both paths"] + H["Assign each chain to the highest-priority source that supports it"] + end + + subgraph v5sub ["v5"] + direction TB + V5A["Subscribe each source with accounts + assigned chains"] + V5WS["Account Activity · merge"] + V5API["Accounts API · merge"] + V5SNAP["Snap · merge"] + V5RPC["RPC · merge"] + V5R["Extra RPC poll: customAssetsOnly
pins on chains another source already owns · merge"] + V5A --> V5WS + V5A --> V5API + V5A --> V5SNAP + V5A --> V5RPC + V5A --> V5R + end + + subgraph v6sub ["v6"] + direction TB + V6A["Subscribe each source with accounts + assigned chains only"] + V6WS["Account Activity · event deltas · merge"] + V6API["Accounts API · visibility from state · full"] + V6SNAP["Snap · visible assets from state · full snapshot
balance events stay merge"] + V6RPC["RPC · visible assets from state · full"] + V6A --> V6WS + V6A --> V6API + V6A --> V6SNAP + V6A --> V6RPC + end + + shared --> v5sub + shared --> v6sub +``` + +- **v5** adds a separate RPC poll (`customAssetsOnly`) for pins that sit on a + chain another source already owns. +- **v6** subscribe only assigns accounts and chains. There is no + `customAssetsOnly` supplement: RPC already polls the visible set on the + chains it was assigned, and Accounts API / Snap cover pins on their chains + through visibility. Account Activity events, RPC token detection, staked + balance updates, and Snap `accountBalancesUpdated` events stay `merge`. + +## 3. Handling an incoming update + +```mermaid +flowchart TB + subgraph v5upd ["v5 · #handleAssetsUpdateV5"] + direction TB + V5I["Incoming update"] --> V5G["CustomAssetGraduation
Accounts API + Account Activity"] --> V5R["RpcFallback when source is Accounts API"] --> V5D[Detection] --> V5T["Token + Price"] --> V5S["State: applyV5
merge, honor replaceCoveredChainBalances"] + end + subgraph v6upd ["v6 · #handleAssetsUpdateV6"] + direction TB + V6I["Incoming update"] --> V6F["RpcFallback when basic functionality is on
errored chains"] --> V6D[Detection] --> V6T["Token + Price"] --> V6S["State: applyV6
full replace or merge overlay"] + end +``` + +Both paths still run the Account Activity occurrence-floor filter when basic +functionality is on. + +- **v5** (`#handleAssetsUpdateV5`): CustomAssetGraduation (Accounts API and + Account Activity) → RpcFallback when the source is Accounts API → Detection → + Token + Price → `applyV5AccountBalanceUpdate` (`effectiveAccountBalancesV5`). +- **v6** (`#handleAssetsUpdateV6`): never CustomAssetGraduation. RpcFallback + when basic functionality is on (failed chains; native + visible assets from + state, still `full`) → Detection → Token + Price → + `applyV6AccountBalanceUpdate` (`effectiveAccountBalancesV6`). + +`#updateState` picks the writer from the flag, not from `updateMode`: + +- flag off → v5 always +- flag on → v6 always. Inside v6, `updateMode: 'merge'` overlays + `{ ...previous, ...incoming }`; `updateMode: 'full'` replaces each covered + chain slice. Visible assets the client asked for survive a `full` replace + via `shouldKeepAsset` (natives, pins, and default tracked assets, plus staking) + vaults until Accounts API returns ETH staked balances). Hidden assets are + excluded from the snapshot and dropped; unhide is expected to fetch again. + +## What a v6 `full` snapshot may drop + +A `full` write is authoritative for the **chains it covers**. It must never +drop a token because one RPC read failed. It still drops tokens whose absence +means “gone”: a detected ERC-20 successfully read as `0`, a hidden token, or a +detected token a complete Accounts API snapshot never returned. + +### RPC v6: build the snapshot + +RPC fetch and poll (`#fetchV6` / `#handleBalanceUpdateV6`) share this tree. +Requested assets are the visible set plus already-tracked ERC-20s, minus +hidden and staking vault IDs. + +```mermaid +flowchart TD + start["RPC v6 · one account-chain"] --> any{"Did any requested asset resolve?"} + any -->|"No — every balanceOf / decimals failed"| fail["Put chain in errors
write no balances for this chain
existing state unchanged until fallback"] + any -->|Yes| asset["For each requested asset"] + asset --> bo{"balanceOf succeeded?"} + bo -->|No| prev{"Previous amount in state?"} + prev -->|No| omit["Leave out of snapshot"] + prev -->|Yes| zeroPrev{"Previous is 0 and asset is not visible?"} + zeroPrev -->|Yes| omit + zeroPrev -->|No| carry["Copy previous amount into snapshot"] + bo -->|Yes| vis{"Raw balance is 0 and asset is not visible?"} + vis -->|Yes| omit + vis -->|No| dec{"Decimals resolved?"} + dec -->|No| prev + dec -->|Yes| write["Write fresh amount"] + omit --> more{"More requested assets?"} + carry --> more + write --> more + more -->|Yes| asset + more -->|No| full["Stamp updateMode: full"] +``` + +A failed `balanceOf` or unknown decimals therefore keeps the last good amount +and still lets every other token on that chain update. Only a chain that +resolved **nothing** is treated as a failed chain. + +### Controller: apply the snapshot + +`effectiveAccountBalancesV6` is `{ ...keptPrevious, ...incoming }`. Incoming +always overwrites a kept row. Visible / staking only **survives omission** — +they do not freeze the previous amount when the snapshot includes them. + +```mermaid +flowchart TD + in["Incoming v6 balance update"] --> mode{"updateMode"} + mode -->|merge| overlay["{ ...previous, ...incoming }"] + mode -->|full| covered["Covered chains = chains present in incoming"] + covered --> prev["For each previous asset"] + prev --> onChain{"On a covered chain?"} + onChain -->|No| keep["Keep previous"] + onChain -->|Yes| listed{"Present in incoming snapshot?"} + listed -->|Yes| useIn["Use incoming amount"] + listed -->|No| undel{"Visible native / pin / default tracked
or staking vault?"} + undel -->|Yes| keep + undel -->|No| drop["Drop"] + keep --> done["{ ...keptPrevious, ...incoming }"] + useIn --> done + drop --> done +``` + +Detected tokens omitted from a `full` payload are dropped — that is how an +empty detected token leaves the list. + +| Situation | Result | +| ----------------------------------------------------------------------------- | ----------------------------------------------------------------------------- | +| One RPC `balanceOf` or decimals lookup fails, others succeed | Keep previous amount for the unread token; rest of the chain updates (`full`) | +| RPC resolved nothing on the chain | No write for that chain; state unchanged until RpcFallback | +| Detected ERC-20 successfully read as `0` (not native / pin / default tracked) | Omit from snapshot → `full` drops it | +| Visible asset (native, pin, default tracked) at `0` | Write a zero row so `full` cannot drop it | +| Hidden asset | Not fetched; dropped on `full` | +| Staking vault omitted from the snapshot | Kept (`shouldKeepAsset`) | +| Accounts API `full` omits a detected token | Dropped unless it is visible or staking | +| `merge` event (detection, Account Activity, Snap balance event, staking) | Overlay only; nothing is removed | +| Snap `getAccountBalances` throws or returns `{}` | No write for that account; last-known amounts stay | +| Snap returns some balances but omits a visible asset that already has a row | Copy that row into the snapshot (`full`) | +| Snap omits a visible asset with no previous row | Seed a chain-specific zero (`getZeroAssetBalance`) so `full` cannot drop it | + +## When v6 keeps a stale balance + +A `full` snapshot has three ways to treat an asset it did not freshly read: + +| Choice | What the snapshot does | What the user sees | +| ----------- | ----------------------------------------------- | ------------------------------------------------------------------------------------------------------------- | +| Wipe | Write `0` (or an empty Stellar metadata object) | The token is still listed, but the amount looks spent or gone | +| Omit | Leave it out of a `full` payload | Detected tokens disappear from the list. Visible natives / pins / defaults survive only via `shouldKeepAsset` | +| Carry stale | Copy the last good row into the snapshot | The token stays listed at the last successful read; the next poll retries | + +v6 chooses **carry stale** whenever a read failed but we already had an amount, and **omit the whole chain** (no write) when the source resolved nothing. It writes zero only when the source successfully reported zero, or when we must seed a visible asset that has never been seen. + +Wipe is worse than a slightly old number. `0` reads as “you have nothing”: send and swap use that amount, and a Snap or RPC outage looks like a drained account. Last-known-good is still an observation we actually made. + +Omit is worse on a `full` replace than it was on v5 `merge`. v5 overlaid whatever arrived and left the rest of state alone, so a missed token lingered. v6 replaces the covered chain slice: an omitted detected ERC-20 is dropped, so a single failed `balanceOf` would hide a token the user still holds. Copying the previous amount into the snapshot keeps that token on the list and still lets every other token on the chain update. Visible assets would survive omission through `shouldKeepAsset`, but putting the previous row in the snapshot makes the carry explicit and keeps Snap / RPC aligned with how the controller applies `full`. + +Doing nothing for the **entire chain** because one token failed is also worse. RPC would then skip a successful ETH read just because one ERC-20 timed out. Carry the unread token, write fresh amounts for the rest. Only when **every** requested read failed (RPC) or the Snap returned no balances at all do we leave the chain untouched — there is no partial snapshot we trust, so we do not cover that slice. + +### Snap v6: build the snapshot + +```mermaid +flowchart TD + start["Snap v6 · one account"] --> fetchBalances{"getAccountBalances"} + fetchBalances -->|throws or empty map| skip["Contribute nothing for this account
last-known amounts stay"] + fetchBalances -->|at least one balance| vis["For each omitted visible asset"] + vis --> prev{"Previous row in state?"} + prev -->|Yes| carry["Copy previous amount and metadata"] + prev -->|No| seed["Seed getZeroAssetBalance
native vs token, Stellar-aware"] + carry --> full["Stamp updateMode: full"] + seed --> full +``` + +### Stale balance scenarios (`full`) + +| # | What happened | What we write | What the user sees | +| --- | --------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- | -------------------------------------------------------------------------- | +| 1 | Snap is unreachable (throws or returns nothing) | Nothing for that account | Last amounts stay on screen | +| 2 | Snap answers but skips a native coin, a token the user added, or a default such as mUSD | Copy the last known amount into the snapshot | That token stays listed at the last successful read | +| 3 | Same skip, but we have never stored an amount for it | A placeholder zero so the asset appears in the list | The asset shows at `0`. This is **not** stale — there is no earlier amount | +| 4 | One RPC `balanceOf` / decimals call fails; the rest of the chain succeeds | Keep the last amount for that token only; write fresh amounts for the others | The unread token stays listed; everything else updates | +| 5 | Same RPC failure, but the token was already `0` and is not one we track on purpose | Leave it out of the snapshot | The unused detected token disappears from the list | +| 6 | Every RPC call on the chain fails | Nothing for that chain; retry via RPC fallback | Last amounts on that chain stay until fallback | +| 7 | Accounts API could not fully answer a chain | Mark the chain failed and strip its balances; retry via RPC fallback | Last amounts on that chain stay until fallback | + +### Stale balance scenarios (`merge`) + +| # | What happened | What we write | What the user sees | +| --- | ------------------------------------------------------------------------------------------- | ------------------------------------------- | ------------------------------------------------------------------ | +| 8 | Account Activity event after a tx omits a token that actually changed | Overlay only the tokens named in the event | Omitted token stays at the last poll | +| 9 | Account Activity row has no `asset`, `postBalance`, or `decimals`; or the event is empty | Nothing for those tokens; no write if empty | Last amounts stay | +| 10 | Account Activity sends an amount that is wrong | Overlay overwrites that token | Wrong number until the next Accounts API `full` poll | +| 11 | Snap live event names some holdings but skips another token that changed | Overlay only the tokens in the event | Omitted token stays until the next Snap `full` poll | +| 12 | Staking source updates one vault and not another, or a token poll never includes staked ETH | Do not overwrite the omitted vault | Staked ETH stays at the last staking read until that vault is sent | + +## Failed chains + +A v6 source either contributes a complete chain snapshot or nothing for that +chain: + +- Accounts API: unprocessed networks and unresolved `includeAssetIds` go in + `errors` +- RPC: a chain that resolved **nothing** goes in `errors`. A single failed + `balanceOf` or unknown decimals does **not** fail the chain; the previous + amount is carried so the snapshot stays complete. v5 still overlays only + the tokens that succeeded (`merge`) +- Snap: a thrown `getAccountBalances`, an empty map, or a snap that does not + own the account contributes nothing for that account; remaining chains stay + on the request for the next middleware +- `filterFailedChainBalances` strips balances on failed chains before state is + written, so an empty failed chain cannot cover (and wipe) that slice + +RpcFallback then retries only `errors` keys. + +## Behavior differences at a glance + +| Concern | v5 | v6 | +| ----------------------- | ----------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- | +| Accounts API endpoint | `fetchV5MultiAccountBalances` | `fetchV6MultiAccountBalances` | +| Fetch update mode | `merge` for Accounts API, Snap, and RPC | `full` for Accounts API, Snap, and RPC; RPC carries previous amounts for unread tokens | +| Subscribe update mode | `merge` for every source | `full` for Accounts API, Snap snapshots, and RPC (carry previous on unread tokens); `merge` for events (Snap, Account Activity, staking, detection) | +| State writer | `effectiveAccountBalancesV5` | `effectiveAccountBalancesV6(updateMode)` | +| Covered-chain replace | `replaceCoveredChainBalances` on force refresh; restore custom + staked | `full` replaces the covered slice; keep visible natives/pins/defaults + staking via `shouldKeepAsset` | +| Pinned assets | `request.customAssets` + merge restore | Visibility → `includeAssetIds` / Snap fill from state-or-0 / RPC fetch list | +| Hidden assets | Not sent to the endpoint | `excludeAssetIds`; skipped by RPC and Snap; omitted balances dropped on `full` (unhide fetches) | +| Default tracked assets | Survive only if already in state or returned | Always in visibility, so a `full` refresh keeps them at zero when unheld | +| Token detection filter | Drop unknown tokens when detection is off | Not applied; v6 snapshot is kept in full | +| RPC token list | Request `customAssets` + tracked balances | Visible natives, pins, and default tracked from state, plus already-tracked ERC-20s | +| RPC pin supplement | Extra `customAssetsOnly` poll | None | +| Custom-asset graduation | Fast lane and live Accounts API / Account Activity | Never | +| Failed chain | Overlay whatever tokens succeeded | Accounts API / Snap drop the chain; RPC drops only chains that resolved nothing, and carries previous amounts for unread tokens on a `full` snapshot | + +v6 `updateMode` by source: + +| Source | Fetch / force refresh | Subscribe / live updates | +| ---------------- | ----------------------------------------------- | ------------------------------------------------------------------------ | +| Accounts API | `full` | `full` (polls `fetch`) | +| Snap | `full` | `merge` | +| RPC | `full` (carry previous amount on unread tokens) | `full` (carry previous amount on unread tokens); detection stays `merge` | +| Account Activity | n/a (event-only) | `merge` | +| Staked balances | `merge` | `merge` | + +Fast-lane Accounts API `full` merged with staking `merge` still stamps `full` +(`mergeDataResponses` promotes it). Staking rows that are present in that +merged payload are written with the snapshot; if they are omitted, `shouldKeepAsset` +keeps the prior vault balance. + +## Where the code lives + +| Concern | v5 | v6 | +| --------------------- | ------------------------------------------ | ------------------------------------------------------------------ | +| Flag | `#isBalanceV6Enabled` (shared) | injected into Accounts API, Snap, RPC, BalanceFetcher, RpcFallback | +| Visibility | n/a (request `customAssets`) | `getAssetVisibility` / `#getAssetVisibility` | +| Force-update request | `#buildForceUpdateRequestV5` | `#buildDataRequest` (pins read from state) | +| Force-update pipeline | `#forceUpdateAssetsV5`, `#runFastFetchV5` | `#forceUpdateAssetsV6`, `#runFastFetchV6` | +| Fast lane composition | `buildFastFetchSources` (graduation on) | `buildFastFetchSources` (graduation off) | +| Subscribe | `#subscribeAssetsBalance` + RPC supplement | `#subscribeAssetsBalance` only | +| Update handling | `#handleAssetsUpdateV5` | `#handleAssetsUpdateV6` | +| State apply | `applyV5AccountBalanceUpdate` | `applyV6AccountBalanceUpdate` | +| Balance merge | `effectiveAccountBalancesV5` | `effectiveAccountBalancesV6` | +| Accounts API fetch | `#fetchV5Balances` | `#fetchV6Balances` | +| Snap fetch | `#fetchV5` | `#fetchV6` | +| RPC fetch | `#fetchV5` | `#fetchV6` | +| RPC fetch list | `#getAssetsToFetchV5` | `#getAssetsToFetchV6` | +| RPC fallback | `#recoverV5` | `#recoverV6` | +| Failed-chain filter | `filterFailedChainBalances` (shared) | same helper; v6 uses it to keep snapshots atomic | + +## Deleting v5 after rollout + +Once v6 is accepted as the only behavior: + +1. Remove the v5 methods (`#fetchV5`, `#fetchV5Balances`, `#recoverV5`, + `#forceUpdateAssetsV5`, `#handleAssetsUpdateV5`, + `applyV5AccountBalanceUpdate`, `effectiveAccountBalancesV5`, + `#subscribeRpcCustomAssetsSupplement`) and every + `!this.#isBalanceV6Enabled()` branch. +2. Keep the v6 methods as the single orchestration path. Drop + `isStakingContractAssetId` from `shouldKeepAsset` when Accounts API returns ETH + staked balances, and drop its native branch when Accounts API always returns + an explicit zero row for every requested native. +3. Remove this document's v5/v6 comparison tables. diff --git a/packages/assets-controller/docs/decisions/README.md b/packages/assets-controller/docs/decisions/README.md new file mode 100644 index 00000000000..7afc5eecae8 --- /dev/null +++ b/packages/assets-controller/docs/decisions/README.md @@ -0,0 +1,7 @@ +# Local decisions + +Package-scoped records for `@metamask/assets-controller`. + +Name each file `YYYY-MM-DD-short-kebab-title.md`. These sit next to the +code they describe; org-wide ADRs still go in +[MetaMask/decisions](https://github.com/MetaMask/decisions). diff --git a/packages/assets-controller/jest.config.cjs b/packages/assets-controller/jest.config.cjs index 795b4b105c0..1782d6db10e 100644 --- a/packages/assets-controller/jest.config.cjs +++ b/packages/assets-controller/jest.config.cjs @@ -18,7 +18,7 @@ module.exports = merge(baseConfig, { coverageThreshold: { global: { branches: 82.17, - functions: 90.06, + functions: 90.01, lines: 90.56, statements: 90.56, }, diff --git a/packages/assets-controller/src/AssetsController-method-action-types.ts b/packages/assets-controller/src/AssetsController-method-action-types.ts index 0e0479ff245..68590b2ce8a 100644 --- a/packages/assets-controller/src/AssetsController-method-action-types.ts +++ b/packages/assets-controller/src/AssetsController-method-action-types.ts @@ -122,7 +122,8 @@ export type AssetsControllerGetStateForTransactionPayAction = { /** * Add a custom asset for an account. * Custom assets are included in subscription and fetch operations. - * Adding a custom asset also unhides it if it was previously hidden. + * Adding a custom asset also unhides it if it was previously hidden, + * and force-fetches that asset's chain (including every visible pin). * * When `pendingMetadata` is provided (e.g. from the extension's pending-tokens * flow), the token metadata is persisted immediately into `assetsInfo` so the @@ -173,6 +174,9 @@ export type AssetsControllerHideAssetAction = { /** * Unhide an asset globally. + * Force-fetches that asset's chain (including every visible pin) so a + * `full` snapshot can restore the balance immediately, then re-evaluates + * subscriptions so later polls stop excluding it. * * @param assetId - The CAIP-19 asset ID to unhide. */ diff --git a/packages/assets-controller/src/AssetsController.bsc-spam-token-filtering.integration.test.ts b/packages/assets-controller/src/AssetsController.bsc-spam-token-filtering.integration.test.ts index 66048972372..1e9eb69b023 100644 --- a/packages/assets-controller/src/AssetsController.bsc-spam-token-filtering.integration.test.ts +++ b/packages/assets-controller/src/AssetsController.bsc-spam-token-filtering.integration.test.ts @@ -28,9 +28,9 @@ import type { AssetsControllerState } from './AssetsController.js'; * * Boots the real controller, answers the same captured APIs as * `buildFastFetchSources.bsc-spam-token-filtering.integration.test.ts`, and - * asserts CDOGE never lands in persisted state. - * - * Integration Expectation - CDOGE is correctly filtered out of controller state. + * asserts CDOGE never lands in persisted state — unless the user imported it + * as a custom asset, in which case it must survive (see the custom-asset + * suite below). */ type StateSurface = { @@ -63,14 +63,17 @@ async function withController( { state = buildEmptyAssetsState(), queryApiClient = createTestApiClient(), + remoteFeatureFlags = {}, }: { state?: Partial; queryApiClient?: ApiPlatformClient; + remoteFeatureFlags?: Record; }, fn: WithControllerCallback, ): Promise { const { rootMessenger, assetsControllerMessenger } = createMockMessengers({ - registerCustomRootActions: registerBscSpamControllerActions, + registerCustomRootActions: (messenger) => + registerBscSpamControllerActions(messenger, { remoteFeatureFlags }), }); const controller = new AssetsController({ @@ -90,25 +93,31 @@ async function withController( async function fetchWallet( state: Partial = buildEmptyAssetsState(), + remoteFeatureFlags: Record = {}, ): Promise { const { accountsSupportedNetworks } = mockBscSpamApis(); - return await withController({ state }, async ({ controller }) => { - // wait for `AccountsApiDataSource` to ask `/v2/supportedNetworks` to indicate the fast-lane is ready - await waitFor(() => expect(accountsSupportedNetworks.isDone()).toBe(true)); + return await withController( + { state, remoteFeatureFlags }, + async ({ controller }) => { + // wait for `AccountsApiDataSource` to ask `/v2/supportedNetworks` to indicate the fast-lane is ready + await waitFor(() => + expect(accountsSupportedNetworks.isDone()).toBe(true), + ); - await controller.getAssets([buildBscSpamAccount()], { - chainIds: [BSC_CHAIN_ID], - forceUpdate: true, - }); + await controller.getAssets([buildBscSpamAccount()], { + chainIds: [BSC_CHAIN_ID], + forceUpdate: true, + }); - // `getAssets` awaits the fast lane only; the slow lane is fire-and-forget - // and can still be writing. Let state settle so the assertions about CDOGE - // being absent cannot pass just because nothing has landed yet. - await waitUntilStable(() => controller.state); + // `getAssets` awaits the fast lane only; the slow lane is fire-and-forget + // and can still be writing. Let state settle so the assertions about CDOGE + // being absent cannot pass just because nothing has landed yet. + await waitUntilStable(() => controller.state); - return controller.state; - }); + return controller.state; + }, + ); } const WALLET_PASSES = [ @@ -168,3 +177,40 @@ describe('AssetsController: BNB Chain spam token (CDOGE)', () => { }); }); }); + +describe('AssetsController: BNB Chain spam token (CDOGE) imported as a custom asset', () => { + afterEach(() => { + cleanAll(); + }); + + function buildCustomAssetWalletState(): AssetsControllerState { + return buildEmptyAssetsState({ + customAssets: { [BSC_SPAM_ACCOUNT_ID]: [CDOGE_ASSET_ID_CHECKSUM] }, + assetsBalance: { + [BSC_SPAM_ACCOUNT_ID]: { + [CDOGE_ASSET_ID_CHECKSUM]: { amount: '0' }, + }, + }, + assetsInfo: { + [CDOGE_ASSET_ID_CHECKSUM]: { + type: 'erc20', + symbol: 'CDOGE', + name: '$$$DOGECHAIN', + decimals: 9, + }, + }, + }); + } + + it('keeps the imported token in customAssets, balances and metadata', async () => { + const state = await fetchWallet(buildCustomAssetWalletState(), { + assetsAccountsApiV6: true, + }); + + expect(state.customAssets[BSC_SPAM_ACCOUNT_ID]).toContain( + CDOGE_ASSET_ID_CHECKSUM, + ); + expect(BALANCES.lookUp(state, CDOGE_ASSET_ID_LOWERCASE)).toBeDefined(); + expect(METADATA.lookUp(state, CDOGE_ASSET_ID_LOWERCASE)).toBeDefined(); + }); +}); diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 5df26b2945a..9b512f390a9 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -20,12 +20,17 @@ import type { AssetsControllerMessenger, AssetsControllerState, } from './AssetsController.js'; +import { AccountActivityDataSource } from './data-sources/AccountActivityDataSource.js'; import type { AccountsApiDataSourceConfig } from './data-sources/AccountsApiDataSource.js'; +import { AccountsApiDataSource } from './data-sources/AccountsApiDataSource.js'; import type { PriceDataSourceConfig } from './data-sources/PriceDataSource.js'; import { PriceDataSource } from './data-sources/PriceDataSource.js'; import { RpcDataSource } from './data-sources/RpcDataSource.js'; import { TokenDataSource } from './data-sources/TokenDataSource.js'; -import { buildDefaultAssetsInfo } from './defaults.js'; +import { + buildDefaultAssetsInfo, + getDefaultTrackedAssetsForChain, +} from './defaults.js'; import type { Assets3346MigrationState } from './migrations/healAssetsInfoMetadata.js'; import type { Caip19AssetId, @@ -34,6 +39,7 @@ import type { DataRequest, DataResponse, FungibleAssetMetadata, + Middleware, } from './types.js'; import { formatExchangeRatesForBridge, @@ -100,6 +106,9 @@ const MOCK_ASSET_ID = const MOCK_ASSET_ID_LOWERCASE = 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48' as Caip19AssetId; const MOCK_NATIVE_ASSET_ID = 'eip155:1/slip44:60' as Caip19AssetId; +/** mUSD on mainnet — a controller-managed default tracked asset. */ +const [MOCK_DEFAULT_TRACKED_ASSET_ID] = + getDefaultTrackedAssetsForChain('eip155:1'); /** * Activate asset tracking by marking the UI open, the keyring unlocked, and @@ -108,14 +117,15 @@ const MOCK_NATIVE_ASSET_ID = 'eip155:1/slip44:60' as Caip19AssetId; * * @param messenger - The root messenger used to publish lifecycle events. */ +type LifecyclePublish = (topic: string, payload?: unknown) => void; + async function activateTracking(messenger: RootMessenger): Promise { - ( - messenger as unknown as { - publish: (topic: string, payload?: unknown) => void; - } - ).publish('ClientController:stateChanged', { isUiOpen: true }); + (messenger as unknown as { publish: LifecyclePublish }).publish( + 'ClientController:stateChanged', + { isUiOpen: true }, + ); messenger.publish('KeyringController:unlock'); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); @@ -524,7 +534,7 @@ describe('AssetsController', () => { // Action handlers should be registered expect(() => { - (messenger.call as CallableFunction)( + (messenger as unknown as { call: LifecyclePublish }).call( 'AssetsController:getCustomAssets', MOCK_ACCOUNT_ID, ); @@ -703,6 +713,75 @@ describe('AssetsController', () => { }); }); + it('seeds a Stellar custom asset with empty trustline metadata', async () => { + const stellarAssetId = + 'stellar:pubnet/asset:USDC-GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN' as Caip19AssetId; + + await withController(async ({ controller }) => { + await controller.addCustomAsset(MOCK_ACCOUNT_ID, stellarAssetId); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[stellarAssetId], + ).toStrictEqual({ + amount: '0', + metadata: { + authorized: false, + limit: '0', + sponsored: false, + }, + }); + }); + }); + + it('force-fetches the asset chain with every pinned token as includeAssetIds', async () => { + // Use a valid checksummed address (DAI token address) + const secondAssetId = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await flushPromises(); + + await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); + fetchV6MultiAccountBalances.mockClear(); + await controller.addCustomAsset(MOCK_ACCOUNT_ID, secondAssetId); + }, + ); + + // The v6 snapshot is applied as `updateMode: 'full'`, so every pin on + // the chain must be requested — otherwise other custom tokens drop. + expect(fetchV6MultiAccountBalances).toHaveBeenCalled(); + for (const [, params] of fetchV6MultiAccountBalances.mock.calls) { + expect(params?.includeAssetIds ?? []).toStrictEqual( + expect.arrayContaining([MOCK_ASSET_ID, secondAssetId]), + ); + } + }); + it('does not overwrite an existing balance when re-adding a custom asset', async () => { await withController( { @@ -892,6 +971,60 @@ describe('AssetsController', () => { }); }); + describe('unhideAsset', () => { + it('force-fetches the unhidden token as an includeAssetId on v6', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + state: { + customAssets: { + [MOCK_ACCOUNT_ID]: [MOCK_ASSET_ID], + }, + assetPreferences: { + [MOCK_ASSET_ID]: { hidden: true }, + }, + }, + }, + async ({ controller }) => { + await flushPromises(); + fetchV6MultiAccountBalances.mockClear(); + + await controller.unhideAsset(MOCK_ASSET_ID); + + expect( + controller.state.assetPreferences[MOCK_ASSET_ID]?.hidden, + ).toBeUndefined(); + expect(fetchV6MultiAccountBalances).toHaveBeenCalled(); + for (const [, params] of fetchV6MultiAccountBalances.mock.calls) { + expect(params?.includeAssetIds ?? []).toContain(MOCK_ASSET_ID); + expect(params?.excludeAssetIds ?? []).not.toContain(MOCK_ASSET_ID); + } + }, + ); + }); + }); + describe('getAssetMetadata', () => { it('returns metadata for existing asset', async () => { const initialState: Partial = { @@ -1447,7 +1580,7 @@ describe('AssetsController', () => { await withController( { queryApiClient, - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, }, async ({ controller }) => { await flushPromises(); @@ -1463,109 +1596,494 @@ describe('AssetsController', () => { ); }); - describe('pipeline splitting', () => { - it('returns from getAssets before background pipelines complete', async () => { - // Spy on handleAssetsUpdate to count how many times state is written. - // Fast pipeline commits once; background pipelines each commit once more. - await withController(async ({ controller }) => { - const updateSpy = jest.spyOn(controller, 'handleAssetsUpdate'); + it('forwards user-pinned custom assets to the Accounts API v6 endpoint as includeAssetIds', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); - await controller.getAssets([createMockInternalAccount()], { - forceUpdate: true, - }); + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; - // getAssets has returned — fast pipeline committed to state. - // Background pipelines are still in flight (fire-and-forget). - expect(updateSpy).not.toHaveBeenCalled(); // internal #updateState, not handleAssetsUpdate + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; - // Let all pending microtasks resolve so background pipelines finish. + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { await flushPromises(); - updateSpy.mockRestore(); - }); - }); + await controller.addCustomAsset(MOCK_ACCOUNT_ID, customToken); - it('getAssets resolves without error in basic functionality mode', async () => { - await withController(async ({ controller }) => { - const assets = await controller.getAssets( - [createMockInternalAccount()], - { forceUpdate: true }, + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + + expect(fetchV6MultiAccountBalances).toHaveBeenCalledWith( + expect.any(Array), + expect.objectContaining({ + includeAssetIds: expect.arrayContaining([customToken]), + }), + expect.anything(), ); - expect(assets).toBeDefined(); + }, + ); + }); - await flushPromises(); - }); + it('stops forwarding a pinned custom asset as includeAssetIds once it is hidden', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], }); - it('background pipelines merge state without overwriting fast-pipeline results', async () => { - const initialState: Partial = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [MOCK_NATIVE_ASSET_ID]: { amount: '1' }, - }, - }, - }; + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; - await withController( - { state: initialState }, - async ({ controller }) => { - await controller.getAssets([createMockInternalAccount()], { - forceUpdate: true, - }); + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; - await flushPromises(); + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await flushPromises(); - // Background pipelines overlay balances without wiping fast-pipeline results. - expect( - controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ - MOCK_NATIVE_ASSET_ID - ], - ).toBeDefined(); - }, - ); - }); + await controller.addCustomAsset(MOCK_ACCOUNT_ID, customToken); + controller.hideAsset(customToken); - it('getAssets resolves without error when isBasicFunctionality is false', async () => { - await withController( - { isBasicFunctionality: () => false }, - async ({ controller }) => { - const assets = await controller.getAssets( - [createMockInternalAccount()], - { forceUpdate: true }, - ); - expect(assets).toBeDefined(); + // The pin is preserved so the import is remembered. + expect(controller.getCustomAssets(MOCK_ACCOUNT_ID)).toContain( + customToken, + ); - await flushPromises(); - }, - ); - }); + fetchV6MultiAccountBalances.mockClear(); + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); - it('does not run token or price middleware in getAssets pipelines when isBasicFunctionality is false', async () => { - const tokenMiddlewareGetter = jest.spyOn( - TokenDataSource.prototype, - 'assetsMiddleware', - // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. - 'get', - ) as unknown as jest.SpyInstance; - const priceMiddlewareGetter = jest.spyOn( - PriceDataSource.prototype, + expect(fetchV6MultiAccountBalances).toHaveBeenCalledWith( + expect.any(Array), + expect.objectContaining({ + excludeAssetIds: expect.arrayContaining([customToken]), + }), + expect.anything(), + ); + for (const [, params] of fetchV6MultiAccountBalances.mock.calls) { + expect(params?.includeAssetIds ?? []).not.toContain(customToken); + } + }, + ); + }); + + it('includes every pinned asset on the v5 force-update request, including other chains', async () => { + const mainnetToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const polygonToken = + 'eip155:137/erc20:0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174' as Caip19AssetId; + + const capturedCustomAssets: (Caip19AssetId[] | undefined)[] = []; + const accountsApiMiddleware: jest.MockedFunction = jest.fn( + async (ctx, next) => { + capturedCustomAssets.push(ctx.request.customAssets); + return next(ctx); + }, + ); + const middlewareGetter = jest + .spyOn( + AccountsApiDataSource.prototype, 'assetsMiddleware', // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. 'get', - ) as unknown as jest.SpyInstance; - - await withController( - { isBasicFunctionality: () => false }, - async ({ controller }) => { - tokenMiddlewareGetter.mockClear(); - priceMiddlewareGetter.mockClear(); + ) + .mockReturnValue(accountsApiMiddleware) as unknown as jest.SpyInstance; - await controller.getAssets([createMockInternalAccount()], { - forceUpdate: true, - }); - await flushPromises(); - }, - ); + await withController(async ({ controller }) => { + await controller.addCustomAsset(MOCK_ACCOUNT_ID, mainnetToken); + await controller.addCustomAsset(MOCK_ACCOUNT_ID, polygonToken); + + capturedCustomAssets.length = 0; + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + }); + + expect(capturedCustomAssets.length).toBeGreaterThan(0); + for (const customAssets of capturedCustomAssets) { + expect(customAssets).toStrictEqual( + expect.arrayContaining([mainnetToken, polygonToken]), + ); + } + + middlewareGetter.mockRestore(); + }); + + it('includes only pins on the requested chains as includeAssetIds', async () => { + const mainnetToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const polygonToken = + 'eip155:137/erc20:0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174' as Caip19AssetId; + + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1, 137], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await flushPromises(); + + await controller.addCustomAsset(MOCK_ACCOUNT_ID, mainnetToken); + await controller.addCustomAsset(MOCK_ACCOUNT_ID, polygonToken); + + fetchV6MultiAccountBalances.mockClear(); + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + }, + ); + + expect(fetchV6MultiAccountBalances).toHaveBeenCalled(); + for (const [, params] of fetchV6MultiAccountBalances.mock.calls) { + expect(params?.includeAssetIds ?? []).toContain(mainnetToken); + expect(params?.includeAssetIds ?? []).not.toContain(polygonToken); + } + }); + + it('forwards user-hidden assets to the Accounts API v6 endpoint as excludeAssetIds', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + const hiddenToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await flushPromises(); + + controller.hideAsset(hiddenToken); + + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + + expect(fetchV6MultiAccountBalances).toHaveBeenCalledWith( + expect.any(Array), + expect.objectContaining({ + excludeAssetIds: expect.arrayContaining([hiddenToken]), + }), + expect.anything(), + ); + }, + ); + }); + + it('forwards bypassServerCache on Accounts API v6 force updates', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await flushPromises(); + + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + bypassServerCache: true, + }); + + expect(fetchV6MultiAccountBalances).toHaveBeenCalledWith( + expect.any(Array), + expect.objectContaining({ + includeAssetIds: expect.arrayContaining([ + MOCK_DEFAULT_TRACKED_ASSET_ID, + ]), + }), + expect.objectContaining({ bypassServerCache: true }), + ); + }, + ); + }); + + describe('pipeline splitting', () => { + it('returns from getAssets before background pipelines complete', async () => { + // Spy on handleAssetsUpdate to count how many times state is written. + // Fast pipeline commits once; background pipelines each commit once more. + await withController(async ({ controller }) => { + const updateSpy = jest.spyOn(controller, 'handleAssetsUpdate'); + + await controller.getAssets([createMockInternalAccount()], { + forceUpdate: true, + }); + + // getAssets has returned — fast pipeline committed to state. + // Background pipelines are still in flight (fire-and-forget). + expect(updateSpy).not.toHaveBeenCalled(); // internal #updateState, not handleAssetsUpdate + + // Let all pending microtasks resolve so background pipelines finish. + await flushPromises(); + + updateSpy.mockRestore(); + }); + }); + + it('getAssets resolves without error in basic functionality mode', async () => { + await withController(async ({ controller }) => { + const assets = await controller.getAssets( + [createMockInternalAccount()], + { forceUpdate: true }, + ); + expect(assets).toBeDefined(); + + await flushPromises(); + }); + }); + + it('background pipelines merge state without overwriting fast-pipeline results', async () => { + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '1' }, + }, + }, + }; + + await withController( + { state: initialState }, + async ({ controller }) => { + await controller.getAssets([createMockInternalAccount()], { + forceUpdate: true, + }); + + await flushPromises(); + + // Background pipelines overlay balances without wiping fast-pipeline results. + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_NATIVE_ASSET_ID + ], + ).toBeDefined(); + }, + ); + }); + + it('getAssets resolves without error when isBasicFunctionality is false', async () => { + await withController( + { isBasicFunctionality: () => false }, + async ({ controller }) => { + const assets = await controller.getAssets( + [createMockInternalAccount()], + { forceUpdate: true }, + ); + expect(assets).toBeDefined(); + + await flushPromises(); + }, + ); + }); + + it('falls back to RPC on the force-update fast path for unprocessed include asset ids', async () => { + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [customToken], + }); + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + const rpcRequests: { + chainIds: ChainId[]; + customAssets: Caip19AssetId[] | undefined; + }[] = []; + const rpcMiddleware: jest.MockedFunction = jest.fn( + async (ctx, next) => { + rpcRequests.push({ + chainIds: ctx.request.chainIds, + customAssets: ctx.request.customAssets, + }); + return next(ctx); + }, + ); + const rpcMiddlewareGetter = jest + .spyOn( + RpcDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) + .mockReturnValue(rpcMiddleware) as unknown as jest.SpyInstance; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await flushPromises(); + + await controller.addCustomAsset(MOCK_ACCOUNT_ID, customToken); + + rpcMiddleware.mockClear(); + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + }, + ); + + // RpcFallbackMiddleware retries the pin's chain; RPC reads native, + // balances, and pins from state rather than a scoped customAssets list. + expect(rpcMiddleware).toHaveBeenCalled(); + expect( + rpcRequests.some(({ chainIds }) => chainIds.includes('eip155:1')), + ).toBe(true); + expect( + rpcRequests.every(({ customAssets }) => customAssets === undefined), + ).toBe(true); + + rpcMiddlewareGetter.mockRestore(); + }); + + it('does not run token or price middleware in getAssets pipelines when isBasicFunctionality is false', async () => { + const tokenMiddlewareGetter = jest.spyOn( + TokenDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + const priceMiddlewareGetter = jest.spyOn( + PriceDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + + await withController( + { isBasicFunctionality: () => false }, + async ({ controller }) => { + tokenMiddlewareGetter.mockClear(); + priceMiddlewareGetter.mockClear(); + + await controller.getAssets([createMockInternalAccount()], { + forceUpdate: true, + }); + await flushPromises(); + }, + ); expect(tokenMiddlewareGetter).not.toHaveBeenCalled(); expect(priceMiddlewareGetter).not.toHaveBeenCalled(); @@ -1573,6 +2091,67 @@ describe('AssetsController', () => { tokenMiddlewareGetter.mockRestore(); priceMiddlewareGetter.mockRestore(); }); + + it('does not call Accounts API v6 on force update when isBasicFunctionality is false', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + const tokenMiddlewareGetter = jest.spyOn( + TokenDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + const priceMiddlewareGetter = jest.spyOn( + PriceDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + await withController( + { + isBasicFunctionality: () => false, + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + tokenMiddlewareGetter.mockClear(); + priceMiddlewareGetter.mockClear(); + fetchV6MultiAccountBalances.mockClear(); + + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + await flushPromises(); + }, + ); + + expect(fetchV6MultiAccountBalances).not.toHaveBeenCalled(); + expect(tokenMiddlewareGetter).not.toHaveBeenCalled(); + expect(priceMiddlewareGetter).not.toHaveBeenCalled(); + + tokenMiddlewareGetter.mockRestore(); + priceMiddlewareGetter.mockRestore(); + }); }); it('filters by chainIds option', async () => { @@ -2016,17 +2595,89 @@ describe('AssetsController', () => { }, }, }, - 'RpcDataSource', - request, - ); - }, - ); + 'RpcDataSource', + request, + ); + }, + ); + + expect(tokenMiddlewareGetter).not.toHaveBeenCalled(); + expect(priceMiddlewareGetter).not.toHaveBeenCalled(); + + tokenMiddlewareGetter.mockRestore(); + priceMiddlewareGetter.mockRestore(); + }); + + it('falls back to RPC for chains a subscription update flagged as errored (e.g. unprocessedNetworks)', async () => { + const rpcMiddlewareGetter = jest.spyOn( + RpcDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + + const request: DataRequest = { + accountsWithSupportedChains: [], + chainIds: ['eip155:1'], + dataTypes: ['balance'], + }; + + await withController( + { remoteFeatureFlags: { assetsAccountsApiV6: true } }, + async ({ controller }) => { + rpcMiddlewareGetter.mockClear(); + + await controller.handleAssetsUpdate( + { + assetsBalance: {}, + errors: { 'eip155:1': 'Unprocessed networks' }, + }, + 'AccountsApiDataSource', + request, + ); + }, + ); + + // The RpcFallbackMiddleware pulls the RPC data source middleware only when + // there are errored chains to recover. + expect(rpcMiddlewareGetter).toHaveBeenCalled(); + + rpcMiddlewareGetter.mockRestore(); + }); + + it('does not run the RPC fallback when a subscription update has no errored chains', async () => { + const rpcMiddlewareGetter = jest.spyOn( + RpcDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + + const request: DataRequest = { + accountsWithSupportedChains: [], + chainIds: ['eip155:1'], + dataTypes: ['balance'], + }; + + await withController(async ({ controller }) => { + rpcMiddlewareGetter.mockClear(); + + await controller.handleAssetsUpdate( + { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '1' }, + }, + }, + }, + 'AccountsApiDataSource', + request, + ); + }); - expect(tokenMiddlewareGetter).not.toHaveBeenCalled(); - expect(priceMiddlewareGetter).not.toHaveBeenCalled(); + expect(rpcMiddlewareGetter).not.toHaveBeenCalled(); - tokenMiddlewareGetter.mockRestore(); - priceMiddlewareGetter.mockRestore(); + rpcMiddlewareGetter.mockRestore(); }); }); @@ -2188,6 +2839,56 @@ describe('AssetsController', () => { }); }); + describe('two-axis subscription handoff (chains + custom assets)', () => { + it('does not attach pinned assets to the account-activity subscription', async () => { + // Account activity owns eip155:1. Pins stay in controller state; the + // websocket does not take include/exclude lists. + jest + .spyOn(AccountActivityDataSource.prototype, 'getActiveChainsSync') + .mockReturnValue(['eip155:1']); + const wsSubscribeSpy = jest + .spyOn(AccountActivityDataSource.prototype, 'subscribe') + .mockResolvedValue(undefined); + const rpcSubscribeSpy = jest + .spyOn(RpcDataSource.prototype, 'subscribe') + .mockResolvedValue(undefined); + + await withController( + { remoteFeatureFlags: { assetsAccountsApiV6: true } }, + async ({ controller }) => { + await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); + + const wsRequest = wsSubscribeSpy.mock.calls.at(-1)?.[0].request; + expect(wsRequest?.chainIds).toStrictEqual(['eip155:1']); + expect(wsRequest?.customAssets).toBeUndefined(); + + expect(rpcSubscribeSpy).not.toHaveBeenCalled(); + }, + ); + }); + + it('does not create an RPC subscription for pinned assets on unassigned chains', async () => { + // Account activity is not active on the pin's chain... + jest + .spyOn(AccountActivityDataSource.prototype, 'getActiveChainsSync') + .mockReturnValue([]); + jest + .spyOn(AccountActivityDataSource.prototype, 'subscribe') + .mockResolvedValue(undefined); + const rpcSubscribeSpy = jest + .spyOn(RpcDataSource.prototype, 'subscribe') + .mockResolvedValue(undefined); + + await withController(async ({ controller }) => { + // ...and RPC has no provider for the chain (no networks configured in + // the mocked NetworkController), so it is assigned no chains. + await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); + + expect(rpcSubscribeSpy).not.toHaveBeenCalled(); + }); + }); + }); + describe('handleAssetsUpdate - state updates', () => { it('updates state with balance data', async () => { await withController(async ({ controller }) => { @@ -2485,79 +3186,374 @@ describe('AssetsController', () => { [MOCK_ASSET_ID]: { amount: '1000000' }, }, }, - }, - 'TestSource', - ); + }, + 'TestSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toStrictEqual({ amount: '1000000' }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_NATIVE_ASSET_ID + ], + ).toStrictEqual({ amount: '0' }); + }); + }); + + it('does not add default native balance for chains without a registered identifier', async () => { + await withController(async ({ controller }) => { + await controller.handleAssetsUpdate( + { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '1' }, + }, + }, + }, + 'TestSource', + ); + + const accountBalances = + controller.state.assetsBalance[MOCK_ACCOUNT_ID] ?? {}; + const nativeIds = Object.keys(accountBalances).filter((id) => + id.includes('/slip44:'), + ); + expect(nativeIds).toStrictEqual([MOCK_NATIVE_ASSET_ID]); + }); + }); + + it('preserves existing balances when merge update adds new chain data', async () => { + const polygonNative = 'eip155:137/slip44:966' as Caip19AssetId; + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '1' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + }, + }, + }; + + await withController({ state: initialState }, async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'merge', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [polygonNative]: { amount: '10' }, + }, + }, + }, + 'TestSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toStrictEqual({ amount: '1' }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_NATIVE_ASSET_ID + ], + ).toStrictEqual({ amount: '0.5' }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[polygonNative], + ).toStrictEqual({ amount: '10' }); + }); + }); + + it('keeps pins a full update omits, since an omission means the source could not report them', async () => { + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '1' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + }, + }, + customAssets: { + [MOCK_ACCOUNT_ID]: [MOCK_ASSET_ID], + }, + }; + + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '2' }, + }, + }, + }, + 'TestSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toStrictEqual({ amount: '1' }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_NATIVE_ASSET_ID + ], + ).toStrictEqual({ amount: '2' }); + }, + ); + }); + + it('drops a detected token a full update omits', async () => { + const detectedAssetId = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [detectedAssetId]: { amount: '1' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + }, + }, + }; + + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '2' }, + }, + }, + }, + 'TestSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[detectedAssetId], + ).toBeUndefined(); + }, + ); + }); + + it('drops an unknown-native fallback row a full update omits', async () => { + // A chain absent from the native asset map resolves its native to the + // zero-address ERC-20 fallback. That fabricated ID must never be + // undeletable: a full snapshot covering the chain that omits it drops + // the row instead of shielding it forever. + const unknownChainNativeId = + 'eip155:999999/erc20:0x0000000000000000000000000000000000000000' as Caip19AssetId; + const unknownChainTokenId = + 'eip155:999999/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [unknownChainNativeId]: { amount: '5' }, + }, + }, + }; + + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [unknownChainTokenId]: { amount: '1' }, + }, + }, + }, + 'TestSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + unknownChainNativeId + ], + ).toBeUndefined(); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + unknownChainTokenId + ], + ).toStrictEqual({ amount: '1' }); + }, + ); + }); + + it('keeps a known native balance a full update omits', async () => { + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '1' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + }, + }, + }; + + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '3' }, + }, + }, + }, + 'TestSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_NATIVE_ASSET_ID + ], + ).toStrictEqual({ amount: '0.5' }); + }, + ); + }); + + it('drops a hidden balance a full update omits, since unhide fetches a new snapshot', async () => { + const initialState: Partial = { + assetsInfo: { + [MOCK_ASSET_ID]: { + type: 'erc20', + symbol: 'TEST', + name: 'Hidden Test Token', + decimals: 18, + }, + }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '1' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + }, + }, + customAssets: { + [MOCK_ACCOUNT_ID]: [MOCK_ASSET_ID], + }, + assetPreferences: { + [MOCK_ASSET_ID]: { hidden: true }, + }, + }; + + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '2' }, + }, + }, + }, + 'TestSource', + ); - expect( - controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], - ).toStrictEqual({ amount: '1000000' }); - expect( - controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ - MOCK_NATIVE_ASSET_ID - ], - ).toStrictEqual({ amount: '0' }); - }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toBeUndefined(); + expect( + controller.getAccountAssetByID(MOCK_ACCOUNT_ID, MOCK_ASSET_ID), + ).toBeUndefined(); + }, + ); }); - it('does not add default native balance for chains without a registered identifier', async () => { - await withController(async ({ controller }) => { - await controller.handleAssetsUpdate( - { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [MOCK_ASSET_ID]: { amount: '1' }, + it('keeps default tracked assets a full update returns at zero', async () => { + // A force refresh (e.g. "Refresh list") replaces the chain slice, so + // mUSD only survives for accounts holding none because v6 requests it + // as an includeAssetId and the response carries it back at zero. + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_DEFAULT_TRACKED_ASSET_ID]: { amount: '0' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + }, + }, + }; + + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '2' }, + [MOCK_DEFAULT_TRACKED_ASSET_ID]: { amount: '0' }, + }, }, }, - }, - 'TestSource', - ); + 'TestSource', + ); - const accountBalances = - controller.state.assetsBalance[MOCK_ACCOUNT_ID] ?? {}; - const nativeIds = Object.keys(accountBalances).filter((id) => - id.includes('/slip44:'), - ); - expect(nativeIds).toStrictEqual([MOCK_NATIVE_ASSET_ID]); - }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_DEFAULT_TRACKED_ASSET_ID + ], + ).toStrictEqual({ amount: '0' }); + }, + ); }); - it('preserves existing balances when merge update adds new chain data', async () => { - const polygonNative = 'eip155:137/slip44:966' as Caip19AssetId; + it('does not overwrite a default tracked asset balance returned by a full update', async () => { const initialState: Partial = { assetsBalance: { [MOCK_ACCOUNT_ID]: { - [MOCK_ASSET_ID]: { amount: '1' }, - [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + [MOCK_DEFAULT_TRACKED_ASSET_ID]: { amount: '5' }, }, }, }; - await withController({ state: initialState }, async ({ controller }) => { - await controller.handleAssetsUpdate( - { - updateMode: 'merge', - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [polygonNative]: { amount: '10' }, + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_DEFAULT_TRACKED_ASSET_ID]: { amount: '7' }, + }, }, }, - }, - 'TestSource', - ); + 'TestSource', + ); - expect( - controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], - ).toStrictEqual({ amount: '1' }); - expect( - controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ - MOCK_NATIVE_ASSET_ID - ], - ).toStrictEqual({ amount: '0.5' }); - expect( - controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[polygonNative], - ).toStrictEqual({ amount: '10' }); - }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_DEFAULT_TRACKED_ASSET_ID + ], + ).toStrictEqual({ amount: '7' }); + }, + ); }); it('replaces covered-chain balances in merge mode when replaceCoveredChainBalances is set', async () => { @@ -2636,6 +3632,52 @@ describe('AssetsController', () => { }); }); + it('preserves existing staked balances when a full update omits them', async () => { + const stakingAssetId = + 'eip155:1/erc20:0x4FEF9D741011476750A243aC70b9789a63dd47Df' as Caip19AssetId; + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '1' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + [stakingAssetId]: { amount: '1.5' }, + }, + }, + }; + + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '2' }, + }, + }, + }, + 'AccountsApiDataSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toBeUndefined(); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_NATIVE_ASSET_ID + ], + ).toStrictEqual({ amount: '2' }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[stakingAssetId], + ).toStrictEqual({ amount: '1.5' }); + }, + ); + }); + it('replaces state when full update has authoritative data', async () => { const initialState: Partial = { assetsBalance: { @@ -2705,6 +3747,46 @@ describe('AssetsController', () => { }); }); + it('overlays balances without removing tokens when v6 merge mode is used', async () => { + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '6.185173' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.000390285791392' }, + }, + }, + }; + + await withController( + { + state: initialState, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'merge', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '0.000389261286724' }, + }, + }, + }, + 'TestSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toStrictEqual({ amount: '6.185173' }); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_NATIVE_ASSET_ID + ], + ).toStrictEqual({ amount: '0.000389261286724' }); + }, + ); + }); + it('seeds missing metadata in merge mode for RPC-only chains', async () => { const avaxNative = 'eip155:43114/slip44:9005' as Caip19AssetId; @@ -2853,7 +3935,7 @@ describe('AssetsController', () => { } ).publish('ClientController:stateChanged', { isUiOpen: true }); messenger.publish('KeyringController:unlock'); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); @@ -3138,15 +4220,17 @@ describe('AssetsController', () => { }); it('does not emit stateChange when a data source re-reports the same balance', async () => { - // Seed the native asset too: the controller backfills a zero-balance - // native entry for any account missing one, and that first backfill is - // itself a real (one-time) state change. Seeding it here isolates what - // this test actually checks — a repeated, otherwise-identical response. + // Seed the native and default tracked assets too: the controller + // backfills zero-balance entries for any account missing them, and that + // first backfill is itself a real (one-time) state change. Seeding them + // here isolates what this test actually checks — a repeated, + // otherwise-identical response. const initialState: Partial = { assetsBalance: { [MOCK_ACCOUNT_ID]: { [MOCK_ASSET_ID]: { amount: '1000000' }, [MOCK_NATIVE_ASSET_ID]: { amount: '0' }, + [MOCK_DEFAULT_TRACKED_ASSET_ID]: { amount: '0' }, }, }, }; @@ -3409,7 +4493,7 @@ describe('AssetsController', () => { } ).publish('ClientController:stateChanged', { isUiOpen: true }); messenger.publish('KeyringController:unlock'); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); @@ -3475,6 +4559,76 @@ describe('AssetsController', () => { ); }); + it('replaces pre-lock balances on unlock with a full Accounts API v6 snapshot', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + balances: [ + { + accountId: 'eip155:1:0x1234567890123456789012345678901234567890', + object: 'token', + type: 'native', + assetId: MOCK_NATIVE_ASSET_ID, + balance: '2', + }, + ], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + await withController( + { + clientControllerState: { isUiOpen: true }, + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + state: { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '100' }, + [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' }, + }, + }, + }, + }, + async ({ controller, messenger }) => { + ( + messenger as unknown as { + publish: (topic: string, payload?: unknown) => void; + } + ).publish('ClientController:stateChanged', { isUiOpen: true }); + messenger.publish('KeyringController:unlock'); + (messenger as unknown as { publish: LifecyclePublish }).publish( + 'AccountTreeController:initialized', + {}, + ); + + await flushPromises(); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toBeUndefined(); + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[ + MOCK_NATIVE_ASSET_ID + ], + ).toStrictEqual({ amount: '2' }); + }, + ); + }); + it('replaces pre-lock balances on unlock via merge with covered-chain replacement', async () => { const fetchV5MultiAccountBalances = jest.fn().mockResolvedValue({ balances: [ @@ -3518,7 +4672,7 @@ describe('AssetsController', () => { } ).publish('ClientController:stateChanged', { isUiOpen: true }); messenger.publish('KeyringController:unlock'); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); @@ -3559,7 +4713,7 @@ describe('AssetsController', () => { } ).publish('ClientController:stateChanged', { isUiOpen: true }); messenger.publish('KeyringController:unlock'); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); @@ -3620,7 +4774,7 @@ describe('AssetsController', () => { // #runStartupRefresh(), whose forced getAssets() call is still // pending (it never resolves until we call resolveGetAssets below). messenger.publish('KeyringController:unlock'); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); @@ -3714,7 +4868,7 @@ describe('AssetsController', () => { // Action handlers should be unregistered expect(() => { // The handler is unregistered, so calling it should throw - (messenger.call as CallableFunction)( + (messenger as unknown as { call: LifecyclePublish }).call( 'AssetsController:getAssets', createMockInternalAccount(), ); @@ -3728,7 +4882,7 @@ describe('AssetsController', () => { describe('network changes', () => { it('handles enabled networks change', async () => { await withController(async ({ messenger }) => { - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'NetworkEnablementController:stateChange', { enabledNetworkMap: { @@ -3770,7 +4924,9 @@ describe('AssetsController', () => { }), ]); - (controller.messenger.publish as CallableFunction)( + ( + controller.messenger as unknown as { publish: LifecyclePublish } + ).publish( 'NetworkEnablementController:stateChange', { enabledNetworkMap: { @@ -3808,7 +4964,9 @@ describe('AssetsController', () => { }), ]); - (controller.messenger.publish as CallableFunction)( + ( + controller.messenger as unknown as { publish: LifecyclePublish } + ).publish( 'NetworkEnablementController:stateChange', { enabledNetworkMap: { @@ -3838,7 +4996,7 @@ describe('AssetsController', () => { it('handles network being disabled', async () => { await withController(async ({ messenger }) => { - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'NetworkEnablementController:stateChange', { enabledNetworkMap: { @@ -3857,7 +5015,7 @@ describe('AssetsController', () => { await new Promise(process.nextTick); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'NetworkEnablementController:stateChange', { enabledNetworkMap: { @@ -3881,11 +5039,11 @@ describe('AssetsController', () => { it('refreshes assets when a network is added or removed', async () => { await withController(async ({ messenger }) => { - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'NetworkController:networkAdded', { chainId: '0x89' }, ); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'NetworkController:networkRemoved', { chainId: '0x89' }, ); @@ -4002,7 +5160,7 @@ describe('AssetsController', () => { } ).publish('ClientController:stateChanged', { isUiOpen: true }); messenger.publish('KeyringController:unlock'); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); @@ -4012,7 +5170,7 @@ describe('AssetsController', () => { fetchV2SupportedNetworks.mockClear(); selectedNetworkClientId = 'mainnet'; - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'NetworkController:networkDidChange', getNetworkState(), ); @@ -4047,7 +5205,7 @@ describe('AssetsController', () => { getAssetsSpy.mockClear(); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:selectedAccountGroupChange', 'entropy:mock-keyring-id-1/1', 'entropy:mock-keyring-id-1/0', @@ -4064,7 +5222,7 @@ describe('AssetsController', () => { await withController(async ({ controller, messenger }) => { const getAssetsSpy = jest.spyOn(controller, 'getAssets'); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:selectedAccountGroupChange', '', 'entropy:mock-keyring-id-1/0', @@ -4089,7 +5247,7 @@ describe('AssetsController', () => { messenger.publish('KeyringController:unlock'); await flushPromises(); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:selectedAccountGroupChange', 'entropy:mock-keyring-id-1/0', '', @@ -4098,7 +5256,7 @@ describe('AssetsController', () => { expect(getAssetsSpy).not.toHaveBeenCalled(); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); @@ -4119,7 +5277,7 @@ describe('AssetsController', () => { getAssetsSpy.mockClear(); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:selectedAccountGroupChange', 'entropy:mock-keyring-id-1/0', 'entropy:mock-keyring-id-1/0', @@ -4218,7 +5376,7 @@ describe('AssetsController', () => { // Intermediate tree mutations during init must not start tracking. getAccountsMock.mockReturnValue([createMockInternalAccount()]); - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:stateChange', {}, [], @@ -4228,7 +5386,7 @@ describe('AssetsController', () => { expect(getAssetsSpy).not.toHaveBeenCalled(); // Step 2: AccountTreeController.init() completes — tree is ready - (messenger.publish as CallableFunction)( + (messenger as unknown as { publish: LifecyclePublish }).publish( 'AccountTreeController:initialized', {}, ); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 8a864bedc3c..8ed6db23720 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -130,8 +130,7 @@ import { } from './pipeline/index.js'; import type { AccountId, - AssetPreferences, - AssetsControllerStateInternal, + AssetsControllerState, AssetsDataSource, AssetsUpdateMode, ChainId, @@ -149,7 +148,12 @@ import type { SubscriptionResponse, Asset, } from './types.js'; -import { ZERO_ADDRESS } from './utils/constants.js'; +import type { AssetVisibility } from './utils/assetVisibility.js'; +import { getAssetVisibility } from './utils/assetVisibility.js'; +import { + UNKNOWN_EVM_NATIVE_ASSET_REFERENCE, + ZERO_ADDRESS, +} from './utils/constants.js'; import { pickRpcCustomAssetsSupplement } from './utils/customAssetsRpcSupplement.js'; import { normalizeAmountString, @@ -159,12 +163,14 @@ import { buildNativeAssetsFromConstant, buildNativeAssetsFromApi, getDefaultNativeAssetBalance, + getZeroAssetBalance, } from './utils/index.js'; import type { BridgeExchangeRatesFormat, TransactionPayLegacyFormat, } from './utils/index.js'; import { emitTrace, withTrace } from './utils/trace.js'; +import type { TraceSpanData } from './utils/trace.js'; const NATIVE_ASSETS_QUERY_KEY = ['nativeAssets']; @@ -239,28 +245,7 @@ const log = createModuleLogger(projectLogger, CONTROLLER_NAME); // STATE TYPES // ============================================================================ -/** - * State structure for AssetsController. - * - * All values are JSON-serializable. UI preferences (e.g. hidden) are in - * assetPreferences, not in metadata. - * - * @see AssetsControllerStateInternal for the semantic type structure - */ -export type AssetsControllerState = { - /** Shared metadata for all assets (stored once per asset) */ - assetsInfo: { [assetId: string]: AssetMetadata }; - /** Per-account balance data */ - assetsBalance: { [accountId: string]: { [assetId: string]: AssetBalance } }; - /** Price data for assets */ - assetsPrice: { [assetId: string]: AssetPrice }; - /** Custom assets added by users per account (CAIP-19 asset IDs) */ - customAssets: { [accountId: string]: Caip19AssetId[] }; - /** UI preferences per asset (e.g. hidden) */ - assetPreferences: { [assetId: string]: AssetPreferences }; - /** Currently-active ISO 4217 currency code */ - selectedCurrency: SupportedCurrency; -}; +export type { AssetsControllerState }; /** * Returns the default state for AssetsController. @@ -583,16 +568,17 @@ function normalizeResponse(response: DataResponse): DataResponse { } /** - * Merge account balances from a data-source response into prior state. + * Compute the effective account balances from a data-source response and prior state. * * @param previousBalances - Balances already in state for this account. * @param accountBalances - Balances from the incoming response. - * @param customAssetIds - Custom assets to preserve when replacing covered chains. - * @param replaceCoveredChains - When true, drop prior balances on chains present - * in the response before applying it (authoritative chain slice). - * @returns The merged balance map for the account. + * @param customAssetIds - Custom assets to preserve when replacing covered chains + * on the Accounts API v5 path (`replaceCoveredChainBalances`). + * @param replaceCoveredChains - When true (v5 force refresh), drop prior balances + * on chains present in the response before applying it, then restore custom assets. + * @returns The effective balance map for the account. */ -function mergeAccountBalances( +function effectiveAccountBalancesV5( previousBalances: Record, accountBalances: Record, customAssetIds: Caip19AssetId[], @@ -618,7 +604,7 @@ function mergeAccountBalances( for (const customId of customAssetIds) { if (!Object.prototype.hasOwnProperty.call(next, customId)) { const prev = previousBalances[customId]; - next[customId] = prev ?? { amount: '0' }; + next[customId] = prev ?? getZeroAssetBalance(customId); } } @@ -637,6 +623,56 @@ function mergeAccountBalances( return next; } +/** + * Compute the effective account balances for an Accounts API v6 update. + * + * `merge` overlays the incoming balances on prior state (Account Activity, + * Snap events, staking). `full` replaces each covered chain slice. v6 asks + * for natives, pins, and default tracked assets through `includeAssetIds` + * and either returns them (at zero when the account holds none) or names + * them in `unprocessedIncludeAssetIds`. Assets the client tracks on purpose + * survive a `full` replace even when the response omits them, so an upstream + * gap cannot wipe them — see `shouldKeepAsset`. + * + * @param previousBalances - Balances already in state for this account. + * @param updatedBalances - Balances from the incoming response. + * @param undeletableAssetIds - Visible assets the client asked for (native, + * pin, or default tracked). They survive a `full` replace if the response + * omits them. Hidden assets are excluded from snapshots and dropped here; + * unhide is expected to fetch a new snapshot. + * @param updateMode - `'merge'` overlays incoming balances; `'full'` replaces + * each updated chain slice except undeletable and staking assets. + * @returns The effective balance map for the account. + */ +function effectiveAccountBalancesV6( + previousBalances: Record, + updatedBalances: Record, + undeletableAssetIds: Set, + updateMode: AssetsUpdateMode, +): Record { + if (updateMode === 'merge') { + return { ...previousBalances, ...updatedBalances }; + } + + const updatedChains = new Set( + Object.keys(updatedBalances).map((assetId) => + extractChainId(assetId as Caip19AssetId), + ), + ); + const shouldKeepAsset = (assetId: Caip19AssetId): boolean => + !updatedChains.has(extractChainId(assetId)) || + undeletableAssetIds.has(assetId.toLowerCase()) || + isStakingContractAssetId(assetId); + + const keptBalances: Record = {}; + for (const [assetId, balance] of Object.entries(previousBalances)) { + if (shouldKeepAsset(assetId as Caip19AssetId)) { + keptBalances[assetId] = balance; + } + } + return { ...keptBalances, ...updatedBalances }; +} + // ============================================================================ // CONTROLLER IMPLEMENTATION // ============================================================================ @@ -963,19 +999,27 @@ export class AssetsController extends BaseController< queryApiClient, onActiveChainsUpdated: this.#onActiveChainsUpdated, ...accountsApiDataSourceConfig, + isBalanceV6Enabled: (): boolean => this.#isBalanceV6Enabled(), + getAssetsState: (): AssetsControllerState => this.state, + getAssetVisibility: this.#getAssetVisibility.bind(this), }); this.#snapDataSource = new SnapDataSource({ messenger: this.messenger, onActiveChainsUpdated: this.#onActiveChainsUpdated, onAssetsUpdate: (response): Promise => this.handleAssetsUpdate(response, 'SnapDataSource'), + isBalanceV6Enabled: (): boolean => this.#isBalanceV6Enabled(), + getAssetVisibility: this.#getAssetVisibility.bind(this), + getAssetsState: (): AssetsControllerState => this.state, }); this.#rpcDataSource = new RpcDataSource({ messenger: this.messenger, + getAssetsState: (): AssetsControllerState => this.state, + isBalanceV6Enabled: (): boolean => this.#isBalanceV6Enabled(), + getAssetVisibility: this.#getAssetVisibility.bind(this), onActiveChainsUpdated: this.#onActiveChainsUpdated, - getNativeAssetForChain: (chainId: ChainId): Caip19AssetId => - this.#getNativeAssetMap()[chainId] ?? - `${chainId}/erc20:${ZERO_ADDRESS}`, + getNativeAssetForChain: (chainId: ChainId): Caip19AssetId | undefined => + this.#getNativeAssetForChain(chainId), // Share the API platform's TanStack Query client so the RPC token // detector caches/dedupes its top-token-list fetches alongside the rest // of the package's API calls. Caller-provided rpcConfig.queryClient @@ -998,13 +1042,17 @@ export class AssetsController extends BaseController< }, getAssetType: (assetId: Caip19AssetId): 'native' | 'erc20' | 'spl' => this.#getAssetType(assetId), + getAssetsState: (): AssetsControllerState => this.state, }); this.#priceDataSource = new PriceDataSource({ queryApiClient, getSelectedCurrency: (): SupportedCurrency => this.state.selectedCurrency, + getAssetsState: (): AssetsControllerState => this.state, ...priceDataSourceConfig, }); - this.#detectionMiddleware = new DetectionMiddleware(); + this.#detectionMiddleware = new DetectionMiddleware({ + getAssetsState: (): AssetsControllerState => this.state, + }); this.#customAssetGraduationMiddleware = new CustomAssetGraduationMiddleware( { getSelectedAccountId: (): AccountId | undefined => { @@ -1016,10 +1064,13 @@ export class AssetsController extends BaseController< }, removeCustomAsset: (accountId, assetId): void => this.removeCustomAsset(accountId, assetId), + getAssetsState: (): AssetsControllerState => this.state, }, ); this.#rpcFallbackMiddleware = new RpcFallbackMiddleware({ rpcDataSource: this.#rpcDataSource, + isBalanceV6Enabled: (): boolean => this.#isBalanceV6Enabled(), + getAssetsState: (): AssetsControllerState => this.state, }); log('Initializing AssetsController', { @@ -1453,7 +1504,7 @@ export class AssetsController extends BaseController< * Returns response and exclusive duration per source (sum ≈ wall time). * * Thin wrapper over {@link executeAssetsPipeline} that supplies the - * controller-owned state accessor and exception reporter. + * controller-owned exception reporter. * * @param params - Middleware execution options. * @param params.sources - Data sources or middlewares with getName() and assetsMiddleware. @@ -1476,7 +1527,6 @@ export class AssetsController extends BaseController< }> { return executeAssetsPipeline({ ...params, - getAssetsState: () => this.state, captureException: this.#captureException, }); } @@ -1500,8 +1550,6 @@ export class AssetsController extends BaseController< bypassServerCache?: boolean; dataTypes?: DataType[]; assetsForPriceUpdate?: Caip19AssetId[]; - /** When set to `'merge'`, fetch result is merged with existing state instead of replacing. Use for partial fetches (e.g. newly added chains). */ - updateMode?: AssetsUpdateMode; }, ): Promise>> { const chainIds = options?.chainIds ?? [...this.#enabledChains]; @@ -1512,166 +1560,370 @@ export class AssetsController extends BaseController< return this.#getAssetsFromState(accounts, chainIds, assetTypes); } - // Collect custom assets for all requested accounts - const customAssets: Caip19AssetId[] = []; - for (const account of accounts) { - const accountCustomAssets = this.getCustomAssets(account.id); - customAssets.push(...accountCustomAssets); - } - if (options?.forceUpdate) { - // Pipeline spans only on unlock/first-init fetch; later forceUpdates pass - // `undefined` so `emitTrace` / `withTrace` no-op without call-site if/else. const pipelineTrace = this.#firstInitFetchReported ? undefined : this.#trace; - - const request = this.#buildDataRequest(accounts, chainIds, { + const requestOptions = { assetTypes, dataTypes, - customAssets: customAssets.length > 0 ? customAssets : undefined, - forceUpdate: true, + forceUpdate: true as const, bypassServerCache: options?.bypassServerCache, assetsForPriceUpdate: options?.assetsForPriceUpdate, - }); + }; + + if (this.#isBalanceV6Enabled()) { + await this.#forceUpdateAssetsV6({ + accounts, + chainIds, + request: this.#buildDataRequest(accounts, chainIds, requestOptions), + pipelineTrace, + }); + } else { + await this.#forceUpdateAssetsV5({ + accounts, + chainIds, + request: this.#buildForceUpdateRequestV5( + accounts, + chainIds, + requestOptions, + ), + pipelineTrace, + }); + } + } - // Fast pipeline: accountsApi + stakedBalance → detection → token + price. - // Snap and RPC are excluded here due to their latency (snap triggers account - // creation, RPC is slow on many chains). Results are committed to state - // immediately so the UI can display balances without waiting for them. - // - // Fast/slow pipelines use merge so partial API snapshots cannot wipe - // tokens missing from the response (e.g. USDC when only native balance - // is returned). Balances present in the response are still refreshed. - const fastSources = buildFastFetchSources( - { - accountsApiDataSource: this.#accountsApiDataSource, - stakedBalanceDataSource: this.#stakedBalanceDataSource, - customAssetGraduationMiddleware: - this.#customAssetGraduationMiddleware, - rpcFallbackMiddleware: this.#rpcFallbackMiddleware, - detectionMiddleware: this.#detectionMiddleware, - tokenDataSource: this.#tokenDataSource, - priceDataSource: this.#priceDataSource, - }, - { isBasicFunctionality: this.#isBasicFunctionality() }, - ); + const result = this.#getAssetsFromState(accounts, chainIds, assetTypes); + return result; + } - const { response } = await withTrace({ - name: TRACE_FETCH_PIPELINE, - trace: pipelineTrace, - data: { - chain_count: chainIds.length, - account_count: accounts.length, - basic_functionality: this.#isBasicFunctionality(), - }, - fn: async (parentContext) => { - const startTime = performance.now(); - const result = await this.#executeMiddlewares({ - sources: fastSources, - request, - parentContext, - trace: pipelineTrace, - }); - await this.#updateState({ - ...result.response, - updateMode: 'merge', - replaceCoveredChainBalances: true, - }); + /** + * v5 force-update request: every pin of the requested accounts, unscoped. + * Delete with the rest of the v5 path when `assetsAccountsApiV6` is the default. + * + * @param accounts - Accounts in this fetch. + * @param chainIds - Chains in this fetch. + * @param requestOptions - Shared force-update request fields. + * @param requestOptions.assetTypes - Asset types to fetch. + * @param requestOptions.dataTypes - Data types to fetch. + * @param requestOptions.forceUpdate - Always `true` to bypass caches. + * @param requestOptions.bypassServerCache - Also bypass server-side HTTP caches. + * @param requestOptions.assetsForPriceUpdate - Assets to refresh prices for. + * @returns The v5 data request. + */ + #buildForceUpdateRequestV5( + accounts: InternalAccount[], + chainIds: ChainId[], + requestOptions: { + assetTypes: AssetType[]; + dataTypes: DataType[]; + forceUpdate: true; + bypassServerCache?: boolean; + assetsForPriceUpdate?: Caip19AssetId[]; + }, + ): DataRequest { + const customAssets: Caip19AssetId[] = []; + for (const account of accounts) { + customAssets.push(...this.getCustomAssets(account.id)); + } - const durationMs = performance.now() - startTime; - - // Summary fields for Assets Health (nested under the parent span). - emitTrace({ - name: TRACE_FULL_FETCH, - trace: pipelineTrace, - data: { - duration_ms: durationMs, - chain_count: chainIds.length, - account_count: accounts.length, - basic_functionality: this.#isBasicFunctionality(), - asset_count: result.response.assetsBalance - ? Object.values(result.response.assetsBalance).reduce( - (sum, acct) => sum + Object.keys(acct).length, - 0, - ) - : 0, - price_count: result.response.assetsPrice - ? Object.keys(result.response.assetsPrice).length - : 0, - ...result.durationByDataSource, - }, - parentContext, - }); + return this.#buildDataRequest(accounts, chainIds, { + ...requestOptions, + customAssets: customAssets.length > 0 ? customAssets : undefined, + }); + } - emitTrace({ - name: TRACE_FIRST_INIT_FETCH, - trace: pipelineTrace, - data: { - duration_ms: durationMs, - chain_ids: JSON.stringify(chainIds), - ...result.durationByDataSource, - }, - parentContext, - }); + async #forceUpdateAssetsV5({ + accounts, + chainIds, + request, + pipelineTrace, + }: { + accounts: InternalAccount[]; + chainIds: ChainId[]; + request: DataRequest; + pipelineTrace?: TraceCallback; + }): Promise { + const isBasicFunctionality = this.#isBasicFunctionality(); + const fastSources = buildFastFetchSources( + { + accountsApiDataSource: this.#accountsApiDataSource, + stakedBalanceDataSource: this.#stakedBalanceDataSource, + customAssetGraduationMiddleware: this.#customAssetGraduationMiddleware, + rpcFallbackMiddleware: this.#rpcFallbackMiddleware, + detectionMiddleware: this.#detectionMiddleware, + tokenDataSource: this.#tokenDataSource, + priceDataSource: this.#priceDataSource, + }, + { isBasicFunctionality, includeCustomAssetGraduation: true }, + ); - return result; - }, - }); + const fastResponse = await this.#runFastFetchV5({ + accounts, + chainIds, + request, + pipelineTrace, + fastSources, + isBasicFunctionality, + }); + this.#firstInitFetchReported = true; - // Mark after the unlock/first forceUpdate so later polls skip pipeline spans. - this.#firstInitFetchReported = true; + this.#runBackgroundFetch({ + accounts, + request, + pipelineTrace, + chainIds: this.#getSlowPipelineChainIds(chainIds, fastResponse), + isBasicFunctionality, + }); + } - // Background (slow) lane — flattened sibling of the fast lane (not nested - // inside its withTrace callback). Still fire-and-forget so UI is not blocked. - const slowPipelineChainIds = this.#getSlowPipelineChainIds( - chainIds, - response, - ); + async #forceUpdateAssetsV6({ + accounts, + chainIds, + request, + pipelineTrace, + }: { + accounts: InternalAccount[]; + chainIds: ChainId[]; + request: DataRequest; + pipelineTrace?: TraceCallback; + }): Promise { + const isBasicFunctionality = this.#isBasicFunctionality(); + const fastSources = buildFastFetchSources( + { + accountsApiDataSource: this.#accountsApiDataSource, + stakedBalanceDataSource: this.#stakedBalanceDataSource, + customAssetGraduationMiddleware: this.#customAssetGraduationMiddleware, + rpcFallbackMiddleware: this.#rpcFallbackMiddleware, + detectionMiddleware: this.#detectionMiddleware, + tokenDataSource: this.#tokenDataSource, + priceDataSource: this.#priceDataSource, + }, + { + isBasicFunctionality, + includeCustomAssetGraduation: false, + }, + ); - if (slowPipelineChainIds.length > 0) { - const slowSources = this.#isBasicFunctionality() - ? [this.#snapDataSource, this.#rpcDataSource] - : [this.#rpcDataSource]; + const fastResponse = await this.#runFastFetchV6({ + accounts, + chainIds, + request, + pipelineTrace, + fastSources, + isBasicFunctionality, + }); + this.#firstInitFetchReported = true; + + this.#runBackgroundFetch({ + accounts, + request, + pipelineTrace, + chainIds: this.#getSlowPipelineChainIds(chainIds, fastResponse), + isBasicFunctionality, + }); + } - const slowRequest = { ...request, chainIds: slowPipelineChainIds }; + async #runFastFetchV5({ + accounts, + chainIds, + request, + pipelineTrace, + fastSources, + isBasicFunctionality, + }: { + accounts: InternalAccount[]; + chainIds: ChainId[]; + request: DataRequest; + pipelineTrace?: TraceCallback; + fastSources: AssetsDataSource[]; + isBasicFunctionality: boolean; + }): Promise { + const { response } = await withTrace({ + name: TRACE_FETCH_PIPELINE, + trace: pipelineTrace, + data: { + chain_count: chainIds.length, + account_count: accounts.length, + basic_functionality: isBasicFunctionality, + }, + fn: async (parentContext) => { + const startTime = performance.now(); + const result = await this.#executeMiddlewares({ + sources: fastSources, + request, + parentContext, + trace: pipelineTrace, + }); + await this.#updateState({ + ...result.response, + updateMode: 'merge', + replaceCoveredChainBalances: true, + }); + this.#emitFastFetchTraces({ + accounts, + chainIds, + result, + durationMs: performance.now() - startTime, + isBasicFunctionality, + parentContext, + pipelineTrace, + }); + return result; + }, + }); + return response; + } - withTrace({ - name: TRACE_BACKGROUND_FETCH, + async #runFastFetchV6({ + accounts, + chainIds, + request, + pipelineTrace, + fastSources, + isBasicFunctionality, + }: { + accounts: InternalAccount[]; + chainIds: ChainId[]; + request: DataRequest; + pipelineTrace?: TraceCallback; + fastSources: AssetsDataSource[]; + isBasicFunctionality: boolean; + }): Promise { + const { response } = await withTrace({ + name: TRACE_FETCH_PIPELINE, + trace: pipelineTrace, + data: { + chain_count: chainIds.length, + account_count: accounts.length, + basic_functionality: isBasicFunctionality, + }, + fn: async (parentContext) => { + const startTime = performance.now(); + const result = await this.#executeMiddlewares({ + sources: fastSources, + request, + parentContext, trace: pipelineTrace, - data: { - chain_count: slowPipelineChainIds.length, - account_count: accounts.length, - }, - fn: async (slowParentContext) => { - const { response: slowResponse } = await this.#executeMiddlewares({ - sources: [ - createParallelBalanceMiddleware(slowSources), - this.#detectionMiddleware, - ...(this.#isBasicFunctionality() - ? [ - createParallelMiddleware([ - this.#tokenDataSource, - this.#priceDataSource, - ]), - ] - : []), - ], - request: slowRequest, - parentContext: slowParentContext, - trace: pipelineTrace, - }); - await this.#updateState({ - ...slowResponse, - updateMode: 'merge', - }); - }, - }).catch((error) => log('Background pipeline failed', { error })); - } + }); + await this.#updateState(result.response); + this.#emitFastFetchTraces({ + accounts, + chainIds, + result, + durationMs: performance.now() - startTime, + isBasicFunctionality, + parentContext, + pipelineTrace, + }); + return result; + }, + }); + return response; + } + + #runBackgroundFetch({ + accounts, + chainIds, + request, + pipelineTrace, + isBasicFunctionality, + }: { + accounts: InternalAccount[]; + chainIds: ChainId[]; + request: DataRequest; + pipelineTrace?: TraceCallback; + isBasicFunctionality: boolean; + }): void { + if (chainIds.length === 0) { + return; } + const slowSources = isBasicFunctionality + ? [this.#snapDataSource, this.#rpcDataSource] + : [this.#rpcDataSource]; + const slowRequest = { ...request, chainIds }; - const result = this.#getAssetsFromState(accounts, chainIds, assetTypes); - return result; + withTrace({ + name: TRACE_BACKGROUND_FETCH, + trace: pipelineTrace, + data: { + chain_count: chainIds.length, + account_count: accounts.length, + }, + fn: async (parentContext) => { + const { response } = await this.#executeMiddlewares({ + sources: [ + createParallelBalanceMiddleware(slowSources), + this.#detectionMiddleware, + ...(isBasicFunctionality + ? [ + createParallelMiddleware([ + this.#tokenDataSource, + this.#priceDataSource, + ]), + ] + : []), + ], + request: slowRequest, + parentContext, + trace: pipelineTrace, + }); + await this.#updateState(response); + }, + }).catch((error) => log('Background pipeline failed', { error })); + } + + #emitFastFetchTraces({ + accounts, + chainIds, + result, + durationMs, + isBasicFunctionality, + parentContext, + pipelineTrace, + }: { + accounts: InternalAccount[]; + chainIds: ChainId[]; + result: { + response: DataResponse; + durationByDataSource: Record; + }; + durationMs: number; + isBasicFunctionality: boolean; + parentContext?: TraceContext; + pipelineTrace?: TraceCallback; + }): void { + emitTrace({ + name: TRACE_FULL_FETCH, + trace: pipelineTrace, + data: { + duration_ms: durationMs, + chain_count: chainIds.length, + account_count: accounts.length, + basic_functionality: isBasicFunctionality, + asset_count: result.response.assetsBalance + ? Object.values(result.response.assetsBalance).reduce( + (sum, account) => sum + Object.keys(account).length, + 0, + ) + : 0, + price_count: result.response.assetsPrice + ? Object.keys(result.response.assetsPrice).length + : 0, + ...result.durationByDataSource, + }, + parentContext, + }); + emitTrace({ + name: TRACE_FIRST_INIT_FETCH, + trace: pipelineTrace, + data: { + duration_ms: durationMs, + chain_ids: JSON.stringify(chainIds), + ...result.durationByDataSource, + }, + parentContext, + }); } async getAssetsBalance( @@ -1953,7 +2205,8 @@ export class AssetsController extends BaseController< /** * Add a custom asset for an account. * Custom assets are included in subscription and fetch operations. - * Adding a custom asset also unhides it if it was previously hidden. + * Adding a custom asset also unhides it if it was previously hidden, + * and force-fetches that asset's chain (including every visible pin). * * When `pendingMetadata` is provided (e.g. from the extension's pending-tokens * flow), the token metadata is persisted immediately into `assetsInfo` so the @@ -2026,24 +2279,21 @@ export class AssetsController extends BaseController< Record >; balances[accountId] ??= {}; - balances[accountId][normalizedAssetId] ??= { amount: '0' }; + balances[accountId][normalizedAssetId] ??= + getZeroAssetBalance(normalizedAssetId); }); - // Fetch data for the newly added custom asset (merge to preserve other chains) const account = this.#getSelectedAccounts().find((a) => a.id === accountId); if (account) { - const chainId = extractChainId(normalizedAssetId); await this.getAssets([account], { - chainIds: [chainId], + chainIds: [extractChainId(normalizedAssetId)], dataTypes: ['balance', 'metadata', 'price'], assetTypes: ['fungible'], forceUpdate: true, - updateMode: 'merge', }); } - // Re-evaluate subscriptions so the supplemental RPC poll picks up the - // new customAsset on chains another data source already owns. + // Re-evaluate subscriptions so polls pick up the new pin. this.#subscribeAssets(); } @@ -2071,8 +2321,7 @@ export class AssetsController extends BaseController< } }); - // Re-evaluate subscriptions so the supplemental RPC poll for that chain - // is torn down when no more customAssets remain there. + // Re-evaluate subscriptions so polls drop the removed pin. this.#subscribeAssets(); } @@ -2108,14 +2357,20 @@ export class AssetsController extends BaseController< } state.assetPreferences[normalizedAssetId].hidden = true; }); + + // Re-evaluate subscriptions so polls exclude the newly hidden asset. + this.#subscribeAssets(); } /** * Unhide an asset globally. + * Force-fetches that asset's chain (including every visible pin) so a + * `full` snapshot can restore the balance immediately, then re-evaluates + * subscriptions so later polls stop excluding it. * * @param assetId - The CAIP-19 asset ID to unhide. */ - unhideAsset(assetId: Caip19AssetId): void { + async unhideAsset(assetId: Caip19AssetId): Promise { const normalizedAssetId = normalizeAssetId(assetId); log('Unhiding asset', { assetId: normalizedAssetId }); @@ -2129,6 +2384,41 @@ export class AssetsController extends BaseController< } } }); + + const chainId = extractChainId(normalizedAssetId); + const account = this.#getSelectedAccounts().find((selectedAccount) => + this.#getEnabledChainsForAccount(selectedAccount).includes(chainId), + ); + if (account) { + await this.getAssets([account], { + chainIds: [chainId], + dataTypes: ['balance', 'metadata', 'price'], + assetTypes: ['fungible'], + forceUpdate: true, + }); + } + + this.#subscribeAssets(); + } + + /** + * Whether Accounts API v6 (and the v6 custom-asset path) is enabled. + * Injected into AccountsApiDataSource and RpcFallbackMiddleware. + * + * `RemoteFeatureFlagController` resolves a threshold-scoped flag to the + * selected group's `value`, so processed state is a boolean. + * + * @returns `true` when the v6 remote flag is on. + */ + #isBalanceV6Enabled(): boolean { + try { + const { remoteFeatureFlags } = this.messenger.call( + 'RemoteFeatureFlagController:getState', + ); + return remoteFeatureFlags?.assetsAccountsApiV6 === true; + } catch { + return false; + } } // ============================================================================ @@ -2260,7 +2550,6 @@ export class AssetsController extends BaseController< isUpdate, onAssetsUpdate: (response) => this.handleAssetsUpdate(response, 'PriceDataSource'), - getAssetsState: () => this.state, }; this.#priceDataSource.subscribe(subscribeReq).catch(console.error); @@ -2312,6 +2601,79 @@ export class AssetsController extends BaseController< ); } + /** + * Canonical native CAIP-19 ID for a chain. Uses the native asset map when + * the chain is registered. An unregistered EVM chain falls back to the + * zero-address ERC-20 encoding used for EVM natives without a SLIP-44 id; + * an unregistered non-EVM chain resolves no native at all — fabricating an + * `erc20:` asset on a non-EVM namespace would produce an ID no data source + * ever writes. + * + * @param chainId - CAIP-2 chain ID. + * @returns The native asset ID for the chain, or `undefined` when the + * chain is unregistered and non-EVM. + */ + #getNativeAssetForChain(chainId: ChainId): Caip19AssetId | undefined { + const registered = this.#getNativeAssetMap()[chainId]; + if (registered) { + return registered; + } + if (parseCaipChainId(chainId).namespace === KnownCaipNamespace.Eip155) { + return `${chainId}/${UNKNOWN_EVM_NATIVE_ASSET_REFERENCE}` as Caip19AssetId; + } + return undefined; + } + + #getAssetVisibility( + accountIds: AccountId[], + chainIds: ChainId[], + ): AssetVisibility { + return getAssetVisibility({ + state: this.state, + accountIds, + chainIds, + getNativeAssetForChain: (chainId) => + this.#getNativeAssetForChain(chainId), + }); + } + + /** + * Lowercased asset IDs a v6 `full` snapshot must keep for an account, even + * when the response omits them: natives, pins, and default tracked assets + * the client asked for. Hidden assets are sent as `excludeAssetIds` and are + * dropped here; unhide is expected to fetch a new snapshot. Everything else + * on a covered chain was discovered by the source, so its absence means the + * balance is gone. + * + * @param state - State being updated (the in-flight draft, not `this.state`). + * @param accountId - Account whose pins apply. + * @param accountBalances - Incoming balances, which define the covered chains. + * @returns Lowercased CAIP-19 IDs that must survive the replace. + */ + #getUndeletableAssetIds( + state: AssetsControllerState, + accountId: AccountId, + accountBalances: Record, + ): Set { + const coveredChainIds = [ + ...new Set( + Object.keys(accountBalances).map( + (assetId) => assetId.split('/')[0] as ChainId, + ), + ), + ]; + const { visibleAssetIds } = getAssetVisibility({ + state, + accountIds: [accountId], + chainIds: coveredChainIds, + // Registered natives only: the unknown-native fallback is a fabricated + // ID, so it must never shield a row from a `full` replace. + getNativeAssetForChain: (chainId) => this.#getNativeAssetMap()[chainId], + }); + + return new Set(visibleAssetIds.map((assetId) => assetId.toLowerCase())); + } + /** * Checks whether the given CAIP-19 asset ID represents a native asset * according to the cached native asset map, the asset namespace, and the asset reference. @@ -2553,6 +2915,7 @@ export class AssetsController extends BaseController< async #updateState(response: DataResponse): Promise { const normalizedResponse = normalizeResponse(response); const mode: AssetsUpdateMode = normalizedResponse.updateMode ?? 'merge'; + const isBalanceV6Enabled = this.#isBalanceV6Enabled(); const releaseLock = await this.#controllerMutex.acquire(); @@ -2640,11 +3003,11 @@ export class AssetsController extends BaseController< } if (normalizedResponse.assetsBalance) { - for (const [accountId, accountBalances] of Object.entries( - normalizedResponse.assetsBalance, - )) { - const previousBalances = - previousState.assetsBalance[accountId] ?? {}; + const applyV5AccountBalanceUpdate = ( + accountId: string, + accountBalances: Record, + previousBalances: Record, + ): void => { const customAssetIds = (state.customAssets as Record)[ accountId @@ -2654,7 +3017,7 @@ export class AssetsController extends BaseController< mode === 'full' || normalizedResponse.replaceCoveredChainBalances === true; - const effective = mergeAccountBalances( + const effectiveAccountBalances = effectiveAccountBalancesV5( previousBalances, accountBalances, customAssetIds, @@ -2670,14 +3033,19 @@ export class AssetsController extends BaseController< : this.#getNativeAssetIdsForEnabledChains(); for (const nativeAssetId of nativeAssetIdsForAccount) { if ( - !Object.prototype.hasOwnProperty.call(effective, nativeAssetId) + !Object.prototype.hasOwnProperty.call( + effectiveAccountBalances, + nativeAssetId, + ) ) { - effective[nativeAssetId] = + effectiveAccountBalances[nativeAssetId] = getDefaultNativeAssetBalance(nativeAssetId); } } - for (const [assetId, balance] of Object.entries(effective)) { + for (const [assetId, balance] of Object.entries( + effectiveAccountBalances, + )) { const previousBalance = previousBalances[ assetId as Caip19AssetId ] as AssetBalance | undefined; @@ -2698,16 +3066,18 @@ export class AssetsController extends BaseController< const newMetadata = (balance as FungibleAssetBalance).metadata ?? (previousBalance as FungibleAssetBalance | undefined)?.metadata; - effective[assetId] = { + effectiveAccountBalances[assetId] = { amount: newAmount, ...(newMetadata === undefined ? {} : { metadata: newMetadata }), }; const oldAmount = previousBalance?.amount; - const isNewDefaultNativeZero = + const isNewSeededZero = oldAmount === undefined && newAmount === '0' && - nativeAssetIdsForAccount.includes(assetId as Caip19AssetId); - if (oldAmount !== newAmount && !isNewDefaultNativeZero) { + (nativeAssetIdsForAccount.includes(assetId as Caip19AssetId) || + getDefaultAssetMetadata(assetId as Caip19AssetId) !== + undefined); + if (oldAmount !== newAmount && !isNewSeededZero) { changedBalances.push({ accountId, assetId, @@ -2717,8 +3087,87 @@ export class AssetsController extends BaseController< } } - if (!isEqual(previousBalances, effective)) { - balances[accountId] = effective; + if (!isEqual(previousBalances, effectiveAccountBalances)) { + balances[accountId] = effectiveAccountBalances; + } + }; + + // Delete with the rest of the v6 path if `assetsAccountsApiV6` is + // rolled back; v5 keeps using `applyV5AccountBalanceUpdate`. + const applyV6AccountBalanceUpdate = ( + accountId: string, + accountBalances: Record, + previousBalances: Record, + ): void => { + const effectiveAccountBalances = effectiveAccountBalancesV6( + previousBalances, + accountBalances, + this.#getUndeletableAssetIds( + state as unknown as AssetsControllerState, + accountId, + accountBalances, + ), + mode, + ); + + for (const [assetId, balance] of Object.entries( + effectiveAccountBalances, + )) { + const previousBalance = previousBalances[ + assetId as Caip19AssetId + ] as AssetBalance | undefined; + const assetDecimals = ( + metadata[assetId] as { decimals?: number } | undefined + )?.decimals; + const newAmount = normalizeAmountString( + (balance as { amount: unknown }).amount, + assetDecimals, + ); + const newMetadata = + (balance as FungibleAssetBalance).metadata ?? + (previousBalance as FungibleAssetBalance | undefined)?.metadata; + effectiveAccountBalances[assetId] = { + amount: newAmount, + ...(newMetadata === undefined ? {} : { metadata: newMetadata }), + }; + const oldAmount = previousBalance?.amount; + // A v6 snapshot reports assets the account holds none of (its + // `includeAssetIds`) at zero. Publishing those as a change would + // emit `0 -> 0`, since a missing previous amount reads as '0'. + const isNewZero = oldAmount === undefined && newAmount === '0'; + if (oldAmount !== newAmount && !isNewZero) { + changedBalances.push({ + accountId, + assetId, + oldAmount, + newAmount, + }); + } + } + + if (!isEqual(previousBalances, effectiveAccountBalances)) { + balances[accountId] = effectiveAccountBalances; + } + }; + + for (const [accountId, accountBalances] of Object.entries( + normalizedResponse.assetsBalance, + )) { + const previousBalances = + previousState.assetsBalance[accountId] ?? {}; + + if (isBalanceV6Enabled) { + applyV6AccountBalanceUpdate( + accountId, + accountBalances, + previousBalances, + ); + } else { + applyV5AccountBalanceUpdate( + accountId, + accountBalances, + previousBalances, + ); } } } @@ -2757,12 +3206,10 @@ export class AssetsController extends BaseController< log('State updated', { changedBalances: changedBalances.length > 0 ? changedBalances : undefined, - changedMetadataCount: - changedMetadata.length > 0 ? changedMetadata.length : undefined, - changedPricesCount: - changedPriceAssets.length > 0 - ? changedPriceAssets.length - : undefined, + changedMetadata: + changedMetadata.length > 0 ? changedMetadata : undefined, + changedPrices: + changedPriceAssets.length > 0 ? changedPriceAssets : undefined, newAssets: Object.keys(detectedAssets).length > 0 ? Object.entries(detectedAssets).map(([accountId, assets]) => ({ @@ -3034,16 +3481,10 @@ export class AssetsController extends BaseController< if (!subscriptionKey.startsWith('ds:')) { continue; } - // Subscription keys take the form `ds:` for the regular - // subscription or `ds::` for supplemental - // subscriptions (e.g. `ds:RpcDataSource:custom`). Split on `:` and - // pick the source-name segment so both shapes resolve correctly. + // Subscription keys take the form `ds:`. const [, sourceId] = subscriptionKey.split(':'); const source = allSources.find((ds) => ds.getName() === sourceId); if (source) { - // Unsubscribe by the actual key — `#unsubscribeDataSource` only - // knows the regular `ds:` shape and would miss - // supplemental subscriptions, leaking their polling timers. this.#unsubscribeBySubscriptionKey(source, subscriptionKey); } } @@ -3093,10 +3534,12 @@ export class AssetsController extends BaseController< * Strategy to minimize data source calls: * 1. Collect all chains to subscribe based on enabled networks * 2. Map chains to accounts based on their scopes - * 3. Split by data source (ordered by priority) - each data source gets ONE subscription - * - * This ensures we make minimal subscriptions to each data source while covering - * all accounts and chains. + * 3. Split by data source (priority order) — each source gets one + * subscription for the chains it was assigned (accounts + chains only). + * v6 sources read pins/hides/defaults from state themselves (Accounts API + * sends `includeAssetIds` / `excludeAssetIds`; RPC polls native, balances, + * pins, and default tracked assets). v5 uses a separate RPC + * `customAssetsOnly` supplement for pins on chains another source owns. * * @param accounts - Accounts to subscribe balance updates for. * @param chainIds - Chain IDs to subscribe for. @@ -3113,13 +3556,9 @@ export class AssetsController extends BaseController< new Set(chainIds), ); const remainingChains = new Set(chainToAccounts.keys()); - // When basic functionality is on, use all balance data sources; when off, - // RPC only. - const isBasicFunctionality = this.#isBasicFunctionality(); - const balanceDataSources = isBasicFunctionality + const balanceDataSources = this.#isBasicFunctionality() ? this.#allBalanceDataSources : [this.#rpcDataSource]; - let rpcAssignedChains: Set = new Set(); for (const source of balanceDataSources) { @@ -3132,16 +3571,13 @@ export class AssetsController extends BaseController< remainingChains.delete(chainId); } } - if (assignedChains.length === 0) { this.#unsubscribeDataSource(source); continue; } - if (source === this.#rpcDataSource) { rpcAssignedChains = new Set(assignedChains); } - const seenIds = new Set(); const accountsForSource = assignedChains .flatMap((chainId) => chainToAccounts.get(chainId) ?? []) @@ -3162,27 +3598,20 @@ export class AssetsController extends BaseController< } } - // Supplemental RPC subscription for customAssets on chains another data - // source claimed during regular handoff. RPC is the sole balance fetcher - // for customAssets, so we must always poll them — even when (e.g.) - // AccountsApi is already covering the chain for normal balances. The - // supplemental subscription runs in `customAssetsOnly` mode so it does - // NOT double-poll the regular tracked balances. - this.#subscribeRpcCustomAssetsSupplement( - accounts, - chainToAccounts, - rpcAssignedChains, - ); + if (!this.#isBalanceV6Enabled()) { + this.#subscribeRpcCustomAssetsSupplement( + accounts, + chainToAccounts, + rpcAssignedChains, + ); + } } /** * Guarantee that customAssets are **always** polled by RPC, even when * AccountsApi or another data source has claimed the chain in the * regular handoff. RPC is the sole balance fetcher for user-imported - * tokens (see `pickRpcCustomAssetsSupplement` for the full rationale), - * so we run a dedicated subscription in `customAssetsOnly` mode under a - * distinct subscription key (`ds:RpcDataSource:custom`) that does not - * interfere with the regular RPC subscription. + * tokens on the Accounts API v5 path (see `pickRpcCustomAssetsSupplement`). * * @param accounts - Accounts to consider for customAssets. * @param chainToAccounts - Map of chain → accounts (built by caller). @@ -3307,7 +3736,8 @@ export class AssetsController extends BaseController< * @param chains - Array of chain IDs to subscribe for. * @param options - Optional subscription overrides. * @param options.subscriptionKey - Custom subscription key (default: `ds:`). - * @param options.customAssetsOnly - When true, only poll customAssets for these chains. + * @param options.customAssetsOnly - When true, only poll customAssets for these + * chains (Accounts API v5 supplemental RPC subscription). * @param options.skipInitialFetch - When true, skip the data source's subscribe-time fetch. */ #subscribeDataSource( @@ -3348,7 +3778,6 @@ export class AssetsController extends BaseController< isUpdate, onAssetsUpdate: (response, request) => this.handleAssetsUpdate(response, sourceId, request), - getAssetsState: () => this.state, ...(options.skipInitialFetch === true ? { skipInitialFetch: true } : {}), }; @@ -3577,13 +4006,13 @@ export class AssetsController extends BaseController< // Refresh subscriptions for new chain set this.#subscribeAssets(); - // Do one-time fetch for newly enabled chains; merge so we keep existing chain balances + // One-time fetch for newly enabled chains. v6 `full` replace is scoped to + // those chains; existing balances on other chains are left untouched. const accounts = this.#getSelectedAccounts(); if (addedChains.length > 0 && accounts.length > 0) { await this.getAssets(accounts, { chainIds: addedChains, forceUpdate: true, - updateMode: 'merge', }); } @@ -3770,57 +4199,35 @@ export class AssetsController extends BaseController< ): Promise { log('Assets updated from data source', { sourceId, - hasBalance: Boolean(response.assetsBalance), - hasPrice: Boolean(response.assetsPrice), + assetsBalance: response.assetsBalance, + assetsPrice: response.assetsPrice, }); - // Enrichment spans only before unlock/first-init fetch completes. + if (this.#isBalanceV6Enabled()) { + await this.#handleAssetsUpdateV6(response, sourceId, request); + return; + } + await this.#handleAssetsUpdateV5(response, sourceId, request); + } + + async #handleAssetsUpdateV5( + response: DataResponse, + sourceId: string, + request?: DataRequest, + ): Promise { const pipelineTrace = this.#firstInitFetchReported ? undefined : this.#trace; - await withTrace({ name: TRACE_UPDATE_PARENT, trace: pipelineTrace, - data: { - source: sourceId, - has_balance: Boolean(response.assetsBalance), - has_price: Boolean(response.assetsPrice), - balance_account_count: response.assetsBalance - ? Object.keys(response.assetsBalance).length - : 0, - }, + data: this.#getUpdateTraceData(response, sourceId), fn: async (parentContext) => { const updateStart = performance.now(); - - const resolvedRequest: DataRequest = request ?? { - accountsWithSupportedChains: [], - chainIds: [], - dataTypes: ['balance', 'metadata', 'price'], - }; - - // RPC-only mode (basic functionality off): never run token/price APIs. Strip - // those data types so downstream middleware cannot treat them as requested. - const pipelineRequest: DataRequest = this.#isBasicFunctionality() - ? resolvedRequest - : { - ...resolvedRequest, - dataTypes: resolvedRequest.dataTypes.filter( - (dt) => dt !== 'metadata' && dt !== 'price', - ), - }; - - // Graduate custom assets only when AccountsAPI / AccountActivity reports - // them. RPC already fetches custom assets on purpose, and Snap handles - // non-EVM chains the rule does not apply to, so skip the middleware for - // those. + const pipelineRequest = this.#getUpdatePipelineRequest(request); const shouldGraduateCustomAssets = sourceId === 'AccountsApiDataSource' || sourceId === 'AccountActivityDataSource'; - - // Websocket updates can carry brand-new spam airdrops: enrich them - // with Token API occurrences and drop below-floor tokens BEFORE - // detection, so spam is never detected, enriched, priced or persisted. const shouldFilterOccurrences = sourceId === 'AccountActivityDataSource' && this.#isBasicFunctionality(); @@ -3863,27 +4270,141 @@ export class AssetsController extends BaseController< ...enrichedResponse, replaceCoveredChainBalances: response.replaceCoveredChainBalances, }); + this.#emitUpdateTrace({ + response, + sourceId, + enrichedResponse, + updateStart, + pipelineTrace, + parentContext, + }); + }, + }); + } + + async #handleAssetsUpdateV6( + response: DataResponse, + sourceId: string, + request?: DataRequest, + ): Promise { + const pipelineTrace = this.#firstInitFetchReported + ? undefined + : this.#trace; + await withTrace({ + name: TRACE_UPDATE_PARENT, + trace: pipelineTrace, + data: this.#getUpdateTraceData(response, sourceId), + fn: async (parentContext) => { + const updateStart = performance.now(); + const pipelineRequest = this.#getUpdatePipelineRequest(request); + const shouldFilterOccurrences = + sourceId === 'AccountActivityDataSource' && + this.#isBasicFunctionality(); + const enrichmentSources: AssetsDataSource[] = [ + ...(shouldFilterOccurrences + ? [ + { + getName: () => 'OccurrenceFloorFilter', + assetsMiddleware: + this.#tokenDataSource.occurrenceFilterMiddleware, + }, + ] + : []), + ]; + if (this.#isBasicFunctionality()) { + enrichmentSources.push(this.#rpcFallbackMiddleware); + } + enrichmentSources.push(this.#detectionMiddleware); + if (this.#isBasicFunctionality()) { + enrichmentSources.push( + createParallelMiddleware([ + this.#tokenDataSource, + this.#priceDataSource, + ]), + ); + } - // Summary fields for Assets Health (nested under the parent span). - emitTrace({ - name: TRACE_UPDATE_PIPELINE, + const { response: enrichedResponse } = await this.#executeMiddlewares({ + sources: enrichmentSources, + request: pipelineRequest, + initialResponse: response, + parentContext, trace: pipelineTrace, - data: { - source: sourceId, - duration_ms: performance.now() - updateStart, - has_balance: Boolean(response.assetsBalance), - has_price: Boolean(response.assetsPrice), - has_metadata: Boolean(enrichedResponse.assetsInfo), - balance_account_count: response.assetsBalance - ? Object.keys(response.assetsBalance).length - : 0, - }, + }); + + await this.#updateState(enrichedResponse); + this.#emitUpdateTrace({ + response, + sourceId, + enrichedResponse, + updateStart, + pipelineTrace, parentContext, }); }, }); } + #getUpdatePipelineRequest(request?: DataRequest): DataRequest { + const resolvedRequest: DataRequest = request ?? { + accountsWithSupportedChains: [], + chainIds: [], + dataTypes: ['balance', 'metadata', 'price'], + }; + if (this.#isBasicFunctionality()) { + return resolvedRequest; + } + return { + ...resolvedRequest, + dataTypes: resolvedRequest.dataTypes.filter( + (dataType) => dataType !== 'metadata' && dataType !== 'price', + ), + }; + } + + #getUpdateTraceData(response: DataResponse, sourceId: string): TraceSpanData { + return { + source: sourceId, + has_balance: Boolean(response.assetsBalance), + has_price: Boolean(response.assetsPrice), + balance_account_count: response.assetsBalance + ? Object.keys(response.assetsBalance).length + : 0, + }; + } + + #emitUpdateTrace({ + response, + sourceId, + enrichedResponse, + updateStart, + pipelineTrace, + parentContext, + }: { + response: DataResponse; + sourceId: string; + enrichedResponse: DataResponse; + updateStart: number; + pipelineTrace?: TraceCallback; + parentContext?: TraceContext; + }): void { + emitTrace({ + name: TRACE_UPDATE_PIPELINE, + trace: pipelineTrace, + data: { + source: sourceId, + duration_ms: performance.now() - updateStart, + has_balance: Boolean(response.assetsBalance), + has_price: Boolean(response.assetsPrice), + has_metadata: Boolean(enrichedResponse.assetsInfo), + balance_account_count: response.assetsBalance + ? Object.keys(response.assetsBalance).length + : 0, + }, + parentContext, + }); + } + // ============================================================================ // CLEANUP // ============================================================================ diff --git a/packages/assets-controller/src/__fixtures__/bsc-spam-token/api-responses/index.ts b/packages/assets-controller/src/__fixtures__/bsc-spam-token/api-responses/index.ts index 66476c232ec..90b4dd7319c 100644 --- a/packages/assets-controller/src/__fixtures__/bsc-spam-token/api-responses/index.ts +++ b/packages/assets-controller/src/__fixtures__/bsc-spam-token/api-responses/index.ts @@ -62,13 +62,16 @@ function mockAccountsSupportedNetworks(): nock.Scope { * Intercept `GET {ACCOUNTS}/v5/multiaccount/balances` with the wallet's * captured 38 holdings. * + * @param omitAssetIds - Asset IDs (any casing) to drop from the captured + * balances, simulating a token the Accounts API does not index. * @returns The nock scope and the account IDs each request asked about. */ -function mockV5MultiAccountBalances(): { +function mockV5MultiAccountBalances(omitAssetIds: string[] = []): { scope: nock.Scope; requestedAccountIds: string[][]; } { const requestedAccountIds: string[][] = []; + const omitted = new Set(omitAssetIds.map((assetId) => assetId.toLowerCase())); const scope = nock(API_URLS.ACCOUNTS) .persist() @@ -78,12 +81,39 @@ function mockV5MultiAccountBalances(): { requestedAccountIds.push( readListParam(uri, 'accountIds', API_URLS.ACCOUNTS), ); - return v5MultiAccountBalances; + const balances = v5MultiAccountBalances.balances.filter( + (entry) => !omitted.has(entry.assetId.toLowerCase()), + ); + return { ...v5MultiAccountBalances, count: balances.length, balances }; }); return { scope, requestedAccountIds }; } +/** + * Intercept `GET {ACCOUNTS}/v6/multiaccount/balances` with the wallet's + * captured token balances. + * + * @param omitAssetIds - Asset IDs (any casing) to drop from the captured + * balances. + * @returns The nock scope. + */ +function mockV6MultiAccountBalances(omitAssetIds: string[] = []): nock.Scope { + const omitted = new Set(omitAssetIds.map((assetId) => assetId.toLowerCase())); + + return nock(API_URLS.ACCOUNTS) + .persist() + .get('/v6/multiaccount/balances') + .query(true) + .reply(200, { + balances: v5MultiAccountBalances.balances.filter( + (entry) => !omitted.has(entry.assetId.toLowerCase()), + ), + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); +} + /** * Intercept `GET {TOKENS}/v2/supportedNetworks`. `eip155:56` is in the captured * `fullSupport` list, so BNB Chain assets genuinely reach the occurrence filter @@ -190,9 +220,16 @@ function mockV3SpotPrices(): BatchRecordingMock { * All interceptors persist, so batch composition and cache misses cannot make a * test fail for want of an interceptor. * + * @param options - Options for shaping the intercepted responses. + * @param options.omitBalanceAssetIds - Asset IDs (any casing) to drop from + * the captured balances, simulating tokens the Accounts API does not index. * @returns The recording mocks, and other utils */ -export function mockBscSpamApis(): { +export function mockBscSpamApis({ + omitBalanceAssetIds = [], +}: { + omitBalanceAssetIds?: string[]; +} = {}): { accountsSupportedNetworks: nock.Scope; balances: { requestedAccountIds: string[][] }; assets: BatchRecordingMock; @@ -204,7 +241,8 @@ export function mockBscSpamApis(): { mockPricesSupportedNetworks(); mockChainIdNetwork(); - const balances = mockV5MultiAccountBalances(); + const balances = mockV5MultiAccountBalances(omitBalanceAssetIds); + mockV6MultiAccountBalances(omitBalanceAssetIds); const assets = mockV3Assets(); const prices = mockV3SpotPrices(); diff --git a/packages/assets-controller/src/__fixtures__/bsc-spam-token/bscSpamWallet.ts b/packages/assets-controller/src/__fixtures__/bsc-spam-token/bscSpamWallet.ts index 25979bb0477..96db9780d31 100644 --- a/packages/assets-controller/src/__fixtures__/bsc-spam-token/bscSpamWallet.ts +++ b/packages/assets-controller/src/__fixtures__/bsc-spam-token/bscSpamWallet.ts @@ -1,6 +1,6 @@ import type { InternalAccount } from '@metamask/keyring-internal-api'; -import type { AssetsControllerStateInternal } from '../../types.js'; +import type { AssetsControllerState } from '../../types.js'; import { BSC_CHAIN_ID, BSC_SPAM_ACCOUNT_ID, @@ -44,8 +44,8 @@ export function buildBscSpamAccount( * @returns The starting state. */ export function buildEmptyAssetsState( - overrides?: Partial, -): AssetsControllerStateInternal { + overrides?: Partial, +): AssetsControllerState { return { assetsInfo: {}, assetsBalance: {}, diff --git a/packages/assets-controller/src/__fixtures__/scam-token-cleanup/scamWalletState.ts b/packages/assets-controller/src/__fixtures__/scam-token-cleanup/scamWalletState.ts index 142cc738453..be1d4bf1114 100644 --- a/packages/assets-controller/src/__fixtures__/scam-token-cleanup/scamWalletState.ts +++ b/packages/assets-controller/src/__fixtures__/scam-token-cleanup/scamWalletState.ts @@ -1,6 +1,6 @@ import type { AssetMetadata, - AssetsControllerStateInternal, + AssetsControllerState, Caip19AssetId, } from '../../types.js'; @@ -1620,8 +1620,8 @@ export const SCAM_WALLET_SURVIVING_ASSET_IDS = ( * @returns Full internal controller state. */ export function buildScamWalletState( - overrides: Partial = {}, -): AssetsControllerStateInternal { + overrides: Partial = {}, +): AssetsControllerState { const assetsInfo = SCAM_WALLET_ASSETS_INFO as Record< Caip19AssetId, AssetMetadata diff --git a/packages/assets-controller/src/__fixtures__/spamWalletState.ts b/packages/assets-controller/src/__fixtures__/spamWalletState.ts index e11af79bcfe..62c057d0125 100644 --- a/packages/assets-controller/src/__fixtures__/spamWalletState.ts +++ b/packages/assets-controller/src/__fixtures__/spamWalletState.ts @@ -4,7 +4,7 @@ import { getDefaultTrackedAssetsForChain } from '../defaults.js'; import type { AssetMetadata, AssetPrice, - AssetsControllerStateInternal, + AssetsControllerState, Caip19AssetId, ChainId, } from '../types.js'; @@ -286,8 +286,8 @@ export const SPAM_WALLET_BALANCES = { * @returns Full internal controller state. */ export function buildSpamWalletState( - overrides: Partial = {}, -): AssetsControllerStateInternal { + overrides: Partial = {}, +): AssetsControllerState { return { assetsInfo: { ...SPAM_WALLET_ASSETS_INFO }, assetsBalance: structuredClone(SPAM_WALLET_BALANCES), @@ -308,7 +308,7 @@ export function buildSpamWalletState( * * @returns Full internal controller state. */ -export function buildLowercasedSpamWalletState(): AssetsControllerStateInternal { +export function buildLowercasedSpamWalletState(): AssetsControllerState { const state = buildSpamWalletState(); return { @@ -352,7 +352,7 @@ function lowercaseKeys( */ export function buildManyTokensState(count: number): { assetIds: Caip19AssetId[]; - state: AssetsControllerStateInternal; + state: AssetsControllerState; } { const assetIds = Array.from({ length: count }, (_, index) => { const address = getChecksumAddress( diff --git a/packages/assets-controller/src/data-sources/AbstractDataSource.ts b/packages/assets-controller/src/data-sources/AbstractDataSource.ts index 6023b18dfb2..b64503552d9 100644 --- a/packages/assets-controller/src/data-sources/AbstractDataSource.ts +++ b/packages/assets-controller/src/data-sources/AbstractDataSource.ts @@ -1,9 +1,4 @@ -import type { - ChainId, - DataRequest, - DataResponse, - AssetsControllerStateInternal, -} from '../types.js'; +import type { ChainId, DataRequest, DataResponse } from '../types.js'; // ============================================================================ // DATA SOURCE BASE TYPES @@ -21,11 +16,6 @@ export type SubscriptionRequest = { response: DataResponse, request?: DataRequest, ) => void | Promise; - /** - * Optional state access (e.g. for price/token data sources that need assetsBalance). - * Provided by the controller when subscribing. - */ - getAssetsState?: () => AssetsControllerStateInternal; /** * When true, skip the one-shot fetch that normally runs when a subscription * is created. Used after the controller has already force-fetched balances diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts index da95936479e..e7d0cde3470 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts @@ -4,13 +4,18 @@ import type { InternalAccount } from '@metamask/keyring-internal-api'; import { Messenger, MOCK_ANY_NAMESPACE } from '@metamask/messenger'; import type { MockAnyNamespace } from '@metamask/messenger'; +import { getDefaultTrackedAssetsForChain } from '../defaults.js'; import type { ChainId, Caip19AssetId, DataRequest, Context, - AssetsControllerStateInternal, + AssetsControllerState, } from '../types.js'; +import type { AssetVisibility } from '../utils/assetVisibility.js'; +import { getAssetVisibility } from '../utils/assetVisibility.js'; +import { ZERO_ADDRESS } from '../utils/constants.js'; +import { NATIVE_ASSETS } from '../utils/native-assets.js'; import { SNAPS_ASSETS_MIGRATION_FLAG_KEYS, SnapsAssetsMigrationStage, @@ -33,6 +38,19 @@ const CHAIN_MAINNET = 'eip155:1' as ChainId; const CHAIN_POLYGON = 'eip155:137' as ChainId; const CHAIN_ARBITRUM = 'eip155:42161' as ChainId; const MOCK_ADDRESS = '0x1234567890123456789012345678901234567890'; +const [MAINNET_MUSD] = getDefaultTrackedAssetsForChain(CHAIN_MAINNET); +const MAINNET_NATIVE = 'eip155:1/slip44:60' as Caip19AssetId; +const POLYGON_NATIVE = 'eip155:137/slip44:966' as Caip19AssetId; + +function getNativeAssetForChain(chainId: ChainId): Caip19AssetId { + return NATIVE_ASSETS[chainId] ?? `${chainId}/erc20:${ZERO_ADDRESS}`; +} + +function isBalanceV6EnabledFromFlags( + remoteFeatureFlags: Record, +): boolean { + return remoteFeatureFlags.assetsAccountsApiV6 === true; +} type MockApiClient = { accounts: { @@ -67,6 +85,7 @@ function createMockApiClient( balances: V5BalanceItem[] = [], unprocessedNetworks: string[] = [], v6Balances: V6BalanceItem[] = [], + unprocessedIncludeAssetIds: string[] = [], partialSupport: (number | string)[] = [], ): MockApiClient { return { @@ -82,7 +101,7 @@ function createMockApiClient( fetchV6MultiAccountBalances: jest.fn().mockResolvedValue({ balances: v6Balances, unprocessedNetworks, - unprocessedIncludeAssetIds: [], + unprocessedIncludeAssetIds, }), }, }; @@ -137,11 +156,24 @@ function createDataRequest( }; } +function createAssetsState( + overrides: Partial = {}, +): AssetsControllerState { + return { + assetsInfo: {}, + assetsBalance: {}, + assetsPrice: {}, + customAssets: {}, + assetPreferences: {}, + selectedCurrency: 'usd', + ...overrides, + }; +} + function createMiddlewareContext(overrides?: Partial): Context { return { request: createDataRequest(), response: {}, - getAssetsState: jest.fn(), ...overrides, }; } @@ -160,9 +192,12 @@ async function setupController( partialSupport?: (number | string)[]; balances?: V5BalanceItem[]; unprocessedNetworks?: string[]; + unprocessedIncludeAssetIds?: string[]; fetchTimeoutMs?: number; v6Balances?: V6BalanceItem[]; remoteFeatureFlags?: Record; + getAssetsState?: () => AssetsControllerState; + tokenDetectionEnabled?: () => boolean; } = {}, ): Promise { const { @@ -170,9 +205,12 @@ async function setupController( partialSupport = [], balances = [], unprocessedNetworks = [], + unprocessedIncludeAssetIds = [], fetchTimeoutMs, v6Balances = [], remoteFeatureFlags = {}, + getAssetsState = (): AssetsControllerState => createAssetsState(), + tokenDetectionEnabled, } = options; const rootMessenger = new Messenger({ @@ -213,6 +251,7 @@ async function setupController( balances, unprocessedNetworks, v6Balances, + unprocessedIncludeAssetIds, partialSupport, ); @@ -223,7 +262,18 @@ async function setupController( apiClient as unknown as AccountsApiDataSourceOptions['queryApiClient'], onActiveChainsUpdated: (dataSourceName, chains, previousChains): void => activeChainsUpdateHandler(dataSourceName, chains, previousChains), + isBalanceV6Enabled: (): boolean => + isBalanceV6EnabledFromFlags(remoteFeatureFlags), + getAssetsState, + getAssetVisibility: (accountIds, chainIds): AssetVisibility => + getAssetVisibility({ + state: getAssetsState(), + accountIds, + chainIds, + getNativeAssetForChain, + }), ...(fetchTimeoutMs === undefined ? {} : { fetchTimeoutMs }), + ...(tokenDetectionEnabled === undefined ? {} : { tokenDetectionEnabled }), }); // Wait for async initialization @@ -794,7 +844,7 @@ describe('AccountsApiDataSource', () => { it('uses the v6 endpoint when the assetsAccountsApiV6 remote flag is enabled', async () => { const { controller, apiClient } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, }); await controller.fetch(createDataRequest()); @@ -809,26 +859,21 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('uses the v5 endpoint when the assetsAccountsApiV6 remote flag is disabled', async () => { - const { controller, apiClient } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: false } }, + it('sets updateMode to full for v6 fetches', async () => { + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, }); - await controller.fetch(createDataRequest()); + const response = await controller.fetch(createDataRequest()); - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(1); - expect( - apiClient.accounts.fetchV6MultiAccountBalances, - ).not.toHaveBeenCalled(); + expect(response.updateMode).toBe('full'); controller.destroy(); }); - it('uses the v5 endpoint when the flag is a plain boolean instead of the JSON value shape', async () => { + it('uses the v5 endpoint when the assetsAccountsApiV6 remote flag is disabled', async () => { const { controller, apiClient } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: true }, + remoteFeatureFlags: { assetsAccountsApiV6: false }, }); await controller.fetch(createDataRequest()); @@ -843,9 +888,9 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('calls the v6 endpoint without extra params when enabled', async () => { + it('calls the v6 endpoint with default tracked assets as includeAssetIds when enabled', async () => { const { controller, apiClient } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, }); await controller.fetch(createDataRequest()); @@ -854,7 +899,7 @@ describe('AccountsApiDataSource', () => { apiClient.accounts.fetchV6MultiAccountBalances, ).toHaveBeenCalledWith( [`eip155:1:${MOCK_ADDRESS}`], - undefined, + { includeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD] }, undefined, ); expect( @@ -866,7 +911,7 @@ describe('AccountsApiDataSource', () => { it('reads the flag per fetch so it can revert to v5 at runtime', async () => { const remoteFeatureFlags = { - assetsAccountsApiV6: { value: true }, + assetsAccountsApiV6: true, }; const { controller, apiClient } = await setupController({ remoteFeatureFlags, @@ -877,7 +922,7 @@ describe('AccountsApiDataSource', () => { apiClient.accounts.fetchV6MultiAccountBalances, ).toHaveBeenCalledTimes(1); - remoteFeatureFlags.assetsAccountsApiV6 = { value: false }; + remoteFeatureFlags.assetsAccountsApiV6 = false; await controller.fetch(createDataRequest()); expect( apiClient.accounts.fetchV5MultiAccountBalances, @@ -889,7 +934,7 @@ describe('AccountsApiDataSource', () => { it('processes v6 token balances grouped by account', async () => { const accountId = `eip155:1:${MOCK_ADDRESS}`; const { controller } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, v6Balances: [ createMockV6BalanceItem( accountId, @@ -933,7 +978,7 @@ describe('AccountsApiDataSource', () => { const { controller } = await setupController({ supportedChains: [1, STELLAR_CHAIN_ID], remoteFeatureFlags: { - assetsAccountsApiV6: { value: true }, + assetsAccountsApiV6: true, [SNAPS_ASSETS_MIGRATION_FLAG_KEYS.stellar]: { stage: SnapsAssetsMigrationStage.ReadAssetsControllerWithFallback, }, @@ -994,7 +1039,7 @@ describe('AccountsApiDataSource', () => { it('ignores v6 defi positions', async () => { const accountId = `eip155:1:${MOCK_ADDRESS}`; const { controller } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, v6Balances: [ createMockV6BalanceItem( accountId, @@ -1026,7 +1071,7 @@ describe('AccountsApiDataSource', () => { const stakingAssetId = 'eip155:1/erc20:0x4fef9d741011476750a243ac70b9789a63dd47df'; const { controller } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, v6Balances: [ createMockV6BalanceItem( accountId, @@ -1054,24 +1099,45 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('marks v6 unprocessed networks as errors', async () => { + it('omits v6 balances from unprocessed networks', async () => { const { controller } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, unprocessedNetworks: ['eip155:1'], + v6Balances: [ + createMockV6BalanceItem( + `eip155:1:${MOCK_ADDRESS}`, + MAINNET_NATIVE, + '1000000000000000000', + 'token', + 'native', + ), + createMockV6BalanceItem( + `eip155:137:${MOCK_ADDRESS}`, + 'eip155:137/slip44:966', + '2000000000000000000', + 'token', + 'native', + ), + ], }); - const response = await controller.fetch(createDataRequest()); + const response = await controller.fetch( + createDataRequest({ chainIds: [CHAIN_MAINNET, CHAIN_POLYGON] }), + ); expect(response.errors?.[CHAIN_MAINNET]).toBe( 'Unprocessed by Accounts API', ); + expect(response.assetsBalance?.['mock-account-id']).toStrictEqual({ + 'eip155:137/slip44:966': { amount: '2000000000000000000' }, + }); controller.destroy(); }); it('handles v6 API errors', async () => { const { controller, apiClient } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, }); apiClient.accounts.fetchV6MultiAccountBalances.mockRejectedValueOnce( @@ -1085,28 +1151,530 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('does not pass includeAssetIds to v6 even when custom assets are present', async () => { + it('passes EVM custom assets on requested chains to v6 as includeAssetIds', async () => { + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; const { controller, apiClient } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [customToken] }, + }), }); + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { + includeAssetIds: expect.arrayContaining([customToken, MAINNET_MUSD]), + }, + undefined, + ); + + controller.destroy(); + }); + + it('omits custom assets that are not on a requested chain from includeAssetIds', async () => { + const polygonToken = + 'eip155:137/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [polygonToken] }, + }), + }); + + await controller.fetch(createDataRequest({ chainIds: [CHAIN_MAINNET] })); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD] }, + undefined, + ); + + controller.destroy(); + }); + + it('errors a chain that came back without every requested include asset id', async () => { + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + unprocessedIncludeAssetIds: [customToken], + v6Balances: [ + createMockV6BalanceItem( + `eip155:1:${MOCK_ADDRESS}`, + MAINNET_NATIVE, + '1000000000000000000', + 'token', + 'native', + ), + ], + }); + + const response = await controller.fetch(createDataRequest()); + + // A snapshot missing one of its `includeAssetIds` is not a full + // snapshot, so the chain fails and its partial balances are dropped. + expect(response.errors?.[CHAIN_MAINNET]).toBe( + 'Unresolved includeAssetIds', + ); + expect(response.assetsBalance).toBeUndefined(); + + controller.destroy(); + }); + + it('errors only the chain of the unresolved include asset id', async () => { const customToken = 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + unprocessedIncludeAssetIds: [customToken], + v6Balances: [ + createMockV6BalanceItem( + `eip155:137:${MOCK_ADDRESS}`, + POLYGON_NATIVE, + '5000000000000000000', + 'token', + 'native', + ), + ], + }); + + const response = await controller.fetch( + createDataRequest({ chainIds: [CHAIN_MAINNET, CHAIN_POLYGON] }), + ); + + expect(response.errors?.[CHAIN_POLYGON]).toBeUndefined(); + expect( + response.assetsBalance?.['mock-account-id']?.[POLYGON_NATIVE], + ).toStrictEqual({ amount: '5000000000000000000' }); + + controller.destroy(); + }); + + it('omits pins not on a requested chain from includeAssetIds', async () => { + const solanaToken = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v' as Caip19AssetId; + const malformed = 'not-a-caip-asset' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [solanaToken, malformed] }, + }), + }); + + await controller.fetch(createDataRequest()); + + // No custom asset on a requested chain — still include defaults. + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD] }, + undefined, + ); + + controller.destroy(); + }); + + it('ignores malformed unprocessed include asset ids', async () => { + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + unprocessedIncludeAssetIds: ['not-a-caip-asset'], + }); + + const response = await controller.fetch(createDataRequest()); + + // The malformed unprocessed id cannot be parsed, so it is dropped from + // both axes (no error, no asset-axis entry). + expect(response.errors).toBeUndefined(); + expect(response.unprocessedCustomAssets).toBeUndefined(); + + controller.destroy(); + }); + + it('passes EVM hidden assets on requested chains to v6 as excludeAssetIds', async () => { + const hiddenToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + assetPreferences: { [hiddenToken]: { hidden: true } }, + }), + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { + includeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD], + excludeAssetIds: [hiddenToken], + }, + undefined, + ); + + controller.destroy(); + }); + + it('omits hidden assets that are not on a requested chain from excludeAssetIds', async () => { + const polygonToken = + 'eip155:137/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + assetPreferences: { [polygonToken]: { hidden: true } }, + }), + }); + + await controller.fetch(createDataRequest({ chainIds: [CHAIN_MAINNET] })); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD] }, + undefined, + ); + + controller.destroy(); + }); + + it('omits hidden assets not on a requested chain from excludeAssetIds', async () => { + const solanaToken = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v' as Caip19AssetId; + const malformed = 'not-a-caip-asset' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + assetPreferences: { + [solanaToken]: { hidden: true }, + [malformed]: { hidden: true }, + }, + }), + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD] }, + undefined, + ); + + controller.destroy(); + }); + + it('passes Solana pins on requested Solana chains to v6 as includeAssetIds', async () => { + const solanaChain = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId; + const solanaNative = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501' as Caip19AssetId; + const solanaToken = + `${solanaChain}/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v` as Caip19AssetId; + const { controller, apiClient } = await setupController({ + supportedChains: [1, solanaChain], + remoteFeatureFlags: { + assetsAccountsApiV6: true, + [SNAPS_ASSETS_MIGRATION_FLAG_KEYS.solana]: { + stage: SnapsAssetsMigrationStage.ReadAssetsControllerWithFallback, + }, + }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [solanaToken] }, + }), + }); + + await controller.fetch(createDataRequest({ chainIds: [solanaChain] })); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`${solanaChain}:${MOCK_ADDRESS}`], + { includeAssetIds: [solanaNative, solanaToken] }, + undefined, + ); + + controller.destroy(); + }); + + it('passes hidden Solana assets on requested Solana chains to v6 as excludeAssetIds', async () => { + const solanaChain = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId; + const solanaNative = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501' as Caip19AssetId; + const solanaToken = + `${solanaChain}/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v` as Caip19AssetId; + const { controller, apiClient } = await setupController({ + supportedChains: [1, solanaChain], + remoteFeatureFlags: { + assetsAccountsApiV6: true, + [SNAPS_ASSETS_MIGRATION_FLAG_KEYS.solana]: { + stage: SnapsAssetsMigrationStage.ReadAssetsControllerWithFallback, + }, + }, + getAssetsState: () => + createAssetsState({ + assetPreferences: { [solanaToken]: { hidden: true } }, + }), + }); + + await controller.fetch(createDataRequest({ chainIds: [solanaChain] })); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`${solanaChain}:${MOCK_ADDRESS}`], + { + includeAssetIds: [solanaNative], + excludeAssetIds: [solanaToken], + }, + undefined, + ); + + controller.destroy(); + }); + + it('reads natives, visible pins, and default tracked assets as includeAssetIds, not detected ERC-20 balances', async () => { + const pinnedToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const trackedToken = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + const native = 'eip155:1/slip44:60' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [pinnedToken] }, + assetsBalance: { + 'mock-account-id': { + [native]: { amount: '1' }, + [trackedToken]: { amount: '2' }, + }, + }, + }), + }); + + await controller.fetch(createDataRequest()); + + const [, params] = + apiClient.accounts.fetchV6MultiAccountBalances.mock.calls[0]; + expect(params?.includeAssetIds ?? []).toStrictEqual( + expect.arrayContaining([pinnedToken, native, MAINNET_MUSD]), + ); + expect(params?.includeAssetIds ?? []).not.toContain(trackedToken); + + controller.destroy(); + }); + + it('includes map-encoded erc20 natives in includeAssetIds, not synthetic slip44 ids', async () => { + const gnosisNative = + 'eip155:100/erc20:0x0000000000000000000000000000000000000000' as Caip19AssetId; + const syntheticSlip44 = 'eip155:1/slip44:bandwidth' as Caip19AssetId; + const gnosisChain = 'eip155:100' as ChainId; + const { controller, apiClient } = await setupController({ + supportedChains: [1, 100], + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + assetsBalance: { + 'mock-account-id': { + [gnosisNative]: { amount: '1' }, + [syntheticSlip44]: { amount: '1' }, + }, + }, + }), + }); + + await controller.fetch( + createDataRequest({ chainIds: [CHAIN_MAINNET, gnosisChain] }), + ); + + const [, params] = + apiClient.accounts.fetchV6MultiAccountBalances.mock.calls[0]; + expect(params?.includeAssetIds ?? []).toContain(gnosisNative); + expect(params?.includeAssetIds ?? []).not.toContain(syntheticSlip44); + + controller.destroy(); + }); + + it('ignores request.customAssets when building v6 includeAssetIds', async () => { + const statePin = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const requestOnlyToken = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [statePin] }, + }), + }); + await controller.fetch( - createDataRequest({ customAssets: [customToken] }), + createDataRequest({ customAssets: [requestOnlyToken] }), + ); + + const [, params] = + apiClient.accounts.fetchV6MultiAccountBalances.mock.calls[0]; + expect(params?.includeAssetIds).toStrictEqual( + expect.arrayContaining([statePin, MAINNET_MUSD]), + ); + expect(params?.includeAssetIds).not.toContain(requestOnlyToken); + + controller.destroy(); + }); + + it('does not include a state pin that is also hidden', async () => { + const token = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [token] }, + assetPreferences: { [token]: { hidden: true } }, + }), + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { + includeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD], + excludeAssetIds: [token], + }, + undefined, + ); + + controller.destroy(); + }); + + it('excludes hidden native and default tracked assets from includeAssetIds', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + assetPreferences: { + [MAINNET_NATIVE]: { hidden: true }, + [MAINNET_MUSD]: { hidden: true }, + }, + }), + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { + excludeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD], + }, + undefined, ); + controller.destroy(); + }); + + it('sends both includeAssetIds and excludeAssetIds when pins and hidden assets differ', async () => { + const pinned = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const hidden = + 'eip155:1/erc20:0xdAC17F958D2ee523a2206206994597C13D831ec7' as Caip19AssetId; + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [pinned] }, + assetPreferences: { [hidden]: { hidden: true } }, + }), + }); + + await controller.fetch(createDataRequest()); + expect( apiClient.accounts.fetchV6MultiAccountBalances, ).toHaveBeenCalledWith( [`eip155:1:${MOCK_ADDRESS}`], + { + includeAssetIds: expect.arrayContaining([pinned, MAINNET_MUSD]), + excludeAssetIds: [hidden], + }, undefined, + ); + + controller.destroy(); + }); + + it('includes default tracked assets as includeAssetIds even with no pins or native in state', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [MAINNET_NATIVE, MAINNET_MUSD] }, undefined, ); controller.destroy(); }); + + it('sends the same includeAssetIds on subscribe polls as on fetch', async () => { + const pinnedToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const { controller, apiClient, assetsUpdateHandler } = + await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + getAssetsState: () => + createAssetsState({ + customAssets: { 'mock-account-id': [pinnedToken] }, + }), + }); + + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest(), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { + includeAssetIds: expect.arrayContaining([pinnedToken, MAINNET_MUSD]), + }, + { staleTime: 0, gcTime: 0 }, + ); + + controller.destroy(); + }); }); it('fetch marks every requested chain as errored when the call exceeds the configured timeout', async () => { @@ -1185,6 +1753,21 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); + it('middleware forwards full updateMode from v6 fetches', async () => { + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }); + + const next = jest.fn().mockResolvedValue(undefined); + const context = createMiddlewareContext(); + + await controller.assetsMiddleware(context, next); + + expect(context.response.updateMode).toBe('full'); + + controller.destroy(); + }); + it('middleware skips Accounts API when balance is not requested', async () => { const { controller, apiClient } = await setupController({ balances: [ @@ -1366,6 +1949,8 @@ describe('AccountsApiDataSource', () => { balances?: V5BalanceItem[]; unprocessedNetworks?: string[]; tokenDetectionEnabled?: boolean; + getAssetsState?: () => AssetsControllerState; + remoteFeatureFlags?: Record; } = {}, ): Promise { const { @@ -1373,6 +1958,8 @@ describe('AccountsApiDataSource', () => { balances = [], unprocessedNetworks = [], tokenDetectionEnabled, + getAssetsState = (): AssetsControllerState => createAssetsState(), + remoteFeatureFlags = {}, } = options; const rootMessenger = new Messenger< @@ -1398,7 +1985,7 @@ describe('AccountsApiDataSource', () => { registerActionHandler: (a: string, h: () => unknown) => void; } ).registerActionHandler('RemoteFeatureFlagController:getState', () => ({ - remoteFeatureFlags: {}, + remoteFeatureFlags, cacheTimestamp: 0, })); @@ -1424,6 +2011,16 @@ describe('AccountsApiDataSource', () => { apiClient as unknown as AccountsApiDataSourceOptions['queryApiClient'], onActiveChainsUpdated: (dataSourceName, chains, previousChains): void => activeChainsUpdateHandler(dataSourceName, chains, previousChains), + isBalanceV6Enabled: (): boolean => + isBalanceV6EnabledFromFlags(remoteFeatureFlags), + getAssetsState, + getAssetVisibility: (accountIds, chainIds): AssetVisibility => + getAssetVisibility({ + state: getAssetsState(), + accountIds, + chainIds, + getNativeAssetForChain, + }), }; if (tokenDetectionEnabled !== undefined) { @@ -1473,17 +2070,6 @@ describe('AccountsApiDataSource', () => { request: createDataRequest(), isUpdate: false, onAssetsUpdate: assetsUpdateHandler, - getAssetsState: () => ({ - assetsInfo: {}, - assetsBalance: { - [ACCOUNT_ID]: { - [KNOWN_ASSET]: { amount: '500' }, - }, - }, - assetsPrice: {}, - customAssets: {}, - assetPreferences: {}, - }), }); expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); @@ -1512,6 +2098,14 @@ describe('AccountsApiDataSource', () => { '2000', ), ], + getAssetsState: () => + createAssetsState({ + assetsBalance: { + [ACCOUNT_ID]: { + [KNOWN_ASSET]: { amount: '500' }, + }, + }, + }), }); await controller.subscribe({ @@ -1519,17 +2113,6 @@ describe('AccountsApiDataSource', () => { request: createDataRequest(), isUpdate: false, onAssetsUpdate: assetsUpdateHandler, - getAssetsState: () => ({ - assetsInfo: {}, - assetsBalance: { - [ACCOUNT_ID]: { - [KNOWN_ASSET]: { amount: '500' }, - }, - }, - assetsPrice: {}, - customAssets: {}, - assetPreferences: {}, - }), }); expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); @@ -1554,6 +2137,7 @@ describe('AccountsApiDataSource', () => { '2000', ), ], + getAssetsState: () => createAssetsState(), }); await controller.subscribe({ @@ -1561,13 +2145,6 @@ describe('AccountsApiDataSource', () => { request: createDataRequest(), isUpdate: false, onAssetsUpdate: assetsUpdateHandler, - getAssetsState: () => ({ - assetsInfo: {}, - assetsBalance: {}, - assetsPrice: {}, - customAssets: {}, - assetPreferences: {}, - }), }); expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); @@ -1593,25 +2170,14 @@ describe('AccountsApiDataSource', () => { '2000', ), ], - }); - - // Set up state accessor via subscribe (middleware uses the stored getAssetsState) - await controller.subscribe({ - subscriptionId: 'sub-setup', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - getAssetsState: () => ({ - assetsInfo: {}, - assetsBalance: { - [ACCOUNT_ID]: { - [KNOWN_ASSET]: { amount: '500' }, + getAssetsState: () => + createAssetsState({ + assetsBalance: { + [ACCOUNT_ID]: { + [KNOWN_ASSET]: { amount: '500' }, + }, }, - }, - assetsPrice: {}, - customAssets: {}, - assetPreferences: {}, - }), + }), }); const middleware = controller.assetsMiddleware; @@ -1642,40 +2208,93 @@ describe('AccountsApiDataSource', () => { '1000000000000000000', ), ], + getAssetsState: () => createAssetsState(), }); - // Simulate new account: subscribe with empty state so #getAssetsState is set. - // Filter will then remove all API balance data (no assets in state). - await controller.subscribe({ - subscriptionId: 'sub-setup', + const nextFn = jest.fn().mockResolvedValue(undefined); + const context = createMiddlewareContext({ request: createDataRequest({ chainIds: [CHAIN_MAINNET, CHAIN_POLYGON], }), - isUpdate: false, - onAssetsUpdate: jest.fn(), - getAssetsState: () => ({ - assetsInfo: {}, - assetsBalance: {}, - assetsPrice: {}, - customAssets: {}, - assetPreferences: {}, + }); + + await controller.assetsMiddleware(context, nextFn); + + // All chains must still be passed to next middleware so RPC can fetch native balances + expect(nextFn).toHaveBeenCalledWith( + expect.objectContaining({ + request: expect.objectContaining({ + chainIds: [CHAIN_MAINNET, CHAIN_POLYGON], + }), }), + ); + + controller.destroy(); + }); + + it('does not filter unknown tokens on the v6 path when tokenDetectionEnabled is false', async () => { + const known = createMockV6BalanceItem( + `eip155:1:${MOCK_ADDRESS}`, + KNOWN_ASSET, + '1000', + ); + const unknown = createMockV6BalanceItem( + `eip155:1:${MOCK_ADDRESS}`, + UNKNOWN_ASSET, + '2000', + ); + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + tokenDetectionEnabled: (): boolean => false, + v6Balances: [known, unknown], + getAssetsState: () => + createAssetsState({ + assetsBalance: { + [ACCOUNT_ID]: { + [KNOWN_ASSET]: { amount: '500' }, + }, + }, + }), + }); + + const response = await controller.fetch(createDataRequest()); + const accountBalances = response.assetsBalance?.[ACCOUNT_ID] ?? {}; + expect(Object.keys(accountBalances)).toHaveLength(2); + expect(accountBalances[KNOWN_ASSET as Caip19AssetId]).toBeDefined(); + expect(accountBalances[UNKNOWN_ASSET as Caip19AssetId]).toBeDefined(); + + controller.destroy(); + }); + + it('claims handled chains on v6 when tokenDetectionEnabled is false', async () => { + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + tokenDetectionEnabled: (): boolean => false, + v6Balances: [ + createMockV6BalanceItem( + `eip155:1:${MOCK_ADDRESS}`, + 'eip155:1/slip44:60', + '1000000000000000000', + 'token', + 'native', + ), + ], + getAssetsState: () => createAssetsState(), }); const nextFn = jest.fn().mockResolvedValue(undefined); const context = createMiddlewareContext({ request: createDataRequest({ - chainIds: [CHAIN_MAINNET, CHAIN_POLYGON], + chainIds: [CHAIN_MAINNET], }), }); await controller.assetsMiddleware(context, nextFn); - // All chains must still be passed to next middleware so RPC can fetch native balances expect(nextFn).toHaveBeenCalledWith( expect.objectContaining({ request: expect.objectContaining({ - chainIds: [CHAIN_MAINNET, CHAIN_POLYGON], + chainIds: [], }), }), ); @@ -1698,14 +2317,8 @@ describe('filterResponseToKnownAssets', () => { function buildState( balances: Record>, - ): AssetsControllerStateInternal { - return { - assetsInfo: {}, - assetsBalance: balances, - assetsPrice: {}, - customAssets: {}, - assetPreferences: {}, - }; + ): AssetsControllerState { + return createAssetsState({ assetsBalance: balances }); } it('returns response unchanged when assetsBalance is undefined', () => { diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts index f3ca6acee06..4a9e0f67c94 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts @@ -8,6 +8,7 @@ import type { Json } from '@metamask/utils'; import { isCaipChainId, KnownCaipNamespace, + parseCaipAssetType, toCaipChainId, } from '@metamask/utils'; @@ -20,8 +21,10 @@ import type { DataRequest, DataResponse, Middleware, - AssetsControllerStateInternal, + AssetsControllerState, } from '../types.js'; +import type { GetAssetVisibility } from '../utils/assetVisibility.js'; +import { filterFailedChainBalances } from '../utils/filterFailedChainBalances.js'; import { fetchWithTimeout, normalizeAssetId } from '../utils/index.js'; import { getMigrationStages, @@ -50,9 +53,9 @@ const log = createModuleLogger(projectLogger, CONTROLLER_NAME); // Allowed actions that AccountsApiDataSource can call. Balances are fetched via // ApiPlatformClient directly (no BackendApiClient actions needed); the messenger -// is used to read the Accounts API v6 balances feature flag and to subscribe to -// `RemoteFeatureFlagController:stateChange` (see constructor) so migration flag -// changes refresh the active chains. +// is used to subscribe to `RemoteFeatureFlagController:stateChange` (see +// constructor) so migration flag changes refresh the active chains, and to +// read those flags when listing active chains. export type AccountsApiDataSourceAllowedActions = RemoteFeatureFlagControllerGetStateAction; @@ -96,7 +99,7 @@ export type AccountsApiDataSourceConfig = { export type AccountsApiDataSourceOptions = AccountsApiDataSourceConfig & { /** * The AssetsController messenger (shared by all data sources). Used to read - * the `assetsAccountsApiV6` remote feature flag. + * Snaps → AssetsController migration flags and subscribe to flag changes. */ messenger: AssetsControllerMessenger; /** ApiPlatformClient for API calls with caching */ @@ -107,6 +110,19 @@ export type AccountsApiDataSourceOptions = AccountsApiDataSourceConfig & { chains: ChainId[], previousChains: ChainId[], ) => void; + /** + * Whether to use the Accounts API v6 balances endpoint (and `includeAssetIds`). + * Injected by AssetsController so the flag is read in one place. Defaults to + * `false` (v5). Read on demand, not cached. + */ + isBalanceV6Enabled?: () => boolean; + /** + * Current AssetsController state. Used for v6 include/exclude asset IDs and + * for filtering when token detection is off. + */ + getAssetsState: () => AssetsControllerState; + /** Returns shared visible/hidden assets for an account/chain scope. */ + getAssetVisibility: GetAssetVisibility; state?: Partial; }; @@ -182,7 +198,7 @@ function caipChainIdToChainId(chainIdStr: string): ChainId { */ export function filterResponseToKnownAssets( response: DataResponse, - assetsState: AssetsControllerStateInternal, + assetsState: AssetsControllerState, ): DataResponse { if (!response.assetsBalance) { return response; @@ -252,15 +268,19 @@ export class AccountsApiDataSource extends AbstractDataSource< /** Shared AssetsController messenger, used to read remote feature flags. */ readonly #messenger: AssetsControllerMessenger; + /** Injected by AssetsController; `true` when the v6 balances endpoint should be used. */ + readonly #isBalanceV6Enabled: () => boolean; + + readonly #getAssetsState: () => AssetsControllerState; + + readonly #getAssetVisibility: GetAssetVisibility; + /** ApiPlatformClient for cached API calls */ readonly #apiClient: ApiPlatformClient; /** Chains refresh timer */ #chainsRefreshTimer: ReturnType | null = null; - /** State accessor from subscriptions (for filtering when tokenDetectionEnabled is false) */ - #getAssetsState?: () => AssetsControllerStateInternal; - constructor(options: AccountsApiDataSourceOptions) { super(CONTROLLER_NAME, { ...defaultState, @@ -273,6 +293,10 @@ export class AccountsApiDataSource extends AbstractDataSource< this.#tokenDetectionEnabled = options.tokenDetectionEnabled ?? ((): boolean => true); this.#messenger = options.messenger; + this.#isBalanceV6Enabled = + options.isBalanceV6Enabled ?? ((): boolean => false); + this.#getAssetsState = options.getAssetsState; + this.#getAssetVisibility = options.getAssetVisibility; this.#apiClient = options.queryApiClient; // The Snaps → AssetsController migration flags gate which migration networks @@ -310,35 +334,6 @@ export class AccountsApiDataSource extends AbstractDataSource< } } - /** - * Whether the Accounts API v6 balances endpoint is enabled, read from the - * RemoteFeatureFlagController (`assetsAccountsApiV6`). Read on demand (per - * fetch) rather than cached so the value can be toggled at runtime, and behaves - * the same across clients (extension, mobile). Defaults to `false` when the - * flag is unset or the controller is unavailable. - * - * The flag is a LaunchDarkly JSON variation shaped `{ value: boolean }` (same - * shape as `backendWebSocketConnection`), so the nested `value` is read rather - * than treating the flag as a plain boolean. - * - * @returns `true` when the v6 balances endpoint should be used. - */ - #isBalanceV6Enabled(): boolean { - try { - const { remoteFeatureFlags } = this.#messenger.call( - 'RemoteFeatureFlagController:getState', - ); - const flag = remoteFeatureFlags?.assetsAccountsApiV6; - return ( - typeof flag === 'object' && - flag !== null && - Boolean((flag as { value?: unknown }).value) - ); - } catch { - return false; - } - } - // ============================================================================ // INITIALIZATION // ============================================================================ @@ -422,57 +417,44 @@ export class AccountsApiDataSource extends AbstractDataSource< // ============================================================================ async fetch(request: DataRequest): Promise { - let response: DataResponse = {}; + if (this.#isBalanceV6Enabled()) { + return this.#fetchV6(request); + } - // Filter to only chains supported by Accounts API - const supportedChains = new Set(this.state.activeChains); - const chainsToFetch = request.chainIds.filter((chainId) => - supportedChains.has(chainId), - ); + return this.#fetchV5(request); + } + + /** + * v5 fetch. Unchanged from the pre-v6 handler; delete the v6 sibling first + * if `assetsAccountsApiV6` is rolled back. + * + * @param request - The data request. + * @returns Balances stamped `merge`, plus errors for chains the API could + * not process. + */ + async #fetchV5(request: DataRequest): Promise { + let response: DataResponse = {}; + const chainsToFetch = this.#getChainsToFetch(request); if (chainsToFetch.length === 0) { - // Mark unsupported chains as errors so they pass to next middleware - for (const chainId of request.chainIds) { - if (!supportedChains.has(chainId)) { - response.errors = response.errors ?? {}; - response.errors[chainId] = 'Chain not supported by Accounts API'; - } - } + this.#markUnsupportedChains(response, request); return response; } try { - // Build CAIP-10 account IDs (e.g., "eip155:1:0x1234...") - // Use pre-computed supportedChains per account from the request - const accountIds = request.accountsWithSupportedChains.flatMap( - ({ account, supportedChains: accountChains }) => - chainsToFetch - .filter((chainId) => accountChains.includes(chainId)) - .map((chainId) => `${chainId}:${account.address}`), - ); + const accountIds = this.#buildAccountIds(request, chainsToFetch); // Skip API call if no valid account-chain combinations if (accountIds.length === 0) { return response; } - const fetchOptions = - request.forceUpdate || request.bypassServerCache - ? { - staleTime: 0, - gcTime: 0, - // Also defeats the API's server-side cache (via a random - // bypassServerCache query param) so a post-transaction refresh cannot - // be answered with a pre-transaction snapshot. - ...(request.bypassServerCache ? { bypassServerCache: true } : {}), - } - : undefined; - - // Feature-flagged: v6 endpoint with a fallback to legacy v5. The flag is - // read here (not cached) so a runtime toggle can revert v6 -> v5. - const { unprocessedNetworks, assetsBalance } = this.#isBalanceV6Enabled() - ? await this.#fetchV6Balances(accountIds, fetchOptions, request) - : await this.#fetchV5Balances(accountIds, fetchOptions, request); + const { unprocessedNetworks, assetsBalance } = + await this.#fetchV5Balances( + accountIds, + this.#buildFetchOptions(request), + request, + ); // Handle unprocessed networks - these will be passed to next middleware if (unprocessedNetworks.length > 0) { @@ -489,30 +471,191 @@ export class AccountsApiDataSource extends AbstractDataSource< response.assetsBalance = assetsBalance; response.updateMode = 'merge'; } catch (error) { - log('Fetch FAILED', { error, chains: chainsToFetch }); + this.#markFetchFailure(response, chainsToFetch, error); + } + + this.#markUnsupportedChains(response, request); + + // When token detection is disabled, filter out tokens not already in state + if (!this.#tokenDetectionEnabled()) { + response = filterResponseToKnownAssets(response, this.#getAssetsState()); + } + + return response; + } + + /** + * v6 fetch. Reads visibility from state, sends it as `includeAssetIds` / + * `excludeAssetIds`, and stamps `full`. A chain the API left unprocessed, or + * answered without every requested `includeAssetId`, is reported in `errors` + * and contributes no balances, so the RPC fallback can recover it. + * + * @param request - The data request. + * @returns An authoritative snapshot for the chains that succeeded. + */ + async #fetchV6(request: DataRequest): Promise { + const response: DataResponse = {}; + const chainsToFetch = this.#getChainsToFetch(request); + + if (chainsToFetch.length === 0) { + this.#markUnsupportedChains(response, request); + return response; + } + + try { + const accountIds = this.#buildAccountIds(request, chainsToFetch); - // On error, mark all chains as errors so they can be handled by next middleware - response.errors = response.errors ?? {}; - for (const chainId of chainsToFetch) { - response.errors[chainId] = - `Fetch failed: ${error instanceof Error ? error.message : String(error)}`; + if (accountIds.length === 0) { + return response; } + + const { visibleAssetIds, hiddenAssetIds } = this.#getAssetVisibility( + request.accountsWithSupportedChains.map(({ account }) => account.id), + chainsToFetch, + ); + + const { unprocessedNetworks, unprocessedIncludeAssetIds, assetsBalance } = + await this.#fetchV6Balances( + accountIds, + this.#buildFetchOptions(request), + request, + visibleAssetIds.length > 0 ? visibleAssetIds : undefined, + hiddenAssetIds.length > 0 ? hiddenAssetIds : undefined, + ); + response.updateMode = 'full'; + + const unprocessedChainIds = unprocessedNetworks.map(caipChainIdToChainId); + // Chains answered without every requested `includeAssetIds`. The + // snapshot is incomplete, so the chain counts as failed rather than + // silently dropping the assets it left out. + const incompleteChainIds = this.#getChainIdsForAssetIds( + unprocessedIncludeAssetIds, + ); + const failedChainIds = new Set([ + ...unprocessedChainIds, + ...incompleteChainIds, + ]); + + // Errors hand these chains to the next middleware (the RPC fallback). + if (failedChainIds.size > 0) { + response.errors = response.errors ?? {}; + for (const chainId of unprocessedChainIds) { + response.errors[chainId] = 'Unprocessed by Accounts API'; + } + for (const chainId of incompleteChainIds) { + response.errors[chainId] ??= 'Unresolved includeAssetIds'; + } + } + + // A failed chain contributes nothing, so its balances stay as they are. + response.assetsBalance = filterFailedChainBalances( + assetsBalance, + failedChainIds, + ); + } catch (error) { + this.#markFetchFailure(response, chainsToFetch, error); } - // Mark unsupported chains as errors so they pass to next middleware + this.#markUnsupportedChains(response, request); + + return response; + } + + /** + * Requested chains the Accounts API currently supports. + * + * @param request - The data request being fetched. + * @returns The subset of `request.chainIds` this source can answer. + */ + #getChainsToFetch(request: DataRequest): ChainId[] { + const supportedChains = new Set(this.state.activeChains); + return request.chainIds.filter((chainId) => supportedChains.has(chainId)); + } + + /** + * Mark unsupported chains as errors so they pass to next middleware. + * + * @param response - Response being built; mutated in place. + * @param request - The data request being fetched. + */ + #markUnsupportedChains(response: DataResponse, request: DataRequest): void { + const supportedChains = new Set(this.state.activeChains); for (const chainId of request.chainIds) { if (!supportedChains.has(chainId)) { response.errors = response.errors ?? {}; response.errors[chainId] = 'Chain not supported by Accounts API'; } } + } - // When token detection is disabled, filter out tokens not already in state - if (!this.#tokenDetectionEnabled() && this.#getAssetsState) { - response = filterResponseToKnownAssets(response, this.#getAssetsState()); + /** + * Build CAIP-10 account IDs (e.g., "eip155:1:0x1234...") from the + * pre-computed supportedChains per account on the request. + * + * @param request - The data request being fetched. + * @param chainsToFetch - Chains this source will ask the API for. + * @returns One CAIP-10 ID per account-chain combination. + */ + #buildAccountIds(request: DataRequest, chainsToFetch: ChainId[]): string[] { + return request.accountsWithSupportedChains.flatMap( + ({ account, supportedChains: accountChains }) => + chainsToFetch + .filter((chainId) => accountChains.includes(chainId)) + .map((chainId) => `${chainId}:${account.address}`), + ); + } + + #buildFetchOptions( + request: DataRequest, + ): + | { staleTime: number; gcTime: number; bypassServerCache?: boolean } + | undefined { + if (!request.forceUpdate && !request.bypassServerCache) { + return undefined; } - return response; + return { + staleTime: 0, + gcTime: 0, + // Also defeats the API's server-side cache (via a random + // bypassServerCache query param) so a post-transaction refresh cannot + // be answered with a pre-transaction snapshot. + ...(request.bypassServerCache ? { bypassServerCache: true } : {}), + }; + } + + /** + * On error, mark all chains as errors so the next middleware handles them. + * + * @param response - Response being built; mutated in place. + * @param chainsToFetch - Chains the failed request covered. + * @param error - The thrown error. + */ + #markFetchFailure( + response: DataResponse, + chainsToFetch: ChainId[], + error: unknown, + ): void { + log('Fetch FAILED', { error, chains: chainsToFetch }); + + response.errors = response.errors ?? {}; + for (const chainId of chainsToFetch) { + response.errors[chainId] = + `Fetch failed: ${error instanceof Error ? error.message : String(error)}`; + } + } + + #getChainIdsForAssetIds(assetIds: string[]): Set { + const chainIds = new Set(); + for (const assetId of assetIds) { + try { + chainIds.add(parseCaipAssetType(assetId as Caip19AssetId).chainId); + } catch { + // An unparseable ID cannot be attributed to a chain, so it cannot + // invalidate one either. + } + } + return chainIds; } /** @@ -531,6 +674,7 @@ export class AccountsApiDataSource extends AbstractDataSource< request: DataRequest, ): Promise<{ unprocessedNetworks: string[]; + unprocessedIncludeAssetIds: string[]; assetsBalance: Record>; }> { const apiResponse = await fetchWithTimeout( @@ -550,6 +694,8 @@ export class AccountsApiDataSource extends AbstractDataSource< return { unprocessedNetworks: apiResponse.unprocessedNetworks, + // v5 has no `includeAssetIds` support. + unprocessedIncludeAssetIds: [], assetsBalance, }; } @@ -560,7 +706,10 @@ export class AccountsApiDataSource extends AbstractDataSource< * @param accountIds - CAIP-10 account IDs to fetch balances for. * @param fetchOptions - Cache/fetch options (e.g. force update settings). * @param request - The original data request containing accounts to map. - * @returns Unprocessed networks and processed asset balances by account. + * @param includeAssetIds - Pinned asset IDs the backend must always return. + * @param excludeAssetIds - Hidden asset IDs the backend must drop. + * @returns Unprocessed networks, unprocessed pinned assets, and processed + * asset balances by account. */ async #fetchV6Balances( accountIds: string[], @@ -568,15 +717,26 @@ export class AccountsApiDataSource extends AbstractDataSource< | { staleTime: number; gcTime: number; bypassServerCache?: boolean } | undefined, request: DataRequest, + includeAssetIds: Caip19AssetId[] | undefined, + excludeAssetIds: Caip19AssetId[] | undefined, ): Promise<{ unprocessedNetworks: string[]; + unprocessedIncludeAssetIds: string[]; assetsBalance: Record>; }> { + const params = + includeAssetIds || excludeAssetIds + ? { + ...(includeAssetIds && { includeAssetIds }), + ...(excludeAssetIds && { excludeAssetIds }), + } + : undefined; + const apiResponse = await fetchWithTimeout( () => this.#apiClient.accounts.fetchV6MultiAccountBalances( accountIds, - undefined, + params, fetchOptions, ), this.#fetchTimeoutMs, @@ -589,6 +749,7 @@ export class AccountsApiDataSource extends AbstractDataSource< return { unprocessedNetworks: apiResponse.unprocessedNetworks, + unprocessedIncludeAssetIds: apiResponse.unprocessedIncludeAssetIds, assetsBalance, }; } @@ -759,77 +920,145 @@ export class AccountsApiDataSource extends AbstractDataSource< * @returns The middleware function for the assets pipeline. */ get assetsMiddleware(): Middleware { - return async (context, next) => { - const { request } = context; + if (this.#isBalanceV6Enabled()) { + return this.#assetsMiddlewareV6; + } + + return this.#assetsMiddlewareV5; + } + + readonly #assetsMiddlewareV5: Middleware = async (context, next) => { + const { request } = context; + + // Price/metadata-only requests must not hit the Accounts API. + if (!request.dataTypes.includes('balance')) { + return next(context); + } + + // If no chains requested, skip to next middleware + if (request.chainIds.length === 0) { + return next(context); + } + + let successfullyHandledChains: ChainId[] = []; - // Price/metadata-only requests must not hit the Accounts API. - if (!request.dataTypes.includes('balance')) { - return next(context); + try { + const response = await this.fetch(request); + + // Merge response into context + if (response.assetsBalance) { + context.response.assetsBalance ??= {}; + for (const [accountId, accountBalances] of Object.entries( + response.assetsBalance, + )) { + context.response.assetsBalance[accountId] = { + ...context.response.assetsBalance[accountId], + ...accountBalances, + }; + } } - // If no chains requested, skip to next middleware - if (request.chainIds.length === 0) { - return next(context); + // Determine successfully handled chains (exclude unprocessed/error chains) + const unprocessedChains = new Set(Object.keys(response.errors ?? {})); + successfullyHandledChains = request.chainIds.filter( + (chainId) => !unprocessedChains.has(chainId), + ); + + // When token detection is off and we filtered out all balance data (e.g. new + // account with empty state), do not claim any chain as handled so that RPC + // middleware can still process them and fetch native balances (ETH, MATIC, etc.). + if ( + !this.#tokenDetectionEnabled() && + (!response.assetsBalance || + Object.keys(response.assetsBalance).length === 0) + ) { + successfullyHandledChains = []; } + } catch (error) { + log('Middleware fetch failed', { error }); + successfullyHandledChains = []; + } - let successfullyHandledChains: ChainId[] = []; + // Remove successfully handled chains from request for next middleware + if (successfullyHandledChains.length > 0) { + const remainingChains = request.chainIds.filter( + (chainId) => !successfullyHandledChains.includes(chainId), + ); - try { - const response = await this.fetch(request); - - // Merge response into context - if (response.assetsBalance) { - context.response.assetsBalance ??= {}; - for (const [accountId, accountBalances] of Object.entries( - response.assetsBalance, - )) { - context.response.assetsBalance[accountId] = { - ...context.response.assetsBalance[accountId], - ...accountBalances, - }; - } - } + return next({ + ...context, + request: { + ...request, + chainIds: remainingChains, + }, + }); + } - // Determine successfully handled chains (exclude unprocessed/error chains) - const unprocessedChains = new Set(Object.keys(response.errors ?? {})); - successfullyHandledChains = request.chainIds.filter( - (chainId) => !unprocessedChains.has(chainId), - ); + // No chains handled - pass context unchanged + return next(context); + }; + + readonly #assetsMiddlewareV6: Middleware = async (context, next) => { + const { request } = context; + + // Price/metadata-only requests must not hit the Accounts API. + if (!request.dataTypes.includes('balance')) { + return next(context); + } - // When token detection is off and we filtered out all balance data (e.g. new - // account with empty state), do not claim any chain as handled so that RPC - // middleware can still process them and fetch native balances (ETH, MATIC, etc.). - if ( - !this.#tokenDetectionEnabled() && - (!response.assetsBalance || - Object.keys(response.assetsBalance).length === 0) - ) { - successfullyHandledChains = []; + // If no chains requested, skip to next middleware + if (request.chainIds.length === 0) { + return next(context); + } + + let remainingChains = request.chainIds; + + try { + const response = await this.fetch(request); + + // Merge response into context + if (response.assetsBalance) { + context.response.assetsBalance ??= {}; + for (const [accountId, accountBalances] of Object.entries( + response.assetsBalance, + )) { + context.response.assetsBalance[accountId] = { + ...context.response.assetsBalance[accountId], + ...accountBalances, + }; } - } catch (error) { - log('Middleware fetch failed', { error }); - successfullyHandledChains = []; } - // Remove successfully handled chains from request for next middleware - if (successfullyHandledChains.length > 0) { - const remainingChains = request.chainIds.filter( - (chainId) => !successfullyHandledChains.includes(chainId), - ); + context.response = { + ...context.response, + updateMode: 'full', + }; - return next({ - ...context, - request: { - ...request, - chainIds: remainingChains, - }, - }); + const errors = response.errors ?? {}; + remainingChains = request.chainIds.filter((chainId) => errors[chainId]); + + // Remaining chains are exactly the failed ones; RPC fallback retries them. + if (remainingChains.length > 0) { + context.response.errors = { + ...context.response.errors, + ...Object.fromEntries( + remainingChains.map((chainId) => [chainId, errors[chainId]]), + ), + }; } + } catch (error) { + log('Middleware fetch failed', { error }); + remainingChains = request.chainIds; + } - // No chains handled - pass context unchanged - return next(context); - }; - } + return next({ + ...context, + request: { + ...request, + chainIds: remainingChains, + }, + }); + }; // ============================================================================ // SUBSCRIBE @@ -839,11 +1068,6 @@ export class AccountsApiDataSource extends AbstractDataSource< const { request, subscriptionId, isUpdate, skipInitialFetch } = subscriptionRequest; - // Store state accessor for filtering when tokenDetectionEnabled is false - if (subscriptionRequest.getAssetsState) { - this.#getAssetsState = subscriptionRequest.getAssetsState; - } - // Try all requested chains - API will handle unsupported ones via unprocessedNetworks const chainsToSubscribe = request.chainIds; diff --git a/packages/assets-controller/src/data-sources/PriceDataSource.test.ts b/packages/assets-controller/src/data-sources/PriceDataSource.test.ts index a62dd70bfef..3d8fa7eb401 100644 --- a/packages/assets-controller/src/data-sources/PriceDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/PriceDataSource.test.ts @@ -6,7 +6,7 @@ import type { DataRequest, Context, Caip19AssetId, - AssetsControllerStateInternal, + AssetsControllerState, } from '../types.js'; import { normalizeAssetId } from '../utils/index.js'; import type { PriceDataSourceOptions } from './PriceDataSource.js'; @@ -31,7 +31,7 @@ type MockApiClient = { type SetupResult = { controller: PriceDataSource; apiClient: MockApiClient; - getAssetsState: () => AssetsControllerStateInternal; + getAssetsState: () => AssetsControllerState; assetsUpdateHandler: jest.Mock; }; @@ -76,7 +76,6 @@ function createMiddlewareContext(overrides?: Partial): Context { return { request: createDataRequest(), response: {}, - getAssetsState: jest.fn().mockReturnValue({ assetsPrice: {} }), ...overrides, }; } @@ -119,6 +118,7 @@ function setupController( getSelectedCurrency?: () => SupportedCurrency; pollInterval?: number; supportedNetworks?: string[]; + getAssetsState?: () => AssetsControllerState; } = {}, ): SetupResult { const { @@ -127,14 +127,21 @@ function setupController( getSelectedCurrency = (): SupportedCurrency => 'usd', pollInterval, supportedNetworks = ['eip155:1'], + getAssetsState: getAssetsStateOverride, } = options; const apiClient = createMockApiClient(priceResponse, supportedNetworks); + const getAssetsState = + getAssetsStateOverride ?? + ((): AssetsControllerState => + ({ assetsBalance: balanceState }) as AssetsControllerState); + const controllerOptions: PriceDataSourceOptions = { queryApiClient: apiClient as unknown as PriceDataSourceOptions['queryApiClient'], getSelectedCurrency, + getAssetsState, }; if (pollInterval) { @@ -143,8 +150,6 @@ function setupController( const controller = new PriceDataSource(controllerOptions); - const getAssetsState = (): AssetsControllerStateInternal => - ({ assetsBalance: balanceState }) as AssetsControllerStateInternal; const assetsUpdateHandler = jest.fn().mockResolvedValue(undefined); return { @@ -181,14 +186,11 @@ describe('PriceDataSource', () => { }); it('fetch returns empty response when no assets in balance state', async () => { - const { controller, getAssetsState } = setupController({ + const { controller } = setupController({ balanceState: {}, }); - const response = await controller.fetch( - createDataRequest(), - getAssetsState, - ); + const response = await controller.fetch(createDataRequest()); expect(response).toStrictEqual({}); @@ -196,7 +198,7 @@ describe('PriceDataSource', () => { }); it('fetch retrieves prices for assets in balance state', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -207,10 +209,7 @@ describe('PriceDataSource', () => { }, }); - const response = await controller.fetch( - createDataRequest(), - getAssetsState, - ); + const response = await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledWith( [MOCK_NATIVE_ASSET], @@ -230,7 +229,7 @@ describe('PriceDataSource', () => { }); it('fetch skips malformed asset IDs in balance state and still fetches prices for valid assets', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { 'not-a-valid-caip19': { amount: '999' }, @@ -242,10 +241,7 @@ describe('PriceDataSource', () => { }, }); - const response = await controller.fetch( - createDataRequest(), - getAssetsState, - ); + const response = await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledWith( [MOCK_NATIVE_ASSET], @@ -277,14 +273,13 @@ describe('PriceDataSource', () => { assetIds.map((id) => [id, { amount: '1' }]), ); - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': balanceState }, priceResponse, }); const response = await controller.fetch( createDataRequest({ chainIds: [] }), - getAssetsState, ); // With 120 assets and a batch size of 50, the API should be called three times. @@ -305,7 +300,7 @@ describe('PriceDataSource', () => { }); it('fetch uses custom currency', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ getSelectedCurrency: () => 'eur', balanceState: { 'mock-account-id': { @@ -317,7 +312,7 @@ describe('PriceDataSource', () => { }, }); - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledWith( expect.anything(), @@ -328,7 +323,7 @@ describe('PriceDataSource', () => { }); it('fetch filters by account ID', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -342,7 +337,7 @@ describe('PriceDataSource', () => { }, }); - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledWith( [MOCK_NATIVE_ASSET], @@ -355,7 +350,7 @@ describe('PriceDataSource', () => { it('fetch filters by chain ID', async () => { const polygonAsset = 'eip155:137/erc20:0x0000000000000000000000000000000000001010' as Caip19AssetId; - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -368,10 +363,7 @@ describe('PriceDataSource', () => { supportedNetworks: ['eip155:1', 'eip155:137'], }); - await controller.fetch( - createDataRequest({ chainIds: [CHAIN_POLYGON] }), - getAssetsState, - ); + await controller.fetch(createDataRequest({ chainIds: [CHAIN_POLYGON] })); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledWith( [polygonAsset], @@ -388,7 +380,7 @@ describe('PriceDataSource', () => { const tronStakedAsset = 'tron:0x2b6653dc/slip44:195-staked-for-bandwidth' as Caip19AssetId; - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -402,7 +394,7 @@ describe('PriceDataSource', () => { }, }); - await controller.fetch(createDataRequest({ chainIds: [] }), getAssetsState); + await controller.fetch(createDataRequest({ chainIds: [] })); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledWith( [MOCK_NATIVE_ASSET], @@ -413,7 +405,7 @@ describe('PriceDataSource', () => { }); it('fetch skips assets with invalid market data', async () => { - const { controller, getAssetsState } = setupController({ + const { controller } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -428,7 +420,6 @@ describe('PriceDataSource', () => { const response = await controller.fetch( createDataRequest({ chainIds: [] }), - getAssetsState, ); expect(response.assetsPrice?.[MOCK_NATIVE_ASSET]).toBeDefined(); @@ -438,7 +429,7 @@ describe('PriceDataSource', () => { }); it('fetch handles API errors gracefully', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -450,10 +441,7 @@ describe('PriceDataSource', () => { new Error('API Error'), ); - const response = await controller.fetch( - createDataRequest(), - getAssetsState, - ); + const response = await controller.fetch(createDataRequest()); expect(response).toStrictEqual({}); @@ -461,16 +449,13 @@ describe('PriceDataSource', () => { }); it('fetch handles getState error gracefully', async () => { - const { controller } = setupController(); - - const getAssetsStateThatThrows = (): never => { - throw new Error('State Error'); - }; + const { controller } = setupController({ + getAssetsState: (): never => { + throw new Error('State Error'); + }, + }); - const response = await controller.fetch( - createDataRequest(), - getAssetsStateThatThrows, - ); + const response = await controller.fetch(createDataRequest()); expect(response).toStrictEqual({}); @@ -478,25 +463,22 @@ describe('PriceDataSource', () => { }); it('subscribe performs initial fetch', async () => { - const { controller, assetsUpdateHandler, getAssetsState } = setupController( - { - balanceState: { - 'mock-account-id': { - [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, - }, - }, - priceResponse: { - [MOCK_NATIVE_ASSET]: createMockPriceData(2500), + const { controller, assetsUpdateHandler } = setupController({ + balanceState: { + 'mock-account-id': { + [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, }, }, - ); + priceResponse: { + [MOCK_NATIVE_ASSET]: createMockPriceData(2500), + }, + }); await controller.subscribe({ subscriptionId: 'sub-1', request: createDataRequest(), isUpdate: false, onAssetsUpdate: assetsUpdateHandler, - getAssetsState, }); expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); @@ -512,7 +494,7 @@ describe('PriceDataSource', () => { }); it('subscribe polls at specified interval', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ pollInterval: 5000, balanceState: { 'mock-account-id': { @@ -527,7 +509,6 @@ describe('PriceDataSource', () => { await controller.subscribe({ subscriptionId: 'sub-1', request: createDataRequest(), - getAssetsState, isUpdate: false, onAssetsUpdate: jest.fn(), }); @@ -541,7 +522,7 @@ describe('PriceDataSource', () => { }); it('subscribe uses request updateInterval when provided', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ pollInterval: 60000, balanceState: { 'mock-account-id': { @@ -558,7 +539,6 @@ describe('PriceDataSource', () => { request: createDataRequest({ updateInterval: 10000 }), isUpdate: false, onAssetsUpdate: jest.fn(), - getAssetsState, }); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); @@ -576,7 +556,7 @@ describe('PriceDataSource', () => { }); it('subscribe update refreshes request and fetches missing prices', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -592,7 +572,6 @@ describe('PriceDataSource', () => { request: createDataRequest(), isUpdate: false, onAssetsUpdate: jest.fn(), - getAssetsState, }); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); @@ -603,7 +582,6 @@ describe('PriceDataSource', () => { request: createDataRequest({ chainIds: [CHAIN_POLYGON] }), isUpdate: true, onAssetsUpdate: jest.fn(), - getAssetsState, }); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); @@ -613,14 +591,6 @@ describe('PriceDataSource', () => { it('subscribe update fetches newly held assets that have no price yet', async () => { let balanceState: Record> = {}; - const { controller, apiClient, assetsUpdateHandler } = setupController({ - balanceState: {}, - priceResponse: { - [MOCK_NATIVE_ASSET]: createMockPriceData(2500), - [MOCK_TOKEN_ASSET]: createMockPriceData(1), - }, - }); - const getAssetsState = jest.fn(() => ({ assetsBalance: balanceState, assetsPrice: {}, @@ -629,13 +599,19 @@ describe('PriceDataSource', () => { assetPreferences: {}, selectedCurrency: 'usd' as const, })); + const { controller, apiClient, assetsUpdateHandler } = setupController({ + getAssetsState, + priceResponse: { + [MOCK_NATIVE_ASSET]: createMockPriceData(2500), + [MOCK_TOKEN_ASSET]: createMockPriceData(1), + }, + }); await controller.subscribe({ subscriptionId: 'sub-1', request: createDataRequest(), isUpdate: false, onAssetsUpdate: assetsUpdateHandler, - getAssetsState, }); // Initial subscribe had no balances, so no price API call. @@ -655,7 +631,6 @@ describe('PriceDataSource', () => { request: createDataRequest(), isUpdate: true, onAssetsUpdate: assetsUpdateHandler, - getAssetsState, }); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalled(); @@ -673,13 +648,6 @@ describe('PriceDataSource', () => { it('subscribe update swallows onAssetsUpdate errors without throwing', async () => { let balanceState: Record> = {}; - const { controller } = setupController({ - balanceState: {}, - priceResponse: { - [MOCK_NATIVE_ASSET]: createMockPriceData(2500), - }, - }); - const getAssetsState = jest.fn(() => ({ assetsBalance: balanceState, assetsPrice: {}, @@ -688,13 +656,18 @@ describe('PriceDataSource', () => { assetPreferences: {}, selectedCurrency: 'usd' as const, })); + const { controller } = setupController({ + getAssetsState, + priceResponse: { + [MOCK_NATIVE_ASSET]: createMockPriceData(2500), + }, + }); await controller.subscribe({ subscriptionId: 'sub-1', request: createDataRequest(), isUpdate: false, onAssetsUpdate: jest.fn(), - getAssetsState, }); balanceState = { @@ -711,7 +684,6 @@ describe('PriceDataSource', () => { onAssetsUpdate: jest .fn() .mockRejectedValue(new Error('handler failed')), - getAssetsState, }), ).toBeUndefined(); @@ -719,19 +691,16 @@ describe('PriceDataSource', () => { }); it('subscribe does not report when no prices fetched', async () => { - const { controller, assetsUpdateHandler, getAssetsState } = setupController( - { - balanceState: {}, - priceResponse: {}, - }, - ); + const { controller, assetsUpdateHandler } = setupController({ + balanceState: {}, + priceResponse: {}, + }); await controller.subscribe({ subscriptionId: 'sub-1', request: createDataRequest(), isUpdate: false, onAssetsUpdate: assetsUpdateHandler, - getAssetsState, }); expect(assetsUpdateHandler).not.toHaveBeenCalled(); @@ -740,7 +709,7 @@ describe('PriceDataSource', () => { }); it('unsubscribe stops polling', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ pollInterval: 5000, balanceState: { 'mock-account-id': { @@ -757,7 +726,6 @@ describe('PriceDataSource', () => { request: createDataRequest(), isUpdate: false, onAssetsUpdate: jest.fn(), - getAssetsState, }); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); @@ -971,7 +939,7 @@ describe('PriceDataSource', () => { }); it('skips fetching prices for assets fetched within the freshness TTL', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -983,18 +951,18 @@ describe('PriceDataSource', () => { }); // First fetch — asset is stale, API is called - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); // Second fetch immediately after — asset is fresh, API is NOT called again - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); controller.destroy(); }); it('re-fetches prices after the freshness TTL expires', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ pollInterval: 10_000, balanceState: { 'mock-account-id': { @@ -1006,13 +974,13 @@ describe('PriceDataSource', () => { }, }); - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); // Advance past the TTL (pollInterval = 10s is used as freshness TTL) await jest.advanceTimersByTimeAsync(11_000); - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(2); controller.destroy(); @@ -1029,7 +997,7 @@ describe('PriceDataSource', () => { jest.setSystemTime(fetchLatencyMs); - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ pollInterval, balanceState: { 'mock-account-id': { @@ -1048,7 +1016,6 @@ describe('PriceDataSource', () => { request: createDataRequest(), isUpdate: false, onAssetsUpdate: jest.fn(), - getAssetsState, }); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); @@ -1057,7 +1024,7 @@ describe('PriceDataSource', () => { // the poll interval but above the capped TTL (9000ms), so the asset is // stale and must be re-fetched rather than skipped. jest.setSystemTime(pollInterval); - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(2); controller.destroy(); @@ -1152,7 +1119,7 @@ describe('PriceDataSource', () => { }); it('freshness is per-asset — stale assets are fetched while fresh ones are skipped', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -1165,7 +1132,7 @@ describe('PriceDataSource', () => { }); // Fetch only MOCK_NATIVE_ASSET via balance state - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledWith( [MOCK_NATIVE_ASSET], @@ -1193,7 +1160,7 @@ describe('PriceDataSource', () => { }); it('destroy clears the freshness cache', async () => { - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ balanceState: { 'mock-account-id': { [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, @@ -1204,7 +1171,7 @@ describe('PriceDataSource', () => { }, }); - await controller.fetch(createDataRequest(), getAssetsState); + await controller.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(1); controller.destroy(); @@ -1215,18 +1182,17 @@ describe('PriceDataSource', () => { queryApiClient: apiClient as unknown as PriceDataSourceOptions['queryApiClient'], getSelectedCurrency: (): SupportedCurrency => 'usd', - }); - - const getAssetsState2 = (): AssetsControllerStateInternal => - ({ - assetsBalance: { - 'mock-account-id': { - [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, + getAssetsState: (): AssetsControllerState => + ({ + assetsBalance: { + 'mock-account-id': { + [MOCK_NATIVE_ASSET]: { amount: '1000000000000000000' }, + }, }, - }, - }) as unknown as AssetsControllerStateInternal; + }) as AssetsControllerState, + }); - await controller2.fetch(createDataRequest(), getAssetsState2); + await controller2.fetch(createDataRequest()); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(2); controller2.destroy(); @@ -1236,7 +1202,7 @@ describe('PriceDataSource', () => { const polygonAsset = 'eip155:137/erc20:0x0000000000000000000000000000000000001010' as Caip19AssetId; - const { controller, apiClient, getAssetsState } = setupController({ + const { controller, apiClient } = setupController({ pollInterval: 5000, balanceState: { 'mock-account-id': { @@ -1256,7 +1222,6 @@ describe('PriceDataSource', () => { request: createDataRequest(), isUpdate: false, onAssetsUpdate: jest.fn(), - getAssetsState, }); await controller.subscribe({ @@ -1264,7 +1229,6 @@ describe('PriceDataSource', () => { request: createDataRequest({ chainIds: [CHAIN_POLYGON] }), isUpdate: false, onAssetsUpdate: jest.fn(), - getAssetsState, }); expect(apiClient.prices.fetchV3SpotPrices).toHaveBeenCalledTimes(2); diff --git a/packages/assets-controller/src/data-sources/PriceDataSource.ts b/packages/assets-controller/src/data-sources/PriceDataSource.ts index 75309142c07..3d10d12d758 100644 --- a/packages/assets-controller/src/data-sources/PriceDataSource.ts +++ b/packages/assets-controller/src/data-sources/PriceDataSource.ts @@ -14,7 +14,7 @@ import type { DataResponse, FungibleAssetPrice, Middleware, - AssetsControllerStateInternal, + AssetsControllerState, } from '../types.js'; import { DedupingBatchFetcher } from '../utils/dedupingBatchFetcher.js'; import { fetchWithTimeout, safeNormalizeAssetId } from '../utils/index.js'; @@ -69,6 +69,8 @@ export type PriceDataSourceOptions = PriceDataSourceConfig & { queryApiClient: ApiPlatformClient; /** Function returning the currently-active ISO 4217 currency code */ getSelectedCurrency: () => SupportedCurrency; + /** Current AssetsController state. Used for balance-based pricing and freshness checks. */ + getAssetsState: () => AssetsControllerState; }; // ============================================================================ @@ -134,9 +136,9 @@ function isValidMarketData(data: unknown): data is SpotPriceMarketData { * This data source: * - Fetches prices from Price API v3 spot-prices endpoint * - Supports one-time fetch and subscription-based polling - * - In subscribe mode, uses getAssetsState from SubscriptionRequest to read assetsBalance and fetch prices + * - In subscribe mode, uses constructor `getAssetsState` to read assetsBalance and fetch prices * - * Usage: Create with queryApiClient; subscribe() requires getAssetsState in the request for balance-based pricing. + * Usage: Create with queryApiClient and getAssetsState; subscribe() polls prices for held assets. */ export class PriceDataSource { static readonly controllerName = CONTROLLER_NAME; @@ -162,6 +164,8 @@ export class PriceDataSource { */ readonly #deduper: DedupingBatchFetcher; + readonly #getAssetsState: () => AssetsControllerState; + /** Active subscriptions by ID */ readonly #activeSubscriptions: Map< string, @@ -169,7 +173,6 @@ export class PriceDataSource { cleanup: () => void; request: DataRequest; onAssetsUpdate: (response: DataResponse) => void | Promise; - getAssetsState?: () => AssetsControllerStateInternal; } > = new Map(); @@ -178,6 +181,7 @@ export class PriceDataSource { this.#pollInterval = options.pollInterval ?? DEFAULT_POLL_INTERVAL; this.#apiClient = options.queryApiClient; this.#fetchTimeoutMs = options.fetchTimeoutMs ?? DEFAULT_FETCH_TIMEOUT_MS; + this.#getAssetsState = options.getAssetsState; this.#deduper = new DedupingBatchFetcher({ fetchBatch: ( assetIds, @@ -211,7 +215,7 @@ export class PriceDataSource { // Extract response from context const { response, request } = ctx; - const statePrices = (ctx.getAssetsState()?.assetsPrice ?? {}) as Record< + const statePrices = (this.#getAssetsState()?.assetsPrice ?? {}) as Record< string, FungibleAssetPrice >; @@ -414,18 +418,11 @@ export class PriceDataSource { * Filters by accounts and chains from the request. * * @param request - Data request with accounts and chainIds filters. - * @param getAssetsState - State access; when omitted, returns []. * @returns Array of CAIP-19 asset IDs from balance state. */ - #getAssetIdsFromBalanceState( - request: DataRequest, - getAssetsState?: () => AssetsControllerStateInternal, - ): Caip19AssetId[] { - if (!getAssetsState) { - return []; - } + #getAssetIdsFromBalanceState(request: DataRequest): Caip19AssetId[] { try { - const state = getAssetsState(); + const state = this.#getAssetsState(); const assetIds = new Set(); const accountIds = request.accountsWithSupportedChains.map( @@ -527,23 +524,16 @@ export class PriceDataSource { /** * Fetch prices for assets held by the accounts and chains in the request. - * When getAssetsState is provided, gets asset IDs from balance state; otherwise returns empty. + * Reads asset IDs from balance state via the constructor `getAssetsState`. * * @param request - The data request specifying accounts and chains. - * @param getAssetsState - Optional state access (e.g. from SubscriptionRequest). * @returns DataResponse containing asset prices. */ - async fetch( - request: DataRequest, - getAssetsState?: () => AssetsControllerStateInternal, - ): Promise { + async fetch(request: DataRequest): Promise { const response: DataResponse = {}; - // Get asset IDs from balance state when state access is provided - const rawAssetIds = this.#getAssetIdsFromBalanceState( - request, - getAssetsState, - ); + // Get asset IDs from balance state + const rawAssetIds = this.#getAssetIdsFromBalanceState(request); // Filter out non-priceable assets (e.g., Tron bandwidth/energy resources) const priceableAssetIds = rawAssetIds.filter(isPriceableAsset); @@ -597,13 +587,9 @@ export class PriceDataSource { if (existing) { existing.request = request; existing.onAssetsUpdate = subscriptionRequest.onAssetsUpdate; - existing.getAssetsState = subscriptionRequest.getAssetsState; try { - const fetchResponse = await this.fetch( - request, - subscriptionRequest.getAssetsState, - ); + const fetchResponse = await this.fetch(request); if ( fetchResponse.assetsPrice && Object.keys(fetchResponse.assetsPrice).length > 0 @@ -635,7 +621,7 @@ export class PriceDataSource { Math.floor(pollInterval * FRESHNESS_TTL_POLL_RATIO), ); - // Create poll function - fetches prices using getAssetsState from subscription. + // Create poll function - fetches prices using constructor getAssetsState. // The freshness TTL naturally gates re-fetches: assets fetched less than // `priceFreshnessTtlMs` ago are skipped, preventing duplicates when middleware // or other triggers already fetched the same assets between polls. @@ -648,10 +634,7 @@ export class PriceDataSource { return; } - const fetchResponse = await this.fetch( - subscription.request, - subscription.getAssetsState, - ); + const fetchResponse = await this.fetch(subscription.request); // Only report if we got prices if ( @@ -675,14 +658,13 @@ export class PriceDataSource { pollFn().catch(console.error); }, pollInterval); - // Store subscription (getAssetsState from request for balance-based pricing) + // Store subscription this.#activeSubscriptions.set(subscriptionId, { cleanup: () => { clearInterval(timer); }, request, onAssetsUpdate: subscriptionRequest.onAssetsUpdate, - getAssetsState: subscriptionRequest.getAssetsState, }); // Initial fetch diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts index 7254faa6102..5699013e784 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts @@ -8,7 +8,6 @@ import type { TransactionMeta } from '@metamask/transaction-controller'; import { createMockMessengers, MockRootMessenger, - registerAssetsControllerStateMock, registerRpcDataSourceActions, } from '../__fixtures__/MockAssetControllerMessenger.js'; import { getDefaultAssetsControllerState } from '../AssetsController.js'; @@ -16,11 +15,13 @@ import type { AssetsControllerMessenger, AssetsControllerState, } from '../AssetsController.js'; +import { getDefaultTrackedAssetsForChain } from '../defaults.js'; import type { Caip19AssetId, ChainId, DataRequest, Context } from '../types.js'; +import type { AssetVisibility } from '../utils/assetVisibility.js'; +import { getAssetVisibility } from '../utils/assetVisibility.js'; import { normalizeAssetId } from '../utils/index.js'; import { BalanceFetcher, TokenDetector } from './evm-rpc-services/index.js'; import type { - Address, BalanceFetchResult, TokenDetectionResult, } from './evm-rpc-services/index.js'; @@ -36,6 +37,7 @@ const MOCK_CHAIN_ID_HEX = '0x1'; const MOCK_CHAIN_ID_CAIP = 'eip155:1' as ChainId; const MOCK_ACCOUNT_ID = 'mock-account-id'; const MOCK_ADDRESS = '0x1234567890123456789012345678901234567890'; +const [MAINNET_MUSD] = getDefaultTrackedAssetsForChain(MOCK_CHAIN_ID_CAIP); type EthereumProvider = { request: jest.Mock }; function createBalanceFetchResult( @@ -164,14 +166,6 @@ async function withController( const { rootMessenger, assetsControllerMessenger } = createMockMessengers(); const defaultNetworkState = networkState ?? createMockNetworkState(); - // TODO - code smell, why is our internal logic trying to call its own methods via messenger? - registerAssetsControllerStateMock( - assetsControllerMessenger, - actionHandlerOverrides?.['AssetsController:getState'] as - | (() => AssetsControllerState) - | undefined, - ); - if (actionHandlerOverrides) { for (const [action, handler] of Object.entries(actionHandlerOverrides)) { if (handler && action !== 'AssetsController:getState') { @@ -231,13 +225,29 @@ async function withController( const defaultNativeAssetMap: Record = { [MOCK_CHAIN_ID_CAIP]: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, }; + const getAssetsState = + options.getAssetsState ?? + ((): AssetsControllerState => { + try { + return ( + (actionHandlerOverrides?.['AssetsController:getState']?.() as + | AssetsControllerState + | undefined) ?? getDefaultAssetsControllerState() + ); + } catch { + return getDefaultAssetsControllerState(); + } + }); const onActiveChainsUpdated = options.onActiveChainsUpdated ?? jest.fn(); const controller = new RpcDataSource({ messenger: assetsControllerMessenger, onActiveChainsUpdated, + // Mirrors AssetsController.#getNativeAssetForChain: registered natives + // from the map, zero-address ERC-20 fallback for unregistered EVM chains. getNativeAssetForChain: (chainId: ChainId): Caip19AssetId => - defaultNativeAssetMap[chainId], + defaultNativeAssetMap[chainId] ?? + (`${chainId}/erc20:0x0000000000000000000000000000000000000000` as Caip19AssetId), getAssetType: (assetId: Caip19AssetId): 'native' | 'erc20' | 'spl' => { const isNative = Object.values(defaultNativeAssetMap).some( @@ -245,7 +255,15 @@ async function withController( ) || assetId.includes('/slip44:'); return isNative ? 'native' : 'erc20'; }, + getAssetVisibility: (accountIds, chainIds): AssetVisibility => + getAssetVisibility({ + state: getAssetsState(), + accountIds, + chainIds, + getNativeAssetForChain: (chainId) => defaultNativeAssetMap[chainId], + }), ...options, + getAssetsState, }); try { @@ -296,9 +314,11 @@ describe('createRpcDataSource', () => { const { assetsControllerMessenger } = createMockMessengers(); const source = createRpcDataSource({ messenger: assetsControllerMessenger, + getAssetsState: getDefaultAssetsControllerState, onActiveChainsUpdated: jest.fn(), getNativeAssetForChain: jest.fn(), getAssetType: jest.fn(), + getAssetVisibility: jest.fn(), }); expect(source).toBeInstanceOf(RpcDataSource); source.destroy(); @@ -491,6 +511,7 @@ describe('RpcDataSource', () => { const response = await controller.fetch(createDataRequest()); expect(response).toBeDefined(); expect(response.assetsBalance).toBeDefined(); + expect(response.updateMode).toBe('merge'); }); }); @@ -544,7 +565,9 @@ describe('RpcDataSource', () => { }); it('returns empty balances when Multicall aggregate3 fails after retries', async () => { - const { Web3Provider } = jest.requireMock('@ethersproject/providers'); + const { Web3Provider } = jest.requireMock('@ethersproject/providers') as { + Web3Provider: jest.Mock; + }; jest.mocked(shouldSkipNativeForCaipChainId).mockReturnValue(true); const mockCall = jest @@ -554,7 +577,7 @@ describe('RpcDataSource', () => { const mockGetBalance = jest .fn() .mockResolvedValue({ toString: () => '1000000000000000000' }); - (Web3Provider as jest.Mock).mockImplementationOnce(() => ({ + Web3Provider.mockImplementationOnce(() => ({ call: mockCall, getBalance: mockGetBalance, })); @@ -569,7 +592,9 @@ describe('RpcDataSource', () => { it('uses getBalance when Multicall aggregate3 fails after retries', async () => { const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; - const { Web3Provider } = jest.requireMock('@ethersproject/providers'); + const { Web3Provider } = jest.requireMock('@ethersproject/providers') as { + Web3Provider: jest.Mock; + }; const mockCall = jest .fn() .mockRejectedValueOnce(new Error('aggregate3 unavailable')) @@ -578,7 +603,7 @@ describe('RpcDataSource', () => { const mockGetBalance = jest .fn() .mockResolvedValue({ toString: () => '1000000000000000000' }); - (Web3Provider as jest.Mock).mockImplementationOnce(() => ({ + Web3Provider.mockImplementationOnce(() => ({ call: mockCall, getBalance: mockGetBalance, })); @@ -641,6 +666,185 @@ describe('RpcDataSource', () => { }); }); + it('keeps the previous balance and stays full when a balanceOf fails on the v6 path', async () => { + const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; + const trackedAssetId = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + assetsInfo: { + [trackedAssetId]: { + type: 'erc20' as const, + symbol: 'USDC', + name: 'USD Coin', + decimals: 6, + }, + }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [trackedAssetId]: { amount: '42' } }, + }, + }), + }, + }, + async ({ controller }) => { + jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue( + createBalanceFetchResult({ + balances: [ + { + assetId: nativeAssetId, + accountId: MOCK_ACCOUNT_ID, + chainId: MOCK_CHAIN_ID_HEX, + balance: '1000000000000000000', + formattedBalance: '1', + decimals: 18, + timestamp: Date.now(), + }, + ], + failedAddresses: ['0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48'], + }), + ); + const response = await controller.fetch(createDataRequest()); + + // The chain still updates on one failed read: native takes its fresh + // amount and the unreadable token keeps the amount already in state. + expect(response.errors).toBeUndefined(); + expect(response.updateMode).toBe('full'); + expect( + response.assetsBalance?.[MOCK_ACCOUNT_ID]?.[nativeAssetId], + ).toStrictEqual({ amount: '1' }); + expect( + response.assetsBalance?.[MOCK_ACCOUNT_ID]?.[trackedAssetId], + ).toStrictEqual({ amount: '42' }); + }, + ); + }); + + it('keeps the previous balance and stays full when decimals cannot be resolved on the v6 path', async () => { + const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; + const erc20AssetId = + 'eip155:1/erc20:0xAbc0000000000000000000000000000000000001' as Caip19AssetId; + const { Web3Provider } = jest.requireMock('@ethersproject/providers') as { + Web3Provider: jest.Mock; + }; + Web3Provider.mockImplementationOnce(() => ({ + getBalance: jest + .fn() + .mockResolvedValue({ toString: () => '1000000000000000000' }), + call: jest.fn().mockResolvedValue('0x'), + })); + + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + // No metadata for the token, so its decimals stay unresolved. + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [erc20AssetId]: { amount: '7' } }, + }, + }), + }, + }, + async ({ controller }) => { + jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue( + createBalanceFetchResult({ + balances: [ + { + assetId: nativeAssetId, + accountId: MOCK_ACCOUNT_ID, + chainId: MOCK_CHAIN_ID_HEX, + balance: '1000000000000000000', + formattedBalance: '1', + decimals: 18, + timestamp: Date.now(), + }, + { + assetId: erc20AssetId, + accountId: MOCK_ACCOUNT_ID, + chainId: MOCK_CHAIN_ID_HEX, + balance: '1000000', + formattedBalance: '1', + timestamp: Date.now(), + }, + ], + }), + ); + const response = await controller.fetch(createDataRequest()); + expect(response.errors).toBeUndefined(); + expect(response.updateMode).toBe('full'); + expect( + response.assetsBalance?.[MOCK_ACCOUNT_ID]?.[nativeAssetId], + ).toStrictEqual({ amount: '1' }); + expect( + response.assetsBalance?.[MOCK_ACCOUNT_ID]?.[ + normalizeAssetId(erc20AssetId) + ], + ).toStrictEqual({ amount: '7' }); + }, + ); + }); + + it('overlays tokens that resolve and skips unresolved decimals on the v5 path', async () => { + const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; + const erc20AssetId = + 'eip155:1/erc20:0xAbc0000000000000000000000000000000000001' as Caip19AssetId; + const normalizedErc20Id = normalizeAssetId(erc20AssetId); + const { Web3Provider } = jest.requireMock('@ethersproject/providers') as { + Web3Provider: jest.Mock; + }; + Web3Provider.mockImplementationOnce(() => ({ + getBalance: jest + .fn() + .mockResolvedValue({ toString: () => '1000000000000000000' }), + call: jest.fn().mockResolvedValue('0x'), + })); + + await withController(async ({ controller }) => { + jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue( + createBalanceFetchResult({ + balances: [ + { + assetId: nativeAssetId, + accountId: MOCK_ACCOUNT_ID, + chainId: MOCK_CHAIN_ID_HEX, + balance: '1000000000000000000', + formattedBalance: '1', + decimals: 18, + timestamp: Date.now(), + }, + { + assetId: erc20AssetId, + accountId: MOCK_ACCOUNT_ID, + chainId: MOCK_CHAIN_ID_HEX, + balance: '1000000', + formattedBalance: '1', + timestamp: Date.now(), + }, + ], + }), + ); + const response = await controller.fetch(createDataRequest()); + expect(response.errors).toBeUndefined(); + expect( + response.assetsBalance?.[MOCK_ACCOUNT_ID]?.[nativeAssetId], + ).toStrictEqual({ amount: '1' }); + expect( + response.assetsBalance?.[MOCK_ACCOUNT_ID]?.[normalizedErc20Id], + ).toBeUndefined(); + }); + }); + it('initializes assetsBalance[accountId] with no native in catch when first fetch for account throws on native skip chain', async () => { await withController(async ({ controller }) => { jest @@ -809,7 +1013,7 @@ describe('RpcDataSource', () => { ); }); - it('handles AssetsController:getState throw in metadata lookup', async () => { + it('handles AssetsController:getState throw without failing fetch metadata', async () => { await withController( { actionHandlerOverrides: { @@ -869,120 +1073,509 @@ describe('RpcDataSource', () => { .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') .mockResolvedValue(createBalanceFetchResult()); - await withController(async ({ controller }) => { - const request = createDataRequest({ - customAssets: [matchingAsset, otherChainAsset], - }); - await controller.fetch(request); - - expect(fetchSpy).toHaveBeenCalledWith( - MOCK_CHAIN_ID_HEX, - MOCK_ACCOUNT_ID, - MOCK_ADDRESS, - [ - { - assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, - address: '0x0000000000000000000000000000000000000000', - }, - expect.objectContaining({ - assetId: matchingAsset, - address: '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { + [MOCK_ACCOUNT_ID]: [matchingAsset, otherChainAsset], + }, }), - ], - ); - }); + }, + }, + async ({ controller }) => { + const request = createDataRequest({ + customAssets: [matchingAsset, otherChainAsset], + }); + await controller.fetch(request); + + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + [ + { + assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, + address: '0x0000000000000000000000000000000000000000', + }, + expect.objectContaining({ + assetId: matchingAsset, + address: '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', + }), + ], + ); + }, + ); fetchSpy.mockRestore(); }); - }); - describe('detectTokens', () => { - it('returns empty when token detection disabled', async () => { + it('fetches existing balances, pins, and default tracked assets from state on the v6 path when the request carries no customAssets', async () => { + const pinnedAsset = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const trackedBalanceAsset = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue(createBalanceFetchResult()); + await withController( { - options: { - tokenDetectionEnabled: () => false, - useExternalService: () => true, + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { [MOCK_ACCOUNT_ID]: [pinnedAsset] }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [trackedBalanceAsset]: { amount: '1' }, + }, + }, + }), }, }, async ({ controller }) => { - const result = await controller.detectTokens( - MOCK_CHAIN_ID_CAIP, - createMockInternalAccount(), + const response = await controller.fetch(createDataRequest()); + + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + expect.arrayContaining([ + { + assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, + address: '0x0000000000000000000000000000000000000000', + }, + expect.objectContaining({ assetId: pinnedAsset }), + expect.objectContaining({ assetId: trackedBalanceAsset }), + expect.objectContaining({ assetId: MAINNET_MUSD }), + ]), ); - expect(result).toStrictEqual({}); + expect(response.updateMode).toBe('full'); }, ); + + fetchSpy.mockRestore(); }); - it('returns empty when useExternalService false', async () => { + it('omits an unheld asset outside the visible set from the v6 snapshot', async () => { + const nativeAsset = `${MOCK_CHAIN_ID_CAIP}/slip44:60` as Caip19AssetId; + const pinnedAsset = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const detectedAsset = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue( + createBalanceFetchResult({ + balances: [ + { + assetId: nativeAsset, + balance: '1000000000000000000', + }, + { assetId: pinnedAsset, balance: '0' }, + { assetId: detectedAsset, balance: '0' }, + ] as BalanceFetchResult['balances'], + }), + ); + await withController( { - options: { - tokenDetectionEnabled: () => true, - useExternalService: () => false, + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + assetsInfo: { + [pinnedAsset]: { + type: 'erc20' as const, + symbol: 'USDC', + name: 'USD Coin', + decimals: 6, + }, + [detectedAsset]: { + type: 'erc20' as const, + symbol: 'DAI', + name: 'Dai', + decimals: 18, + }, + }, + customAssets: { [MOCK_ACCOUNT_ID]: [pinnedAsset] }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [detectedAsset]: { amount: '3' } }, + }, + }), }, }, async ({ controller }) => { - const result = await controller.detectTokens( - MOCK_CHAIN_ID_CAIP, - createMockInternalAccount(), - ); - expect(result).toStrictEqual({}); + const response = await controller.fetch(createDataRequest()); + + // The pin keeps its zero row; the detected token is left out, so the + // `full` write drops it instead of persisting an empty balance. + expect(response.assetsBalance?.[MOCK_ACCOUNT_ID]).toStrictEqual({ + [nativeAsset]: { amount: '1' }, + [pinnedAsset]: { amount: '0' }, + }); + expect(response.updateMode).toBe('full'); }, ); + + fetchSpy.mockRestore(); }); - it('returns detected tokens when detection succeeds', async () => { - const assetId = 'eip155:1/erc20:0xabc' as Caip19AssetId; - const detectTokensSpy = jest - .spyOn(TokenDetector.prototype, 'detectTokens') - .mockResolvedValue({ - chainId: MOCK_CHAIN_ID_HEX, - accountId: MOCK_ACCOUNT_ID, - accountAddress: MOCK_ADDRESS, - detectedAssets: [ - { - assetId, - symbol: 'TST', - name: 'Test', - decimals: 18, - } as TokenDetectionResult['detectedAssets'][0], - ], - detectedBalances: [ - { - assetId, - balance: '1000000000000000000', - } as TokenDetectionResult['detectedBalances'][0], - ], - zeroBalanceAddresses: [], - failedAddresses: [], - timestamp: Date.now(), - }); + it('still emits a v6 snapshot when every read resolved to an omitted zero', async () => { + const nativeAsset = `${MOCK_CHAIN_ID_CAIP}/slip44:60` as Caip19AssetId; + const detectedAsset = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; - await withController(async ({ controller }) => { - const result = await controller.detectTokens( - MOCK_CHAIN_ID_CAIP, - createMockInternalAccount(), + // Only the detected token resolves — at zero — so nothing is written. + // The zero read is still a success: the snapshot must go out so the + // `full` write can drop the stale detected token from state. + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue( + createBalanceFetchResult({ + balances: [ + { assetId: detectedAsset, balance: '0' }, + ] as BalanceFetchResult['balances'], + failedAddresses: ['0x0000000000000000000000000000000000000000'], + }), ); - expect(result.detectedAssets).toBeDefined(); - expect(result.assetsBalance).toBeDefined(); - expect(Object.keys(result.assetsInfo ?? {})).toHaveLength(1); - }); - detectTokensSpy.mockRestore(); + + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + assetsInfo: { + [detectedAsset]: { + type: 'erc20' as const, + symbol: 'DAI', + name: 'Dai', + decimals: 18, + }, + }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [nativeAsset]: { amount: '2' }, + [detectedAsset]: { amount: '3' }, + }, + }, + }), + }, + }, + async ({ controller }) => { + const response = await controller.fetch(createDataRequest()); + + // The unreadable native keeps its previous amount; the detected + // token is omitted so the `full` write drops it. + expect(response.assetsBalance?.[MOCK_ACCOUNT_ID]).toStrictEqual({ + [nativeAsset]: { amount: '2' }, + }); + expect(response.updateMode).toBe('full'); + }, + ); + + fetchSpy.mockRestore(); }); - it('returns detected asset ids but omits assetsBalance when detected token has no decimals', async () => { - const assetId = - 'eip155:1/erc20:0xabcdefabcdefabcdefabcdefabcdefabcdefabcd' as Caip19AssetId; - const detectTokensSpy = jest - .spyOn(TokenDetector.prototype, 'detectTokens') - .mockResolvedValue({ - chainId: MOCK_CHAIN_ID_HEX, - accountId: MOCK_ACCOUNT_ID, - accountAddress: MOCK_ADDRESS, - detectedAssets: [ - { + it('does not fetch hidden native, pinned, default, or existing assets on the v6 path', async () => { + const nativeAsset = `${MOCK_CHAIN_ID_CAIP}/slip44:60` as Caip19AssetId; + const pinnedAsset = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const trackedBalanceAsset = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue(createBalanceFetchResult()); + + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { [MOCK_ACCOUNT_ID]: [pinnedAsset] }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [trackedBalanceAsset]: { amount: '1' }, + }, + }, + assetPreferences: { + [nativeAsset]: { hidden: true }, + [pinnedAsset]: { hidden: true }, + [trackedBalanceAsset]: { hidden: true }, + [MAINNET_MUSD]: { hidden: true }, + }, + }), + }, + }, + async ({ controller }) => { + await controller.fetch(createDataRequest()); + + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + [], + ); + }, + ); + + fetchSpy.mockRestore(); + }); + + it('fetches only native on the v5 path when the request carries no customAssets', async () => { + const pinnedAsset = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const trackedBalanceAsset = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue(createBalanceFetchResult()); + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { [MOCK_ACCOUNT_ID]: [pinnedAsset] }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [trackedBalanceAsset]: { amount: '1' }, + }, + }, + }), + }, + }, + async ({ controller }) => { + await controller.fetch(createDataRequest()); + + // v5 reads ERC-20s from the scoped request only; state pins and + // tracked balances are polled by the balance fetcher instead. + const assets = fetchSpy.mock.calls[0][3] as { assetId: string }[]; + const assetIds = assets.map((entry) => entry.assetId); + expect(assetIds).toStrictEqual([`${MOCK_CHAIN_ID_CAIP}/slip44:60`]); + expect(assetIds).not.toContain(pinnedAsset); + expect(assetIds).not.toContain(trackedBalanceAsset); + }, + ); + + fetchSpy.mockRestore(); + }); + + it('fetches only the request custom assets when the request is scoped', async () => { + const requestedAsset = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const otherPinnedAsset = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue(createBalanceFetchResult()); + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { + [MOCK_ACCOUNT_ID]: [requestedAsset, otherPinnedAsset], + }, + }), + }, + }, + async ({ controller }) => { + // A scoped fetch must not widen into the account's other pins. + await controller.fetch( + createDataRequest({ customAssets: [requestedAsset] }), + ); + + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + [ + { + assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, + address: '0x0000000000000000000000000000000000000000', + }, + expect.objectContaining({ assetId: requestedAsset }), + ], + ); + }, + ); + + fetchSpy.mockRestore(); + }); + + it('skips hidden custom assets from state', async () => { + const hiddenAsset = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue(createBalanceFetchResult()); + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { [MOCK_ACCOUNT_ID]: [hiddenAsset] }, + assetPreferences: { [hiddenAsset]: { hidden: true } }, + }), + }, + }, + async ({ controller }) => { + await controller.fetch(createDataRequest()); + + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + [ + { + assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, + address: '0x0000000000000000000000000000000000000000', + }, + ], + ); + }, + ); + + fetchSpy.mockRestore(); + }); + + it('does not fetch custom assets belonging to another account', async () => { + const otherAccountAsset = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue(createBalanceFetchResult()); + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { 'other-account-id': [otherAccountAsset] }, + }), + }, + }, + async ({ controller }) => { + await controller.fetch(createDataRequest()); + + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + [ + { + assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, + address: '0x0000000000000000000000000000000000000000', + }, + ], + ); + }, + ); + + fetchSpy.mockRestore(); + }); + }); + + describe('detectTokens', () => { + it('returns empty when token detection disabled', async () => { + await withController( + { + options: { + tokenDetectionEnabled: () => false, + useExternalService: () => true, + }, + }, + async ({ controller }) => { + const result = await controller.detectTokens( + MOCK_CHAIN_ID_CAIP, + createMockInternalAccount(), + ); + expect(result).toStrictEqual({}); + }, + ); + }); + + it('returns empty when useExternalService false', async () => { + await withController( + { + options: { + tokenDetectionEnabled: () => true, + useExternalService: () => false, + }, + }, + async ({ controller }) => { + const result = await controller.detectTokens( + MOCK_CHAIN_ID_CAIP, + createMockInternalAccount(), + ); + expect(result).toStrictEqual({}); + }, + ); + }); + + it('returns detected tokens when detection succeeds', async () => { + const assetId = 'eip155:1/erc20:0xabc' as Caip19AssetId; + const detectTokensSpy = jest + .spyOn(TokenDetector.prototype, 'detectTokens') + .mockResolvedValue({ + chainId: MOCK_CHAIN_ID_HEX, + accountId: MOCK_ACCOUNT_ID, + accountAddress: MOCK_ADDRESS, + detectedAssets: [ + { + assetId, + symbol: 'TST', + name: 'Test', + decimals: 18, + } as TokenDetectionResult['detectedAssets'][0], + ], + detectedBalances: [ + { + assetId, + balance: '1000000000000000000', + } as TokenDetectionResult['detectedBalances'][0], + ], + zeroBalanceAddresses: [], + failedAddresses: [], + timestamp: Date.now(), + }); + + await withController(async ({ controller }) => { + const result = await controller.detectTokens( + MOCK_CHAIN_ID_CAIP, + createMockInternalAccount(), + ); + expect(result.detectedAssets).toBeDefined(); + expect(result.assetsBalance).toBeDefined(); + expect(Object.keys(result.assetsInfo ?? {})).toHaveLength(1); + }); + detectTokensSpy.mockRestore(); + }); + + it('returns detected asset ids but omits assetsBalance when detected token has no decimals', async () => { + const assetId = + 'eip155:1/erc20:0xabcdefabcdefabcdefabcdefabcdefabcdefabcd' as Caip19AssetId; + const detectTokensSpy = jest + .spyOn(TokenDetector.prototype, 'detectTokens') + .mockResolvedValue({ + chainId: MOCK_CHAIN_ID_HEX, + accountId: MOCK_ACCOUNT_ID, + accountAddress: MOCK_ADDRESS, + detectedAssets: [ + { assetId, symbol: 'ND', name: 'No Decimals', @@ -1188,6 +1781,24 @@ describe('RpcDataSource', () => { expect(next).toHaveBeenCalledWith(context); }); }); + + it('forwards full updateMode on the v6 path', async () => { + await withController( + { options: { isBalanceV6Enabled: (): boolean => true } }, + async ({ controller }) => { + const context: Context = { + request: createDataRequest(), + response: {}, + getAssetsState: jest.fn(), + }; + const next = jest + .fn() + .mockImplementation((ctx: Context) => Promise.resolve(ctx)); + await controller.assetsMiddleware(context, next); + expect(context.response.updateMode).toBe('full'); + }, + ); + }); }); describe('subscribe', () => { @@ -1382,6 +1993,30 @@ describe('RpcDataSource', () => { expect(detectionStopSpy).toHaveBeenCalled(); }); }); + + it('does not poll when no chains are assigned even if customAssets are present', async () => { + const balanceStartSpy = jest.spyOn( + BalanceFetcher.prototype, + 'startPolling', + ); + const customAssetId = + `${MOCK_CHAIN_ID_CAIP}/erc20:0x1111111111111111111111111111111111111111` as Caip19AssetId; + + await withController(async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest({ + chainIds: [], + customAssets: [customAssetId], + }), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate: jest.fn(), + }); + + expect(balanceStartSpy).not.toHaveBeenCalled(); + await controller.unsubscribe('test-sub'); + }); + }); }); describe('unsubscribe', () => { @@ -1501,10 +2136,7 @@ describe('RpcDataSource', () => { ).toStrictEqual({ amount: '1' }); }); - it('omits unknown ERC-20 from assetsInfo when not in existing state', async () => { - const tokenAddress = '0xAbc0000000000000000000000000000000000001'; - const erc20AssetId = `eip155:1/erc20:${tokenAddress}` as Caip19AssetId; - const normalizedId = normalizeAssetId(erc20AssetId); + it('stamps full updateMode on the v6 poll path', async () => { let balanceUpdateCallback: | ((result: BalanceFetchResult) => void | Promise) | null = null; @@ -1516,13 +2148,445 @@ describe('RpcDataSource', () => { const onAssetsUpdate = jest.fn(); await withController( - { - actionHandlerOverrides: { - 'AssetsController:getState': () => ({ - ...getDefaultAssetsControllerState(), - assetsInfo: {}, - }), - }, + { options: { isBalanceV6Enabled: (): boolean => true } }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate, + }); + await balanceUpdateCallback?.( + createBalanceFetchResult({ + balances: [ + { + assetId: 'eip155:1/slip44:60' as Caip19AssetId, + balance: '1000000000000000000', + } as BalanceFetchResult['balances'][0], + ], + }), + ); + }, + ); + + expect(onAssetsUpdate).toHaveBeenCalledWith( + expect.objectContaining({ updateMode: 'full' }), + expect.any(Object), + ); + }); + + it('keeps the previous balance and stays full when a balanceOf failed on the v6 poll path', async () => { + const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; + const trackedAssetId = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + let balanceUpdateCallback: + | ((result: BalanceFetchResult) => void | Promise) + | null = null; + jest + .spyOn(BalanceFetcher.prototype, 'setOnBalanceUpdate') + .mockImplementation(function (this: BalanceFetcher, callback) { + balanceUpdateCallback = callback; + }); + + const onAssetsUpdate = jest.fn(); + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + assetsInfo: { + [trackedAssetId]: { + type: 'erc20' as const, + symbol: 'USDC', + name: 'USD Coin', + decimals: 6, + }, + }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [trackedAssetId]: { amount: '42' } }, + }, + }), + }, + }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate, + }); + await balanceUpdateCallback?.( + createBalanceFetchResult({ + balances: [ + { + assetId: nativeAssetId, + balance: '1000000000000000000', + } as BalanceFetchResult['balances'][0], + ], + failedAddresses: ['0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48'], + }), + ); + }, + ); + + expect(onAssetsUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [nativeAssetId]: { amount: '1' }, + [trackedAssetId]: { amount: '42' }, + }, + }, + }), + expect.any(Object), + ); + }); + + it('keeps the previous balance and stays full when an asset has unknown decimals on the v6 poll path', async () => { + const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; + const erc20AssetId = + 'eip155:1/erc20:0xAbc0000000000000000000000000000000000001' as Caip19AssetId; + let balanceUpdateCallback: + | ((result: BalanceFetchResult) => void | Promise) + | null = null; + jest + .spyOn(BalanceFetcher.prototype, 'setOnBalanceUpdate') + .mockImplementation(function (this: BalanceFetcher, callback) { + balanceUpdateCallback = callback; + }); + + const onAssetsUpdate = jest.fn(); + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + // No metadata for the token, so its decimals stay unresolved. + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [erc20AssetId]: { amount: '7' } }, + }, + }), + }, + }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate, + }); + await balanceUpdateCallback?.( + createBalanceFetchResult({ + balances: [ + { + assetId: nativeAssetId, + balance: '1000000000000000000', + } as BalanceFetchResult['balances'][0], + { + assetId: erc20AssetId, + balance: '500', + } as BalanceFetchResult['balances'][0], + ], + }), + ); + }, + ); + + expect(onAssetsUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [nativeAssetId]: { amount: '1' }, + [normalizeAssetId(erc20AssetId)]: { amount: '7' }, + }, + }, + }), + expect.any(Object), + ); + }); + + it('leaves an unheld asset outside the visible set out of the v6 poll snapshot', async () => { + const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; + // Only tracked because detection found it earlier: not native, pinned, + // or default tracked. At zero it is not worth a row. + const detectedAssetId = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + let balanceUpdateCallback: + | ((result: BalanceFetchResult) => void | Promise) + | null = null; + jest + .spyOn(BalanceFetcher.prototype, 'setOnBalanceUpdate') + .mockImplementation(function (this: BalanceFetcher, callback) { + balanceUpdateCallback = callback; + }); + + const onAssetsUpdate = jest.fn(); + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + assetsInfo: { + [detectedAssetId]: { + type: 'erc20' as const, + symbol: 'DAI', + name: 'Dai', + decimals: 18, + }, + }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [detectedAssetId]: { amount: '3' } }, + }, + }), + }, + }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate, + }); + await balanceUpdateCallback?.( + createBalanceFetchResult({ + balances: [ + { + assetId: nativeAssetId, + balance: '1000000000000000000', + } as BalanceFetchResult['balances'][0], + { + assetId: detectedAssetId, + balance: '0', + } as BalanceFetchResult['balances'][0], + ], + }), + ); + }, + ); + + expect(onAssetsUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [nativeAssetId]: { amount: '1' }, + }, + }, + }), + expect.any(Object), + ); + }); + + it('still emits a v6 poll snapshot when every read resolved to an omitted zero', async () => { + const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; + const detectedAssetId = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + let balanceUpdateCallback: + | ((result: BalanceFetchResult) => void | Promise) + | null = null; + jest + .spyOn(BalanceFetcher.prototype, 'setOnBalanceUpdate') + .mockImplementation(function (this: BalanceFetcher, callback) { + balanceUpdateCallback = callback; + }); + + const onAssetsUpdate = jest.fn(); + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + assetsInfo: { + [detectedAssetId]: { + type: 'erc20' as const, + symbol: 'DAI', + name: 'Dai', + decimals: 18, + }, + }, + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [nativeAssetId]: { amount: '2' }, + [detectedAssetId]: { amount: '3' }, + }, + }, + }), + }, + }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate, + }); + // The only read that resolved is a zero for a detected token, which + // is deliberately left out of the snapshot. The update must still + // publish so the `full` write drops the stale token from state. + await balanceUpdateCallback?.( + createBalanceFetchResult({ + balances: [ + { + assetId: detectedAssetId, + balance: '0', + } as BalanceFetchResult['balances'][0], + ], + }), + ); + }, + ); + + expect(onAssetsUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [nativeAssetId]: { amount: '2' }, + }, + }, + }), + expect.any(Object), + ); + }); + + it('keeps a zero balance for a pinned asset on the v6 poll snapshot', async () => { + const nativeAssetId = 'eip155:1/slip44:60' as Caip19AssetId; + const pinnedAssetId = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + let balanceUpdateCallback: + | ((result: BalanceFetchResult) => void | Promise) + | null = null; + jest + .spyOn(BalanceFetcher.prototype, 'setOnBalanceUpdate') + .mockImplementation(function (this: BalanceFetcher, callback) { + balanceUpdateCallback = callback; + }); + + const onAssetsUpdate = jest.fn(); + await withController( + { + options: { isBalanceV6Enabled: (): boolean => true }, + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + assetsInfo: { + [pinnedAssetId]: { + type: 'erc20' as const, + symbol: 'USDC', + name: 'USD Coin', + decimals: 6, + }, + }, + customAssets: { [MOCK_ACCOUNT_ID]: [pinnedAssetId] }, + }), + }, + }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate, + }); + await balanceUpdateCallback?.( + createBalanceFetchResult({ + balances: [ + { + assetId: nativeAssetId, + balance: '1000000000000000000', + } as BalanceFetchResult['balances'][0], + { + assetId: pinnedAssetId, + balance: '0', + } as BalanceFetchResult['balances'][0], + ], + }), + ); + }, + ); + + expect(onAssetsUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + updateMode: 'full', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [nativeAssetId]: { amount: '1' }, + [pinnedAssetId]: { amount: '0' }, + }, + }, + }), + expect.any(Object), + ); + }); + + it('skips the v6 poll snapshot when no asset can be converted', async () => { + const erc20AssetId = + 'eip155:1/erc20:0xAbc0000000000000000000000000000000000001' as Caip19AssetId; + let balanceUpdateCallback: + | ((result: BalanceFetchResult) => void | Promise) + | null = null; + jest + .spyOn(BalanceFetcher.prototype, 'setOnBalanceUpdate') + .mockImplementation(function (this: BalanceFetcher, callback) { + balanceUpdateCallback = callback; + }); + + const onAssetsUpdate = jest.fn(); + await withController( + { options: { isBalanceV6Enabled: (): boolean => true } }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate, + }); + // Non-zero, so the token is not omitted as an unheld invisible + // zero; its decimals are unknown, so it cannot be converted. + await balanceUpdateCallback?.( + createBalanceFetchResult({ + balances: [ + { + assetId: erc20AssetId, + balance: '500', + } as BalanceFetchResult['balances'][0], + ], + }), + ); + }, + ); + + expect(onAssetsUpdate).not.toHaveBeenCalled(); + }); + + it('omits unknown ERC-20 from assetsInfo when not in existing state', async () => { + const tokenAddress = '0xAbc0000000000000000000000000000000000001'; + const erc20AssetId = `eip155:1/erc20:${tokenAddress}` as Caip19AssetId; + const normalizedId = normalizeAssetId(erc20AssetId); + let balanceUpdateCallback: + | ((result: BalanceFetchResult) => void | Promise) + | null = null; + jest + .spyOn(BalanceFetcher.prototype, 'setOnBalanceUpdate') + .mockImplementation(function (this: BalanceFetcher, callback) { + balanceUpdateCallback = callback; + }); + + const onAssetsUpdate = jest.fn(); + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + assetsInfo: {}, + }), + }, }, async ({ controller }) => { await controller.subscribe({ @@ -1969,6 +3033,50 @@ describe('RpcDataSource', () => { } }); + it('refreshes balance with full mode on the v6 path when transaction confirmed', async () => { + const onAssetsUpdate = jest.fn().mockResolvedValue(undefined); + const fetchSpy = jest + .spyOn(RpcDataSource.prototype, 'fetch') + .mockResolvedValue({ + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + 'eip155:1/slip44:60': { amount: '2' }, + }, + }, + updateMode: 'full', + }); + + try { + await withController( + { options: { isBalanceV6Enabled: (): boolean => true } }, + async ({ controller, rootMessenger }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate, + }); + + rootMessenger.publish( + 'TransactionController:transactionConfirmed', + { + chainId: MOCK_CHAIN_ID_HEX, + txParams: { from: MOCK_ADDRESS }, + } as unknown as TransactionMeta, + ); + await new Promise(process.nextTick); + + expect(onAssetsUpdate).toHaveBeenCalledWith( + expect.objectContaining({ updateMode: 'full' }), + expect.objectContaining({ dataTypes: ['balance'] }), + ); + }, + ); + } finally { + fetchSpy.mockRestore(); + } + }); + it('does not refresh when transaction confirmed has no chainId', async () => { await withController(async ({ rootMessenger }) => { rootMessenger.publish( @@ -2047,9 +3155,11 @@ describe('RpcDataSource', () => { }); const controller = new RpcDataSource({ messenger: assetsControllerMessenger, + getAssetsState: getDefaultAssetsControllerState, onActiveChainsUpdated: jest.fn(), getNativeAssetForChain: jest.fn(), getAssetType: jest.fn().mockReturnValue('erc20'), + getAssetVisibility: jest.fn(), }); controller.destroy(); expect(controller).toBeDefined(); diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.ts b/packages/assets-controller/src/data-sources/RpcDataSource.ts index 71344bd9792..f7a52693eca 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.ts @@ -15,30 +15,30 @@ import type { import { isStrictHexString, isCaipChainId, + KnownCaipNamespace, parseCaipAssetType, parseCaipChainId, hexToNumber, toCaipChainId, - KnownCaipNamespace, } from '@metamask/utils'; import type { Hex } from '@metamask/utils'; import { BigNumber as BigNumberJS } from 'bignumber.js'; -import type { - AssetsControllerGetStateAction, - AssetsControllerMessenger, -} from '../AssetsController.js'; +import type { AssetsControllerMessenger } from '../AssetsController.js'; import { projectLogger, createModuleLogger } from '../logger.js'; import type { ChainId, Caip19AssetId, AssetBalance, AssetMetadata, + AssetsControllerState, DataRequest, DataResponse, Middleware, } from '../types.js'; +import type { GetAssetVisibility } from '../utils/assetVisibility.js'; import { ZERO_ADDRESS } from '../utils/constants.js'; +import { filterFailedChainBalances } from '../utils/filterFailedChainBalances.js'; import { normalizeAssetId } from '../utils/index.js'; import { AbstractDataSource } from './AbstractDataSource.js'; import type { @@ -47,6 +47,7 @@ import type { } from './AbstractDataSource.js'; import { BalanceFetcher, + isStakingContractAssetId, MulticallClient, TokenDetector, TokensApiClient, @@ -74,8 +75,7 @@ const log = createModuleLogger(projectLogger, CONTROLLER_NAME); // Allowed actions that RpcDataSource can call export type RpcDataSourceAllowedActions = | NetworkControllerGetStateAction - | NetworkControllerGetNetworkClientByIdAction - | AssetsControllerGetStateAction; + | NetworkControllerGetNetworkClientByIdAction; // Allowed events that RpcDataSource can subscribe to export type RpcDataSourceAllowedEvents = @@ -117,6 +117,11 @@ export type RpcDataSourceConfig = { export type RpcDataSourceOptions = { /** The AssetsController messenger (shared by all data sources). */ messenger: AssetsControllerMessenger; + /** + * Current AssetsController state. Used to include the account's visible + * `customAssets` on fetch, and to read metadata for converting balances. + */ + getAssetsState: () => AssetsControllerState; /** Called when active chains are updated. Pass dataSourceName so the controller knows the source. */ onActiveChainsUpdated: ( dataSourceName: string, @@ -124,7 +129,7 @@ export type RpcDataSourceOptions = { previousChains: ChainId[], ) => void; /** Resolves CAIP-2 chain ID to CAIP-19 native asset ID from the cached native asset map. */ - getNativeAssetForChain: (chainId: ChainId) => Caip19AssetId; + getNativeAssetForChain: (chainId: ChainId) => Caip19AssetId | undefined; /** Request timeout in ms */ timeout?: number; /** Balance polling interval in ms (default: 30s) */ @@ -145,6 +150,14 @@ export type RpcDataSourceOptions = { /** Returns the asset type ('native' | 'erc20' | 'spl') for the given CAIP-19 asset ID */ getAssetType: (assetId: Caip19AssetId) => 'native' | 'erc20' | 'spl'; + /** + * Whether Accounts API v6 is enabled. Unscoped fetches then include native, + * existing balances, pins, and default tracked assets. v5 unscoped fetches + * stay pins-from-state only. Omitted means v5. + */ + isBalanceV6Enabled?: () => boolean; + /** Returns shared visible/hidden assets for an account/chain scope. */ + getAssetVisibility: GetAssetVisibility; }; /** @@ -208,13 +221,17 @@ export class RpcDataSource extends AbstractDataSource< > { readonly #messenger: AssetsControllerMessenger; + readonly #getAssetsState: () => AssetsControllerState; + readonly #onActiveChainsUpdated: ( dataSourceName: string, chains: ChainId[], previousChains: ChainId[], ) => void; - readonly #getNativeAssetForChain: (chainId: ChainId) => Caip19AssetId; + readonly #getNativeAssetForChain: ( + chainId: ChainId, + ) => Caip19AssetId | undefined; readonly #timeout: number; @@ -249,9 +266,14 @@ export class RpcDataSource extends AbstractDataSource< assetId: Caip19AssetId, ) => 'native' | 'erc20' | 'spl'; + readonly #isBalanceV6Enabled: () => boolean; + + readonly #getAssetVisibility: GetAssetVisibility; + constructor(options: RpcDataSourceOptions) { super(CONTROLLER_NAME, { activeChains: [] }); this.#messenger = options.messenger; + this.#getAssetsState = options.getAssetsState; this.#onActiveChainsUpdated = options.onActiveChainsUpdated; this.#getNativeAssetForChain = options.getNativeAssetForChain; this.#getAssetType = options.getAssetType; @@ -261,6 +283,9 @@ export class RpcDataSource extends AbstractDataSource< this.#useExternalService = options.useExternalService ?? ((): boolean => true); this.#isOnboarded = options.isOnboarded ?? ((): boolean => true); + this.#isBalanceV6Enabled = + options.isBalanceV6Enabled ?? ((): boolean => false); + this.#getAssetVisibility = options.getAssetVisibility; const balanceInterval = options.balanceInterval ?? DEFAULT_BALANCE_INTERVAL; const detectionInterval = @@ -291,37 +316,18 @@ export class RpcDataSource extends AbstractDataSource< )?.contracts?.multicall3, }); - // Create messenger adapters for BalanceFetcher and TokenDetector - const balanceFetcherMessenger = { - call: ( - _action: 'AssetsController:getState', - ): { - assetsBalance: Record>; - customAssets?: Record; - } => { - // The messenger is not being called from a constructor, so this is safe. - // eslint-disable-next-line no-restricted-syntax - const state = this.#messenger.call('AssetsController:getState'); - return { - assetsBalance: state.assetsBalance ?? {}, - customAssets: state.customAssets ?? {}, - }; - }, - }; - // Initialize BalanceFetcher with polling interval - this.#balanceFetcher = new BalanceFetcher( - this.#multicallClient, - balanceFetcherMessenger, - { - pollingInterval: balanceInterval, - isNativeAsset: (assetId: Caip19AssetId): boolean => { - const { chainId } = parseCaipAssetType(assetId); - const nativeId = this.#getNativeAssetForChain(chainId); - return nativeId?.toLowerCase() === assetId.toLowerCase(); - }, + this.#balanceFetcher = new BalanceFetcher(this.#multicallClient, { + pollingInterval: balanceInterval, + getAssetsState: this.#getAssetsState, + isNativeAsset: (assetId: Caip19AssetId): boolean => { + const { chainId } = parseCaipAssetType(assetId); + const nativeId = this.#getNativeAssetForChain(chainId); + return nativeId?.toLowerCase() === assetId.toLowerCase(); }, - ); + isBalanceV6Enabled: (): boolean => this.#isBalanceV6Enabled(), + getAssetVisibility: this.#getAssetVisibility, + }); // Polling controller awaits this callback; rejections must not become unhandled. this.#balanceFetcher.setOnBalanceUpdate(async (result) => { try { @@ -490,6 +496,21 @@ export class RpcDataSource extends AbstractDataSource< * @param result - The balance fetch result. */ async #handleBalanceUpdate(result: BalanceFetchResult): Promise { + if (this.#isBalanceV6Enabled()) { + await this.#handleBalanceUpdateV6(result); + return; + } + + await this.#handleBalanceUpdateV5(result); + } + + /** + * v5 poll path. Unchanged from the pre-v6 handler; delete the v6 sibling + * first if `assetsAccountsApiV6` is rolled back. + * + * @param result - The balance fetch result. + */ + async #handleBalanceUpdateV5(result: BalanceFetchResult): Promise { const newBalances: Record = {}; // Convert hex chain ID to CAIP-2 format @@ -552,9 +573,126 @@ export class RpcDataSource extends AbstractDataSource< }); for (const subscription of this.#activeSubscriptions.values()) { - subscription.onAssetsUpdate(response, request)?.catch((error) => { - log('Failed to update assets', { error }); - }); + subscription + .onAssetsUpdate(response, request) + ?.catch((error: unknown) => { + log('Failed to update assets', { error }); + }); + } + } + + /** + * v6 poll path. Always stamps `full`: assets that resolve get their fresh + * amount, and a failed `balanceOf` or missing metadata keeps the amount + * already in state instead of dropping the asset. Delete with the rest of + * the v6 path if `assetsAccountsApiV6` is rolled back. + * + * @param result - The balance fetch result. + */ + async #handleBalanceUpdateV6(result: BalanceFetchResult): Promise { + const newBalances: Record = {}; + + // Convert hex chain ID to CAIP-2 format + const chainIdDecimal = parseInt(result.chainId, 16); + const caipChainId = `eip155:${chainIdDecimal}` as ChainId; + + // Normalize asset IDs from BalanceFetcher (lowercase) to checksummed form + const normalizedBalances = result.balances.map((b) => ({ + ...b, + assetId: normalizeAssetId(b.assetId), + })); + + // Collect metadata for all balances + const assetsInfo = this.#collectMetadataForBalances( + normalizedBalances, + caipChainId, + ); + + // Convert balances that have decimals. Skip the rest so a missing + // metadata lookup cannot drop the tokens that did resolve. + const visibleAssetKeys = this.#getVisibleAssetKeys( + result.accountId, + caipChainId, + ); + const handledAssetKeys = new Set(); + const existingMetadata = this.#getExistingAssetsMetadata(); + for (const balance of normalizedBalances) { + if ( + this.#isUnheldInvisibleAsset( + balance.assetId, + balance.balance, + visibleAssetKeys, + ) + ) { + handledAssetKeys.add(balance.assetId.toLowerCase()); + continue; + } + + const stateMetadata = existingMetadata[balance.assetId]; + const pipelineMetadata = assetsInfo[balance.assetId]; + const decimals = this.#pickValidDecimals(stateMetadata, pipelineMetadata); + + if (decimals === undefined) { + continue; + } + + const humanReadableAmount = this.#convertToHumanReadable( + balance.balance, + decimals, + ); + + newBalances[balance.assetId] = { + amount: humanReadableAmount, + }; + handledAssetKeys.add(balance.assetId.toLowerCase()); + } + + if (handledAssetKeys.size === 0) { + return; + } + + // Everything requested that did not resolve keeps its previous amount, so + // the snapshot stays complete for the chain and can be applied as `full`. + const unresolved = this.#getAssetsToFetchV6( + result.accountId, + caipChainId, + ).filter( + (entry) => + !handledAssetKeys.has(normalizeAssetId(entry.assetId).toLowerCase()), + ); + this.#carryOverUnresolvedBalances({ + unresolved, + accountId: result.accountId, + chainId: caipChainId, + balances: newBalances, + visibleAssetKeys, + }); + + const response: DataResponse = { + assetsBalance: { + [result.accountId]: newBalances, + }, + assetsInfo, + updateMode: 'full', + }; + + const request: DataRequest = { + accountsWithSupportedChains: [], + chainIds: [caipChainId], + dataTypes: ['balance'], + }; + + log('Balance update response', { + accountId: result.accountId, + newBalanceCount: Object.keys(newBalances).length, + }); + + for (const subscription of this.#activeSubscriptions.values()) { + subscription + .onAssetsUpdate(response, request) + ?.catch((error: unknown) => { + log('Failed to update assets', { error }); + }); } } @@ -629,9 +767,11 @@ export class RpcDataSource extends AbstractDataSource< }; for (const subscription of this.#activeSubscriptions.values()) { - subscription.onAssetsUpdate(response, request)?.catch((error) => { - log('Failed to update detected assets', { error }); - }); + subscription + .onAssetsUpdate(response, request) + ?.catch((error: unknown) => { + log('Failed to update detected assets', { error }); + }); } } @@ -662,7 +802,7 @@ export class RpcDataSource extends AbstractDataSource< } const caipChainId = `eip155:${parseInt(hexChainId, 16)}` as ChainId; this.#refreshBalanceForChains([caipChainId], 'transactionConfirmed').catch( - (error) => { + (error: unknown) => { log('Failed to refresh balance after transaction confirmed', { error }); }, ); @@ -987,22 +1127,113 @@ export class RpcDataSource extends AbstractDataSource< log('Skipping fetch - onboarding not complete'); return {}; } + if (this.#isBalanceV6Enabled()) { + return this.#fetchV6(request); + } + return this.#fetchV5(request); + } - const response: DataResponse = {}; + /** + * v5 fetch. Overlays tokens that succeed, stamps `merge`. Delete with the + * rest of the v5 path if `assetsAccountsApiV6` is rolled back. + * + * @param request - The data request. + * @returns Balance and metadata for the requested chains. + */ + async #fetchV5(request: DataRequest): Promise { + const chainsToFetch = this.#getChainsToFetch(request); + if (chainsToFetch.length === 0) { + log('No active chains to fetch'); + return {}; + } - const chainsToFetch = request.chainIds.filter((chainId) => - this.#activeChains.includes(chainId), - ); + const assetsBalance: Record< + string, + Record + > = {}; + const assetsInfo: Record = {}; + const failedChains: ChainId[] = []; - log('Fetch requested', { - accounts: request.accountsWithSupportedChains.map((a) => a.account.id), - requestedChains: request.chainIds, + for (const { + account, + supportedChains, + } of request.accountsWithSupportedChains) { + const chainsForAccount = chainsToFetch.filter((chain) => + supportedChains.includes(chain), + ); + if (chainsForAccount.length === 0) { + continue; + } + + const { address, id: accountId } = account; + + for (const chainId of chainsForAccount) { + const hexChainId = caipChainIdToHex(chainId); + const nativeAssetId = this.#getNativeAssetForChain(chainId); + const shouldSkipNative = shouldSkipNativeForCaipChainId(chainId); + const assetsToFetch = this.#getAssetsToFetchV5( + request, + chainId, + shouldSkipNative, + ); + + try { + const result = await this.#balanceFetcher.fetchBalancesForAssets( + hexChainId, + accountId, + address as Address, + assetsToFetch, + ); + await this.#ingestFetchedBalances( + result, + accountId, + chainId, + assetsBalance, + assetsInfo, + ); + } catch (error) { + this.#recordFetchChainFailure({ + address, + chainId, + error, + accountId, + nativeAssetId, + shouldSkipNative, + assetsBalance, + assetsInfo, + failedChains, + }); + } + } + } + + return this.#completeFetchResponse({ + assetsBalance, + assetsInfo, + failedChains, chainsToFetch, + updateMode: 'merge', }); + } + /** + * v6 fetch. Visibility-scoped snapshot, always stamped `full`. + * + * A single failed `balanceOf` or unknown decimals no longer holds back the + * rest of the chain: the assets that resolved are written, and the ones that + * did not keep the amount already in state (see + * `#carryOverUnresolvedBalances`). A chain that resolved nothing at all is + * reported in `errors` and contributes no balances. + * + * @param request - The data request. + * @returns An authoritative snapshot for the chains that resolved at least + * one balance. Fully failed chains are omitted. + */ + async #fetchV6(request: DataRequest): Promise { + const chainsToFetch = this.#getChainsToFetch(request); if (chainsToFetch.length === 0) { log('No active chains to fetch'); - return response; + return {}; } const assetsBalance: Record< @@ -1012,7 +1243,6 @@ export class RpcDataSource extends AbstractDataSource< const assetsInfo: Record = {}; const failedChains: ChainId[] = []; - // Fetch balances for each account and its supported chains (pre-computed in request) for (const { account, supportedChains, @@ -1029,43 +1259,14 @@ export class RpcDataSource extends AbstractDataSource< for (const chainId of chainsForAccount) { const hexChainId = caipChainIdToHex(chainId); const nativeAssetId = this.#getNativeAssetForChain(chainId); - const shouldSkipNative = shouldSkipNativeForCaipChainId(chainId); - const assetsToFetch: AssetFetchEntry[] = []; - if (!shouldSkipNative) { - // Build a single AssetFetchEntry[] for native + custom ERC-20s - assetsToFetch.push({ assetId: nativeAssetId, address: ZERO_ADDRESS }); - } - - if (request.customAssets) { - const existingMetadata = this.#getExistingAssetsMetadata(); - - for (const assetId of request.customAssets) { - try { - const parsed = parseCaipAssetType(assetId); - const assetChainId = `${parsed.chain.namespace}:${parsed.chain.reference}`; - if ( - assetChainId === chainId && - this.#getAssetType(assetId) === 'erc20' - ) { - const tokenAddress = - parsed.assetReference.toLowerCase() as Address; - const normalizedId = normalizeAssetId(assetId); - const decimals = existingMetadata[normalizedId]?.decimals; - - assetsToFetch.push({ - assetId, - address: tokenAddress, - decimals, - }); - } - } catch { - // Skip unparseable asset IDs - } - } - } try { + const assetsToFetch = this.#getAssetsToFetchV6(accountId, chainId); + const visibleAssetKeys = this.#getVisibleAssetKeys( + accountId, + chainId, + ); const result = await this.#balanceFetcher.fetchBalancesForAssets( hexChainId, accountId, @@ -1073,93 +1274,347 @@ export class RpcDataSource extends AbstractDataSource< assetsToFetch, ); - if (!assetsBalance[accountId]) { - assetsBalance[accountId] = {}; + const handledAssetKeys = await this.#ingestFetchedBalances( + result, + accountId, + chainId, + assetsBalance, + assetsInfo, + visibleAssetKeys, + ); + const unresolved = assetsToFetch.filter( + (entry) => + !handledAssetKeys.has( + normalizeAssetId(entry.assetId).toLowerCase(), + ), + ); + + if (unresolved.length === 0) { + continue; + } + + if (handledAssetKeys.size === 0) { + // Nothing at all resolved, so there is no trustworthy slice to + // write. Leave the chain to the RPC fallback. + if (!failedChains.includes(chainId)) { + failedChains.push(chainId); + } + continue; } - // Normalize asset IDs from BalanceFetcher (which uses lowercase - // addresses) to checksummed form so they match assetsInfo state keys. - const normalizedBalances = result.balances.map((b) => ({ - ...b, - assetId: normalizeAssetId(b.assetId), - })); + log('Partial v6 RPC fetch', { + accountId, + chainId, + failedAddresses: result.failedAddresses, + unresolved: unresolved.map((entry) => entry.assetId), + }); + this.#carryOverUnresolvedBalances({ + unresolved, + accountId, + chainId, + balances: assetsBalance[accountId], + visibleAssetKeys, + }); + } catch (error) { + this.#recordFetchChainFailure({ + address, + chainId, + error, + accountId, + nativeAssetId, + shouldSkipNative, + assetsBalance, + assetsInfo, + failedChains, + }); + } + } + } + + return this.#completeFetchResponse({ + // A fully failed chain contributes nothing, so its balances stay as they are. + assetsBalance: filterFailedChainBalances( + assetsBalance, + new Set(failedChains), + ), + assetsInfo, + failedChains, + chainsToFetch, + updateMode: 'full', + }); + } + + #getChainsToFetch(request: DataRequest): ChainId[] { + const chainsToFetch = request.chainIds.filter((chainId) => + this.#activeChains.includes(chainId), + ); + + log('Fetch requested', { + accounts: request.accountsWithSupportedChains.map((a) => a.account.id), + requestedChains: request.chainIds, + chainsToFetch, + }); + + return chainsToFetch; + } + + /** + * Lowercased visible asset IDs (native, pins, default tracked) for one + * account-chain scope, for membership checks against normalized IDs. + * + * @param accountId - Account being fetched. + * @param chainId - Chain being fetched. + * @returns Lookup keys for the visible set. + */ + #getVisibleAssetKeys(accountId: string, chainId: ChainId): Set { + const { visibleAssetIds } = this.#getAssetVisibility( + [accountId], + [chainId], + ); + return new Set(visibleAssetIds.map((assetId) => assetId.toLowerCase())); + } + + /** + * Whether a v6 snapshot should leave a fetched asset out. + * + * RPC also reads tokens that are only tracked because they were detected + * earlier. A zero balance for one of those is not worth a row: writing it + * keeps an empty token in state indefinitely, and on a `full` snapshot the + * omission drops it instead. Visible assets (native, pins, default tracked) + * are always written, at zero when the account holds none. + * + * @param assetId - Normalized CAIP-19 asset ID. + * @param rawBalance - Raw on-chain amount as returned by the fetcher. + * @param visibleAssetKeys - Lowercased visible IDs for this account-chain. + * @returns True when the asset should be omitted from the snapshot. + */ + #isUnheldInvisibleAsset( + assetId: Caip19AssetId, + rawBalance: string, + visibleAssetKeys: Set, + ): boolean { + return rawBalance === '0' && !visibleAssetKeys.has(assetId.toLowerCase()); + } + + /** + * Convert fetched raw balances into human-readable amounts. + * + * Tokens that resolve are overlaid; tokens whose decimals cannot be resolved + * are skipped so callers can preserve their previous amount. + * + * @param result - The balance fetch result. + * @param accountId - The account the balances belong to. + * @param chainId - The CAIP-2 chain ID. + * @param assetsBalance - Accumulator for converted balances. + * @param assetsInfo - Accumulator for metadata collected from this fetch. + * @param visibleAssetKeys - Lowercased visible IDs (v6 only). When given, + * zero balances for assets outside the set are left out of the snapshot. + * @returns Lowercased IDs this pass settled: written, or deliberately left + * out as an unheld invisible asset. Anything requested but missing from the + * set failed and still needs its previous amount. + */ + async #ingestFetchedBalances( + result: BalanceFetchResult, + accountId: string, + chainId: ChainId, + assetsBalance: Record>, + assetsInfo: Record, + visibleAssetKeys?: Set, + ): Promise> { + assetsBalance[accountId] ??= {}; + + const normalizedBalances = result.balances.map((balance) => ({ + ...balance, + assetId: normalizeAssetId(balance.assetId), + })); + + Object.assign( + assetsInfo, + this.#collectMetadataForBalances(normalizedBalances, chainId), + ); + + const handledAssetKeys = new Set(); + const existingMetadata = this.#getExistingAssetsMetadata(); + for (const balance of normalizedBalances) { + if ( + visibleAssetKeys && + this.#isUnheldInvisibleAsset( + balance.assetId, + balance.balance, + visibleAssetKeys, + ) + ) { + log('Skipping unheld asset outside the visible set', { + accountId, + chainId, + assetId: balance.assetId, + }); + handledAssetKeys.add(balance.assetId.toLowerCase()); + continue; + } + + const stateMetadata = existingMetadata[balance.assetId]; + const pipelineMetadata = assetsInfo[balance.assetId]; + let decimals: number | undefined = this.#pickValidDecimals( + stateMetadata, + pipelineMetadata, + ); - // Collect metadata for all balances - const balanceMetadata = this.#collectMetadataForBalances( - normalizedBalances, + if (decimals === undefined) { + const parsed = parseCaipAssetType(balance.assetId); + if (this.#getAssetType(balance.assetId) === 'erc20') { + decimals = await this.#fetchDecimalsViaRpc( chainId, + parsed.assetReference, ); - Object.assign(assetsInfo, balanceMetadata); - - // Convert balances to human-readable format using decimals from - // assetsInfo state (which includes pendingMetadata from addCustomAsset). - // Resolution: state → pipeline metadata → RPC `decimals()`; omit balance if still unknown. - const existingMetadata = this.#getExistingAssetsMetadata(); - for (const balance of normalizedBalances) { - const stateMetadata = existingMetadata[balance.assetId]; - const pipelineMetadata = assetsInfo[balance.assetId]; - let decimals: number | undefined = this.#pickValidDecimals( - stateMetadata, - pipelineMetadata, - ); - - if (decimals === undefined) { - const parsed = parseCaipAssetType(balance.assetId); - if (this.#getAssetType(balance.assetId) === 'erc20') { - decimals = await this.#fetchDecimalsViaRpc( - chainId, - parsed.assetReference, - ); - } - } + } + } - if (decimals === undefined) { - continue; - } + if (decimals === undefined) { + log('Skipping asset with unresolved decimals on RPC fetch', { + accountId, + chainId, + assetId: balance.assetId, + }); + continue; + } - const humanReadableAmount = this.#convertToHumanReadable( - balance.balance, - decimals, - ); + assetsBalance[accountId][balance.assetId] = { + amount: this.#convertToHumanReadable(balance.balance, decimals), + }; + handledAssetKeys.add(balance.assetId.toLowerCase()); + } - assetsBalance[accountId][balance.assetId] = { - amount: humanReadableAmount, - }; - } - } catch (error) { - log('Failed to fetch balance', { address, chainId, error }); + return handledAssetKeys; + } - if (!assetsBalance[accountId]) { - assetsBalance[accountId] = {}; - } + /** + * Preserve the amount already in state for assets this pass could not read. + * + * A v6 snapshot replaces the chain slice it covers, so an asset left out + * would be dropped. Writing its previous amount back means one failed + * `balanceOf` or unknown decimals no longer blocks the chain from updating: + * every other asset still gets its fresh amount, and the unreadable one is + * retried on the next poll. An unheld asset outside the visible set is not + * preserved — the same rule as a fresh zero read applies. + * + * @param options - Carry-over inputs. + * @param options.unresolved - Requested entries the fetch did not settle. + * @param options.accountId - Account being fetched. + * @param options.chainId - Chain being fetched. + * @param options.balances - Balance map for this account, mutated in place. + * @param options.visibleAssetKeys - Lowercased visible IDs for this scope. + */ + #carryOverUnresolvedBalances({ + unresolved, + accountId, + chainId, + balances, + visibleAssetKeys, + }: { + unresolved: AssetFetchEntry[]; + accountId: string; + chainId: ChainId; + balances: Record; + visibleAssetKeys: Set; + }): void { + const previousBalances = + this.#getAssetsState().assetsBalance[accountId] ?? {}; + const previousByKey = new Map( + Object.entries(previousBalances).map(([assetId, balance]) => [ + assetId.toLowerCase(), + balance, + ]), + ); - if (!shouldSkipNative) { - assetsBalance[accountId][nativeAssetId] = { amount: '0' }; - } - // On error, emit a stub only when no valid metadata exists in state - // yet. Re-emitting existing metadata would overwrite richer entries - // (e.g. image/description added by AccountsAPI) with a simpler stub. - const existingNativeMeta = - this.#getExistingAssetsMetadata()[nativeAssetId]; - if (!this.#hasValidDecimals(existingNativeMeta)) { - const chainStatus = this.#chainStatuses[chainId]; - if (chainStatus) { - assetsInfo[nativeAssetId] = { - type: 'native', - symbol: chainStatus.nativeCurrency, - name: chainStatus.nativeCurrency, - decimals: 18, - }; - } - } + for (const entry of unresolved) { + const assetId = normalizeAssetId(entry.assetId); + const key = assetId.toLowerCase(); + if (balances[assetId] !== undefined) { + continue; + } - if (!failedChains.includes(chainId)) { - failedChains.push(chainId); - } + const previous = previousByKey.get(key); + if (previous === undefined) { + continue; + } + if (!visibleAssetKeys.has(key) && Number(previous.amount) === 0) { + continue; + } + + log('Keeping previous balance for unreadable asset', { + accountId, + chainId, + assetId, + }); + balances[assetId] = { ...previous }; + } + } + + #recordFetchChainFailure({ + address, + chainId, + error, + accountId, + nativeAssetId, + shouldSkipNative, + assetsBalance, + assetsInfo, + failedChains, + }: { + address: string; + chainId: ChainId; + error: unknown; + accountId: string; + nativeAssetId: Caip19AssetId | undefined; + shouldSkipNative: boolean; + assetsBalance: Record>; + assetsInfo: Record; + failedChains: ChainId[]; + }): void { + log('Failed to fetch balance', { address, chainId, error }); + + assetsBalance[accountId] ??= {}; + + if (nativeAssetId) { + if (!shouldSkipNative) { + assetsBalance[accountId][nativeAssetId] = { amount: '0' }; + } + const existingNativeMeta = + this.#getExistingAssetsMetadata()[nativeAssetId]; + if (!this.#hasValidDecimals(existingNativeMeta)) { + const chainStatus = this.#chainStatuses[chainId]; + if (chainStatus) { + assetsInfo[nativeAssetId] = { + type: 'native', + symbol: chainStatus.nativeCurrency, + name: chainStatus.nativeCurrency, + decimals: 18, + }; } } } + if (!failedChains.includes(chainId)) { + failedChains.push(chainId); + } + } + + #completeFetchResponse({ + assetsBalance, + assetsInfo, + failedChains, + chainsToFetch, + updateMode, + }: { + assetsBalance: DataResponse['assetsBalance']; + assetsInfo: Record; + failedChains: ChainId[]; + chainsToFetch: ChainId[]; + updateMode: DataResponse['updateMode']; + }): DataResponse { + const response: DataResponse = { updateMode }; + if (failedChains.length > 0) { log('Fetch PARTIAL - some chains failed', { successChains: chainsToFetch.filter( @@ -1175,19 +1630,16 @@ export class RpcDataSource extends AbstractDataSource< } else { log('Fetch SUCCESS', { chains: chainsToFetch, - accountCount: Object.keys(assetsBalance).length, + accountCount: Object.keys(assetsBalance ?? {}).length, }); } response.assetsBalance = assetsBalance; - // Include metadata for native tokens if we have any if (Object.keys(assetsInfo).length > 0) { response.assetsInfo = assetsInfo; } - response.updateMode = 'merge'; - return response; } @@ -1344,6 +1796,13 @@ export class RpcDataSource extends AbstractDataSource< }; } + if (response.updateMode) { + context.response = { + ...context.response, + updateMode: response.updateMode, + }; + } + const failedChains = new Set(Object.keys(response.errors ?? {})); successfullyHandledChains = supportedChains.filter( (chainId) => !failedChains.has(chainId), @@ -1450,8 +1909,6 @@ export class RpcDataSource extends AbstractDataSource< const balanceToken = this.#balanceFetcher.startPolling(balanceInput); balancePollingTokens.push(balanceToken); - // Token detection is only relevant for "regular" subscriptions — - // a customAssetsOnly subscription should never run detection. if ( request.customAssetsOnly !== true && this.#tokenDetectionEnabled() && @@ -1512,6 +1969,144 @@ export class RpcDataSource extends AbstractDataSource< } } + /** + * v5 entries for one account-chain fetch: native plus the ERC-20s the + * caller scoped via `request.customAssets`. Unchanged from the pre-v6 + * behavior; delete with the rest of the v5 path. + * + * @param request - The data request being fetched. + * @param chainId - Chain being fetched. + * @param shouldSkipNative - When true, omit the native entry. + * @returns Entries to hand to the balance fetcher. + */ + #getAssetsToFetchV5( + request: DataRequest, + chainId: ChainId, + shouldSkipNative: boolean, + ): AssetFetchEntry[] { + const assetsToFetch: AssetFetchEntry[] = []; + + const nativeAssetId = this.#getNativeAssetForChain(chainId); + if (!shouldSkipNative && nativeAssetId) { + // Build a single AssetFetchEntry[] for native + custom ERC-20s + assetsToFetch.push({ + assetId: nativeAssetId, + address: ZERO_ADDRESS, + }); + } + + if (request.customAssets) { + const existingMetadata = this.#getExistingAssetsMetadata(); + + for (const assetId of request.customAssets) { + try { + const parsed = parseCaipAssetType(assetId); + const assetChainId = `${parsed.chain.namespace}:${parsed.chain.reference}`; + if ( + assetChainId === chainId && + this.#getAssetType(assetId) === 'erc20' + ) { + const tokenAddress = parsed.assetReference.toLowerCase() as Address; + const normalizedId = normalizeAssetId(assetId); + const decimals = existingMetadata[normalizedId]?.decimals; + + assetsToFetch.push({ + assetId, + address: tokenAddress, + decimals, + }); + } + } catch { + // Skip unparseable asset IDs + } + } + } + + return assetsToFetch; + } + + /** + * v6 entries for one account-chain fetch: the shared visible set (native, + * pins, default tracked) plus already-tracked ERC-20s that are not hidden. + * Chains with no native token omit native from `visibleAssetIds`. Hide is + * only applied to `assetsBalance` rows. The request is never scoped via + * `customAssets` on v6. + * + * @param accountId - Account being fetched. + * @param chainId - Chain being fetched. + * @returns Entries to hand to the balance fetcher. + */ + #getAssetsToFetchV6(accountId: string, chainId: ChainId): AssetFetchEntry[] { + const { assetsInfo = {}, assetsBalance = {} } = this.#getAssetsState(); + const { visibleAssetIds, hiddenAssetIds } = this.#getAssetVisibility( + [accountId], + [chainId], + ); + const hidden = new Set(hiddenAssetIds); + const assetsToFetch: AssetFetchEntry[] = []; + const seen = new Set(); + + // Fetch visible assets + for (const assetId of visibleAssetIds) { + try { + if (this.#getAssetType(assetId) === 'native') { + assetsToFetch.push({ assetId, address: ZERO_ADDRESS }); + continue; + } + + const parsed = parseCaipAssetType(assetId); + const assetChainId = `${parsed.chain.namespace}:${parsed.chain.reference}`; + if ( + assetChainId !== chainId || + this.#getAssetType(assetId) !== 'erc20' || + seen.has(assetId) + ) { + continue; + } + + seen.add(assetId); + assetsToFetch.push({ + assetId, + address: parsed.assetReference.toLowerCase() as Address, + decimals: assetsInfo[assetId]?.decimals, + }); + } catch { + // Skip unparseable asset IDs + } + } + + // Fetch already-tracked ERC-20s that are not hidden + for (const assetId of Object.keys( + assetsBalance[accountId] ?? {}, + ) as Caip19AssetId[]) { + try { + const parsed = parseCaipAssetType(assetId); + const assetChainId = `${parsed.chain.namespace}:${parsed.chain.reference}`; + const normalizedId = normalizeAssetId(assetId); + if ( + hidden.has(normalizedId) || + isStakingContractAssetId(assetId) || + assetChainId !== chainId || + this.#getAssetType(assetId) !== 'erc20' || + seen.has(normalizedId) + ) { + continue; + } + + seen.add(normalizedId); + assetsToFetch.push({ + assetId, + address: parsed.assetReference.toLowerCase() as Address, + decimals: assetsInfo[normalizedId]?.decimals, + }); + } catch { + // Skip unparseable asset IDs + } + } + + return assetsToFetch; + } + /** * Get existing assets metadata from AssetsController state. * Used to include metadata for ERC20 tokens when returning balance updates. @@ -1519,13 +2114,7 @@ export class RpcDataSource extends AbstractDataSource< * @returns Record of asset IDs to their metadata. */ #getExistingAssetsMetadata(): Record { - try { - const state = this.#messenger.call('AssetsController:getState'); - return state.assetsInfo ?? {}; - } catch (error) { - log('Failed to get existing assets metadata', { error }); - return {}; - } + return this.#getAssetsState().assetsInfo ?? {}; } /** diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts index 26b33f10b67..1f33b4588df 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts @@ -8,8 +8,16 @@ import type { SubjectPermissions, } from '@metamask/permission-controller'; +import { getDefaultAssetsControllerState } from '../AssetsController.js'; import type { AssetsControllerMessenger } from '../AssetsController.js'; -import type { ChainId, DataRequest, Context, Caip19AssetId } from '../types.js'; +import type { + AssetsControllerState, + ChainId, + DataRequest, + Context, + Caip19AssetId, +} from '../types.js'; +import type { GetAssetVisibility } from '../utils/assetVisibility.js'; import type { SnapDataSourceOptions, AccountBalancesUpdatedEventPayload, @@ -43,7 +51,11 @@ type RootMessenger = Messenger; const MOCK_ADDRESS = '0x1234567890123456789012345678901234567890'; const MOCK_SOL_ASSET = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501' as Caip19AssetId; +const MOCK_SOL_PIN = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:So11111111111111111111111111111111111111112' as Caip19AssetId; const MOCK_STELLAR_ASSET = 'stellar:pubnet/slip44:148' as Caip19AssetId; +const MOCK_STELLAR_PIN = + 'stellar:pubnet/asset:USDC-GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN' as Caip19AssetId; const MOCK_BTC_ASSET = 'bip122:000000000019d6689c085ae165831e93/slip44:0' as Caip19AssetId; const MOCK_TRON_ASSET = 'tron:728126428/slip44:195' as Caip19AssetId; @@ -107,7 +119,6 @@ function createMiddlewareContext(overrides?: Partial): Context { return { request: createDataRequest(), response: {}, - getAssetsState: jest.fn().mockReturnValue({ assetsInfo: {} }), ...overrides, }; } @@ -161,16 +172,30 @@ function createMockHandleRequest( accountAssets: string[] = [], balances: Record = {}, ): jest.Mock { - return jest.fn().mockImplementation((params) => { - const { request } = params; - if (request?.method === 'keyring_listAccountAssets') { - return Promise.resolve(accountAssets); - } - if (request?.method === 'keyring_getAccountBalances') { - return Promise.resolve(balances); - } - return Promise.resolve(null); - }); + return jest.fn( + (params: { + request?: { method?: string; params?: { assets?: string[] } }; + }) => { + const { request } = params; + if (request?.method === 'keyring_listAccountAssets') { + return Promise.resolve(accountAssets); + } + if (request?.method === 'keyring_getAccountBalances') { + const requested = request.params?.assets; + if (!requested) { + return Promise.resolve(balances); + } + return Promise.resolve( + Object.fromEntries( + requested + .filter((assetId) => balances[assetId]) + .map((assetId) => [assetId, balances[assetId]]), + ), + ); + } + return Promise.resolve(null); + }, + ); } function setupController( @@ -179,9 +204,22 @@ function setupController( accountAssets?: string[]; balances?: Record; configuredNetworks?: ChainId[]; + isBalanceV6Enabled?: () => boolean; + getAssetVisibility?: GetAssetVisibility; + getAssetsState?: () => AssetsControllerState; } = {}, ): SetupResult { - const { installedSnaps = {}, accountAssets = [], balances = {} } = options; + const { + installedSnaps = {}, + accountAssets = [], + balances = {}, + isBalanceV6Enabled = (): boolean => false, + getAssetVisibility = (): ReturnType => ({ + visibleAssetIds: [], + hiddenAssetIds: [], + }), + getAssetsState = getDefaultAssetsControllerState, + } = options; const rootMessenger = new Messenger({ namespace: MOCK_ANY_NAMESPACE, @@ -254,6 +292,9 @@ function setupController( messenger: controllerMessenger as unknown as AssetsControllerMessenger, onActiveChainsUpdated: activeChainsUpdateHandler, onAssetsUpdate: assetsUpdateHandler, + isBalanceV6Enabled, + getAssetVisibility, + getAssetsState, }; const controller = new SnapDataSource(controllerOptions); @@ -526,6 +567,274 @@ describe('SnapDataSource', () => { cleanup(); }); + it('requests expected visible assets and zero-fills any the snap omitted', async () => { + const getAssetVisibility = jest.fn().mockReturnValue({ + visibleAssetIds: [MOCK_SOL_ASSET, MOCK_SOL_PIN], + hiddenAssetIds: [], + }); + const { controller, mockHandleRequest, cleanup } = setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + accountAssets: [MOCK_SOL_ASSET], + balances: { + [MOCK_SOL_ASSET]: { amount: '1000000000', unit: 'SOL' }, + }, + isBalanceV6Enabled: () => true, + getAssetVisibility, + }); + await new Promise(process.nextTick); + + const response = await controller.fetch(createDataRequest()); + + expect(getAssetVisibility).toHaveBeenCalledWith( + ['mock-account-id'], + [SOLANA_MAINNET], + ); + expect(mockHandleRequest).toHaveBeenCalledWith( + expect.objectContaining({ + request: expect.objectContaining({ + method: 'keyring_getAccountBalances', + params: { + id: 'mock-account-id', + assets: [MOCK_SOL_ASSET, MOCK_SOL_PIN], + }, + }), + }), + ); + expect(response).toStrictEqual({ + assetsBalance: { + 'mock-account-id': { + [MOCK_SOL_ASSET]: { amount: '1000000000' }, + [MOCK_SOL_PIN]: { amount: '0' }, + }, + }, + assetsInfo: {}, + updateMode: 'full', + }); + + cleanup(); + }); + + it('seeds omitted Stellar visible tokens with empty trustline metadata', async () => { + const { controller, cleanup } = setupController({ + installedSnaps: { + [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_MAINNET] }, + }, + accountAssets: [MOCK_STELLAR_ASSET], + balances: { + [MOCK_STELLAR_ASSET]: { amount: '1000000000', unit: 'XLM' }, + }, + isBalanceV6Enabled: () => true, + getAssetVisibility: () => ({ + visibleAssetIds: [MOCK_STELLAR_ASSET, MOCK_STELLAR_PIN], + hiddenAssetIds: [], + }), + }); + await new Promise(process.nextTick); + + const response = await controller.fetch( + createDataRequest({ + chainIds: [STELLAR_MAINNET], + accounts: [ + createMockAccount({ + metadata: { + name: 'Test Account', + keyring: { type: 'HD Key Tree' }, + importTime: Date.now(), + lastSelected: Date.now(), + snap: { + id: STELLAR_SNAP_ID, + name: 'Stellar Snap', + enabled: true, + }, + }, + }), + ], + }), + ); + + expect(response.assetsBalance?.['mock-account-id']).toStrictEqual({ + [MOCK_STELLAR_ASSET]: { amount: '1000000000' }, + [MOCK_STELLAR_PIN]: { + amount: '0', + metadata: { + authorized: false, + limit: '0', + sponsored: false, + }, + }, + }); + + cleanup(); + }); + + it('does not zero-fill visible assets when the snap returns no balances', async () => { + const { controller, cleanup } = setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + accountAssets: [MOCK_SOL_ASSET], + balances: {}, + isBalanceV6Enabled: () => true, + getAssetVisibility: () => ({ + visibleAssetIds: [MOCK_SOL_ASSET, MOCK_SOL_PIN], + hiddenAssetIds: [], + }), + }); + await new Promise(process.nextTick); + + const response = await controller.fetch(createDataRequest()); + + expect(response).toStrictEqual({ + assetsBalance: {}, + assetsInfo: {}, + updateMode: 'full', + }); + + cleanup(); + }); + + it('does not write balances when getAccountBalances throws', async () => { + const { controller, mockHandleRequest, cleanup } = setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + accountAssets: [MOCK_SOL_ASSET], + balances: { + [MOCK_SOL_ASSET]: { amount: '1000000000', unit: 'SOL' }, + }, + isBalanceV6Enabled: () => true, + getAssetVisibility: () => ({ + visibleAssetIds: [MOCK_SOL_ASSET], + hiddenAssetIds: [], + }), + }); + mockHandleRequest.mockImplementation( + (params: { request?: { method: string } }) => { + if (params.request?.method === 'keyring_listAccountAssets') { + return Promise.resolve([MOCK_SOL_ASSET]); + } + if (params.request?.method === 'keyring_getAccountBalances') { + return Promise.reject(new Error('snap failed')); + } + return Promise.resolve(null); + }, + ); + await new Promise(process.nextTick); + + const response = await controller.fetch(createDataRequest()); + + expect(response).toStrictEqual({ + assetsBalance: {}, + assetsInfo: {}, + updateMode: 'full', + }); + + cleanup(); + }); + + it('fills an omitted visible asset from current state when the snap left it out', async () => { + const { controller, cleanup } = setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + accountAssets: [MOCK_SOL_ASSET], + balances: { + [MOCK_SOL_ASSET]: { amount: '1000000000', unit: 'SOL' }, + }, + isBalanceV6Enabled: () => true, + getAssetVisibility: () => ({ + visibleAssetIds: [MOCK_SOL_ASSET, MOCK_SOL_PIN], + hiddenAssetIds: [], + }), + getAssetsState: () => ({ + ...getDefaultAssetsControllerState(), + assetsBalance: { + 'mock-account-id': { + [MOCK_SOL_PIN]: { amount: '42' }, + }, + }, + }), + }); + await new Promise(process.nextTick); + + const response = await controller.fetch(createDataRequest()); + + expect(response.assetsBalance?.['mock-account-id']).toStrictEqual({ + [MOCK_SOL_ASSET]: { amount: '1000000000' }, + [MOCK_SOL_PIN]: { amount: '42' }, + }); + + cleanup(); + }); + + it('uses the snap balance for a visible asset omitted from listAccountAssets', async () => { + const { controller, cleanup } = setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + accountAssets: [MOCK_SOL_ASSET], + balances: { + [MOCK_SOL_ASSET]: { amount: '1000000000', unit: 'SOL' }, + [MOCK_SOL_PIN]: { amount: '42', unit: 'WSOL' }, + }, + isBalanceV6Enabled: () => true, + getAssetVisibility: () => ({ + visibleAssetIds: [MOCK_SOL_ASSET, MOCK_SOL_PIN], + hiddenAssetIds: [], + }), + }); + await new Promise(process.nextTick); + + const response = await controller.fetch(createDataRequest()); + + expect(response.assetsBalance?.['mock-account-id']).toStrictEqual({ + [MOCK_SOL_ASSET]: { amount: '1000000000' }, + [MOCK_SOL_PIN]: { amount: '42' }, + }); + + cleanup(); + }); + + it('does not request hidden assets on the v6 snap fetch', async () => { + const { controller, mockHandleRequest, cleanup } = setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + accountAssets: [MOCK_SOL_ASSET, MOCK_SOL_PIN], + balances: { + [MOCK_SOL_ASSET]: { amount: '1000000000', unit: 'SOL' }, + [MOCK_SOL_PIN]: { amount: '42', unit: 'WSOL' }, + }, + isBalanceV6Enabled: () => true, + getAssetVisibility: () => ({ + visibleAssetIds: [MOCK_SOL_ASSET], + hiddenAssetIds: [MOCK_SOL_PIN], + }), + }); + await new Promise(process.nextTick); + + const response = await controller.fetch(createDataRequest()); + + expect(mockHandleRequest).toHaveBeenCalledWith( + expect.objectContaining({ + request: expect.objectContaining({ + method: 'keyring_getAccountBalances', + params: { + id: 'mock-account-id', + assets: [MOCK_SOL_ASSET], + }, + }), + }), + ); + expect(response.assetsBalance?.['mock-account-id']).toStrictEqual({ + [MOCK_SOL_ASSET]: { amount: '1000000000' }, + }); + + cleanup(); + }); + it('fetch includes balance metadata when provided by the snap', async () => { const balanceMetadata = { spendable: '900000000', @@ -631,8 +940,10 @@ describe('SnapDataSource', () => { installedSnaps: { [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, + isBalanceV6Enabled: () => true, }); await new Promise(process.nextTick); + assetsUpdateHandler.mockClear(); triggerBalancesUpdated({ balances: { @@ -644,15 +955,15 @@ describe('SnapDataSource', () => { await new Promise(process.nextTick); - expect(assetsUpdateHandler).toHaveBeenCalledWith( - expect.objectContaining({ - assetsBalance: { - 'account-1': { - [MOCK_SOL_ASSET]: { amount: '1000000000' }, - }, + expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); + expect(assetsUpdateHandler).toHaveBeenCalledWith({ + updateMode: 'merge', + assetsBalance: { + 'account-1': { + [MOCK_SOL_ASSET]: { amount: '1000000000' }, }, - }), - ); + }, + }); cleanup(); }); @@ -948,6 +1259,37 @@ describe('SnapDataSource', () => { cleanup(); }); + it('middleware forwards full update mode for v6 snapshots', async () => { + const { controller, cleanup } = setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + accountAssets: [MOCK_SOL_ASSET], + balances: { + [MOCK_SOL_ASSET]: { amount: '1000000000', unit: 'SOL' }, + }, + isBalanceV6Enabled: () => true, + getAssetVisibility: () => ({ + visibleAssetIds: [MOCK_SOL_ASSET], + hiddenAssetIds: [], + }), + }); + await new Promise(process.nextTick); + + const next = jest.fn().mockImplementation(async (context) => context); + const context = createMiddlewareContext(); + + await controller.assetsMiddleware(context, next); + + expect(next).toHaveBeenCalledWith( + expect.objectContaining({ + response: expect.objectContaining({ updateMode: 'full' }), + }), + ); + + cleanup(); + }); + it('middleware removes handled chains from next request', async () => { const { controller, cleanup } = setupController({ installedSnaps: { @@ -1072,6 +1414,12 @@ describe('SnapDataSource', () => { messenger: controllerMessenger as unknown as AssetsControllerMessenger, onActiveChainsUpdated: jest.fn(), onAssetsUpdate: jest.fn(), + isBalanceV6Enabled: () => false, + getAssetVisibility: () => ({ + visibleAssetIds: [], + hiddenAssetIds: [], + }), + getAssetsState: getDefaultAssetsControllerState, }); await new Promise(process.nextTick); diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index e022e7fefda..a90c7f2b2d7 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -21,12 +21,15 @@ import type { AssetsControllerMessenger } from '../AssetsController.js'; import { projectLogger, createModuleLogger } from '../logger.js'; import type { AssetBalance, + AssetsControllerState, ChainId, Caip19AssetId, DataRequest, DataResponse, Middleware, } from '../types.js'; +import type { GetAssetVisibility } from '../utils/assetVisibility.js'; +import { getZeroAssetBalance } from '../utils/getZeroAssetBalance.js'; import { AbstractDataSource } from './AbstractDataSource.js'; import type { DataSourceState, @@ -173,6 +176,15 @@ export type SnapDataSourceOptions = { * controller currently has an active subscription for the chain. */ onAssetsUpdate: (response: DataResponse) => void | Promise; + /** Whether authoritative v6 balance snapshots are enabled. */ + isBalanceV6Enabled: () => boolean; + /** Resolves native, pinned, default, and hidden assets for a request scope. */ + getAssetVisibility: GetAssetVisibility; + /** + * Current AssetsController state. Used to keep the last-known amount when + * a v6 snap snapshot omits a visible asset. + */ + getAssetsState: () => AssetsControllerState; /** Configured networks to support (defaults to all snap networks) */ configuredNetworks?: ChainId[]; /** Default polling interval in ms for subscriptions */ @@ -217,6 +229,12 @@ export class SnapDataSource extends AbstractDataSource< readonly #onAssetsUpdate: (response: DataResponse) => void | Promise; + readonly #isBalanceV6Enabled: () => boolean; + + readonly #getAssetVisibility: GetAssetVisibility; + + readonly #getAssetsState: () => AssetsControllerState; + /** Bound handler for snap keyring balance updates, stored for cleanup */ readonly #handleSnapBalancesUpdatedBound: ( payload: AccountBalancesUpdatedEventPayload, @@ -236,6 +254,9 @@ export class SnapDataSource extends AbstractDataSource< this.#messenger = options.messenger; this.#onActiveChainsUpdated = options.onActiveChainsUpdated; this.#onAssetsUpdate = options.onAssetsUpdate; + this.#isBalanceV6Enabled = options.isBalanceV6Enabled; + this.#getAssetVisibility = options.getAssetVisibility; + this.#getAssetsState = options.getAssetsState; // Bind handlers for cleanup in destroy() this.#handleSnapBalancesUpdatedBound = @@ -452,6 +473,22 @@ export class SnapDataSource extends AbstractDataSource< if (!request?.chainIds?.length) { return {}; } + + if (this.#isBalanceV6Enabled()) { + return this.#fetchV6(request); + } + + return this.#fetchV5(request); + } + + /** + * v5 fetch. Unchanged from the pre-v6 handler; delete the v6 sibling first + * if `assetsAccountsApiV6` is rolled back. + * + * @param request - The data request. + * @returns Overlay (`merge`) balances from the snap keyring. + */ + async #fetchV5(request: DataRequest): Promise { if (!request?.accountsWithSupportedChains?.length) { return { assetsBalance: {}, assetsInfo: {}, updateMode: 'merge' }; } @@ -515,6 +552,188 @@ export class SnapDataSource extends AbstractDataSource< return results; } + /** + * v6 fetch: a complete snapshot of the requested account-chain slices. + * Listed holdings and expected visible assets (native, pin, default + * tracked) are requested together. Hidden assets are left out entirely. + * Delete with the rest of the v6 path if `assetsAccountsApiV6` is rolled + * back. + * + * @param request - The data request. + * @returns Authoritative (`full`) balances for the requested slices. + */ + async #fetchV6(request: DataRequest): Promise { + if (!request?.accountsWithSupportedChains?.length) { + return { assetsBalance: {}, assetsInfo: {}, updateMode: 'full' }; + } + + const results: DataResponse = { + assetsBalance: {}, + assetsInfo: {}, + updateMode: 'full', + }; + + // Fetch balances for each account using its snap ID from metadata + for (const { account } of request.accountsWithSupportedChains) { + // Skip accounts without snap metadata (non-snap accounts) + const snapId = account.metadata.snap?.id; + if (!snapId) { + continue; + } + + // Skip accounts whose snap doesn't support any of the requested chains + const supportedChainIds = request.chainIds.filter( + (chainId) => this.state.chainToSnap[chainId] === snapId, + ); + if (supportedChainIds.length === 0) { + continue; + } + + const accountId = account.id; + try { + const client = this.#getKeyringClient(snapId); + const { visibleAssetIds, hiddenAssetIds } = this.#getAssetVisibility( + [accountId], + supportedChainIds, + ); + + // Step 1: Get the list of assets for this account + const accountAssets = await client.listAccountAssets(accountId); + const assetsToFetch = this.#selectAssetsToFetchV6( + accountAssets ?? [], + visibleAssetIds, + supportedChainIds, + hiddenAssetIds, + ); + + // Step 2: Get balances for those specific assets. An empty map is a + // failed fetch (or nothing to ask for), not a zero snapshot. Leave + // this account out so `full` cannot replace last-known amounts with 0. + const balances: Record = assetsToFetch.length + ? await client.getAccountBalances(accountId, assetsToFetch) + : {}; + if ( + !balances || + typeof balances !== 'object' || + Object.keys(balances).length === 0 + ) { + continue; + } + + // Transform keyring response to DataResponse format + const accountBalances: Record = {}; + for (const [assetId, balance] of Object.entries(balances)) { + accountBalances[assetId] = { + amount: balance.amount, + ...(balance.metadata ? { metadata: balance.metadata } : {}), + }; + } + + // Step 3: Guard against an incomplete snap response. A `full` replace + // would drop any expected visible asset that was requested but omitted. + // Keep the last-known amount when we have one; otherwise seed 0. + this.#fillOmittedVisibleAssets( + accountId, + visibleAssetIds, + accountBalances, + ); + + if (results.assetsBalance) { + results.assetsBalance[accountId] = accountBalances; + } + } catch { + // Snap failed or the account does not belong to this snap. Contribute + // nothing so previous balances stay in state. + } + } + + return results; + } + + /** + * Put every omitted visible asset on the snapshot so `full` cannot drop it. + * Reuse the amount already in state when one exists; otherwise seed `0`. + * + * @param accountId - Account whose balances were fetched. + * @param visibleAssetIds - Native, pin, and default tracked IDs for this scope. + * @param accountBalances - Snapshot being built, mutated in place. + */ + #fillOmittedVisibleAssets( + accountId: string, + visibleAssetIds: Caip19AssetId[], + accountBalances: Record, + ): void { + const previousByKey = new Map( + Object.entries(this.#getAssetsState().assetsBalance[accountId] ?? {}).map( + ([assetId, balance]) => [assetId.toLowerCase(), balance], + ), + ); + const presentKeys = new Set( + Object.keys(accountBalances).map((assetId) => assetId.toLowerCase()), + ); + + for (const assetId of visibleAssetIds) { + if (presentKeys.has(assetId.toLowerCase())) { + continue; + } + + const previous = previousByKey.get(assetId.toLowerCase()); + if (previous) { + accountBalances[assetId] = { ...previous }; + continue; + } + + accountBalances[assetId] = getZeroAssetBalance(assetId); + } + } + + /** + * Union the snap's listed holdings with expected visible assets (native, + * pin, default tracked), limited to the requested chains. Hidden assets + * are dropped so they are neither fetched nor written back to state. + * + * @param accountAssets - Asset IDs the snap listed for the account. + * @param visibleAssetIds - Native, pin, and default tracked IDs for this scope. + * @param supportedChainIds - Chains this snap was asked about. + * @param hiddenAssetIds - Asset IDs the user hid in this scope. + * @returns Asset IDs to request balances for. + */ + #selectAssetsToFetchV6( + accountAssets: string[], + visibleAssetIds: Caip19AssetId[], + supportedChainIds: ChainId[], + hiddenAssetIds: Caip19AssetId[], + ): CaipAssetType[] { + const supportedChains = new Set(supportedChainIds); + const hidden = new Set( + hiddenAssetIds.map((assetId) => assetId.toLowerCase()), + ); + const assetsToFetch = new Map(); + + for (const assetId of [...accountAssets, ...visibleAssetIds]) { + const normalizedAssetId = assetId.toLowerCase(); + if ( + assetsToFetch.has(normalizedAssetId) || + hidden.has(normalizedAssetId) + ) { + continue; + } + + try { + if (!supportedChains.has(extractChainFromAssetId(assetId))) { + continue; + } + } catch { + // Skip unparseable asset IDs + continue; + } + + assetsToFetch.set(normalizedAssetId, assetId as CaipAssetType); + } + + return Array.from(assetsToFetch.values()); + } + // ============================================================================ // MIDDLEWARE // ============================================================================ @@ -577,6 +796,12 @@ export class SnapDataSource extends AbstractDataSource< ...response.assetsPrice, }; } + if (response.updateMode) { + context.response = { + ...context.response, + updateMode: response.updateMode, + }; + } // Determine successfully handled chains (exclude errors) const failedChains = new Set(Object.keys(response.errors ?? {})); diff --git a/packages/assets-controller/src/data-sources/StakedBalanceDataSource.test.ts b/packages/assets-controller/src/data-sources/StakedBalanceDataSource.test.ts index 36bb4074c12..77fae538041 100644 --- a/packages/assets-controller/src/data-sources/StakedBalanceDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/StakedBalanceDataSource.test.ts @@ -8,12 +8,7 @@ import { registerStakedMessengerActions, } from '../__fixtures__/MockAssetControllerMessenger.js'; import type { AssetsControllerMessenger } from '../AssetsController.js'; -import type { - AssetsControllerStateInternal, - ChainId, - Context, - DataRequest, -} from '../types.js'; +import type { ChainId, Context, DataRequest } from '../types.js'; import type { StakedBalanceDataSourceOptions } from './StakedBalanceDataSource.js'; import { StakedBalanceDataSource } from './StakedBalanceDataSource.js'; @@ -60,22 +55,10 @@ function createDataRequest( }; } -function getMockAssetsState(): AssetsControllerStateInternal { - return { - assetsInfo: {}, - assetsBalance: {}, - assetsPrice: {}, - customAssets: {}, - assetPreferences: {}, - selectedCurrency: 'usd', - }; -} - function createMiddlewareContext(overrides?: Partial): Context { return { request: createDataRequest(), response: {}, - getAssetsState: getMockAssetsState, ...overrides, }; } @@ -356,7 +339,6 @@ describe('StakedBalanceDataSource', () => { subscriptionId: 'test-sub', isUpdate: false, onAssetsUpdate, - getAssetsState: getMockAssetsState, }); await new Promise((resolve) => { setTimeout(resolve, 100); @@ -378,7 +360,6 @@ describe('StakedBalanceDataSource', () => { subscriptionId: 'test-sub', isUpdate: false, onAssetsUpdate, - getAssetsState: getMockAssetsState, }); expect(onAssetsUpdate).not.toHaveBeenCalled(); }, @@ -394,7 +375,6 @@ describe('StakedBalanceDataSource', () => { subscriptionId: 'test-sub', isUpdate: false, onAssetsUpdate: jest.fn(), - getAssetsState: getMockAssetsState, }); await controller.unsubscribe('test-sub'); const chains = await controller.getActiveChains(); @@ -414,7 +394,6 @@ describe('StakedBalanceDataSource', () => { subscriptionId: 'test-sub', isUpdate: false, onAssetsUpdate, - getAssetsState: getMockAssetsState, }); await new Promise((resolve) => setTimeout(resolve, 100)); onAssetsUpdate.mockClear(); @@ -501,7 +480,6 @@ describe('StakedBalanceDataSource', () => { subscriptionId: 'test-sub', isUpdate: false, onAssetsUpdate, - getAssetsState: getMockAssetsState, }); onAssetsUpdate.mockClear(); expect(await controller.refreshStakedBalance()).toBeUndefined(); diff --git a/packages/assets-controller/src/data-sources/TokenDataSource.test.ts b/packages/assets-controller/src/data-sources/TokenDataSource.test.ts index e9a2ef1723d..c3b511c35f4 100644 --- a/packages/assets-controller/src/data-sources/TokenDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/TokenDataSource.test.ts @@ -8,7 +8,13 @@ import type { import { TokenScanResultType } from '@metamask/phishing-controller'; import type { AssetsControllerMessenger } from '../AssetsController.js'; -import type { Context, DataRequest, Caip19AssetId, ChainId } from '../types.js'; +import type { + AssetsControllerState, + Context, + DataRequest, + Caip19AssetId, + ChainId, +} from '../types.js'; import type { TokenDataSourceOptions } from './TokenDataSource.js'; import { TokenDataSource } from './TokenDataSource.js'; @@ -131,9 +137,6 @@ function createMiddlewareContext(overrides?: Partial): Context { return { request: createDataRequest(), response: {}, - getAssetsState: jest.fn().mockReturnValue({ - assetsInfo: {}, - }), ...overrides, }; } @@ -144,6 +147,7 @@ function setupController(options: { assetsResponse?: V3AssetResponse[]; nativeAssetIds?: string[]; suggestedOccurrenceFloors?: Record; + getAssetsState?: () => Partial; }): SetupResult { const { messenger, @@ -151,6 +155,9 @@ function setupController(options: { assetsResponse = [], nativeAssetIds = [], suggestedOccurrenceFloors = { '1': 3 }, + getAssetsState = (): Partial => ({ + assetsInfo: {}, + }), } = options; const apiClient = createMockApiClient( @@ -176,6 +183,7 @@ function setupController(options: { } return 'erc20'; }, + getAssetsState: getAssetsState as () => AssetsControllerState, }); return { @@ -429,16 +437,7 @@ describe('TokenDataSource', () => { const { controller, apiClient } = setupController({ messenger: createTestMessenger(), supportedNetworks: ['eip155:1'], - }); - - const next = jest.fn().mockResolvedValue(undefined); - const context = createMiddlewareContext({ - response: { - detectedAssets: { - 'mock-account-id': [MOCK_TOKEN_ASSET], - }, - }, - getAssetsState: jest.fn().mockReturnValue({ + getAssetsState: (): Partial => ({ assetsInfo: { [MOCK_TOKEN_ASSET]: { type: 'erc20', @@ -451,6 +450,15 @@ describe('TokenDataSource', () => { }), }); + const next = jest.fn().mockResolvedValue(undefined); + const context = createMiddlewareContext({ + response: { + detectedAssets: { + 'mock-account-id': [MOCK_TOKEN_ASSET], + }, + }, + }); + await controller.assetsMiddleware(context, next); expect(apiClient.tokens.fetchV3Assets).not.toHaveBeenCalled(); @@ -1037,6 +1045,10 @@ describe('TokenDataSource', () => { createMockAssetResponse(spamAsset, { occurrences: 1 }), ], suggestedOccurrenceFloors: { '1': 3 }, + getAssetsState: (): Partial => ({ + assetsBalance: {}, + assetsInfo: {}, + }), }); const next = jest.fn().mockResolvedValue(undefined); @@ -1060,10 +1072,6 @@ describe('TokenDataSource', () => { }, }, }, - getAssetsState: jest.fn().mockReturnValue({ - assetsBalance: {}, - assetsInfo: {}, - }), }); await controller.occurrenceFilterMiddleware(context, next); @@ -1090,6 +1098,14 @@ describe('TokenDataSource', () => { createMockAssetResponse(MOCK_TOKEN_ASSET, { occurrences: 1 }), ], suggestedOccurrenceFloors: { '1': 3 }, + getAssetsState: (): Partial => ({ + assetsBalance: { + 'mock-account-id': { + [MOCK_TOKEN_ASSET_CHECKSUMMED]: { amount: '42' }, + }, + }, + assetsInfo: {}, + }), }); const next = jest.fn().mockResolvedValue(undefined); @@ -1104,14 +1120,6 @@ describe('TokenDataSource', () => { }, }, }, - getAssetsState: jest.fn().mockReturnValue({ - assetsBalance: { - 'mock-account-id': { - [MOCK_TOKEN_ASSET_CHECKSUMMED]: { amount: '42' }, - }, - }, - assetsInfo: {}, - }), }); await controller.occurrenceFilterMiddleware(context, next); @@ -1135,6 +1143,17 @@ describe('TokenDataSource', () => { createMockAssetResponse(MOCK_TOKEN_ASSET, { occurrences: 1 }), ], suggestedOccurrenceFloors: { '1': 3 }, + getAssetsState: (): Partial => ({ + assetsBalance: {}, + assetsInfo: { + [MOCK_TOKEN_ASSET_CHECKSUMMED]: { + type: 'erc20', + name: 'USD Coin', + symbol: 'USDC', + decimals: 6, + }, + }, + }), }); const next = jest.fn().mockResolvedValue(undefined); @@ -1147,17 +1166,6 @@ describe('TokenDataSource', () => { }, }, }, - getAssetsState: jest.fn().mockReturnValue({ - assetsBalance: {}, - assetsInfo: { - [MOCK_TOKEN_ASSET_CHECKSUMMED]: { - type: 'erc20', - name: 'USD Coin', - symbol: 'USDC', - decimals: 6, - }, - }, - }), }); await controller.occurrenceFilterMiddleware(context, next); @@ -1180,6 +1188,15 @@ describe('TokenDataSource', () => { createMockAssetResponse(MOCK_TOKEN_ASSET, { occurrences: 1 }), ], suggestedOccurrenceFloors: { '1': 3 }, + getAssetsState: (): Partial => ({ + assetsBalance: { + 'mock-account-id': { + [MOCK_TOKEN_ASSET_CHECKSUMMED]: { amount: '0' }, + }, + }, + assetsInfo: {}, + customAssets: { 'mock-account-id': [] }, + }), }); const next = jest.fn().mockResolvedValue(undefined); @@ -1195,15 +1212,6 @@ describe('TokenDataSource', () => { }, }, }, - getAssetsState: jest.fn().mockReturnValue({ - assetsBalance: { - 'mock-account-id': { - [MOCK_TOKEN_ASSET_CHECKSUMMED]: { amount: '0' }, - }, - }, - assetsInfo: {}, - customAssets: { 'mock-account-id': [] }, - }), }); await controller.occurrenceFilterMiddleware(context, next); @@ -1226,6 +1234,11 @@ describe('TokenDataSource', () => { createMockAssetResponse(MOCK_TOKEN_ASSET, { occurrences: 1 }), ], suggestedOccurrenceFloors: { '1': 3 }, + getAssetsState: (): Partial => ({ + assetsBalance: {}, + assetsInfo: {}, + customAssets: { 'mock-account-id': [MOCK_TOKEN_ASSET_CHECKSUMMED] }, + }), }); const next = jest.fn().mockResolvedValue(undefined); @@ -1238,11 +1251,6 @@ describe('TokenDataSource', () => { }, }, }, - getAssetsState: jest.fn().mockReturnValue({ - assetsBalance: {}, - assetsInfo: {}, - customAssets: { 'mock-account-id': [MOCK_TOKEN_ASSET_CHECKSUMMED] }, - }), }); await controller.occurrenceFilterMiddleware(context, next); @@ -1359,23 +1367,23 @@ describe('TokenDataSource', () => { assetsResponse: [ createMockAssetResponse(lowercaseCustomAsset, { occurrences: 1 }), ], + getAssetsState: (): Partial => ({ + assetsInfo: {}, + customAssets: { + 'mock-account-id': [checksummedCustomAsset], + }, + }), }); const next = jest.fn().mockResolvedValue(undefined); - const context: Context = { + const context = createMiddlewareContext({ request: createDataRequest(), response: { detectedAssets: { 'mock-account-id': [checksummedCustomAsset], }, }, - getAssetsState: jest.fn().mockReturnValue({ - assetsInfo: {}, - customAssets: { - 'mock-account-id': [checksummedCustomAsset], - }, - }), - }; + }); await controller.assetsMiddleware(context, next); diff --git a/packages/assets-controller/src/data-sources/TokenDataSource.ts b/packages/assets-controller/src/data-sources/TokenDataSource.ts index 013fe6d115d..224dd6eae7f 100644 --- a/packages/assets-controller/src/data-sources/TokenDataSource.ts +++ b/packages/assets-controller/src/data-sources/TokenDataSource.ts @@ -14,6 +14,7 @@ import { forDataTypes } from '../types.js'; import type { Caip19AssetId, AssetMetadata, + AssetsControllerState, Middleware, FungibleAssetMetadata, DataResponse, @@ -71,6 +72,8 @@ export type TokenDataSourceOptions = { * fires, the batch rejects so metadata enrichment proceeds without it. */ fetchTimeoutMs?: number; + /** Current AssetsController state. Used to skip already-known metadata and custom assets. */ + getAssetsState: () => AssetsControllerState; }; /** @@ -226,6 +229,8 @@ export class TokenDataSource { readonly #fetchTimeoutMs: number; + readonly #getAssetsState: () => AssetsControllerState; + constructor( messenger: AssetsControllerMessenger, options: TokenDataSourceOptions, @@ -235,6 +240,7 @@ export class TokenDataSource { this.#getNativeAssetIds = options.getNativeAssetIds; this.#getAssetType = options.getAssetType; this.#fetchTimeoutMs = options.fetchTimeoutMs ?? DEFAULT_FETCH_TIMEOUT_MS; + this.#getAssetsState = options.getAssetsState; } /** @@ -406,7 +412,7 @@ export class TokenDataSource { assetsBalance: stateBalances, assetsInfo: stateMetadata, customAssets, - } = ctx.getAssetsState(); + } = this.#getAssetsState(); const customAssetIds = new Set( Object.values(customAssets ?? {}) @@ -538,7 +544,8 @@ export class TokenDataSource { // Extract response from context const { response } = ctx; - const { assetsInfo: stateMetadata, customAssets } = ctx.getAssetsState(); + const { assetsInfo: stateMetadata, customAssets } = + this.#getAssetsState(); const assetIdsNeedingMetadata = new Set(); // Newly detected asset IDs (lowercase) — subject to spam filtering. const detectedAssetIds = new Set(); diff --git a/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.test.ts b/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.test.ts index 85585253441..119411b105c 100644 --- a/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.test.ts +++ b/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.test.ts @@ -1,5 +1,6 @@ import type { CaipAssetType } from '@metamask/utils'; +import { getDefaultTrackedAssetsForChain } from '../../../defaults.js'; import type { MulticallClient } from '../clients/index.js'; import type { Address, @@ -11,7 +12,6 @@ import type { import { BalanceFetcher } from './BalanceFetcher.js'; import type { BalanceFetcherConfig, - BalanceFetcherMessenger, BalancePollingInput, } from './BalanceFetcher.js'; @@ -75,14 +75,10 @@ function createMockAssetsBalanceState( }; } -function createMockMessenger( +function createMockGetAssetsState( assetsBalanceState?: AssetsBalanceState, -): BalanceFetcherMessenger { - return { - call: (_action: 'AssetsController:getState'): AssetsBalanceState => { - return assetsBalanceState ?? { assetsBalance: {} }; - }, - }; +): () => AssetsBalanceState { + return () => assetsBalanceState ?? { assetsBalance: {} }; } function createMockBalanceResponse( @@ -99,7 +95,16 @@ function createMockBalanceResponse( // ============================================================================= type WithControllerOptions = { - config?: BalanceFetcherConfig; + config?: Omit< + BalanceFetcherConfig, + 'getAssetsState' | 'getAssetVisibility' | 'isBalanceV6Enabled' + > & + Partial< + Pick< + BalanceFetcherConfig, + 'getAssetsState' | 'getAssetVisibility' | 'isBalanceV6Enabled' + > + >; assetsBalanceState?: AssetsBalanceState; }; @@ -125,14 +130,18 @@ async function withController( config = { isNativeAsset: (): boolean => false }, assetsBalanceState, } = options; + const resolvedConfig: BalanceFetcherConfig = { + getAssetsState: createMockGetAssetsState(assetsBalanceState), + isBalanceV6Enabled: () => false, + getAssetVisibility: () => ({ + visibleAssetIds: [], + hiddenAssetIds: [], + }), + ...config, + }; const mockMulticallClient = createMockMulticallClient(); - const mockMessenger = createMockMessenger(assetsBalanceState); - const controller = new BalanceFetcher( - mockMulticallClient, - mockMessenger, - config, - ); + const controller = new BalanceFetcher(mockMulticallClient, resolvedConfig); try { return await fn({ controller, mockMulticallClient }); @@ -312,6 +321,173 @@ describe('BalanceFetcher', () => { ); }); + it('on the v6 path polls default tracked assets even when they have no assetsBalance entry', async () => { + const [mainnetMusd] = getDefaultTrackedAssetsForChain('eip155:1'); + const musdAddress = mainnetMusd.split(':').at(-1)?.toLowerCase() ?? ''; + + const mockState: AssetsBalanceState = { + assetsBalance: { + [TEST_ACCOUNT_ID]: { + [NATIVE_ETH_ASSET_ID]: { amount: '0' }, + }, + }, + }; + + await withController( + { + assetsBalanceState: mockState, + config: { + isNativeAsset: (id: CaipAssetType) => id === NATIVE_ETH_ASSET_ID, + isBalanceV6Enabled: (): boolean => true, + getAssetVisibility: () => ({ + visibleAssetIds: [NATIVE_ETH_ASSET_ID, mainnetMusd], + hiddenAssetIds: [], + }), + }, + }, + async ({ controller, mockMulticallClient }) => { + controller.setOnBalanceUpdate(jest.fn()); + mockMulticallClient.batchBalanceOf.mockResolvedValue([ + createMockBalanceResponse( + ZERO_ADDRESS, + TEST_ACCOUNT, + true, + '1000000000000000000', + ), + createMockBalanceResponse( + musdAddress as Address, + TEST_ACCOUNT, + true, + '0', + ), + ]); + + await controller._executePoll({ + chainId: MAINNET_CHAIN_ID, + accountId: TEST_ACCOUNT_ID, + accountAddress: TEST_ACCOUNT, + }); + + const [, batchedRequests] = + mockMulticallClient.batchBalanceOf.mock.calls[0]; + const requestedTokens = ( + batchedRequests as { tokenAddress: string }[] + ).map((req) => req.tokenAddress.toLowerCase()); + expect(requestedTokens).toContain(ZERO_ADDRESS.toLowerCase()); + expect(requestedTokens).toContain(musdAddress); + }, + ); + }); + + it('on the v6 path stops polling a tracked asset the user hid', async () => { + const mockState: AssetsBalanceState = { + assetsBalance: { + [TEST_ACCOUNT_ID]: { + [NATIVE_ETH_ASSET_ID]: { amount: '1' }, + [TOKEN_1_ASSET_ID]: { amount: '500' }, + }, + }, + }; + + await withController( + { + assetsBalanceState: mockState, + config: { + isNativeAsset: (id: CaipAssetType) => id === NATIVE_ETH_ASSET_ID, + isBalanceV6Enabled: (): boolean => true, + getAssetVisibility: () => ({ + visibleAssetIds: [NATIVE_ETH_ASSET_ID], + hiddenAssetIds: [TOKEN_1_ASSET_ID], + }), + }, + }, + async ({ controller, mockMulticallClient }) => { + controller.setOnBalanceUpdate(jest.fn()); + mockMulticallClient.batchBalanceOf.mockResolvedValue([ + createMockBalanceResponse( + ZERO_ADDRESS, + TEST_ACCOUNT, + true, + '1000000000000000000', + ), + ]); + + await controller._executePoll({ + chainId: MAINNET_CHAIN_ID, + accountId: TEST_ACCOUNT_ID, + accountAddress: TEST_ACCOUNT, + }); + + const [, batchedRequests] = + mockMulticallClient.batchBalanceOf.mock.calls[0]; + const requestedTokens = ( + batchedRequests as { tokenAddress: string }[] + ).map((req) => req.tokenAddress.toLowerCase()); + expect(requestedTokens).toStrictEqual([ZERO_ADDRESS.toLowerCase()]); + }, + ); + }); + + it('on the v6 path does not poll staking vault share tokens from assetsBalance', async () => { + const stakingAssetId = + 'eip155:1/erc20:0x4fef9d741011476750a243ac70b9789a63dd47df' as CaipAssetType; + const stakingAddress = + '0x4fef9d741011476750a243ac70b9789a63dd47df' as Address; + const mockState: AssetsBalanceState = { + assetsBalance: { + [TEST_ACCOUNT_ID]: { + [NATIVE_ETH_ASSET_ID]: { amount: '1' }, + [stakingAssetId]: { amount: '2' }, + [TOKEN_1_ASSET_ID]: { amount: '500' }, + }, + }, + }; + + await withController( + { + assetsBalanceState: mockState, + config: { + isNativeAsset: (id: CaipAssetType) => id === NATIVE_ETH_ASSET_ID, + isBalanceV6Enabled: (): boolean => true, + getAssetVisibility: () => ({ + visibleAssetIds: [NATIVE_ETH_ASSET_ID], + hiddenAssetIds: [], + }), + }, + }, + async ({ controller, mockMulticallClient }) => { + controller.setOnBalanceUpdate(jest.fn()); + mockMulticallClient.batchBalanceOf.mockResolvedValue([ + createMockBalanceResponse( + ZERO_ADDRESS, + TEST_ACCOUNT, + true, + '1000000000000000000', + ), + createMockBalanceResponse(TEST_TOKEN_1, TEST_ACCOUNT, true, '500'), + ]); + + await controller._executePoll({ + chainId: MAINNET_CHAIN_ID, + accountId: TEST_ACCOUNT_ID, + accountAddress: TEST_ACCOUNT, + }); + + const [, batchedRequests] = + mockMulticallClient.batchBalanceOf.mock.calls[0]; + const requestedTokens = ( + batchedRequests as { tokenAddress: string }[] + ) + .map((req) => req.tokenAddress.toLowerCase()) + .sort(); + expect(requestedTokens).toStrictEqual( + [ZERO_ADDRESS.toLowerCase(), TEST_TOKEN_1.toLowerCase()].sort(), + ); + expect(requestedTokens).not.toContain(stakingAddress); + }, + ); + }); + it('in customAssetsOnly mode skips state.assetsBalance and only fetches state.customAssets', async () => { // The supplemental subscription path: another data source covers the // chain for regular balances, but RPC must still poll the user's diff --git a/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.ts b/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.ts index 44d8af36fa1..d86faa982b5 100644 --- a/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.ts +++ b/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.ts @@ -1,6 +1,8 @@ import { StaticIntervalPollingControllerOnly } from '@metamask/polling-controller'; import { parseCaipAssetType } from '@metamask/utils'; +import type { ChainId as Caip2ChainId } from '../../../types.js'; +import type { GetAssetVisibility } from '../../../utils/assetVisibility.js'; import { ZERO_ADDRESS } from '../../../utils/constants.js'; import type { MulticallClient } from '../clients/index.js'; import type { @@ -15,24 +17,29 @@ import type { CaipAssetType, ChainId, } from '../types/index.js'; -import { reduceInBatchesSerially } from '../utils/index.js'; +import { + isStakingContractAssetId, + reduceInBatchesSerially, +} from '../utils/index.js'; const DEFAULT_BALANCE_INTERVAL = 30_000; // 30 seconds -/** - * Minimal messenger interface for BalanceFetcher. - */ -export type BalanceFetcherMessenger = { - call: (action: 'AssetsController:getState') => AssetsBalanceState; -}; - export type BalanceFetcherConfig = { defaultBatchSize?: number; defaultTimeoutMs?: number; /** Polling interval in ms (default: 30s) */ pollingInterval?: number; + /** Current AssetsController balance/custom-asset state. */ + getAssetsState: () => AssetsBalanceState; /** Determines whether a CAIP-19 asset ID represents a native asset. */ isNativeAsset: (assetId: CaipAssetType) => boolean; + /** + * When true, polls default tracked assets (e.g. mUSD) in addition to + * `assetsBalance` and `customAssets`. v5 omits this and stays pins + balances. + */ + isBalanceV6Enabled: () => boolean; + /** Returns shared visible/hidden assets for an account/chain scope. */ + getAssetVisibility: GetAssetVisibility; }; /** @@ -73,29 +80,38 @@ export type OnBalanceUpdateCallback = ( export class BalanceFetcher extends StaticIntervalPollingControllerOnly() { readonly #multicallClient: MulticallClient; - readonly #messenger: BalanceFetcherMessenger; - readonly #config: Required< - Omit + Omit< + BalanceFetcherConfig, + | 'pollingInterval' + | 'getAssetsState' + | 'isNativeAsset' + | 'isBalanceV6Enabled' + | 'getAssetVisibility' + > >; + readonly #getAssetsState: () => AssetsBalanceState; + readonly #isNativeAsset: (assetId: CaipAssetType) => boolean; + readonly #isBalanceV6Enabled: () => boolean; + + readonly #getAssetVisibility: GetAssetVisibility; + #onBalanceUpdate: OnBalanceUpdateCallback | undefined; - constructor( - multicallClient: MulticallClient, - messenger: BalanceFetcherMessenger, - config: BalanceFetcherConfig, - ) { + constructor(multicallClient: MulticallClient, config: BalanceFetcherConfig) { super(); this.#multicallClient = multicallClient; - this.#messenger = messenger; this.#config = { defaultBatchSize: config.defaultBatchSize ?? 300, defaultTimeoutMs: config.defaultTimeoutMs ?? 30000, }; + this.#getAssetsState = config.getAssetsState; this.#isNativeAsset = config.isNativeAsset; + this.#isBalanceV6Enabled = config.isBalanceV6Enabled; + this.#getAssetVisibility = config.getAssetVisibility; // Set the polling interval this.setIntervalLength(config?.pollingInterval ?? DEFAULT_BALANCE_INTERVAL); @@ -143,7 +159,7 @@ export class BalanceFetcher extends StaticIntervalPollingControllerOnly