MultichainTransactionsController.ts

import type {
  AccountsControllerAccountAddedEvent,
  AccountsControllerAccountRemovedEvent,
  AccountsControllerListMultichainAccountsAction,
  AccountsControllerAccountTransactionsUpdatedEvent,
} from '@metamask/accounts-controller';
import { BaseController } from '@metamask/base-controller';
import type {
  ControllerGetStateAction,
  ControllerStateChangeEvent,
} from '@metamask/base-controller';
import {
  AccountActivityServiceTransactionUpdatedEvent,
  Transaction as BackendTransactionUpdate,
} from '@metamask/core-backend';
import { isEvmAccountType, TransactionStatus } from '@metamask/keyring-api';
import type {
  Transaction,
  AccountTransactionsUpdatedEventPayload,
} from '@metamask/keyring-api';
import type { KeyringControllerGetStateAction } from '@metamask/keyring-controller';
import type { InternalAccount } from '@metamask/keyring-internal-api';
import { KeyringClient } from '@metamask/keyring-snap-client';
import type { Messenger } from '@metamask/messenger';
import type { SnapControllerHandleRequestAction } from '@metamask/snaps-controllers';
import type { SnapId } from '@metamask/snaps-sdk';
import { HandlerType } from '@metamask/snaps-utils';
import type { CaipChainId, Json, JsonRpcRequest } from '@metamask/utils';
import type { Draft } from 'immer';

import type { MultichainTransactionsControllerMethodActions } from './MultichainTransactionsController-method-action-types';

const controllerName = 'MultichainTransactionsController';

const MESSENGER_EXPOSED_METHODS = ['updateTransactionsForAccount'] as const;

/**
 * Maps a backend WebSocket status string to a terminal keyring status, or
 * `undefined` for non-terminal / unknown statuses (which are ignored).
 *
 * @param status - The status reported by the account-activity WebSocket.
 * @returns The terminal {@link TransactionStatus}, or `undefined`.
 */
function toTerminalStatus(status: string): TransactionStatus | undefined {
  const normalized = status.toLowerCase();
  if (normalized === TransactionStatus.Confirmed) {
    return TransactionStatus.Confirmed;
  }
  if (normalized === TransactionStatus.Failed) {
    return TransactionStatus.Failed;
  }
  return undefined;
}

/**
 * PaginationOptions
 *
 * Represents options for paginating transaction results
 * limit - The maximum number of transactions to return
 * next - The cursor for the next page of transactions, or null if there is no next page
 */
export type PaginationOptions = {
  limit: number;
  next?: string | null;
};

/**
 * State used by the {@link MultichainTransactionsController} to cache account transactions.
 */
export type MultichainTransactionsControllerState = {
  nonEvmTransactions: {
    [accountId: string]: {
      [chain: CaipChainId]: TransactionStateEntry;
    };
  };
};

/**
 * Constructs the default {@link MultichainTransactionsController} state.
 *
 * @returns The default {@link MultichainTransactionsController} state.
 */
export function getDefaultMultichainTransactionsControllerState(): MultichainTransactionsControllerState {
  return {
    nonEvmTransactions: {},
  };
}

/**
 * Event emitted when a transaction is finalized.
 */
export type MultichainTransactionsControllerTransactionConfirmedEvent = {
  type: `${typeof controllerName}:transactionConfirmed`;
  payload: [Transaction];
};

/**
 * Event emitted when a transaction is submitted.
 */
export type MultichainTransactionsControllerTransactionSubmittedEvent = {
  type: `${typeof controllerName}:transactionSubmitted`;
  payload: [Transaction];
};

/**
 * Returns the state of the {@link MultichainTransactionsController}.
 */
export type MultichainTransactionsControllerGetStateAction =
  ControllerGetStateAction<
    typeof controllerName,
    MultichainTransactionsControllerState
  >;

/**
 * Event emitted when the state of the {@link MultichainTransactionsController} changes.
 */
export type MultichainTransactionsControllerStateChange =
  ControllerStateChangeEvent<
    typeof controllerName,
    MultichainTransactionsControllerState
  >;

/**
 * Actions exposed by the {@link MultichainTransactionsController}.
 */
export type MultichainTransactionsControllerActions =
  | MultichainTransactionsControllerGetStateAction
  | MultichainTransactionsControllerMethodActions;

/**
 * Events emitted by {@link MultichainTransactionsController}.
 */
export type MultichainTransactionsControllerEvents =
  | MultichainTransactionsControllerStateChange
  | MultichainTransactionsControllerTransactionConfirmedEvent
  | MultichainTransactionsControllerTransactionSubmittedEvent;

/**
 * Messenger type for the MultichainTransactionsController.
 */
export type MultichainTransactionsControllerMessenger = Messenger<
  typeof controllerName,
  MultichainTransactionsControllerActions | AllowedActions,
  MultichainTransactionsControllerEvents | AllowedEvents
>;

/**
 * Actions that this controller is allowed to call.
 */
type AllowedActions =
  | AccountsControllerListMultichainAccountsAction
  | KeyringControllerGetStateAction
  | SnapControllerHandleRequestAction;

/**
 * Events that this controller is allowed to subscribe.
 */
type AllowedEvents =
  | AccountsControllerAccountAddedEvent
  | AccountsControllerAccountRemovedEvent
  | AccountsControllerAccountTransactionsUpdatedEvent
  | AccountActivityServiceTransactionUpdatedEvent;

/**
 * {@link MultichainTransactionsController}'s metadata.
 *
 * This allows us to choose if fields of the state should be persisted or not
 * using the `persist` flag; and if they can be sent to Sentry or not, using
 * the `anonymous` flag.
 */
const multichainTransactionsControllerMetadata = {
  nonEvmTransactions: {
    includeInStateLogs: true,
    persist: true,
    includeInDebugSnapshot: false,
    usedInUi: true,
  },
};

/**
 * The state of transactions for a specific chain.
 */
export type TransactionStateEntry = {
  transactions: Transaction[];
  next: string | null;
  lastUpdated: number;
};

/**
 * The MultichainTransactionsController is responsible for fetching and caching account
 * transactions for non-EVM accounts.
 */
export class MultichainTransactionsController extends BaseController<
  typeof controllerName,
  MultichainTransactionsControllerState,
  MultichainTransactionsControllerMessenger
> {
  constructor({
    messenger,
    state,
  }: {
    messenger: MultichainTransactionsControllerMessenger;
    state?: Partial<MultichainTransactionsControllerState>;
  }) {
    super({
      messenger,
      name: controllerName,
      metadata: multichainTransactionsControllerMetadata,
      state: {
        ...getDefaultMultichainTransactionsControllerState(),
        ...state,
      },
    });

    this.messenger.registerMethodActionHandlers(
      this,
      MESSENGER_EXPOSED_METHODS,
    );

    // Fetch initial transactions for all non-EVM accounts
    for (const account of this.#listAccounts()) {
      this.updateTransactionsForAccount(account.id).catch((error) => {
        console.error(
          `Failed to fetch initial transactions for account ${account.id}:`,
          error,
        );
      });
    }

    this.messenger.subscribe(
      'AccountsController:accountAdded',
      (account: InternalAccount) => this.#handleOnAccountAdded(account),
    );
    this.messenger.subscribe(
      'AccountsController:accountRemoved',
      (accountId: string) => this.#handleOnAccountRemoved(accountId),
    );
    this.messenger.subscribe(
      'AccountsController:accountTransactionsUpdated',
      (transactionsUpdate: AccountTransactionsUpdatedEventPayload) =>
        this.#handleOnAccountTransactionsUpdated(transactionsUpdate),
    );
    // Client-owned terminal tracking: flip pending (Submitted) non-EVM entries to
    // Confirmed/Failed from the account-activity WebSocket, replacing the snap's
    // former confirmation tracking (see snap-transactions offload, ticket 5).
    this.messenger.subscribe(
      'AccountActivityService:transactionUpdated',
      (update) => this.#handleOnBackendTransactionUpdated(update),
    );
  }

  /**
   * Lists the multichain accounts coming from the `AccountsController`.
   *
   * @returns A list of multichain accounts.
   */
  #listMultichainAccounts(): InternalAccount[] {
    return this.messenger.call('AccountsController:listMultichainAccounts');
  }

  /**
   * Lists the accounts that we should get transactions for.
   *
   * @returns A list of accounts that we should get transactions for.
   */
  #listAccounts(): InternalAccount[] {
    const accounts = this.#listMultichainAccounts();
    return accounts.filter((account) => this.#isNonEvmAccount(account));
  }

  /**
   * Gets transactions for an account.
   *
   * @param accountId - The ID of the account to get transactions for.
   * @param snapId - The ID of the snap that manages the account.
   * @param pagination - Options for paginating transaction results.
   * @returns A promise that resolves to the transaction data and pagination info.
   */
  async #getTransactions(
    accountId: string,
    snapId: string,
    pagination: PaginationOptions,
  ): Promise<{
    data: Transaction[];
    next: string | null;
  }> {
    return await this.#getClient(snapId).listAccountTransactions(
      accountId,
      pagination,
    );
  }

  /**
   * Updates transactions for a specific account. This is used for the initial fetch
   * when an account is first added.
   *
   * @param accountId - The ID of the account to get transactions for.
   */
  async updateTransactionsForAccount(accountId: string) {
    const { isUnlocked } = this.messenger.call('KeyringController:getState');

    if (!isUnlocked) {
      return;
    }

    try {
      const account = this.#listAccounts().find(
        (accountItem) => accountItem.id === accountId,
      );

      if (account?.metadata.snap) {
        const response = await this.#getTransactions(
          account.id,
          account.metadata.snap.id,
          { limit: 10 },
        );

        const transactionsByChain: Record<CaipChainId, Transaction[]> = {};

        response.data.forEach((transaction) => {
          const { chain } = transaction;

          if (!transactionsByChain[chain]) {
            transactionsByChain[chain] = [];
          }
          transactionsByChain[chain].push(transaction);
        });

        const chainUpdates = Object.entries(transactionsByChain).map(
          ([chain, transactions]) => ({
            chain,
            entry: {
              transactions,
              next: response.next,
              lastUpdated: Date.now(),
            },
          }),
        );

        this.update((state: Draft<MultichainTransactionsControllerState>) => {
          if (!state.nonEvmTransactions[account.id]) {
            state.nonEvmTransactions[account.id] = {};
          }

          chainUpdates.forEach(({ chain, entry }) => {
            state.nonEvmTransactions[account.id][chain as CaipChainId] = entry;
          });
        });
      }
    } catch (error) {
      console.error(
        `Failed to fetch transactions for account ${accountId}:`,
        error,
      );
    }
  }

  /**
   * Checks for non-EVM accounts.
   *
   * @param account - The new account to be checked.
   * @returns True if the account is a non-EVM account, false otherwise.
   */
  #isNonEvmAccount(account: InternalAccount): boolean {
    return (
      !isEvmAccountType(account.type) &&
      // Non-EVM accounts are backed by a Snap for now
      account.metadata.snap !== undefined
    );
  }

  /**
   * Handles changes when a new account has been added.
   *
   * @param account - The new account being added.
   */
  async #handleOnAccountAdded(account: InternalAccount) {
    if (!this.#isNonEvmAccount(account)) {
      return;
    }

    await this.updateTransactionsForAccount(account.id);
  }

  /**
   * Handles changes when a new account has been removed.
   *
   * @param accountId - The account ID being removed.
   */
  async #handleOnAccountRemoved(accountId: string) {
    if (accountId in this.state.nonEvmTransactions) {
      this.update((state: Draft<MultichainTransactionsControllerState>) => {
        delete state.nonEvmTransactions[accountId];
      });
    }
  }

  /**
   * Publishes transaction update events.
   *
   * @param updatedTransaction - The updated transaction.
   */
  #publishTransactionUpdateEvent(updatedTransaction: Transaction) {
    if (updatedTransaction.status === TransactionStatus.Confirmed) {
      this.messenger.publish(
        'MultichainTransactionsController:transactionConfirmed',
        updatedTransaction,
      );
    }

    if (updatedTransaction.status === TransactionStatus.Submitted) {
      this.messenger.publish(
        'MultichainTransactionsController:transactionSubmitted',
        updatedTransaction,
      );
    }
  }

  /**
   * Handles transaction updates received from the AccountsController.
   *
   * @param transactionsUpdate - The transaction update event containing new transactions.
   */
  #handleOnAccountTransactionsUpdated(
    transactionsUpdate: AccountTransactionsUpdatedEventPayload,
  ): void {
    const updatedTransactions: Record<
      string,
      Record<CaipChainId, Transaction[]>
    > = {};
    const transactionsToPublish: Transaction[] = [];

    if (!transactionsUpdate?.transactions) {
      return;
    }

    Object.entries(transactionsUpdate.transactions).forEach(
      ([accountId, newTransactions]) => {
        updatedTransactions[accountId] = {};

        newTransactions.forEach((tx) => {
          const { chain } = tx;

          if (!updatedTransactions[accountId][chain]) {
            updatedTransactions[accountId][chain] = [];
          }

          updatedTransactions[accountId][chain].push(tx);
          transactionsToPublish.push(tx);
        });

        Object.entries(updatedTransactions[accountId]).forEach(
          ([chain, chainTransactions]) => {
            // Account might not have any transactions yet, so use `[]` in that case.
            const oldTransactions =
              this.state.nonEvmTransactions[accountId]?.[chain as CaipChainId]
                ?.transactions ?? [];

            // Uses a `Map` to deduplicate transactions by ID, ensuring we keep the latest version
            // of each transaction while preserving older transactions and transactions from other accounts.
            // Transactions are sorted by timestamp (newest first).
            const transactions = new Map();

            oldTransactions.forEach((tx) => {
              transactions.set(tx.id, tx);
            });

            chainTransactions.forEach((tx) => {
              transactions.set(tx.id, tx);
            });

            // Sorted by timestamp (newest first). If the timestamp is not provided, those
            // transactions will be put in the end of this list.
            updatedTransactions[accountId][chain as CaipChainId] = Array.from(
              transactions.values(),
            ).sort((a, b) => (b.timestamp ?? 0) - (a.timestamp ?? 0));
          },
        );
      },
    );

    this.update((state) => {
      Object.entries(updatedTransactions).forEach(([accountId, chainsData]) => {
        if (!state.nonEvmTransactions[accountId]) {
          state.nonEvmTransactions[accountId] = {};
        }

        Object.entries(chainsData).forEach(([chain, transactions]) => {
          state.nonEvmTransactions[accountId][chain as CaipChainId] = {
            transactions,
            next: null,
            lastUpdated: Date.now(),
          };
        });
      });
    });

    // After we update the state, publish the events for new/updated transactions
    transactionsToPublish.forEach((tx) => {
      this.#publishTransactionUpdateEvent(tx);
    });
  }

  /**
   * Handles a transaction status update from the account-activity WebSocket.
   *
   * Finds the matching pending entry (by chain-specific id) and flips it to the
   * reported terminal status. Non-terminal updates and unknown transactions are
   * ignored.
   *
   * @param update - The backend transaction update.
   */
  #handleOnBackendTransactionUpdated(update: BackendTransactionUpdate): void {
    const terminalStatus = toTerminalStatus(update.status);
    if (!terminalStatus) {
      return;
    }

    let location:
      | { accountId: string; chain: CaipChainId; existing: Transaction }
      | undefined;
    for (const [accountId, chains] of Object.entries(
      this.state.nonEvmTransactions,
    )) {
      for (const [chain, entry] of Object.entries(chains)) {
        const existing = entry.transactions.find((tx) => tx.id === update.id);
        if (existing) {
          location = { accountId, chain: chain as CaipChainId, existing };
          break;
        }
      }
      if (location) {
        break;
      }
    }

    if (!location || location.existing.status === terminalStatus) {
      return;
    }

    const { accountId, chain } = location;
    const updatedTransaction: Transaction = {
      ...location.existing,
      status: terminalStatus,
    };

    this.update((state) => {
      const entry = state.nonEvmTransactions[accountId][chain];
      entry.transactions = entry.transactions.map((tx) =>
        tx.id === update.id ? updatedTransaction : tx,
      );
      entry.lastUpdated = Date.now();
    });

    this.#publishTransactionUpdateEvent(updatedTransaction);
  }

  /**
   * Gets a `KeyringClient` for a Snap.
   *
   * @param snapId - ID of the Snap to get the client for.
   * @returns A `KeyringClient` for the Snap.
   */
  #getClient(snapId: string): KeyringClient {
    return new KeyringClient({
      send: async (request: JsonRpcRequest) =>
        (await this.messenger.call('SnapController:handleRequest', {
          snapId: snapId as SnapId,
          origin: 'metamask',
          handler: HandlerType.OnKeyringRequest,
          request,
        })) as Promise<Json>,
    });
  }
}