feat: isolate strategy-owned KIS balances

trex_core/adapters/broker/ownership.py

@@ -0,0 +1,162 @@
+"""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,
+    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",
+    ) -> bool:
+        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)