Skip to content
This repository was archived by the owner on Jul 31, 2026. It is now read-only.

feat: push special asset snapshots to controller #362

Closed
ulissesferreira wants to merge 2 commits into from
Closed

feat: push special asset snapshots to controller #362

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aaa25ae
feat: gate controller-only asset tracking
ulissesferreira Jul 15, 2026
d236615
feat: push special asset snapshots to controller
ulissesferreira Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/snap/snap.manifest.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
"url": "https://github.com/MetaMask/snap-tron-wallet.git"
},
"source": {
"shasum": "tSxhnKoJe+iTI1HJ13ENq81vPK99Q12rj7owcUCpqxQ=",
"shasum": "KWVrILlsFNeR5YoP7RLalzLse8MDPlTzrXn7tihjabA=",
"location": {
"npm": {
"filePath": "dist/bundle.js",
Expand Down
176 changes: 174 additions & 2 deletions packages/snap/src/services/accounts/AccountsService.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@ import { mockLogger } from '../../utils/mockLogger';
import type { AssetsService } from '../assets/AssetsService';
import type { ConfigProvider } from '../config';
import type { Config } from '../config/ConfigProvider';
import { MigrationStage } from '../migration/stage';
import type { TronAssetsControllerAdapter } from '../migration/TronAssetsControllerAdapter';
import type { TransactionsService } from '../transactions/TransactionsService';

jest.mock('@metamask/keyring-snap-sdk', () => ({
Expand Down Expand Up @@ -120,14 +122,22 @@ type WithAccountsServiceCallback = (payload: {
mockConfigProvider: jest.Mocked<Pick<ConfigProvider, 'get'>>;
mockLogger: ILogger;
mockAssetsService: jest.Mocked<
Pick<AssetsService, 'fetchAssetsAndBalancesForAccount' | 'saveMany'>
Pick<
AssetsService,
| 'fetchAssetsAndBalancesForAccount'
| 'saveMany'
| 'buildControllerSnapshot'
>
>;
mockSnapClient: jest.Mocked<
Pick<SnapClient, 'getBip32Entropy' | 'listEntropySources'>
>;
mockTransactionsService: jest.Mocked<
Pick<TransactionsService, 'fetchNewTransactionsForAccount' | 'saveMany'>
>;
mockTronAssetsControllerAdapter: jest.Mocked<
Pick<TronAssetsControllerAdapter, 'getMigrationStage' | 'pushAssetSnapshot'>
>;
}) => void | Promise<void>;

/**
Expand Down Expand Up @@ -265,10 +275,16 @@ async function withAccountsService(
};

const mockAssetsService: jest.Mocked<
Pick<AssetsService, 'fetchAssetsAndBalancesForAccount' | 'saveMany'>
Pick<
AssetsService,
| 'fetchAssetsAndBalancesForAccount'
| 'saveMany'
| 'buildControllerSnapshot'
>
> = {
fetchAssetsAndBalancesForAccount: jest.fn().mockResolvedValue([]),
saveMany: jest.fn().mockResolvedValue(undefined),
buildControllerSnapshot: jest.fn().mockReturnValue([]),
};

const mockTransactionsService: jest.Mocked<
Expand All @@ -278,13 +294,21 @@ async function withAccountsService(
saveMany: jest.fn().mockResolvedValue(undefined),
};

const mockTronAssetsControllerAdapter: jest.Mocked<
Pick<TronAssetsControllerAdapter, 'getMigrationStage' | 'pushAssetSnapshot'>
> = {
getMigrationStage: jest.fn().mockResolvedValue(MigrationStage.Off),
pushAssetSnapshot: jest.fn().mockResolvedValue(undefined),
};

const accountsService = new AccountsService({
accountsRepository: mockAccountsRepository,
configProvider: mockConfigProvider,
logger: mockLogger,
assetsService: mockAssetsService,
snapClient: mockSnapClient,
transactionsService: mockTransactionsService,
tronAssetsControllerAdapter: mockTronAssetsControllerAdapter,
} as unknown as ConstructorParameters<typeof AccountsService>[0]);

await testFn({
Expand All @@ -295,6 +319,7 @@ async function withAccountsService(
mockAssetsService,
mockSnapClient,
mockTransactionsService,
mockTronAssetsControllerAdapter,
});
}

Expand Down Expand Up @@ -1147,6 +1172,153 @@ describe('AccountsService', () => {
},
);
});

it('does not write controller snapshots at stage Off', async () => {
const account: TronKeyringAccount = {
id: 'off-id',
address: 'TOff123456789012345678901',
type: TrxAccountType.Eoa,
options: {},
methods: [],
scopes: [],
entropySource: 'e1',
derivationPath: "m/44'/195'/0'/0/0",
index: 0,
};
await withAccountsService(
async ({
accountsService,
mockConfigProvider,
mockAssetsService,
mockTronAssetsControllerAdapter,
}) => {
mockConfigProvider.get.mockReturnValue({
...MOCK_CONFIG,
activeNetworks: [Network.Mainnet],
});
mockAssetsService.fetchAssetsAndBalancesForAccount.mockResolvedValue(
[],
);

await accountsService.synchronizeAssets([account]);

expect(
mockTronAssetsControllerAdapter.pushAssetSnapshot,
).not.toHaveBeenCalled();
},
);
});

it('writes snapshots at stage 1 and continues when a write fails', async () => {
const account: TronKeyringAccount = {
id: 'write-id',
address: 'TWrite12345678901234567890',
type: TrxAccountType.Eoa,
options: {},
methods: [],
scopes: [],
entropySource: 'e1',
derivationPath: "m/44'/195'/0'/0/0",
index: 0,
};
const snapshot = [
[
{
assetId: 'asset',
amount: '1',
metadata: { symbol: 'A', name: 'Asset', decimals: 0 },
},
],
];

await withAccountsService(
async ({
accountsService,
mockConfigProvider,
mockAssetsService,
mockTronAssetsControllerAdapter,
}) => {
mockConfigProvider.get.mockReturnValue({
...MOCK_CONFIG,
activeNetworks: [Network.Mainnet],
});
mockTronAssetsControllerAdapter.getMigrationStage.mockResolvedValue(
MigrationStage.ReadAssetsControllerWithFallback,
);
mockAssetsService.buildControllerSnapshot.mockReturnValue(snapshot);
mockTronAssetsControllerAdapter.pushAssetSnapshot.mockRejectedValue(
new Error('controller unavailable'),
);

expect(
await accountsService.synchronizeAssets([account]),
).toBeUndefined();

expect(
mockTronAssetsControllerAdapter.pushAssetSnapshot,
).toHaveBeenCalledWith(account.id, Network.Mainnet, snapshot[0]);
},
);
});

it('does not persist controller-only assets and resumes after rollback', async () => {
const account: TronKeyringAccount = {
id: 'controller-only-id',
address: 'TControllerOnly1234567890',
type: TrxAccountType.Eoa,
options: {},
methods: [],
scopes: [],
entropySource: 'e1',
derivationPath: "m/44'/195'/0'/0/0",
index: 0,
};

await withAccountsService(
async ({
accountsService,
mockConfigProvider,
mockAssetsService,
mockTronAssetsControllerAdapter,
}) => {
const assets: NativeAsset[] = [
{
assetType: `${Network.Mainnet}/slip44:195`,
keyringAccountId: account.id,
network: Network.Mainnet,
symbol: 'TRX',
decimals: 6,
rawAmount: '1000000',
uiAmount: '1',
iconUrl: '',
},
];
mockConfigProvider.get.mockReturnValue({
...MOCK_CONFIG,
activeNetworks: [Network.Mainnet],
});
mockAssetsService.fetchAssetsAndBalancesForAccount.mockResolvedValue(
assets,
);
mockTronAssetsControllerAdapter.getMigrationStage.mockResolvedValueOnce(
MigrationStage.ReadAssetsControllerOnly,
);

await accountsService.synchronizeAssets([account]);

expect(mockAssetsService.saveMany).not.toHaveBeenCalled();

mockAssetsService.saveMany.mockClear();
mockTronAssetsControllerAdapter.getMigrationStage.mockResolvedValueOnce(
MigrationStage.Off,
);

await accountsService.synchronizeAssets([account]);

expect(mockAssetsService.saveMany).toHaveBeenCalledWith(assets);
},
);
});
});

describe('synchronizeTransactions', () => {
Expand Down
101 changes: 98 additions & 3 deletions packages/snap/src/services/accounts/AccountsService.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ import { createPrefixedLogger, type ILogger } from '../../utils/logger';
import { DerivationPathStruct } from '../../validation/structs';
import type { AssetsService } from '../assets/AssetsService';
import type { ConfigProvider } from '../config';
import { MigrationStage, shouldPushSnapAssets } from '../migration/stage';
import type { TronAssetsControllerAdapter } from '../migration/TronAssetsControllerAdapter';
import type { TransactionsService } from '../transactions/TransactionsService';

/**
Expand Down Expand Up @@ -110,27 +112,32 @@ export class AccountsService {

readonly #snapClient: SnapClient;

readonly #tronAssetsControllerAdapter: TronAssetsControllerAdapter;

constructor({
accountsRepository,
configProvider,
logger,
assetsService,
snapClient,
transactionsService,
tronAssetsControllerAdapter,
}: {
accountsRepository: AccountsRepository;
configProvider: ConfigProvider;
logger: ILogger;
assetsService: AssetsService;
snapClient: SnapClient;
transactionsService: TransactionsService;
tronAssetsControllerAdapter: TronAssetsControllerAdapter;
}) {
this.#logger = createPrefixedLogger(logger, '[🔑 AccountsService]');
this.#configProvider = configProvider;
this.#accountsRepository = accountsRepository;
this.#assetsService = assetsService;
this.#transactionsService = transactionsService;
this.#snapClient = snapClient;
this.#tronAssetsControllerAdapter = tronAssetsControllerAdapter;
}

/**
Expand Down Expand Up @@ -509,11 +516,99 @@ export class AccountsService {
}),
);

const assets = assetResponses.flatMap((response) =>
response.status === 'fulfilled' ? response.value : [],
const synchronizedAssets = await Promise.all(
assetResponses.map(async (response, index) => {
if (response.status !== 'fulfilled') {
return { assets: [], shouldPersist: false, combination: undefined };
}

const combination = combinations[index];

if (!combination) {
return { assets: [], shouldPersist: false, combination: undefined };
}

const { scope } = combination;
const stage =
await this.#tronAssetsControllerAdapter.getMigrationStage(scope);

if (stage >= MigrationStage.ReadAssetsControllerOnly) {
this.#logger.info(
'Asset tracking and persistence disabled for controller-only stage',
{ scope, stage },
);
return {
assets: response.value,
shouldPersist: false,
combination: { ...combination, stage },
};
}

return {
assets: response.value,
shouldPersist: true,
combination: { ...combination, stage },
};
}),
);
const assets = synchronizedAssets.flatMap(
({ assets: responseAssets }) => responseAssets,
);
const hasPersistableScope = synchronizedAssets.some(
({ shouldPersist }) => shouldPersist,
);

if (hasPersistableScope || combinations.length === 0) {
await this.#assetsService.saveMany(assets);
}

await this.#assetsService.saveMany(assets);
const writes = synchronizedAssets.flatMap(
({ assets: scopeAssets, combination }) => {
if (!combination || !shouldPushSnapAssets(combination.stage)) {
return [];
}

return this.#assetsService
.buildControllerSnapshot(
combination.account,
combination.scope,
scopeAssets,
)
.map((group) => ({ ...combination, group }));
},
);

if (writes.length > 0) {
const startedAt = Date.now();
this.#logger.info('Assets controller snapshot write attempted', {
attemptCount: writes.length,
});
const results = await Promise.allSettled(
writes.map(async ({ account, scope, group }) =>
this.#tronAssetsControllerAdapter.pushAssetSnapshot(
account.id,
scope,
group,
),
),
);
const failureCount = results.filter(
(result) => result.status === 'rejected',
).length;

this.#logger.info('Assets controller snapshot write completed', {
attemptCount: writes.length,
failureCount,
latencyMs: Date.now() - startedAt,
});
results.forEach((result) => {
if (result.status === 'rejected') {
this.#logger.error('Assets controller snapshot write failed', {
error: result.reason,
});
}
});
}
}

async synchronizeTransactions(accounts: TronKeyringAccount[]): Promise<void> {
Expand Down
Loading