feat(billing): implement balance refresh

refs #247

apps/api/package.json

@@ -32,6 +32,7 @@
     "@akashnetwork/akash-api": "^1.3.0",
     "@akashnetwork/akashjs": "^0.10.0",
     "@akashnetwork/database": "*",
+    "@akashnetwork/http-sdk": "*",
     "@chain-registry/assets": "^0.7.1",
     "@cosmjs/amino": "^0.32.4",
     "@cosmjs/crypto": "^0.32.4",

apps/api/src/billing/providers/http-sdk.provider.ts

@@ -0,0 +1,6 @@
+import { AllowanceHttpService } from "@akashnetwork/http-sdk";
+import { container } from "tsyringe";
+
+import { apiNodeUrl } from "@src/utils/constants";
+
+container.register(AllowanceHttpService, { useValue: new AllowanceHttpService({ baseURL: apiNodeUrl }) });

apps/api/src/billing/providers/index.ts

@@ -1,5 +1,6 @@
 import "./user-wallet-schema.provider";
 import "./config.provider";
+import "./http-sdk.provider";
 
 export * from "./user-wallet-schema.provider";
 export * from "./config.provider";

apps/api/src/billing/repositories/user-wallet/user-wallet.repository.ts

@@ -6,9 +6,9 @@ import { InjectUserWalletSchema, UserWalletSchema } from "@src/billing/providers
 import { ApiPgDatabase, InjectPg } from "@src/core/providers";
 import { TxService } from "@src/core/services";
 
-export type UserInput = Partial<UserWalletSchema["$inferInsert"]>;
-export type DbUserOutput = UserWalletSchema["$inferSelect"];
-export type UserOutput = DbUserOutput & {
+export type UserWalletInput = Partial<UserWalletSchema["$inferInsert"]>;
+export type DbUserWalletOutput = UserWalletSchema["$inferSelect"];
+export type UserWalletOutput = DbUserWalletOutput & {
   creditAmount: number;
 };
 
@@ -29,7 +29,7 @@ export class UserWalletRepository {
     private readonly txManager: TxService
   ) {}
 
-  async create(input: Pick<UserInput, "userId" | "address">) {
+  async create(input: Pick<UserWalletInput, "userId" | "address">) {
     return this.toOutput(
       first(
         await this.cursor
@@ -43,9 +43,9 @@ export class UserWalletRepository {
     );
   }
 
-  async updateById(id: UserOutput["id"], payload: Partial<UserInput>, options?: { returning: true }): Promise<UserOutput>;
-  async updateById(id: UserOutput["id"], payload: Partial<UserInput>): Promise<void>;
-  async updateById(id: UserOutput["id"], payload: Partial<UserInput>, options?: { returning: boolean }): Promise<void | UserOutput> {
+  async updateById(id: UserWalletOutput["id"], payload: Partial<UserWalletInput>, options?: { returning: true }): Promise<UserWalletOutput>;
+  async updateById(id: UserWalletOutput["id"], payload: Partial<UserWalletInput>): Promise<void>;
+  async updateById(id: UserWalletOutput["id"], payload: Partial<UserWalletInput>, options?: { returning: boolean }): Promise<void | UserWalletOutput> {
     const cursor = this.cursor.update(this.userWallet).set(payload).where(eq(this.userWallet.id, id));
 
     if (options?.returning) {
@@ -58,8 +58,8 @@ export class UserWalletRepository {
     return undefined;
   }
 
-  async find(query?: Partial<DbUserOutput>) {
-    const fields = query && (Object.keys(query) as Array<keyof DbUserOutput>);
+  async find(query?: Partial<DbUserWalletOutput>) {
+    const fields = query && (Object.keys(query) as Array<keyof DbUserWalletOutput>);
     const where = fields.length ? and(...fields.map(field => eq(this.userWallet[field], query[field]))) : undefined;
 
     return this.toOutputList(
@@ -69,15 +69,15 @@ export class UserWalletRepository {
     );
   }
 
-  async findByUserId(userId: UserOutput["userId"]) {
-    return await this.cursor.query.userWalletSchema.findFirst({ where: eq(this.userWallet.userId, userId) });
+  async findByUserId(userId: UserWalletOutput["userId"]) {
+    return this.toOutput(await this.cursor.query.userWalletSchema.findFirst({ where: eq(this.userWallet.userId, userId) }));
   }
 
-  private toOutputList(dbOutput: UserWalletSchema["$inferSelect"][]): UserOutput[] {
+  private toOutputList(dbOutput: UserWalletSchema["$inferSelect"][]): UserWalletOutput[] {
     return dbOutput.map(item => this.toOutput(item));
   }
 
-  private toOutput(dbOutput: UserWalletSchema["$inferSelect"]): UserOutput {
+  private toOutput(dbOutput: UserWalletSchema["$inferSelect"]): UserWalletOutput {
     return {
       ...dbOutput,
       creditAmount: parseFloat(dbOutput.deploymentAllowance) + parseFloat(dbOutput.feeAllowance)

apps/api/src/billing/services/balances/balances.service.ts

@@ -0,0 +1,70 @@
+import { AllowanceHttpService } from "@akashnetwork/http-sdk";
+import { singleton } from "tsyringe";
+
+import { BillingConfig, InjectBillingConfig } from "@src/billing/providers";
+import { UserWalletInput, UserWalletOutput, UserWalletRepository } from "@src/billing/repositories";
+import { MasterWalletService } from "@src/billing/services";
+
+@singleton()
+export class BalancesService {
+  constructor(
+    @InjectBillingConfig() private readonly config: BillingConfig,
+    private readonly userWalletRepository: UserWalletRepository,
+    private readonly masterWalletService: MasterWalletService,
+    private readonly allowanceHttpService: AllowanceHttpService
+  ) {}
+
+  async updateUserWalletLimits(userWallet: UserWalletOutput) {
+    const [feeLimit, deploymentLimit] = await Promise.all([this.calculateFeeLimit(userWallet), this.calculateDeploymentLimit(userWallet)]);
+
+    const update: Partial<UserWalletInput> = {};
+
+    const feeLimitStr = feeLimit.toString();
+
+    if (userWallet.feeAllowance !== feeLimitStr) {
+      update.feeAllowance = feeLimitStr;
+    }
+
+    const deploymentLimitStr = deploymentLimit.toString();
+
+    if (userWallet.deploymentAllowance !== deploymentLimitStr) {
+      update.deploymentAllowance = deploymentLimitStr;
+    }
+
+    if (Object.keys(update).length > 0) {
+      await this.userWalletRepository.updateById(userWallet.id, update);
+    }
+  }
+
+  private async calculateFeeLimit(userWallet: UserWalletOutput) {
+    const feeAllowance = await this.allowanceHttpService.getFeeAllowancesForGrantee(userWallet.address);
+    const masterWalletAddress = await this.masterWalletService.getFirstAddress();
+
+    return feeAllowance.reduce((acc, allowance) => {
+      if (allowance.granter !== masterWalletAddress) {
+        return acc;
+      }
+
+      return allowance.allowance.spend_limit.reduce((acc, { denom, amount }) => {
+        if (denom !== this.config.TRIAL_ALLOWANCE_DENOM) {
+          return acc;
+        }
+
+        return acc + parseInt(amount);
+      }, 0);
+    }, 0);
+  }
+
+  private async calculateDeploymentLimit(userWallet: UserWalletOutput) {
+    const deploymentAllowance = await this.allowanceHttpService.getDeploymentAllowancesForGrantee(userWallet.address);
+    const masterWalletAddress = await this.masterWalletService.getFirstAddress();
+
+    return deploymentAllowance.reduce((acc, allowance) => {
+      if (allowance.granter !== masterWalletAddress || allowance.authorization.spend_limit.denom !== this.config.TRIAL_ALLOWANCE_DENOM) {
+        return acc;
+      }
+
+      return acc + parseInt(allowance.authorization.spend_limit.amount);
+    }, 0);
+  }
+}

apps/api/src/billing/services/tx-signer/tx-signer.service.ts

@@ -1,3 +1,4 @@
+import { AllowanceHttpService } from "@akashnetwork/http-sdk";
 import { stringToPath } from "@cosmjs/crypto";
 import { DirectSecp256k1HdWallet, EncodeObject, Registry } from "@cosmjs/proto-signing";
 import { calculateFee, GasPrice, SigningStargateClient } from "@cosmjs/stargate";
@@ -7,8 +8,9 @@ import { singleton } from "tsyringe";
 
 import { BillingConfig, InjectBillingConfig } from "@src/billing/providers";
 import { InjectTypeRegistry } from "@src/billing/providers/type-registry.provider";
-import { UserOutput, UserWalletRepository } from "@src/billing/repositories";
+import { UserWalletOutput, UserWalletRepository } from "@src/billing/repositories";
 import { MasterWalletService } from "@src/billing/services";
+import { BalancesService } from "@src/billing/services/balances/balances.service";
 import { ForbiddenException } from "@src/core";
 
 type StringifiedEncodeObject = Omit<EncodeObject, "value"> & { value: string };
@@ -24,12 +26,14 @@ export class TxSignerService {
 
   constructor(
     @InjectBillingConfig() private readonly config: BillingConfig,
+    @InjectTypeRegistry() private readonly registry: Registry,
     private readonly userWalletRepository: UserWalletRepository,
     private readonly masterWalletService: MasterWalletService,
-    @InjectTypeRegistry() private readonly registry: Registry
+    private readonly allowanceHttpService: AllowanceHttpService,
+    private readonly balancesService: BalancesService
   ) {}
 
-  async signAndBroadcast(userId: UserOutput["userId"], messages: StringifiedEncodeObject[]) {
+  async signAndBroadcast(userId: UserWalletOutput["userId"], messages: StringifiedEncodeObject[]) {
     const userWallet = await this.userWalletRepository.findByUserId(userId);
     const decodedMessages = this.decodeMessages(messages);
 
@@ -38,6 +42,8 @@ export class TxSignerService {
     const client = await this.getClientForAddressIndex(userWallet.id);
     const tx = await client.signAndBroadcast(decodedMessages);
 
+    await this.balancesService.updateUserWalletLimits(userWallet);
+
     return pick(tx, ["code", "transactionHash", "rawLog"]);
   }
 

apps/api/src/billing/services/wallet-initializer/wallet-initializer.service.ts

@@ -1,7 +1,7 @@
 import pick from "lodash/pick";
 import { singleton } from "tsyringe";
 
-import { UserInput, UserWalletRepository } from "@src/billing/repositories";
+import { UserWalletInput, UserWalletRepository } from "@src/billing/repositories";
 import { ManagedUserWalletService } from "@src/billing/services";
 import { WithTransaction } from "@src/core/services";
 
@@ -13,7 +13,7 @@ export class WalletInitializerService {
   ) {}
 
   @WithTransaction()
-  async initialize(userId: UserInput["userId"]) {
+  async initialize(userId: UserWalletInput["userId"]) {
     const { id } = await this.userWalletRepository.create({ userId });
     const wallet = await this.walletManager.createAndAuthorizeTrialSpending({ addressIndex: id });
     const userWallet = await this.userWalletRepository.updateById(

apps/api/test/functional/create-wallet.spec.ts

@@ -1,32 +1,39 @@
-import { faker } from "@faker-js/faker";
 import { eq } from "drizzle-orm";
 import { container } from "tsyringe";
 
 import { app } from "@src/app";
 import { BILLING_CONFIG, BillingConfig, USER_WALLET_SCHEMA, UserWalletSchema } from "@src/billing/providers";
 import { ApiPgDatabase, POSTGRES_DB } from "@src/core";
+import { USER_SCHEMA, UserSchema } from "@src/user/providers";
 
 jest.setTimeout(10000);
 
 describe("wallets", () => {
-  const schema = container.resolve<UserWalletSchema>(USER_WALLET_SCHEMA);
+  const userWalletSchema = container.resolve<UserWalletSchema>(USER_WALLET_SCHEMA);
+  const userSchema = container.resolve<UserSchema>(USER_SCHEMA);
   const config = container.resolve<BillingConfig>(BILLING_CONFIG);
   const db = container.resolve<ApiPgDatabase>(POSTGRES_DB);
   const userWalletsTable = db.query.userWalletSchema;
 
   afterEach(async () => {
-    await db.delete(schema);
+    await Promise.all([db.delete(userWalletSchema), db.delete(userSchema)]);
   });
 
   describe("POST /v1/wallets", () => {
     it("should create a wallet for a user", async () => {
-      const userId = faker.string.uuid();
+      const userResponse = await app.request("/v1/anonymous-users", {
+        method: "POST",
+        headers: new Headers({ "Content-Type": "application/json" })
+      });
+      const {
+        data: { id: userId }
+      } = await userResponse.json();
       const res = await app.request("/v1/wallets", {
         method: "POST",
         body: JSON.stringify({ data: { userId } }),
         headers: new Headers({ "Content-Type": "application/json" })
       });
-      const userWallet = await userWalletsTable.findFirst({ where: eq(schema.userId, userId) });
+      const userWallet = await userWalletsTable.findFirst({ where: eq(userWalletSchema.userId, userId) });
 
       expect(res.status).toBe(200);
       expect(await res.json()).toMatchObject({

apps/api/test/functional/sign-and-broadcast-tx.spec.ts

@@ -1,19 +1,34 @@
 import { certificateManager } from "@akashnetwork/akashjs/build/certificates/certificate-manager";
 import type { Registry } from "@cosmjs/proto-signing";
-import { faker } from "@faker-js/faker";
 import { container } from "tsyringe";
 
 import { app } from "@src/app";
+import { USER_WALLET_SCHEMA, UserWalletSchema } from "@src/billing/providers";
 import { TYPE_REGISTRY } from "@src/billing/providers/type-registry.provider";
+import { ApiPgDatabase, POSTGRES_DB } from "@src/core";
+import { USER_SCHEMA, UserSchema } from "@src/user/providers";
 
 jest.setTimeout(20000);
 
 describe("Tx Sign", () => {
   const registry = container.resolve<Registry>(TYPE_REGISTRY);
+  const db = container.resolve<ApiPgDatabase>(POSTGRES_DB);
+  const userWalletSchema = container.resolve<UserWalletSchema>(USER_WALLET_SCHEMA);
+  const userSchema = container.resolve<UserSchema>(USER_SCHEMA);
+
+  afterEach(async () => {
+    await Promise.all([db.delete(userWalletSchema), db.delete(userSchema)]);
+  });
 
   describe("POST /v1/tx", () => {
     it("should create a wallet for a user", async () => {
-      const userId = faker.string.uuid();
+      const userResponse = await app.request("/v1/anonymous-users", {
+        method: "POST",
+        headers: new Headers({ "Content-Type": "application/json" })
+      });
+      const {
+        data: { id: userId }
+      } = await userResponse.json();
       const walletResponse = await app.request("/v1/wallets", {
         method: "POST",
         body: JSON.stringify({

apps/deploy-web/src/context/WalletProvider/WalletProvider.tsx

@@ -73,7 +73,7 @@ export const WalletProvider = ({ children }) => {
   const { user } = useAnonymousUser();
 
   const userWallet = useSelectedChain();
-  const { wallet: managedWallet, isLoading, create } = useManagedWallet();
+  const { wallet: managedWallet, isLoading, create, refetch } = useManagedWallet();
   const { address: walletAddress, username, isWalletConnected } = useMemo(() => managedWallet || userWallet, [managedWallet, userWallet]);
   const { addEndpoints } = useManager();
   const {
@@ -201,8 +201,6 @@ export const WalletProvider = ({ children }) => {
         txResult = await userWallet.broadcast(txRaw);
 
         setIsBroadcastingTx(false);
-
-        await refreshBalances();
       }
 
       if (txResult.code !== 0) {
@@ -259,6 +257,7 @@ export const WalletProvider = ({ children }) => {
 
       return false;
     } finally {
+      await refreshBalances();
       if (pendingSnackbarKey) {
         closeSnackbar(pendingSnackbarKey);
       }
@@ -289,9 +288,10 @@ export const WalletProvider = ({ children }) => {
 
   async function refreshBalances(address?: string): Promise<{ uakt: number; usdc: number }> {
     if (managedWallet) {
+      const wallet = await refetch();
       const walletBalances = {
         uakt: 0,
-        usdc: managedWallet.creditAmount
+        usdc: wallet.data?.creditAmount || managedWallet.creditAmount
       };
 
       setWalletBalances(walletBalances);

apps/deploy-web/src/hooks/useManagedWallet.ts

@@ -10,7 +10,7 @@ const isBillingEnabled = envConfig.NEXT_PUBLIC_BILLING_ENABLED;
 export const useManagedWallet = () => {
   const { user } = useStoredAnonymousUser();
 
-  const { data: queried, isFetched, isLoading: isFetching } = useManagedWalletQuery(isBillingEnabled && user?.id);
+  const { data: queried, isFetched, isLoading: isFetching, refetch } = useManagedWalletQuery(isBillingEnabled && user?.id);
   const { mutate: create, data: created, isLoading: isCreating, isSuccess: isCreated } = useCreateManagedWalletMutation();
   const wallet = useMemo(() => queried || created, [queried, created]);
   const isLoading = isFetching || isCreating;
@@ -48,7 +48,8 @@ export const useManagedWallet = () => {
             isWalletLoaded: isConfigured
           }
         : undefined,
-      isLoading
+      isLoading,
+      refetch
     };
-  }, [create, isLoading, user, wallet]);
+  }, [create, isLoading, user, wallet, refetch]);
 };

docker-compose.yml

@@ -4,6 +4,7 @@ services:
       target: development
     volumes:
       - ./apps/api:/app/apps/api
+      - ./packages:/app/packages
       - ./package.json:/app/package.json
       - ./package-lock.json:/app/package-lock.json
       - /app/node_modules