GetBalance now returns per-mint balances with the ability to filter

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/dynamodb v1.59.0
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.102.2
 	github.com/code-payments/code-vm-indexer v1.2.0
-	github.com/code-payments/ocp-protobuf-api v1.16.0
+	github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54
 	github.com/emirpasic/gods v1.12.0
 	github.com/envoyproxy/protoc-gen-validate v1.3.3
 	github.com/golang/protobuf v1.5.4

go.sum

@@ -80,8 +80,8 @@ github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
 github.com/code-payments/code-vm-indexer v1.2.0 h1:rSHpBMiT9BKgmKcXg/VIoi/h0t7jNxGx07Qz59m+6Q0=
 github.com/code-payments/code-vm-indexer v1.2.0/go.mod h1:vn91YN2qNqb+gGJeZe2+l+TNxVmEEiRHXXnIn2Y40h8=
-github.com/code-payments/ocp-protobuf-api v1.16.0 h1:j8Z5TRsNd8b02mCjJSq8+e1CoRh0JCnRVDldTlb9+Ds=
-github.com/code-payments/ocp-protobuf-api v1.16.0/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8=
+github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54 h1:idVNVzNMZWunI9RlDid6GNto1+z8HVApQhHqmfRIAPc=
+github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8=
 github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6 h1:NmTXa/uVnDyp0TY5MKi197+3HWcnYWfnHGyaFthlnGw=
 github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6/go.mod h1:GL3xCUCBDV3CZiTSEKksMWbLE66hEyuu9qyDOOqM47Y=
 github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk=

ocp/rpc/balance/server.go

@@ -45,6 +45,19 @@ func (s *server) GetBalance(ctx context.Context, req *balancepb.GetBalanceReques
 	}
 	log = log.With(zap.String("owner_account", owner.PublicKey().ToBase58()))
 
+	var mintFilter map[string]struct{}
+	if len(req.Mints) > 0 {
+		mintFilter = make(map[string]struct{}, len(req.Mints))
+		for i, protoMint := range req.Mints {
+			mint, err := common.NewAccountFromProto(protoMint)
+			if err != nil {
+				log.With(zap.Error(err), zap.Int("index", i)).Warn("invalid mint account")
+				return nil, status.Error(codes.Internal, "")
+			}
+			mintFilter[mint.PublicKey().ToBase58()] = struct{}{}
+		}
+	}
+
 	ownerMetadata, err := common.GetOwnerMetadata(ctx, s.data, owner)
 	if err == common.ErrOwnerNotFound {
 		return &balancepb.GetBalanceResponse{
@@ -61,61 +74,79 @@ func (s *server) GetBalance(ctx context.Context, req *balancepb.GetBalanceReques
 		}, nil
 	}
 
-	coreMintValue, err := s.calculateCoreMintValue(ctx, owner)
+	balancesByMint, err := s.calculateCoreMintValueByMint(ctx, owner, mintFilter)
 	if err != nil {
 		log.With(zap.Error(err)).Warn("failure calculating core mint value")
 		return nil, status.Error(codes.Internal, "")
 	}
 
+	var totalCoreMintValue uint64
+	for _, mintBalance := range balancesByMint {
+		totalCoreMintValue += mintBalance.CoreMintValue
+	}
+
 	return &balancepb.GetBalanceResponse{
-		Result:        balancepb.GetBalanceResponse_OK,
-		CoreMintValue: coreMintValue,
+		Result:         balancepb.GetBalanceResponse_OK,
+		CoreMintValue:  totalCoreMintValue,
+		BalancesByMint: balancesByMint,
 	}, nil
 }
 
-func (s *server) calculateCoreMintValue(ctx context.Context, owner *common.Account) (uint64, error) {
+// calculateCoreMintValueByMint values the owner's holdings in each mint they
+// hold a non-zero balance of. A nil mintFilter includes every mint, otherwise
+// only mints in the filter are included.
+func (s *server) calculateCoreMintValueByMint(ctx context.Context, owner *common.Account, mintFilter map[string]struct{}) (map[string]*balancepb.MintBalance, error) {
 	// The ledger holds a record for every account Code manages for the owner,
 	// and each carries the mint it holds. Accounts that have left the L2 system
 	// don't have a cached balance that can be trusted, so it omits them.
 	balanceByTokenAccount, err := balance.BatchCalculateFromCacheByOwner(ctx, s.data, owner)
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
 
 	quarksByMint := make(map[string]uint64)
 	for _, cached := range balanceByTokenAccount {
+		if mintFilter != nil {
+			if _, ok := mintFilter[cached.MintAccount]; !ok {
+				continue
+			}
+		}
 		quarksByMint[cached.MintAccount] += cached.Quarks
 	}
 
-	var coreMintValue uint64
+	balancesByMint := make(map[string]*balancepb.MintBalance)
 	for mint, quarks := range quarksByMint {
 		if quarks == 0 {
 			continue
 		}
 
-		if mint == common.CoreMintAccount.PublicKey().ToBase58() {
-			coreMintValue += quarks
-			continue
-		}
-
 		mintAccount, err := common.NewAccountFromPublicKeyString(mint)
 		if err != nil {
-			return 0, err
+			return nil, err
 		}
 
-		reserveState, err := s.mintDataProvider.GetLiveReserveState(ctx, mintAccount)
-		if err != nil {
-			return 0, err
+		var coreMintValue uint64
+		if mint == common.CoreMintAccount.PublicKey().ToBase58() {
+			coreMintValue = quarks
+		} else {
+			reserveState, err := s.mintDataProvider.GetLiveReserveState(ctx, mintAccount)
+			if err != nil {
+				return nil, err
+			}
+
+			coreMintValue, _ = currencycreator.EstimateSell(&currencycreator.EstimateSellArgs{
+				CurrentSupplyInQuarks: reserveState.SupplyFromBonding,
+				SellAmountInQuarks:    quarks,
+				ValueMintDecimals:     uint8(common.CoreMintDecimals),
+				SellFeeBps:            0,
+			})
 		}
 
-		sellValue, _ := currencycreator.EstimateSell(&currencycreator.EstimateSellArgs{
-			CurrentSupplyInQuarks: reserveState.SupplyFromBonding,
-			SellAmountInQuarks:    quarks,
-			ValueMintDecimals:     uint8(common.CoreMintDecimals),
-			SellFeeBps:            0,
-		})
-		coreMintValue += sellValue
+		balancesByMint[mint] = &balancepb.MintBalance{
+			Mint:          mintAccount.ToProto(),
+			CoreMintValue: coreMintValue,
+		}
 	}
 
-	return coreMintValue, nil
+	return balancesByMint, nil
 }

ocp/rpc/balance/server_test.go

@@ -97,6 +97,7 @@ func TestGetBalance_HappyPath(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
 	assert.EqualValues(t, 0, resp.CoreMintValue)
+	assert.Empty(t, resp.BalancesByMint)
 
 	setupCachedBalance(t, env, primaryCoreMintAccountRecords, common.ToCoreMintQuarks(42))
 	setupCachedBalance(t, env, primaryLaunchpadMintAccountRecords, currencycreator.ToQuarks(100))
@@ -115,6 +116,77 @@ func TestGetBalance_HappyPath(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
 	assert.EqualValues(t, common.ToCoreMintQuarks(42)+expectedLaunchpadMintValue, resp.CoreMintValue)
+	require.Len(t, resp.BalancesByMint, 2)
+	assertMintBalance(t, resp, common.CoreMintAccount, common.ToCoreMintQuarks(42))
+	assertMintBalance(t, resp, launchpadMint, expectedLaunchpadMintValue)
+}
+
+func TestGetBalance_MintFilter(t *testing.T) {
+	env, cleanup := setup(t)
+	defer cleanup()
+
+	coreVmConfig := testutil.NewRandomVmConfig(t, true)
+	launchpadMint := testutil.SetupLaunchpadCurrency(t, env.data, env.reserveStore, env.holderStore)
+	launchpadVmConfig, err := common.GetVmConfigForMint(env.ctx, env.data, launchpadMint)
+	require.NoError(t, err)
+
+	ownerAccount := testutil.NewRandomAccount(t)
+
+	primaryCoreMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, coreVmConfig, 0, commonpb.AccountType_PRIMARY)
+	primaryLaunchpadMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, launchpadVmConfig, 0, commonpb.AccountType_PRIMARY)
+
+	setupCachedBalance(t, env, primaryCoreMintAccountRecords, common.ToCoreMintQuarks(42))
+	setupCachedBalance(t, env, primaryLaunchpadMintAccountRecords, currencycreator.ToQuarks(100))
+
+	expectedLaunchpadMintValue, _ := currencycreator.EstimateSell(&currencycreator.EstimateSellArgs{
+		CurrentSupplyInQuarks: currencycreator.ToQuarks(1_000),
+		SellAmountInQuarks:    currencycreator.ToQuarks(100),
+		ValueMintDecimals:     uint8(common.CoreMintDecimals),
+		SellFeeBps:            0,
+	})
+	require.NotZero(t, expectedLaunchpadMintValue)
+
+	// Filtering to the core mint excludes the launchpad currency from both the
+	// total and the per-mint breakdown
+	resp, err := env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{
+		Owner: ownerAccount.ToProto(),
+		Mints: []*commonpb.SolanaAccountId{common.CoreMintAccount.ToProto()},
+	})
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
+	assert.EqualValues(t, common.ToCoreMintQuarks(42), resp.CoreMintValue)
+	require.Len(t, resp.BalancesByMint, 1)
+	assertMintBalance(t, resp, common.CoreMintAccount, common.ToCoreMintQuarks(42))
+
+	resp, err = env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{
+		Owner: ownerAccount.ToProto(),
+		Mints: []*commonpb.SolanaAccountId{launchpadMint.ToProto()},
+	})
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
+	assert.EqualValues(t, expectedLaunchpadMintValue, resp.CoreMintValue)
+	require.Len(t, resp.BalancesByMint, 1)
+	assertMintBalance(t, resp, launchpadMint, expectedLaunchpadMintValue)
+
+	// Duplicate mints in the filter don't double count
+	resp, err = env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{
+		Owner: ownerAccount.ToProto(),
+		Mints: []*commonpb.SolanaAccountId{common.CoreMintAccount.ToProto(), launchpadMint.ToProto(), common.CoreMintAccount.ToProto()},
+	})
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
+	assert.EqualValues(t, common.ToCoreMintQuarks(42)+expectedLaunchpadMintValue, resp.CoreMintValue)
+	require.Len(t, resp.BalancesByMint, 2)
+
+	// A mint the owner doesn't hold results in an OK response with no balances
+	resp, err = env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{
+		Owner: ownerAccount.ToProto(),
+		Mints: []*commonpb.SolanaAccountId{testutil.NewRandomAccount(t).ToProto()},
+	})
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
+	assert.EqualValues(t, 0, resp.CoreMintValue)
+	assert.Empty(t, resp.BalancesByMint)
 }
 
 func TestGetBalance_UnmanagedAccountsExcluded(t *testing.T) {
@@ -206,6 +278,13 @@ func setupAccountRecords(t *testing.T, env testEnv, ownerAccount, authorityAccou
 	}
 }
 
+func assertMintBalance(t *testing.T, resp *balancepb.GetBalanceResponse, mint *common.Account, expectedCoreMintValue uint64) {
+	mintBalance, ok := resp.BalancesByMint[mint.PublicKey().ToBase58()]
+	require.True(t, ok)
+	assert.Equal(t, mint.PublicKey().ToBytes(), mintBalance.Mint.Value)
+	assert.EqualValues(t, expectedCoreMintValue, mintBalance.CoreMintValue)
+}
+
 func setupCachedBalance(t *testing.T, env testEnv, accountRecords *common.AccountRecords, quarks uint64) {
 	require.NoError(t, balance_util.ApplyDeltasInTx(env.ctx, env.data, &balance.Delta{
 		TokenAccount: accountRecords.General.TokenAccount,