feat: return estimated balance in account balance endpoints (#2104)

* feat: mempool principal cache

* fix: unify principal activity query

* feat: progress

* chore: mempool etag

* fix: missing column name

* fix: pruned

* fix: tests

* test: feature

* fix: coalesce fee rate

* fix: token amount coalesce

src/api/routes/address.ts

@@ -89,7 +89,7 @@ export const AddressRoutes: FastifyPluginAsync<
   fastify.get(
     '/:principal/stx',
     {
-      preHandler: handlePrincipalCache,
+      preHandler: handlePrincipalMempoolCache,
       schema: {
         operationId: 'get_account_stx_balance',
         summary: 'Get account STX balance',
@@ -120,8 +120,14 @@ export const AddressRoutes: FastifyPluginAsync<
           stxAddress,
           blockHeight
         );
+        let mempoolBalance: bigint | undefined = undefined;
+        if (req.query.until_block === undefined) {
+          const delta = await fastify.db.getPrincipalMempoolStxBalanceDelta(sql, stxAddress);
+          mempoolBalance = stxBalanceResult.balance + delta;
+        }
         const result: AddressStxBalance = {
           balance: stxBalanceResult.balance.toString(),
+          estimated_balance: mempoolBalance?.toString(),
           total_sent: stxBalanceResult.totalSent.toString(),
           total_received: stxBalanceResult.totalReceived.toString(),
           total_fees_sent: stxBalanceResult.totalFeesSent.toString(),
@@ -145,7 +151,7 @@ export const AddressRoutes: FastifyPluginAsync<
   fastify.get(
     '/:principal/balances',
     {
-      preHandler: handlePrincipalCache,
+      preHandler: handlePrincipalMempoolCache,
       schema: {
         operationId: 'get_account_balance',
         summary: 'Get account balances',
@@ -204,9 +210,16 @@ export const AddressRoutes: FastifyPluginAsync<
           };
         });
 
+        let mempoolBalance: bigint | undefined = undefined;
+        if (req.query.until_block === undefined) {
+          const delta = await fastify.db.getPrincipalMempoolStxBalanceDelta(sql, stxAddress);
+          mempoolBalance = stxBalanceResult.balance + delta;
+        }
+
         const result: AddressBalance = {
           stx: {
             balance: stxBalanceResult.balance.toString(),
+            estimated_balance: mempoolBalance?.toString(),
             total_sent: stxBalanceResult.totalSent.toString(),
             total_received: stxBalanceResult.totalReceived.toString(),
             total_fees_sent: stxBalanceResult.totalFeesSent.toString(),

src/api/schemas/entities/balances.ts

@@ -21,6 +21,11 @@ export const NftBalanceSchema = Type.Object(
 export const StxBalanceSchema = Type.Object(
   {
     balance: Type.String(),
+    estimated_balance: Type.Optional(
+      Type.String({
+        description: 'Total STX balance considering pending mempool transactions',
+      })
+    ),
     total_sent: Type.String(),
     total_received: Type.String(),
     total_fees_sent: Type.String(),

src/datastore/pg-store.ts

@@ -2476,6 +2476,36 @@ export class PgStore extends BasePgStore {
     };
   }
 
+  /**
+   * Returns the total STX balance delta affecting a principal from transactions currently in the
+   * mempool.
+   */
+  async getPrincipalMempoolStxBalanceDelta(sql: PgSqlClient, principal: string): Promise<bigint> {
+    const results = await sql<{ delta: string }[]>`
+      WITH sent AS (
+        SELECT SUM(COALESCE(token_transfer_amount, 0) + fee_rate) AS total
+        FROM mempool_txs
+        WHERE pruned = false AND sender_address = ${principal}
+      ),
+      sponsored AS (
+        SELECT SUM(fee_rate) AS total
+        FROM mempool_txs
+        WHERE pruned = false AND sponsor_address = ${principal} AND sponsored = true
+      ),
+      received AS (
+        SELECT SUM(COALESCE(token_transfer_amount, 0)) AS total
+        FROM mempool_txs
+        WHERE pruned = false AND token_transfer_recipient_address = ${principal}
+      )
+      SELECT
+        COALESCE((SELECT total FROM received), 0)
+        - COALESCE((SELECT total FROM sent), 0)
+        - COALESCE((SELECT total FROM sponsored), 0)
+        AS delta
+    `;
+    return BigInt(results[0]?.delta ?? '0');
+  }
+
   async getUnlockedStxSupply(
     args:
       | {

tests/api/address.test.ts

@@ -1544,6 +1544,7 @@ describe('address tests', () => {
     const expectedResp1 = {
       stx: {
         balance: '88679',
+        estimated_balance: '88679',
         total_sent: '6385',
         total_received: '100000',
         total_fees_sent: '4936',
@@ -1587,6 +1588,7 @@ describe('address tests', () => {
     const expectedResp2 = {
       stx: {
         balance: '91',
+        estimated_balance: '91',
         total_sent: '15',
         total_received: '1350',
         total_fees_sent: '1244',
@@ -1622,6 +1624,7 @@ describe('address tests', () => {
     expect(fetchAddrStxBalance1.type).toBe('application/json');
     const expectedStxResp1 = {
       balance: '91',
+      estimated_balance: '91',
       total_sent: '15',
       total_received: '1350',
       total_fees_sent: '1244',
@@ -1652,6 +1655,7 @@ describe('address tests', () => {
     expect(fetchAddrStxBalance1.type).toBe('application/json');
     const expectedStxResp1Sponsored = {
       balance: '3766',
+      estimated_balance: '3766',
       total_sent: '0',
       total_received: '5000',
       total_fees_sent: '1234',

tests/api/mempool.test.ts

@@ -2128,4 +2128,137 @@ describe('mempool tests', () => {
     );
     expect(check2.body.results).toHaveLength(2);
   });
+
+  test('account estimated balance from mempool activity', async () => {
+    const address = 'SP3FXEKSA6D4BW3TFP2BWTSREV6FY863Y90YY7D8G';
+    const url = `/extended/v1/address/${address}/stx`;
+    await db.update(
+      new TestBlockBuilder({
+        block_height: 1,
+        index_block_hash: '0x01',
+        parent_index_block_hash: '0x00',
+      })
+        .addTx({
+          tx_id: '0x0001',
+          token_transfer_recipient_address: address,
+          token_transfer_amount: 2000n,
+        })
+        .addTxStxEvent({ recipient: address, amount: 2000n })
+        .build()
+    );
+
+    // Base balance
+    const balance0 = await supertest(api.server).get(url);
+    expect(balance0.body.balance).toEqual('2000');
+    expect(balance0.body.estimated_balance).toEqual('2000');
+
+    // STX transfer in mempool
+    await db.updateMempoolTxs({
+      mempoolTxs: [
+        testMempoolTx({
+          tx_id: '0x0002',
+          sender_address: address,
+          token_transfer_amount: 100n,
+          fee_rate: 50n,
+        }),
+      ],
+    });
+    const balance1 = await supertest(api.server).get(url);
+    expect(balance1.body.balance).toEqual('2000');
+    expect(balance1.body.estimated_balance).toEqual('1850'); // Minus amount and fee
+
+    // Contract call in mempool
+    await db.updateMempoolTxs({
+      mempoolTxs: [
+        testMempoolTx({
+          tx_id: '0x0002aa',
+          sender_address: address,
+          type_id: DbTxTypeId.ContractCall,
+          token_transfer_amount: 0n,
+          contract_call_contract_id: 'ST27W5M8BRKA7C5MZE2R1S1F4XTPHFWFRNHA9M04Y.hello-world',
+          contract_call_function_args: '',
+          contract_call_function_name: 'test',
+          fee_rate: 50n,
+        }),
+      ],
+    });
+    const balance1b = await supertest(api.server).get(url);
+    expect(balance1b.body.balance).toEqual('2000');
+    expect(balance1b.body.estimated_balance).toEqual('1800'); // Minus fee
+
+    // Sponsored tx in mempool
+    await db.updateMempoolTxs({
+      mempoolTxs: [
+        testMempoolTx({
+          tx_id: '0x0003',
+          sponsor_address: address,
+          sponsored: true,
+          token_transfer_amount: 100n,
+          fee_rate: 50n,
+        }),
+      ],
+    });
+    const balance2 = await supertest(api.server).get(url);
+    expect(balance2.body.balance).toEqual('2000');
+    expect(balance2.body.estimated_balance).toEqual('1750'); // Minus fee
+
+    // STX received in mempool
+    await db.updateMempoolTxs({
+      mempoolTxs: [
+        testMempoolTx({
+          tx_id: '0x0004',
+          token_transfer_recipient_address: address,
+          token_transfer_amount: 100n,
+          fee_rate: 50n,
+        }),
+      ],
+    });
+    const balance3 = await supertest(api.server).get(url);
+    expect(balance3.body.balance).toEqual('2000');
+    expect(balance3.body.estimated_balance).toEqual('1850'); // Plus amount
+
+    // Confirm all txs
+    await db.update(
+      new TestBlockBuilder({
+        block_height: 2,
+        index_block_hash: '0x02',
+        parent_index_block_hash: '0x01',
+      })
+        .addTx({
+          tx_id: '0x0002',
+          sender_address: address,
+          token_transfer_amount: 100n,
+          fee_rate: 50n,
+        })
+        .addTxStxEvent({ sender: address, amount: 100n })
+        .addTx({
+          tx_id: '0x0002aa',
+          sender_address: address,
+          type_id: DbTxTypeId.ContractCall,
+          token_transfer_amount: 0n,
+          contract_call_contract_id: 'ST27W5M8BRKA7C5MZE2R1S1F4XTPHFWFRNHA9M04Y.hello-world',
+          contract_call_function_args: '',
+          contract_call_function_name: 'test',
+          fee_rate: 50n,
+        })
+        .addTx({
+          tx_id: '0x0003',
+          sponsor_address: address,
+          sponsored: true,
+          token_transfer_amount: 100n,
+          fee_rate: 50n,
+        })
+        .addTx({
+          tx_id: '0x0004',
+          token_transfer_recipient_address: address,
+          token_transfer_amount: 100n,
+          fee_rate: 50n,
+        })
+        .addTxStxEvent({ recipient: address, amount: 100n })
+        .build()
+    );
+    const balance4 = await supertest(api.server).get(url);
+    expect(balance4.body.balance).toEqual('1850');
+    expect(balance4.body.estimated_balance).toEqual('1850');
+  });
 });

tests/api/tx.test.ts

@@ -1012,6 +1012,7 @@ describe('tx tests', () => {
 
     const expectedSponsoredRespBefore = {
       balance: '0',
+      estimated_balance: '0',
       total_sent: '0',
       total_received: '0',
       total_fees_sent: '0',
@@ -1119,6 +1120,7 @@ describe('tx tests', () => {
 
     const expectedResp = {
       balance: '0',
+      estimated_balance: '0',
       total_sent: '0',
       total_received: '0',
       total_fees_sent: '0',
@@ -1137,6 +1139,7 @@ describe('tx tests', () => {
     const expectedRespBalance = {
       stx: {
         balance: '0',
+        estimated_balance: '0',
         total_sent: '0',
         total_received: '0',
         total_fees_sent: '0',
@@ -1159,6 +1162,7 @@ describe('tx tests', () => {
 
     const expectedSponsoredRespAfter = {
       balance: '-300',
+      estimated_balance: '-300',
       total_sent: '0',
       total_received: '0',
       total_fees_sent: '300',

tests/utils/test-builders.ts

@@ -302,6 +302,7 @@ interface TestMempoolTxArgs {
   fee_rate?: bigint;
   raw_tx?: string;
   sponsor_address?: string;
+  sponsored?: boolean;
   receipt_time?: number;
 }
 
@@ -322,7 +323,7 @@ export function testMempoolTx(args?: TestMempoolTxArgs): DbMempoolTxRaw {
     status: args?.status ?? DbTxStatus.Pending,
     post_conditions: '0x01f5',
     fee_rate: args?.fee_rate ?? 1234n,
-    sponsored: false,
+    sponsored: args?.sponsored ?? false,
     sponsor_address: args?.sponsor_address,
     origin_hash_mode: 1,
     sender_address: args?.sender_address ?? SENDER_ADDRESS,