"""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 datetime import datetime 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, order_date: str | None = None, submitted_at: datetime | None = None, ) -> FillStatus: result = self._delegate.poll_fill( broker_order_id, side=side, symbol=symbol, order_date=order_date, submitted_at=submitted_at, ) 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 exposes_full_balance_snapshot(self) -> bool: """False: projected balances omit operator-owned/unmanaged broker positions.""" return False def account_scope_hash(self) -> str: return self._delegate.account_scope_hash() def get_available_cash(self, *, price_check_code: str = "005930") -> int: return self._delegate.get_available_cash(price_check_code=price_check_code) def get_current_price(self, symbol: str, market: str = "J"): return self._delegate.get_current_price(symbol, market)