Add tokenBalance api endpoint

src/api/background.ts

@@ -45,6 +45,8 @@ import subscriptionModule from "./modules/subscription";
 import subscription from "./modules/subscription/subscription.background";
 import userTokensModule from "./modules/user_tokens";
 import userTokens from "./modules/user_tokens/user_tokens.background";
+import tokenBalanceModule from "./modules/token_balance";
+import tokenBalance from "./modules/token_balance/token_balance.background";
 
 /** Background modules */
 const modules: BackgroundModule<any>[] = [
@@ -69,6 +71,7 @@ const modules: BackgroundModule<any>[] = [
   { ...signDataItemModule, function: signDataItem },
   { ...subscriptionModule, function: subscription },
   { ...userTokensModule, function: userTokens },
+  { ...tokenBalanceModule, function: tokenBalance },
   { ...batchSignDataItemModule, function: batchSignDataItem }
 ];
 

src/api/foreground.ts

@@ -65,6 +65,8 @@ import signDataItem, {
 } from "./modules/sign_data_item/sign_data_item.foreground";
 import userTokensModule from "./modules/user_tokens";
 import userTokens from "./modules/user_tokens/user_tokens.foreground";
+import tokenBalanceModule from "./modules/token_balance";
+import tokenBalance from "./modules/token_balance/token_balance.foreground";
 
 /** Foreground modules */
 const modules: ForegroundModule[] = [
@@ -101,6 +103,7 @@ const modules: ForegroundModule[] = [
   },
   { ...subscriptionModule, function: subscription },
   { ...userTokensModule, function: userTokens },
+  { ...tokenBalanceModule, function: tokenBalance },
   {
     ...batchSignDataItemModule,
     function: batchSignDataItem,

src/api/modules/token_balance/index.ts

@@ -0,0 +1,11 @@
+import type { PermissionType } from "~applications/permissions";
+import type { ModuleProperties } from "~api/module";
+
+const permissions: PermissionType[] = ["ACCESS_TOKENS"];
+
+const tokenBalanceModule: ModuleProperties = {
+  functionName: "tokenBalance",
+  permissions
+};
+
+export default tokenBalanceModule;

src/api/modules/token_balance/token_balance.background.ts

@@ -0,0 +1,20 @@
+import type { ModuleFunction } from "~api/background";
+import { ExtensionStorage } from "~utils/storage";
+import { getAoTokenBalance, getNativeTokenBalance } from "~tokens/aoTokens/ao";
+import { AO_NATIVE_TOKEN } from "~utils/ao_import";
+import { isAddress } from "~utils/assertions";
+
+const background: ModuleFunction<string> = async (_, id?: string) => {
+  // validate input
+  isAddress(id);
+  const address = await ExtensionStorage.get("active_address");
+
+  const balance =
+    id === AO_NATIVE_TOKEN
+      ? await getNativeTokenBalance(address)
+      : (await getAoTokenBalance(address, id)).toString();
+
+  return balance;
+};
+
+export default background;

src/api/modules/token_balance/token_balance.foreground.ts

@@ -0,0 +1,6 @@
+import type { ModuleFunction } from "~api/module";
+
+// no need to transform anything in the foreground
+const foreground: ModuleFunction<void> = () => {};
+
+export default foreground;

src/tokens/aoTokens/ao.ts

@@ -16,6 +16,8 @@ import {
 import type { Alarms } from "webextension-polyfill";
 import type { KeystoneSigner } from "~wallets/hardware/keystone";
 import browser from "webextension-polyfill";
+import { fetchTokenByProcessId } from "~lib/transactions";
+import { tokenTypeRegistry } from "~tokens/token";
 
 export type AoInstance = ReturnType<typeof connect>;
 
@@ -223,7 +225,7 @@ export async function getAoTokenBalance(
 ): Promise<Quantity> {
   const aoTokens = (await ExtensionStorage.get<TokenInfo[]>("ao_tokens")) || [];
 
-  const aoToken = aoTokens.find((token) => token.processId === process);
+  let aoToken = aoTokens.find((token) => token.processId === process);
 
   const res = await dryrun({
     Id,
@@ -232,13 +234,35 @@ export async function getAoTokenBalance(
     tags: [{ name: "Action", value: "Balance" }]
   });
 
+  const errorMessage = (res as any)?.error || res?.Error;
+
+  if (errorMessage) {
+    throw new Error(errorMessage);
+  }
+
+  if (res.Messages.length === 0) {
+    throw new Error(
+      "Invalid token process: Balance action handler missing or unsupported."
+    );
+  }
+
   for (const msg of res.Messages as Message[]) {
     const balance = getTagValue("Balance", msg.Tags);
 
-    if (balance && aoToken) {
+    if (balance && +balance) {
+      if (!aoToken) {
+        aoToken = await fetchTokenByProcessId(process);
+        if (!aoToken) {
+          throw new Error("Could not load token info.");
+        }
+      }
+
       return new Quantity(BigInt(balance), BigInt(aoToken.Denomination));
     }
   }
+
+  // default return
+  return new Quantity(0n, 12n);
 }
 
 export async function getNativeTokenBalance(address: string): Promise<string> {