trex-core/trex_core/adapters/broker/ownership.py

"""Optional strategy-only balance projection for mixed manual holdings.

The KIS account balance contains both operator-owned shares and shares bought
through this TREX Core instance.  When a per-account ownership path is
configured, reconciliation must see only the latter.  Confirmed TREX fills are
recorded idempotently in an atomic JSON ledger; raw broker totals remain the
upper safety bound.
"""
from __future__ import annotations

import json
import os
from pathlib import Path
from threading import Lock

from trex_core.application.broker_port import (
    BalanceItem,
    BrokerBalanceError,
    CancelOutcome,
    FillStatus,
    OrderAck,
    OrderRequest,
)


class StrategyOwnershipStore:
    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self._lock = Lock()

    def snapshot(self) -> tuple[dict[str, int], set[str]]:
        with self._lock:
            return self._read()

    def apply_fill(
        self,
        *,
        execution_id: str,
        symbol: str,
        side: str,
        quantity: int,
    ) -> bool:
        execution_id = str(execution_id or "").strip()
        symbol = str(symbol or "").strip().zfill(6)
        side = str(side or "").strip().lower()
        quantity = int(quantity)
        if not execution_id or len(symbol) != 6 or not symbol.isdigit():
            raise ValueError("valid execution_id and symbol are required")
        if side not in {"buy", "sell"} or quantity <= 0:
            raise ValueError("fill side/quantity is invalid")
        with self._lock:
            quantities, applied = self._read()
            if execution_id in applied:
                return False
            current = quantities.get(symbol, 0)
            updated = current + quantity if side == "buy" else current - quantity
            if updated < 0:
                raise BrokerBalanceError("ownership_mismatch")
            quantities[symbol] = updated
            applied.add(execution_id)
            self._write(quantities, applied)
            return True

    def project(self, broker_items: list[BalanceItem]) -> list[BalanceItem]:
        raw = {str(item.symbol).zfill(6): int(item.hldg_qty) for item in broker_items}
        quantities, _ = self.snapshot()
        for symbol, owned in quantities.items():
            if owned > raw.get(symbol, 0):
                raise BrokerBalanceError("ownership_mismatch")
        return [
            BalanceItem(symbol=symbol, hldg_qty=quantity)
            for symbol, quantity in sorted(quantities.items())
            if quantity > 0
        ]

    def _read(self) -> tuple[dict[str, int], set[str]]:
        if not self.path.exists():
            return {}, set()
        try:
            payload = json.loads(self.path.read_text(encoding="utf-8"))
            if payload.get("schema_version") != 1:
                raise ValueError("unsupported schema")
            quantities = {
                str(symbol).zfill(6): int(quantity)
                for symbol, quantity in payload["strategy_quantities"].items()
            }
            applied = {str(value) for value in payload["applied_execution_ids"]}
        except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError):
            raise BrokerBalanceError("ownership_mismatch") from None
        if any(quantity < 0 for quantity in quantities.values()):
            raise BrokerBalanceError("ownership_mismatch")
        return quantities, applied

    def _write(self, quantities: dict[str, int], applied: set[str]) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        temporary = self.path.with_name(f".{self.path.name}.tmp")
        temporary.write_text(
            json.dumps(
                {
                    "schema_version": 1,
                    "strategy_quantities": dict(sorted(quantities.items())),
                    "applied_execution_ids": sorted(applied),
                },
                ensure_ascii=False,
                indent=2,
            )
            + "\n",
            encoding="utf-8",
        )
        os.replace(temporary, self.path)


class OwnershipAwareBroker:
    """Broker decorator that exposes only quantities filled by this TREX instance."""

    def __init__(self, delegate, store: StrategyOwnershipStore) -> None:
        self._delegate = delegate
        self._store = store

    def name(self) -> str:
        return self._delegate.name()

    def reinitialize(self, credential) -> None:
        self._delegate.reinitialize(credential)

    def submit_order(self, req: OrderRequest) -> OrderAck:
        return self._delegate.submit_order(req)

    def poll_fill(self, broker_order_id: str, *, side: str, symbol: str) -> FillStatus:
        result = self._delegate.poll_fill(broker_order_id, side=side, symbol=symbol)
        if result.execution_id and result.filled_qty > 0 and (
            result.is_filled or result.is_partial_cancel
        ):
            self._store.apply_fill(
                execution_id=result.execution_id,
                symbol=symbol,
                side=side,
                quantity=result.filled_qty,
            )
        return result

    def cancel_order(
        self,
        broker_order_id: str,
        *,
        side: str,
        symbol: str,
        qty: int,
        venue: str = "KRX",
    ) -> CancelOutcome:
        return self._delegate.cancel_order(
            broker_order_id,
            side=side,
            symbol=symbol,
            qty=qty,
            venue=venue,
        )

    def get_balance(self) -> list[BalanceItem]:
        return self._store.project(self._delegate.get_balance())

    def get_current_price(self, symbol: str, market: str = "J"):
        return self._delegate.get_current_price(symbol, market)