Add configurable balance checking in token account preparation (#971)

* Update TS SDK. Add tests

* Rust SDK. Update tests.

* Changeset

* Format

* Update test

* Revert unwanted changes

Author: calintje

.changeset/flat-toes-post.md

@@ -0,0 +1,6 @@
+---
+"@orca-so/whirlpools-rust": minor
+"@orca-so/whirlpools": minor
+---
+
+Add configurable balance checking in token account preparation, allowing users to disable balance validation to get quotes and instructions even with insufficient token balances.

rust-sdk/whirlpool/src/config.rs

@@ -138,6 +138,25 @@ pub fn set_native_mint_wrapping_strategy(
     Ok(())
 }
 
+/// The default setting for enforcing balance checks during token account preparation.
+pub const DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK: bool = false;
+
+/// The currently selected setting for enforcing balance checks during token account preparation.
+/// When true, the system will assert that token accounts have sufficient balance before proceeding.
+/// When false, balance checks are skipped, allowing users to get quotes and instructions even with insufficient balance.
+pub static ENFORCE_TOKEN_BALANCE_CHECK: Mutex<bool> =
+    Mutex::new(DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK);
+
+/// Sets whether to enforce balance checks during token account preparation.
+///
+/// # Arguments
+///
+/// * `enforce_balance_check` - When true, the system will assert that token accounts have sufficient balance. When false, balance checks are skipped.
+pub fn set_enforce_token_balance_check(enforce_balance_check: bool) -> Result<(), Box<dyn Error>> {
+    *ENFORCE_TOKEN_BALANCE_CHECK.try_lock()? = enforce_balance_check;
+    Ok(())
+}
+
 /// Resets the configuration to its default values.
 pub fn reset_configuration() -> Result<(), Box<dyn Error>> {
     *WHIRLPOOLS_CONFIG_ADDRESS.try_lock()? = SOLANA_MAINNET_WHIRLPOOLS_CONFIG_ADDRESS;
@@ -146,6 +165,7 @@ pub fn reset_configuration() -> Result<(), Box<dyn Error>> {
     *FUNDER.try_lock()? = DEFAULT_FUNDER;
     *NATIVE_MINT_WRAPPING_STRATEGY.try_lock()? = DEFAULT_NATIVE_MINT_WRAPPING_STRATEGY;
     *SLIPPAGE_TOLERANCE_BPS.try_lock()? = DEFAULT_SLIPPAGE_TOLERANCE_BPS;
+    *ENFORCE_TOKEN_BALANCE_CHECK.try_lock()? = DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK;
     Ok(())
 }
 
@@ -214,6 +234,18 @@ mod tests {
         reset_configuration().unwrap();
     }
 
+    #[test]
+    #[serial]
+    fn test_set_enforce_token_balance_check() {
+        set_enforce_token_balance_check(true).unwrap();
+        assert!(*ENFORCE_TOKEN_BALANCE_CHECK.lock().unwrap());
+
+        set_enforce_token_balance_check(false).unwrap();
+        assert!(!(*ENFORCE_TOKEN_BALANCE_CHECK.lock().unwrap()));
+
+        reset_configuration().unwrap();
+    }
+
     #[test]
     #[serial]
     fn test_reset_configuration() {
@@ -231,5 +263,9 @@ mod tests {
             NativeMintWrappingStrategy::Keypair
         );
         assert_eq!(*SLIPPAGE_TOLERANCE_BPS.lock().unwrap(), 100);
+        assert_eq!(
+            *ENFORCE_TOKEN_BALANCE_CHECK.lock().unwrap(),
+            DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK
+        );
     }
 }

rust-sdk/whirlpool/src/increase_liquidity.rs

@@ -992,7 +992,7 @@ mod tests {
 
     #[tokio::test]
     #[serial]
-    async fn test_increase_liquidity_fails_if_deposit_exceeds_user_balance(
+    async fn test_increase_liquidity_succeeds_if_deposit_exceeds_user_balance_when_balance_check_not_enforced(
     ) -> Result<(), Box<dyn Error>> {
         let ctx = RpcContext::new().await;
 
@@ -1015,9 +1015,43 @@ mod tests {
         )
         .await;
 
+        assert!(
+            res.is_ok(),
+            "Should succeed when balance checking is disabled even if deposit exceeds balance"
+        );
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_increase_liquidity_fails_if_deposit_exceeds_user_balance_when_balance_check_enforced(
+    ) -> Result<(), Box<dyn Error>> {
+        let ctx = RpcContext::new().await;
+        crate::set_enforce_token_balance_check(true)?;
+
+        let minted = setup_all_mints(&ctx).await?;
+        let user_atas = setup_all_atas(&ctx, &minted).await?;
+
+        let mint_a_key = minted.get("A").unwrap();
+        let mint_b_key = minted.get("B").unwrap();
+        let pool_pubkey = setup_whirlpool(&ctx, *mint_a_key, *mint_b_key, 64).await?;
+
+        let position_mint = setup_position(&ctx, pool_pubkey, Some((-100, 100)), None).await?;
+
+        // Attempt
+        let res = increase_liquidity_instructions(
+            &ctx.rpc,
+            position_mint,
+            IncreaseLiquidityParam::TokenA(2_000_000_000),
+            Some(100),
+            Some(ctx.signer.pubkey()),
+        )
+        .await;
+
         assert!(
             res.is_err(),
-            "Should fail if user tries depositing more than balance"
+            "Should fail if user tries depositing more than balance when balance checking is enforced"
         );
         let err_str = format!("{:?}", res.err().unwrap());
         assert!(
@@ -1027,6 +1061,7 @@ mod tests {
             err_str
         );
 
+        crate::reset_configuration()?;
         Ok(())
     }
 

rust-sdk/whirlpool/src/token.rs

@@ -20,7 +20,9 @@ use spl_token_2022::ID as TOKEN_2022_PROGRAM_ID;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use std::{collections::HashMap, error::Error};
 
-use crate::{NativeMintWrappingStrategy, NATIVE_MINT_WRAPPING_STRATEGY};
+use crate::{
+    NativeMintWrappingStrategy, ENFORCE_TOKEN_BALANCE_CHECK, NATIVE_MINT_WRAPPING_STRATEGY,
+};
 
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub(crate) enum TokenAccountStrategy {
@@ -124,7 +126,8 @@ pub(crate) async fn prepare_token_accounts_instructions(
             0
         };
 
-        if existing_balance < required_balance {
+        let enforce_balance_check = *ENFORCE_TOKEN_BALANCE_CHECK.try_lock()?;
+        if enforce_balance_check && existing_balance < required_balance {
             return Err(format!("Insufficient balance for mint {}", mint_addresses[i]).into());
         }
     }
@@ -460,10 +463,10 @@ mod tests {
 
     #[tokio::test]
     #[serial]
-    async fn test_insufficient_balance() {
+    async fn test_insufficient_balance_check_not_enforced() {
         let ctx = RpcContext::new().await;
 
-        // Create a mint and token account with small balance using token.rs helpers
+        // Create a mint and token account with small balance
         let mint = setup_mint(&ctx).await.unwrap();
         let initial_amount = 1_000u64;
         let _ata = setup_ata_with_amount(&ctx, mint, initial_amount)
@@ -479,12 +482,87 @@ mod tests {
         )
         .await;
 
-        // Should fail due to insufficient balance
+        // Should succeed because balance checking is disabled
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_insufficient_balance_check_enforced() {
+        let ctx = RpcContext::new().await;
+        crate::set_enforce_token_balance_check(true).unwrap();
+
+        // Create a mint and token account with small balance
+        let mint = setup_mint(&ctx).await.unwrap();
+        let initial_amount = 1_000u64;
+        let _ata = setup_ata_with_amount(&ctx, mint, initial_amount)
+            .await
+            .unwrap();
+
+        // Try to prepare instructions requiring more balance
+        let required_amount = 2_000u64;
+        let result = prepare_token_accounts_instructions(
+            &ctx.rpc,
+            ctx.signer.pubkey(),
+            vec![TokenAccountStrategy::WithBalance(mint, required_amount)],
+        )
+        .await;
+
+        // Should fail due to insufficient balance when checking is enabled
+        assert!(result.is_err());
+        assert!(result
+            .unwrap_err()
+            .to_string()
+            .contains("Insufficient balance"));
+        crate::reset_configuration().unwrap();
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_nonexistent_token_account_balance_check_not_enforced() {
+        let ctx = RpcContext::new().await;
+
+        // Create a mint but no token account
+        let mint = setup_mint(&ctx).await.unwrap();
+
+        // Try to prepare instructions requiring balance for non-existent account
+        let required_amount = 1_000u64;
+        let result = prepare_token_accounts_instructions(
+            &ctx.rpc,
+            ctx.signer.pubkey(),
+            vec![TokenAccountStrategy::WithBalance(mint, required_amount)],
+        )
+        .await;
+
+        // Should succeed because balance checking is disabled
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_nonexistent_token_account_balance_check_enforced() {
+        let ctx = RpcContext::new().await;
+        crate::set_enforce_token_balance_check(true).unwrap();
+
+        // Create a mint but no token account
+        let mint = setup_mint(&ctx).await.unwrap();
+
+        // Try to prepare instructions requiring balance for non-existent account
+        let required_amount = 1_000u64;
+        let result = prepare_token_accounts_instructions(
+            &ctx.rpc,
+            ctx.signer.pubkey(),
+            vec![TokenAccountStrategy::WithBalance(mint, required_amount)],
+        )
+        .await;
+
+        // Should fail due to insufficient balance (0) when checking is enabled
         assert!(result.is_err());
         assert!(result
             .unwrap_err()
             .to_string()
             .contains("Insufficient balance"));
+        crate::reset_configuration().unwrap();
     }
 
     #[tokio::test]

ts-sdk/whirlpool/src/config.ts

@@ -173,6 +173,29 @@ export function setNativeMintWrappingStrategy(
   NATIVE_MINT_WRAPPING_STRATEGY = strategy;
 }
 
+/**
+ * The default setting for enforcing balance checks during token account preparation.
+ */
+export const DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK = false;
+
+/**
+ * The currently selected setting for enforcing balance checks during token account preparation.
+ * When true, the system will assert that token accounts have sufficient balance before proceeding.
+ * When false, balance checks are skipped, allowing users to get quotes and instructions even with insufficient balance.
+ */
+export let ENFORCE_TOKEN_BALANCE_CHECK = DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK;
+
+/**
+ * Sets whether to enforce balance checks during token account preparation.
+ *
+ * @param {boolean} enforceBalanceCheck - When true, the system will assert that token accounts have sufficient balance. When false, balance checks are skipped.
+ */
+export function setEnforceTokenBalanceCheck(
+  enforceBalanceCheck: boolean,
+): void {
+  ENFORCE_TOKEN_BALANCE_CHECK = enforceBalanceCheck;
+}
+
 /**
  * Resets the configuration to its default state.
  *
@@ -185,6 +208,7 @@ export function resetConfiguration() {
   FUNDER = DEFAULT_FUNDER;
   SLIPPAGE_TOLERANCE_BPS = DEFAULT_SLIPPAGE_TOLERANCE_BPS;
   NATIVE_MINT_WRAPPING_STRATEGY = DEFAULT_NATIVE_MINT_WRAPPING_STRATEGY;
+  ENFORCE_TOKEN_BALANCE_CHECK = DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK;
 }
 
 let _payer: KeyPairSigner | undefined;

ts-sdk/whirlpool/src/token.ts

@@ -26,7 +26,10 @@ import {
   getAddressEncoder,
   lamports,
 } from "@solana/kit";
-import { NATIVE_MINT_WRAPPING_STRATEGY } from "./config";
+import {
+  NATIVE_MINT_WRAPPING_STRATEGY,
+  ENFORCE_TOKEN_BALANCE_CHECK,
+} from "./config";
 import {
   getCreateAccountInstruction,
   getCreateAccountWithSeedInstruction,
@@ -143,10 +146,12 @@ export async function prepareTokenAccountsInstructions(
       const existingBalance = tokenAccount.exists
         ? tokenAccount.data.amount
         : 0n;
-      assert(
-        BigInt(spec[mint.address]) <= existingBalance,
-        `Token account for ${mint.address} does not have the required balance`,
-      );
+      if (ENFORCE_TOKEN_BALANCE_CHECK) {
+        assert(
+          BigInt(spec[mint.address]) <= existingBalance,
+          `Token account for ${mint.address} does not have the required balance`,
+        );
+      }
     }
   }
 

ts-sdk/whirlpool/tests/config.test.ts

@@ -15,6 +15,9 @@ import {
   DEFAULT_NATIVE_MINT_WRAPPING_STRATEGY,
   DEFAULT_WHIRLPOOLS_CONFIG_EXTENSION_ADDRESS,
   DEFAULT_WHIRLPOOLS_CONFIG_ADDRESSES,
+  ENFORCE_TOKEN_BALANCE_CHECK,
+  DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK,
+  setEnforceTokenBalanceCheck,
 } from "../src/config";
 import assert from "assert";
 import { address, createKeyPairSignerFromPrivateKeyBytes } from "@solana/kit";
@@ -74,6 +77,14 @@ describe("Configuration", () => {
     assert.strictEqual(NATIVE_MINT_WRAPPING_STRATEGY, "ata");
   });
 
+  it("Should be able to set the enforce token balance check", () => {
+    setEnforceTokenBalanceCheck(true);
+    assert.strictEqual(ENFORCE_TOKEN_BALANCE_CHECK, true);
+
+    setEnforceTokenBalanceCheck(false);
+    assert.strictEqual(ENFORCE_TOKEN_BALANCE_CHECK, false);
+  });
+
   it("Should be able to reset the configuration", () => {
     resetConfiguration();
     assert.strictEqual(
@@ -90,5 +101,9 @@ describe("Configuration", () => {
       NATIVE_MINT_WRAPPING_STRATEGY,
       DEFAULT_NATIVE_MINT_WRAPPING_STRATEGY,
     );
+    assert.strictEqual(
+      ENFORCE_TOKEN_BALANCE_CHECK,
+      DEFAULT_ENFORCE_TOKEN_BALANCE_CHECK,
+    );
   });
 });