trex-core/trex_core/application/account_snapshot.py

"""Bounded account snapshot recovery for transient KIS read failures.

Holdings and cash capacity have independent capture times so a cash-read
failure never invalidates a healthy holdings refresh (and vice versa). Fresh
reconciliation remains live-only; snapshots are consumed solely by the
new-buy admission guard and the pending-sell quantity path.
"""
from __future__ import annotations

import logging
import random
import threading
from contextlib import AbstractContextManager, contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Callable, Optional, Protocol

from trex_core.application.broker_port import BalanceItem, BrokerBalanceError

_log = logging.getLogger(__name__)


@dataclass(frozen=True)
class AccountSnapshot:
    account_hash: str
    holdings: tuple[BalanceItem, ...]
    holdings_captured_at: Optional[datetime]
    available_cash: Optional[int]
    cash_captured_at: Optional[datetime]


class AccountSnapshotStore(Protocol):
    def save(self, snapshot: AccountSnapshot) -> None: ...
    def load(self, account_hash: str) -> Optional[AccountSnapshot]: ...
    def refresh_guard(self, account_hash: str) -> AbstractContextManager[bool]: ...


class AccountSnapshotBroker(Protocol):
    def account_scope_hash(self) -> str: ...
    def get_balance(self) -> list[BalanceItem]: ...
    def get_available_cash(self, *, price_check_code: str = "005930") -> int: ...


class AccountSnapshotUnavailable(RuntimeError):
    """A sanitized fail-closed result for snapshot consumers."""

    _CODES = frozenset({
        "not_ready",
        "account_mismatch",
        "expired",
        "live_read_failed",
        "read_failed",
    })

    def __init__(self, code: str, *, detail_code: str | None = None) -> None:
        if code not in self._CODES:
            raise ValueError(f"unsupported account snapshot error code: {code}")
        self.code = code
        self.detail_code = detail_code
        super().__init__(f"account snapshot unavailable ({code})")


class InMemoryAccountSnapshotStore:
    def __init__(self) -> None:
        self._items: dict[str, AccountSnapshot] = {}
        self._items_lock = threading.Lock()
        self._refresh_lock = threading.Lock()

    def save(self, snapshot: AccountSnapshot) -> None:
        with self._items_lock:
            self._items[snapshot.account_hash] = snapshot

    def load(self, account_hash: str) -> Optional[AccountSnapshot]:
        with self._items_lock:
            return self._items.get(account_hash)

    @contextmanager
    def refresh_guard(self, account_hash: str):
        acquired = self._refresh_lock.acquire(blocking=False)
        try:
            yield acquired
        finally:
            if acquired:
                self._refresh_lock.release()


class AccountSnapshotService:
    """Refresh, age-check, and consume account-scoped broker snapshots."""

    def __init__(
        self,
        *,
        broker: AccountSnapshotBroker,
        store: AccountSnapshotStore,
        buy_max_age_sec: int = 600,
        sell_max_age_sec: int = 3600,
        refresh_interval_sec: float = 300.0,
        clock: Optional[Callable[[], datetime]] = None,
        rng: Optional[Callable[[], float]] = None,
    ) -> None:
        if buy_max_age_sec < 0 or sell_max_age_sec < buy_max_age_sec:
            raise ValueError("invalid account snapshot age policy")
        if sell_max_age_sec > 21_600:
            raise ValueError("account snapshot age exceeds six-hour ceiling")
        if refresh_interval_sec <= 0:
            raise ValueError("account snapshot refresh interval must be positive")
        self._broker = broker
        self._store = store
        self._buy_max_age = buy_max_age_sec
        self._sell_max_age = sell_max_age_sec
        self._refresh_interval = refresh_interval_sec
        self._clock = clock or (lambda: datetime.now(tz=timezone.utc))
        self._rng = rng or random.random
        self._refresh_lock = threading.Lock()
        self._last_holdings_error: Optional[str] = "not_attempted"
        self._last_cash_error: Optional[str] = "not_attempted"

    def refresh(self) -> bool:
        """Refresh holdings and cash independently under one account-scoped owner."""
        account_hash = self._broker.account_scope_hash()
        if not account_hash:
            self._last_holdings_error = "ownership_mismatch"
            self._last_cash_error = "ownership_mismatch"
            return False
        with self._refresh_lock, self._store.refresh_guard(account_hash) as acquired:
            if not acquired:
                return False
            existing = self._store.load(account_hash)
            holdings = existing.holdings if existing is not None else ()
            holdings_captured_at = (
                existing.holdings_captured_at if existing is not None else None
            )
            available_cash = existing.available_cash if existing is not None else None
            cash_captured_at = existing.cash_captured_at if existing is not None else None
            now = self._clock()
            holdings_ok = cash_ok = False
            try:
                holdings = tuple(self._broker.get_balance())
                holdings_captured_at = now
                self._last_holdings_error = None
                holdings_ok = True
            except BrokerBalanceError as exc:
                self._last_holdings_error = exc.code
                _log.warning(
                    "account holdings snapshot refresh failed: code=%s",
                    exc.code,
                )
            try:
                available_cash = int(self._broker.get_available_cash())
                if available_cash < 0:
                    raise BrokerBalanceError("malformed_response")
                cash_captured_at = now
                self._last_cash_error = None
                cash_ok = True
            except BrokerBalanceError as exc:
                self._last_cash_error = exc.code
                _log.warning(
                    "account cash snapshot refresh failed: code=%s",
                    exc.code,
                )
            if not account_hash or self._broker.account_scope_hash() != account_hash:
                self._last_holdings_error = "ownership_mismatch"
                self._last_cash_error = "ownership_mismatch"
                return False
            if not holdings_ok and not cash_ok:
                return False
            snapshot = AccountSnapshot(
                account_hash=account_hash,
                holdings=holdings,
                holdings_captured_at=holdings_captured_at,
                available_cash=available_cash,
                cash_captured_at=cash_captured_at,
            )
            self._store.save(snapshot)
            return holdings_ok and cash_ok

    def run_forever(self, stop: threading.Event) -> None:
        """Refresh continuously with bounded exponential backoff."""
        failures = 0
        while not stop.is_set():
            ok = self.refresh()
            failures = 0 if ok else min(failures + 1, 4)
            if ok:
                wait_sec = self._refresh_interval * (0.9 + 0.2 * self._rng())
            else:
                wait_sec = min(15.0 * (2 ** (failures - 1)), 300.0)
            if stop.wait(wait_sec):
                return

    def _load_current(self) -> AccountSnapshot:
        account_hash = self._broker.account_scope_hash()
        if not account_hash:
            raise AccountSnapshotUnavailable("account_mismatch")
        try:
            snapshot = self._store.load(account_hash)
        except Exception as exc:
            _log.warning(
                "account snapshot load failed: account_hash=%s",
                account_hash,
                exc_info=True,
            )
            raise AccountSnapshotUnavailable("read_failed") from exc
        if snapshot is None or snapshot.account_hash != account_hash:
            raise AccountSnapshotUnavailable("account_mismatch")
        return snapshot

    def _age(self, captured_at: Optional[datetime], max_age_sec: int) -> float:
        if captured_at is None:
            raise AccountSnapshotUnavailable("not_ready")
        age = (self._clock() - captured_at).total_seconds()
        if age < 0 or age > max_age_sec:
            raise AccountSnapshotUnavailable("expired")
        return age

    def holdings_for_sell(self) -> list[BalanceItem]:
        """Prefer a live balance; fallback only on transient transport failure."""
        transient_code: str | None = None
        try:
            return self._broker.get_balance()
        except BrokerBalanceError as exc:
            if exc.code not in {"transport_failed", "timeout", "http_5xx"}:
                raise AccountSnapshotUnavailable(
                    "live_read_failed", detail_code=exc.code
                ) from None
            transient_code = exc.code
        snapshot = self._load_current()
        try:
            age = self._age(snapshot.holdings_captured_at, self._sell_max_age)
        except AccountSnapshotUnavailable as exc:
            raise AccountSnapshotUnavailable(
                exc.code, detail_code=transient_code
            ) from None
        _log.warning(
            "sell balance using account snapshot: age_sec=%.0f",
            age,
        )
        return list(snapshot.holdings)

    def holdings_for_slots_rpc(self) -> tuple[list[BalanceItem], datetime]:
        """Return freshness-validated snapshot holdings for strategy slots reads.

        Unlike ``holdings_for_sell()``, this read path never performs a live broker
        call and therefore cannot mutate broker state or mask snapshot staleness.
        Freshness uses the same ``sell_max_age_sec`` policy (`self._sell_max_age`)
        as pending-sell quantity protection so strategy reads and sell dispatch
        share one stale-snapshot threshold.
        """
        snapshot = self._load_current()
        captured_at = snapshot.holdings_captured_at
        if captured_at is None:
            raise AccountSnapshotUnavailable("not_ready")
        self._age(captured_at, self._sell_max_age)
        return list(snapshot.holdings), captured_at

    def buy_guard(
        self,
        *,
        qty: Optional[float],
        limit_price: Optional[float],
        reserved_cash: float = 0.0,
    ) -> Optional[str]:
        """Return a stable rejection reason, or None when bounded admission passes."""
        if self._last_cash_error not in {None, "transport_failed", "timeout", "http_5xx"}:
            return "account_snapshot_unavailable"
        try:
            snapshot = self._load_current()
            self._age(snapshot.cash_captured_at, self._buy_max_age)
        except AccountSnapshotUnavailable:
            return "account_snapshot_unavailable"
        if qty is None or qty <= 0 or limit_price is None or limit_price <= 0:
            return "account_snapshot_price_required"
        requested = float(qty) * float(limit_price)
        if snapshot.available_cash is None:
            return "account_snapshot_unavailable"
        if requested + max(0.0, reserved_cash) > snapshot.available_cash:
            return "insufficient_snapshot_cash"
        return None