fix(solana): retry Helius 429 once per fetch

blockapi/test/v2/api/test_solana.py

@@ -3,7 +3,7 @@
 from unittest.mock import call, patch
 
 import pytest
-from requests_mock import ANY, Mocker
+from requests_mock import Mocker
 
 from blockapi.test.v2.api.conftest import read_file
 from blockapi.v2.api import SolanaApi, SolscanApi
@@ -77,10 +77,41 @@ def test_use_base_url():
 
 
 @pytest.mark.parametrize(
-    ('rpc_url', 'uses_v2_staking'),
+    ('rpc_url', 'status', 'retry_after'),
     [
-        ('https://mainnet.helius-rpc.com/', True),
-        ('https://proxy/solana/', False),
+        ('https://proxy/solana/', 429, None),
+        ('https://not-helius-rpc.com/', 429, None),
+        ('https://helius-rpc.com.example.org/', 429, None),
+        ('https://mainnet.helius-rpc.com/', 400, None),
+        ('https://mainnet.helius-rpc.com/', 500, None),
+        ('https://mainnet.helius-rpc.com/', 429, '5'),
+        ('https://mainnet.helius-rpc.com/', 429, 'Fri, 31 Dec 9999 23:59:59 GMT'),
+        ('https://mainnet.helius-rpc.com/', 429, 'invalid'),
+    ],
+)
+def test_fetch_does_not_retry_other_errors_or_long_cooldowns(
+    rpc_url, status, retry_after
+):
+    headers = {'Retry-After': retry_after} if retry_after is not None else {}
+    with Mocker() as m, patch('time.sleep') as sleep:
+        m.post(rpc_url, status_code=status, headers=headers)
+        with pytest.raises(ApiException):
+            SolanaApi(base_url=rpc_url).fetch_balances('address')
+
+    assert len(m.request_history) == 1
+    sleep.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ('rpc_url', 'uses_v2_staking', 'retry_index', 'retry_after'),
+    [
+        ('https://mainnet.helius-rpc.com/', True, None, None),
+        ('https://proxy/solana/', False, None, None),
+        ('https://mainnet.helius-rpc.com/', True, 0, None),
+        ('https://mainnet.helius-rpc.com/', True, 1, '0'),
+        ('https://mainnet.helius-rpc.com/', True, 2, '1'),
+        ('https://mainnet.helius-rpc.com/', True, 3, None),
+        ('https://mainnet.helius-rpc.com/', True, 5, None),
     ],
 )
 def test_get_balance_supports_helius_and_legacy_staking_responses(
@@ -90,33 +121,48 @@ def test_get_balance_supports_helius_and_legacy_staking_responses(
     staked_solana_response,
     rpc_url,
     uses_v2_staking,
+    retry_index,
+    retry_after,
 ):
     test_addr = '5PjMxaijeVVQtuEzxK2NxyJeWwUbpTsi2uXuZ653WoHu'
     empty_token_accounts = '{"jsonrpc":"2.0","result":{"context":{"apiVersion":"1.17.34","slot":268207149},"value":[]},"id":1}'
     staking_response = json.loads(staked_solana_response)
     if not uses_v2_staking:
         staking_response['result'] = staking_response['result']['accounts']
 
-    iterator = iter(
-        [
+    responses = [
+        {'text': data}
+        for data in [
             sol_balance_response,
             token_accounts_response,
             empty_token_accounts,
             das_asset_batch_response,
             json.dumps(staking_response),
         ]
-    )
-
-    def get_text(*args, **kwargs):
-        assert args[0].url == rpc_url
-        data = next(iterator)
-        return data
-
-    with Mocker() as m:
-        m.post(ANY, text=get_text)
+    ]
+    if uses_v2_staking:
+        responses.insert(
+            4, {'json': {'result': {'accounts': [], 'paginationKey': 'next-page'}}}
+        )
+    if retry_index is not None:
+        headers = {'Retry-After': retry_after} if retry_after is not None else {}
+        responses.insert(retry_index, {'status_code': 429, 'headers': headers})
+
+    with Mocker() as m, patch('time.sleep') as sleep:
+        m.post(rpc_url, responses)
         api = SolanaApi(base_url=rpc_url)
         balances = api.get_balance(test_addr)
 
+    if retry_index is not None:
+        sleep.assert_called_once_with(1.1)
+        assert len(m.request_history) == len(responses)
+        assert (
+            m.request_history[retry_index].body
+            == m.request_history[retry_index + 1].body
+        )
+    else:
+        sleep.assert_not_called()
+
     staking_balances = {
         balance.asset_type: balance.balance_raw
         for balance in balances
@@ -128,6 +174,47 @@ def get_text(*args, **kwargs):
     }
 
 
+@pytest.mark.parametrize('fail_on_same_rpc', [True, False])
+def test_retry_is_shared_across_fetch_and_resets_for_next_fetch(
+    token_accounts_response, fail_on_same_rpc
+):
+    api = SolanaApi(base_url='https://mainnet.helius-rpc.com/')
+    responses = [{'status_code': 429}]
+    if not fail_on_same_rpc:
+        responses += [
+            {'json': {'result': {'value': 0}}},
+            {'text': token_accounts_response},
+            {'json': {'result': {'value': []}}},
+            {'status_code': 429},  # Optional DAS metadata may fail.
+            {'json': {'result': {'accounts': [], 'paginationKey': 'next-page'}}},
+        ]
+    responses.append({'status_code': 429})
+
+    with Mocker() as m, patch('time.sleep') as sleep:
+        for attempt in range(2):
+            m.post(api.base_url, responses)
+            with pytest.raises(ApiException):
+                api.fetch_balances('address')
+            assert len(m.request_history) == (attempt + 1) * len(responses)
+            if not fail_on_same_rpc:
+                assert (
+                    m.last_request.json()['params'][1]['paginationKey'] == 'next-page'
+                )
+
+    assert sleep.call_args_list == [call(1.1), call(1.1)]
+
+
+def test_get_coin_retries_helius_metadata_once():
+    api = SolanaApi(base_url='https://mainnet.helius-rpc.com/')
+    with Mocker() as m, patch('time.sleep') as sleep:
+        m.post(api.base_url, [{'status_code': 429}, {'json': {'result': []}}])
+        coin = api.get_coin(('mint', 6))
+
+    assert coin.address == 'mint'
+    assert len(m.request_history) == 2
+    sleep.assert_called_once_with(1.1)
+
+
 def test_build_coin_from_das_asset():
     api = SolanaApi()
     asset = {
@@ -317,20 +404,23 @@ def test_fetch_staked_sol_uses_v2_pagination():
         call(
             method='getProgramAccountsV2',
             params=[api.STAKE_PROGRAM_ID, config],
+            retry=None,
         ),
         call(
             method='getProgramAccountsV2',
             params=[
                 api.STAKE_PROGRAM_ID,
                 {**config, 'paginationKey': 'next-page'},
             ],
+            retry=None,
         ),
         call(
             method='getProgramAccountsV2',
             params=[
                 api.STAKE_PROGRAM_ID,
                 {**config, 'paginationKey': 'last-page'},
             ],
+            retry=None,
         ),
     ]
     assert response == {

blockapi/v2/api/solana.py

@@ -1,5 +1,7 @@
 import json
 import logging
+import time
+from dataclasses import dataclass
 from typing import Optional, Union
 from urllib.parse import urlparse
 
@@ -32,6 +34,11 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class _HeliusRetry:
+    used: bool = False
+
+
 class SolanaApi(CustomizableBlockchainApi, BalanceMixin):
     """Solana JSON-RPC client with DAS metadata integration.
 
@@ -67,7 +74,7 @@ class SolanaApi(CustomizableBlockchainApi, BalanceMixin):
     api_options = ApiOptions(
         blockchain=Blockchain.SOLANA,
         base_url='https://api.mainnet-beta.solana.com/',
-        rate_limit=1,
+        rate_limit=2,
         start_offset=0,
         max_items_per_page=1000,
         page_offset_step=1,
@@ -81,6 +88,7 @@ class SolanaApi(CustomizableBlockchainApi, BalanceMixin):
     STAKE_PROGRAM_ID = 'Stake11111111111111111111111111111111111111'
     STAKE_AUTHORITY_OFFSET = 44
     HELIUS_RPC_DOMAIN = 'helius-rpc.com'
+    HELIUS_RETRY_DELAY = 1.1
     HELIUS_PROGRAM_ACCOUNTS_PAGE_SIZE = 5000
     DAS_BATCH_SIZE = 1000
     _JSONRPC_INVALID_PARAMS = -32602
@@ -104,7 +112,8 @@ def __init__(
 
     def fetch_balances(self, address: str) -> FetchResult:
         """Fetch native SOL, token accounts, DAS metadata, and staking data."""
-        sol_response = self._request('getBalance', [address])
+        retry = _HeliusRetry()
+        sol_response = self._request('getBalance', [address], retry=retry)
 
         raw_token_balances = self._request(
             'getTokenAccountsByOwner',
@@ -113,6 +122,7 @@ def fetch_balances(self, address: str) -> FetchResult:
                 {'programId': self.TOKEN_PROGRAM_ID},
                 {'encoding': 'jsonParsed'},
             ],
+            retry=retry,
         )
         raw_token2022_balances = self._request(
             'getTokenAccountsByOwner',
@@ -121,14 +131,15 @@ def fetch_balances(self, address: str) -> FetchResult:
                 {'programId': self.TOKEN_2022_PROGRAM_ID},
                 {'encoding': 'jsonParsed'},
             ],
+            retry=retry,
         )
 
         mint_addresses = self._collect_mint_addresses(
             raw_token_balances, raw_token2022_balances
         )
-        self._fetch_das_assets(mint_addresses)
+        self._fetch_das_assets(mint_addresses, retry=retry)
 
-        raw_staked_sol = self._fetch_staked_sol(address)
+        raw_staked_sol = self._fetch_staked_sol(address, retry=retry)
 
         return FetchResult(
             data=sol_response,
@@ -172,7 +183,7 @@ def parse_balances(self, fetch_result: FetchResult) -> ParseResult:
     def get_coin(self, fetch_params: tuple[str, int]) -> Coin:
         """Fetch and build a Coin for a given contract address and decimals."""
         contract, decimals = fetch_params
-        self._fetch_das_assets([contract])
+        self._fetch_das_assets([contract], retry=_HeliusRetry())
         return self._resolve_coin(contract, decimals)
 
     @staticmethod
@@ -248,7 +259,9 @@ def _extract_token_info(account: dict) -> dict:
 
     # ── DAS integration ────────────────────────────────────────
 
-    def _fetch_das_assets(self, mint_addresses: list[str]) -> None:
+    def _fetch_das_assets(
+        self, mint_addresses: list[str], retry: Optional[_HeliusRetry] = None
+    ) -> None:
         """Batch-fetch token metadata via DAS and populate cache."""
         uncached = list(
             dict.fromkeys(m for m in mint_addresses if m not in self._das_cache)
@@ -262,6 +275,7 @@ def _fetch_das_assets(self, mint_addresses: list[str]) -> None:
                 response = self._request(
                     'getAssetBatch',
                     {'ids': chunk, 'options': {'showFungible': True}},
+                    retry=retry,
                 )
             except (ApiException, RequestException) as e:
                 logger.warning('DAS getAssetBatch failed: %s', e)
@@ -315,7 +329,9 @@ def _build_coin_from_das_asset(self, asset: dict) -> Optional[Coin]:
 
     # ── Staking ────────────────────────────────────────────────
 
-    def _fetch_staked_sol(self, address: str) -> dict:
+    def _fetch_staked_sol(
+        self, address: str, retry: Optional[_HeliusRetry] = None
+    ) -> dict:
         """Fetch staked SOL accounts for a given address."""
         config = {
             'filters': [
@@ -331,11 +347,7 @@ def _fetch_staked_sol(self, address: str) -> dict:
             'commitment': 'finalized',
         }
 
-        hostname = urlparse(self.base_url).hostname or ''
-        is_helius_rpc = hostname == self.HELIUS_RPC_DOMAIN or hostname.endswith(
-            f'.{self.HELIUS_RPC_DOMAIN}'
-        )
-        if not is_helius_rpc:
+        if not self._is_helius_rpc():
             return self._request(
                 method='getProgramAccounts',
                 params=[self.STAKE_PROGRAM_ID, config],
@@ -349,6 +361,7 @@ def _fetch_staked_sol(self, address: str) -> dict:
             response = self._request(
                 method='getProgramAccountsV2',
                 params=[self.STAKE_PROGRAM_ID, config],
+                retry=retry,
             )
             page = response['result']
             accounts.extend(page['accounts'])
@@ -419,7 +432,12 @@ def _parse_rent_reserve(
 
     # ── Infrastructure ─────────────────────────────────────────
 
-    def _request(self, method: str, params: Union[list, dict]) -> dict:
+    def _request(
+        self,
+        method: str,
+        params: Union[list, dict],
+        retry: Optional[_HeliusRetry] = None,
+    ) -> dict:
         """Send a JSON-RPC request to the Solana RPC endpoint."""
         self._request_id += 1
         body = json.dumps(
@@ -430,7 +448,31 @@ def _request(self, method: str, params: Union[list, dict]) -> dict:
                 'params': params,
             }
         )
-        return self.post(body=body, headers={'Content-Type': 'application/json'})
+        headers = {'Content-Type': 'application/json'}
+        if not self._is_helius_rpc():
+            return self.post(body=body, headers=headers)
+
+        response = self._session.post(self.base_url, data=body, headers=headers)
+        if (
+            response.status_code == 429
+            and retry is not None
+            and not retry.used
+            and response.headers.get('Retry-After', '0').strip() in ('0', '1')
+        ):
+            retry.used = True
+            response.close()
+            logger.warning(
+                f'Helius {method} rate limited; retrying in {self.HELIUS_RETRY_DELAY}s'
+            )
+            time.sleep(self.HELIUS_RETRY_DELAY)
+            response = self._session.post(self.base_url, data=body, headers=headers)
+        return self._check_and_get_from_response(response)
+
+    def _is_helius_rpc(self) -> bool:
+        hostname = urlparse(self.base_url).hostname or ''
+        return hostname == self.HELIUS_RPC_DOMAIN or hostname.endswith(
+            f'.{self.HELIUS_RPC_DOMAIN}'
+        )
 
     def _opt_raise_on_other_error(self, response: Response) -> None:
         """Raise ApiException or InvalidAddressException on RPC errors."""