refactor(loaders): extract bounded retry/budget pattern into base.py

ccxt_loader and okx each carried their own copy of the same 25-line
retry+budget loop introduced in HKUDS#121. Hoist the pattern into two
helpers — check_budget() and retry_with_budget() — so any future loader
(yfinance, akshare, tushare) inherits the same guarantees: per-call
transient retries, a wall-clock budget, non-transient exceptions left
unwrapped, original exception preserved as __cause__.

ccxt_loader / okx now declare only what is loader-specific (the
exception class for transient errors and the budget env var); the
scheduling, sleep math, and TimeoutError wrapping live in base.py.

Adds 11 helper-level tests + the 11 pre-existing integration tests
still pass against the refactored loaders.

Follow-up to HKUDS#121.

Author: warren618

agent/backtest/loaders/base.py

@@ -1,8 +1,16 @@
-"""DataLoader Protocol and shared exceptions for all data source loaders."""
+"""DataLoader Protocol, shared exceptions, and bounded-retry helpers.
+
+The retry/budget helpers are the canonical pattern for any loader that calls
+a flaky external API: a wall-clock deadline plus a small backoff schedule
+applied only to a declared transient exception class. New loaders should
+import :func:`check_budget` and :func:`retry_with_budget` rather than
+re-implementing the loop.
+"""
 
 from __future__ import annotations
 
-from typing import Protocol, runtime_checkable
+import time
+from typing import Callable, Protocol, TypeVar, runtime_checkable
 
 import pandas as pd
 
@@ -30,6 +38,91 @@ def validate_date_range(start_date: str, end_date: str) -> None:
         raise ValueError(f"start_date ({start_date}) > end_date ({end_date})")
 
 
+# ---------------------------------------------------------------------------
+# Bounded retry / budget helpers (shared by ccxt_loader, okx, and any future
+# loader calling a flaky external API).
+# ---------------------------------------------------------------------------
+
+DEFAULT_BACKOFF: tuple[float, ...] = (0.5, 1.5, 4.0)
+DEFAULT_MAX_RETRIES = 3
+
+
+def check_budget(deadline: float, label: str, budget_s: float | None = None) -> None:
+    """Raise :class:`TimeoutError` if the monotonic clock has crossed ``deadline``.
+
+    Use this between pages of a paginated fetch to fail fast instead of
+    grinding through more requests once the wall-clock budget is gone.
+
+    Args:
+        deadline: ``time.monotonic()`` instant past which we abort.
+        label: Free-form label used in the exception message
+            (e.g. ``"ccxt fetch for BTC/USDT"``).
+        budget_s: Original budget in seconds, included verbatim in the
+            message when present.
+    """
+    if time.monotonic() > deadline:
+        suffix = f" exceeded {budget_s:.0f}s budget" if budget_s is not None else " exceeded budget"
+        raise TimeoutError(f"{label}{suffix}")
+
+
+_T = TypeVar("_T")
+
+
+def retry_with_budget(
+    fn: Callable[[], _T],
+    *,
+    transient: type[BaseException] | tuple[type[BaseException], ...],
+    deadline: float,
+    label: str,
+    max_retries: int = DEFAULT_MAX_RETRIES,
+    backoff: tuple[float, ...] = DEFAULT_BACKOFF,
+) -> _T:
+    """Call ``fn`` with a bounded retry budget on declared transient errors.
+
+    Between attempts sleeps ``min(backoff[attempt], remaining_budget)`` so a
+    short remaining budget never spends the full backoff. The terminal
+    transient failure — whether ``max_retries`` is exhausted OR the deadline
+    has passed — is wrapped in :class:`TimeoutError`, preserving the original
+    exception as ``__cause__``. Anything not in ``transient`` propagates
+    unchanged on the first occurrence (we never retry an exception class
+    the caller didn't opt in to).
+
+    Args:
+        fn: Zero-arg callable producing the result.
+        transient: Exception class(es) considered transient and retryable.
+        deadline: ``time.monotonic()`` instant past which retries are aborted.
+        label: Free-form label used in the TimeoutError message
+            (e.g. ``"OKX fetch for BTC-USDT"``).
+        max_retries: Additional attempts after the first call. Total
+            attempts = ``max_retries + 1``.
+        backoff: Per-retry sleep seconds. Must have at least
+            ``max_retries`` entries.
+
+    Returns:
+        Whatever ``fn`` returns.
+
+    Raises:
+        ValueError: ``backoff`` is shorter than ``max_retries``.
+        TimeoutError: All retries exhausted or the deadline crossed.
+        Any non-transient exception: Propagated unchanged from ``fn``.
+    """
+    if len(backoff) < max_retries:
+        raise ValueError(
+            f"backoff has {len(backoff)} entries; need >= max_retries ({max_retries})"
+        )
+    for attempt in range(max_retries + 1):
+        try:
+            return fn()
+        except transient as exc:
+            remaining = deadline - time.monotonic()
+            if attempt == max_retries or remaining <= 0:
+                raise TimeoutError(
+                    f"{label} failed after {attempt + 1} attempt(s): {exc}"
+                ) from exc
+            time.sleep(min(backoff[attempt], max(0.0, remaining)))
+    raise AssertionError("unreachable: retry loop must return or raise")  # pragma: no cover
+
+
 @runtime_checkable
 class DataLoaderProtocol(Protocol):
     """Interface that every data source loader must satisfy."""

agent/backtest/loaders/ccxt_loader.py

@@ -14,7 +14,11 @@
 
 import pandas as pd
 
-from backtest.loaders.base import validate_date_range
+from backtest.loaders.base import (
+    check_budget,
+    retry_with_budget,
+    validate_date_range,
+)
 from backtest.loaders.registry import register
 
 logger = logging.getLogger(__name__)
@@ -27,11 +31,10 @@
 # P12-b: ccxt had no request timeout and an unbounded paginated fetch with
 # no retry budget, so a transient disconnect hung get_market_data for 10+
 # minutes. Cap each HTTP call, bound transient retries, and enforce a hard
-# wall-clock budget so the fetch fails fast instead of hanging.
+# wall-clock budget so the fetch fails fast instead of hanging. Retry
+# scheduling is delegated to :mod:`backtest.loaders.base`.
 _CCXT_TIMEOUT_MS = int(os.getenv("CCXT_TIMEOUT_MS", "15000"))
 _CCXT_FETCH_BUDGET_S = float(os.getenv("CCXT_FETCH_BUDGET_S", "60"))
-_CCXT_MAX_RETRIES = 3
-_CCXT_BACKOFF = (0.5, 1.5, 4.0)  # seconds; len == _CCXT_MAX_RETRIES
 
 
 @register
@@ -113,31 +116,19 @@ def _fetch_one(
         cursor = since_ms
         limit = 1000
         deadline = time.monotonic() + _CCXT_FETCH_BUDGET_S
+        label = f"ccxt fetch for {symbol}"
 
         for _ in range(200):
-            if time.monotonic() > deadline:
-                raise TimeoutError(
-                    f"ccxt fetch for {symbol} exceeded "
-                    f"{_CCXT_FETCH_BUDGET_S:.0f}s budget"
-                )
-            ohlcv = None
-            for attempt in range(_CCXT_MAX_RETRIES + 1):
-                try:
-                    ohlcv = exchange.fetch_ohlcv(
-                        symbol, timeframe, since=cursor, limit=limit
-                    )
-                    break
-                except ccxt.NetworkError as exc:
-                    # NetworkError covers RequestTimeout / DDoSProtection /
-                    # ExchangeNotAvailable — the transient family. Anything
-                    # else (e.g. ExchangeError: bad symbol) is not retried.
-                    remaining = deadline - time.monotonic()
-                    if attempt == _CCXT_MAX_RETRIES or remaining <= 0:
-                        raise TimeoutError(
-                            f"ccxt fetch for {symbol} failed after "
-                            f"{attempt + 1} attempt(s): {exc}"
-                        ) from exc
-                    time.sleep(min(_CCXT_BACKOFF[attempt], max(0.0, remaining)))
+            check_budget(deadline, label, budget_s=_CCXT_FETCH_BUDGET_S)
+            # ``ccxt.NetworkError`` covers RequestTimeout / DDoSProtection /
+            # ExchangeNotAvailable — the transient family. Anything else
+            # (e.g. ``ExchangeError`` for a bad symbol) is not retried.
+            ohlcv = retry_with_budget(
+                lambda: exchange.fetch_ohlcv(symbol, timeframe, since=cursor, limit=limit),
+                transient=ccxt.NetworkError,
+                deadline=deadline,
+                label=label,
+            )
             if not ohlcv:
                 break
             all_rows.extend(ohlcv)

agent/backtest/loaders/okx.py

@@ -12,18 +12,21 @@
 import pandas as pd
 import requests
 
-from backtest.loaders.base import validate_date_range
+from backtest.loaders.base import (
+    check_budget,
+    retry_with_budget,
+    validate_date_range,
+)
 from backtest.loaders.registry import register
 
 BASE_URL = "https://www.okx.com/api/v5"
 _MAX_PER_PAGE = 300
 # P12-b parity: OKX already sets a per-request timeout but had no retry
 # budget, so a transient blip dropped the whole symbol and a slow tier
-# could stall ~max_pages*timeout. Bound it like the ccxt loader.
+# could stall ~max_pages*timeout. Bound it like the ccxt loader; retry
+# scheduling is delegated to :mod:`backtest.loaders.base`.
 _OKX_TIMEOUT = int(os.getenv("OKX_TIMEOUT_S", "15"))
 _OKX_FETCH_BUDGET_S = float(os.getenv("OKX_FETCH_BUDGET_S", "60"))
-_OKX_MAX_RETRIES = 3
-_OKX_BACKOFF = (0.5, 1.5, 4.0)  # seconds; len == _OKX_MAX_RETRIES
 
 
 @register
@@ -108,37 +111,31 @@ def _fetch_candles(
         all_rows: list = []
         after = str(end_ts)
         deadline = time.monotonic() + _OKX_FETCH_BUDGET_S
+        label = f"OKX fetch for {inst_id}"
 
         for _ in range(max_pages):
-            if time.monotonic() > deadline:
-                raise TimeoutError(
-                    f"OKX fetch for {inst_id} exceeded "
-                    f"{_OKX_FETCH_BUDGET_S:.0f}s budget"
-                )
+            check_budget(deadline, label, budget_s=_OKX_FETCH_BUDGET_S)
             params = {
                 "instId": inst_id,
                 "bar": bar,
                 "limit": str(_MAX_PER_PAGE),
                 "after": after,
             }
-            data = None
-            for attempt in range(_OKX_MAX_RETRIES + 1):
-                try:
-                    resp = requests.get(
-                        f"{BASE_URL}/market/candles",
-                        params=params,
-                        timeout=_OKX_TIMEOUT,
-                    )
-                    data = resp.json()
-                    break
-                except requests.RequestException as exc:
-                    remaining = deadline - time.monotonic()
-                    if attempt == _OKX_MAX_RETRIES or remaining <= 0:
-                        raise TimeoutError(
-                            f"OKX fetch for {inst_id} failed after "
-                            f"{attempt + 1} attempt(s): {exc}"
-                        ) from exc
-                    time.sleep(min(_OKX_BACKOFF[attempt], max(0.0, remaining)))
+
+            def _do_request() -> dict:
+                resp = requests.get(
+                    f"{BASE_URL}/market/candles",
+                    params=params,
+                    timeout=_OKX_TIMEOUT,
+                )
+                return resp.json()
+
+            data = retry_with_budget(
+                _do_request,
+                transient=requests.RequestException,
+                deadline=deadline,
+                label=label,
+            )
             if data.get("code") != "0" or not data.get("data"):
                 break
 

agent/tests/test_ccxt_loader_bounded.py

@@ -16,6 +16,7 @@
 import ccxt
 
 import backtest.loaders.ccxt_loader as cl
+from backtest.loaders.base import DEFAULT_MAX_RETRIES
 from backtest.loaders.ccxt_loader import DataLoader
 
 SINCE = int(pd.Timestamp("2026-05-01").timestamp() * 1000)
@@ -61,7 +62,7 @@ def test_persistent_disconnect_is_bounded_not_a_hang():
     ex = _FakeEx([ccxt.NetworkError("down")])  # always fails
     with pytest.raises(TimeoutError):
         DataLoader._fetch_one(ex, "BTC/USDT", "1d", SINCE, END)
-    assert ex.calls == cl._CCXT_MAX_RETRIES + 1  # bounded, not range(200)/forever
+    assert ex.calls == DEFAULT_MAX_RETRIES + 1  # bounded, not range(200)/forever
 
 
 def test_non_network_error_is_not_retried():