Coverage for adrs/oms/oms.py: 84%

1import os 

2import json 

3import copy 

4import time 

5import asyncio 

6import logging 

7import signal 

8 

9from typing import Dict, Protocol 

10from decimal import Decimal 

11from pydantic import BaseModel 

12from datetime import datetime, timezone, timedelta 

13 

14from aion import Scheduler, Trigger 

15from nats_client import Msg 

16from adrs.data import MetricStream, MetricBuilder 

17from adrs.data.connector import DEFAULT_METRIC_NAMESPACE 

18from adrs.subjects import portfolio_signal_subject, oms_command_subject 

19from cybotrade import Symbol 

20from cybotrade.io import ExchangeEvent 

21from cybotrade.io.event import Event, EventType 

22from cybotrade.models import Position, OrderSide, OrderStatus 

23 

24from adrs.oms.config import ConfigManager 

25from adrs.oms.price_feed import PriceFeed 

26from adrs.oms.ops.order_executer import OrderExecutor, MAX_CONCURRENT_ORDER_OPS 

27from adrs.oms.ops.order_pool import CancelBacklogs 

28from adrs.oms.position import PositionManager 

29from adrs.oms.risk import RiskEngine 

30from adrs.oms.ops.order_placement_manager import OrderPlacementManager 

31from adrs.oms.rate_limit.rate_limiter import RateLimiter 

32from adrs.oms.rate_limit.exchange_limit_profiles import Endpoints 

33from adrs.oms.rate_limit.error_policy import ExchangeErrorPolicy 

34 

35logger = logging.getLogger(__name__) 

36 

37# A record whose order the exchange history still can't find after this long 

38# is unrecoverable; keeping it would poll the shared budget forever 

39AEGIS_NOT_FOUND_GIVE_UP = timedelta(hours=1) 

40 

41# A frame this much older than our synced clock means the event loop stalled or 

42# our clock drifted. Reported rather than enforced: the staleness guard stays on 

43# the monotonic clock, which is immune to both. 

44PRICE_FEED_EVENT_TIME_DIVERGENCE_WARN_MS = 2_000 

45 

46# How often the divergence summary may be emitted. Detection stays per frame; 

47# only the reporting is throttled. 

48# 

49# A stall delays every frame queued behind it, and bookTicker delivers hundreds 

50# of frames a second, so logging per frame described a single event hundreds of 

51# times: one 9-second startup stall produced 156 near-identical WARNINGs inside 

52# one minute, burying the rest of the log. The aggregated line carries the count 

53# and the worst value, which says more than any single occurrence did. 

54PRICE_FEED_DIVERGENCE_LOG_INTERVAL_SEC = 60.0 

55 

56# How long _handle_shutdown waits before retrying cancels that failed on the way 

57# out. 

58# 

59# This used to be 60 seconds, which is exactly the pods' 

60# terminationGracePeriodSeconds (verified on both a Binance and a Bybit tenant). 

61# The sleep therefore consumed the whole grace period and SIGKILL landed as the 

62# retry began, so a rate-limited cancel was never actually retried and every 

63# shutdown with one took the full minute. The retry was dead code. 

64# 

65# What it is really waiting for is a rate-limit pool to refill, and those recover 

66# in about a second — Bybit's cancel pool is 8/s, Binance's weight window is one 

67# minute but a cancel costs 1 of thousands. A few seconds is therefore ample, and 

68# it leaves the rest of the grace period for the retry itself to complete. 

69SHUTDOWN_CANCEL_RETRY_DELAY_SEC = 5.0 

70 

71 

72class PortfolioSignal(BaseModel): 

73 assets: Dict[str, Decimal] 

74 timestamp: int 

75 

76 

77class OMSEventHandler(Protocol): 

78 """ 

79 Observer for key OMS state transitions. Pass an implementation to OMS.__init__ 

80 to receive structured callbacks instead of parsing log output. 

81 

82 All methods are synchronous — keep them fast (record to a queue, increment a 

83 counter). The default (None) skips every call. 

84 """ 

85 

86 def on_signal_received(self, signal: "PortfolioSignal") -> None: ... 

87 def on_desired_updated(self, symbol: "Symbol", quantity: "Decimal") -> None: ... 

88 def on_signal_skipped(self, symbol: "Symbol", reason: str) -> None: ... 

89 

90 

91def generate_cron(total_seconds: int): 

92 if total_seconds < 1: 

93 return "* * * * * *" 

94 

95 hours = total_seconds // 3600 

96 minutes = (total_seconds % 3600) // 60 

97 seconds = total_seconds % 60 

98 

99 if hours > 0: 

100 return f"{seconds} {minutes} */{hours} * * *" 

101 elif minutes > 0: 

102 return f"{seconds} */{minutes} * * * *" 

103 else: 

104 return f"*/{seconds} * * * * *" 

105 

106 

107def getenv(name: str) -> str: 

108 env = os.getenv(name) 

109 if env is None: 

110 raise ValueError(f"{name} is not present in environment") 

111 return env 

112 

113 

114class OMS: 

115 # Subclasses override these to swap in custom implementations without 

116 # touching __init__. executor_cls is forwarded to OrderPlacementManager. 

117 executor_cls: type[OrderExecutor] = OrderExecutor 

118 position_cls: type[PositionManager] = PositionManager 

119 opm_cls: type[OrderPlacementManager] = OrderPlacementManager 

120 risk_cls: type[RiskEngine] = RiskEngine 

121 

122 def __init__( 

123 self, 

124 config: ConfigManager, 

125 metric_stream: MetricStream, 

126 rate_limiter: RateLimiter, 

127 error_policy: ExchangeErrorPolicy | None = None, 

128 insert_prefix: str = DEFAULT_METRIC_NAMESPACE, 

129 signal_namespace: str | None = None, 

130 observer: OMSEventHandler | None = None, 

131 ): 

132 super().__init__() 

133 self.config = config 

134 self.metric_stream = metric_stream 

135 self.metric_builder = MetricBuilder(self.metric_stream, insert_prefix) 

136 # Must match the PortfolioExecutor's namespace so the OMS subscribes to 

137 # the same `portfolio_signal.<ns>.<portfolio_id>` the executor publishes. 

138 self.signal_namespace = signal_namespace 

139 self.observer = observer 

140 # Per-exchange error classification; derive from config when not injected 

141 error_policy = error_policy or config.config.credentials.to_error_policy() 

142 self.position = self.position_cls( 

143 config=config, 

144 rate_limiter=rate_limiter, 

145 ) 

146 # Owned here, not by the OPM: the feed serves both the executor's reads 

147 # and the signal recompute, and a later Bybit adapter will share it. 

148 self.price_feed = PriceFeed() 

149 # Aggregation state for the event-time divergence report; see 

150 # _note_event_time_divergence. Monotonic, so an NTP step cannot make a 

151 # window look finished or eternal. 

152 self._divergence_count = 0 

153 self._divergence_worst_ms = 0 

154 self._divergence_worst_symbol: str | None = None 

155 self._divergence_window_started_at: float | None = None 

156 self._divergence_last_emit_at: float | None = None 

157 self._divergence_emitted_previous_window = False 

158 self.opm = self.opm_cls( 

159 position=self.position, 

160 config=self.config, 

161 rate_limiter=rate_limiter, 

162 error_policy=error_policy, 

163 executor_cls=self.executor_cls, 

164 price_feed=self.price_feed, 

165 ) 

166 self.risk = self.risk_cls( 

167 config=self.config, 

168 position=self.position, 

169 rate_limiter=rate_limiter, 

170 on_breach=self._trigger_shutdown, 

171 ) 

172 self.scheduler = Scheduler() 

173 self.previous_signal: PortfolioSignal | None = None 

174 self.latest_signal: PortfolioSignal | None = None 

175 self.rate_limiter = rate_limiter 

176 # Initialised in run(); set here so on_refresh_config is safe to call 

177 # before run() starts (e.g. in tests). 

178 self.exchange_events_task: asyncio.Task | None = None 

179 # Initialised in run(); set here so on_refresh_config and shutdown are 

180 # safe to call before run() starts (e.g. in tests). 

181 self.price_feed_task: asyncio.Task | None = None 

182 self.price_feed_ws: ExchangeEvent | None = None 

183 

184 async def init(self): 

185 """Initialise the OMS state when first started.""" 

186 logger.info("[INIT] Initilalizing position values") 

187 

188 # LATEST SIGNAL 

189 self.latest_signal = await self.get_latest_signal( 

190 self.config.config.portfolio_id, 

191 ) 

192 logger.info(f"Latest signal at startup {self.latest_signal}") 

193 

194 # ORDER POOL + POSITION share one open-orders snapshot 

195 snapshot = await self.opm.order_pools.fetch_open_orders_snapshot() 

196 await self.opm.order_pools.resync_order_pool(snapshot) 

197 await self.position.update_exchange() 

198 self.position.update_pending(snapshot) 

199 for s in self.config.config.base_asset_to_symbol_table.values(): 

200 symbol = Symbol(s) 

201 if symbol not in self.position.pending.keys(): 

202 self.position.pending[symbol] = Position( 

203 symbol=symbol, 

204 quantity=Decimal("0"), 

205 entry_price=Decimal("0"), 

206 updated_time=datetime.now(tz=timezone.utc), 

207 ) 

208 if symbol not in self.position.exchange.keys(): 

209 self.position.exchange[symbol] = Position( 

210 symbol=symbol, 

211 quantity=Decimal("0"), 

212 entry_price=Decimal("0"), 

213 updated_time=datetime.now(tz=timezone.utc), 

214 ) 

215 

216 # Init using desired and exchange 

217 quantity = ( 

218 self.position.exchange[symbol].quantity 

219 + self.position.pending[symbol].quantity 

220 ) 

221 self.position.desired[symbol] = Position( 

222 symbol=symbol, 

223 quantity=quantity, 

224 entry_price=Decimal("0"), 

225 updated_time=datetime.now(tz=timezone.utc), 

226 ) 

227 

228 async def get_latest_signal(self, portfolio_id: str) -> PortfolioSignal | None: 

229 """ 

230 Cold-start hook: fetch the last known signal so the OMS does not start 

231 blind. The base has no signal store, so it starts without one and waits 

232 for the first live signal; subclasses override to read their own store. 

233 """ 

234 return None 

235 

236 def _setup_signals(self): 

237 """ 

238 SIGTERM stops the scheduler, not _handle_shutdown() directly: run()'s 

239 finally already calls _handle_shutdown() exactly once whenever 

240 scheduler.start() returns, for any reason. Calling it again here 

241 would race that (double order-cancel pass, and a SystemExit raised 

242 inside this handler's own task wouldn't propagate to end the 

243 process anyway - it would just be a swallowed background-task 

244 exception). Stopping the scheduler reuses the existing, single 

245 shutdown path instead of adding a second one. 

246 """ 

247 loop = asyncio.get_event_loop() 

248 

249 loop.add_signal_handler( 

250 signal.SIGTERM, lambda: asyncio.create_task(self.scheduler.shutdown()) 

251 ) 

252 

253 async def _trigger_shutdown(self): 

254 """ 

255 Single shutdown entry point for risk breaches (and anything else that 

256 must stop the OMS from inside a running handler): stop the scheduler, 

257 which makes scheduler.start() return, and run()'s finally performs the 

258 one order-cancel pass and process exit. Calling _handle_shutdown() 

259 directly from a handler would cancel orders but leave the scheduler 

260 ticking - the next placement tick would put every position right back. 

261 (Same reasoning as the SIGTERM handler, see _setup_signals.) 

262 """ 

263 logger.error("[SHUTDOWN] Triggered - stopping scheduler") 

264 await self.scheduler.shutdown() 

265 

266 async def _handle_shutdown(self): 

267 """To close all pending orders on shutdown""" 

268 logger.info("Shutdown signal received. Cancelling orders...") 

269 

270 self._stop_price_feed() 

271 

272 # Snapshot under the pool lock, then release BEFORE the gather: each 

273 # cancel_single_order re-acquires the same lock, so holding it across 

274 # the gather would self-deadlock. 

275 async with self.opm.order_pools.get_order_pool() as order_pool: 

276 orders_to_cancel = list(order_pool.values()) 

277 

278 sem = asyncio.Semaphore(MAX_CONCURRENT_ORDER_OPS) 

279 

280 async def _bounded_cancel(symbol: str, client_order_id: str): 

281 async with sem: 

282 return await self.opm.executor.cancel_single_order( 

283 Symbol(symbol), client_order_id 

284 ) 

285 

286 cancel_tasks = [ 

287 _bounded_cancel(order.symbol, order.client_order_id) 

288 for order in orders_to_cancel 

289 ] 

290 

291 if not cancel_tasks: 

292 raise SystemExit(0) 

293 

294 cancel_results = await asyncio.gather(*cancel_tasks, return_exceptions=True) 

295 

296 cancel_retries = [] 

297 for result in cancel_results: 

298 if isinstance(result, Exception): 

299 logger.error(f"Failed to cancel order: {result}") 

300 if isinstance(result, CancelBacklogs): 

301 cancel_retries.append(result) 

302 

303 if not cancel_retries: 

304 logger.info("All orders cancelled successfully.") 

305 return 

306 

307 logger.warning( 

308 f"Retrying {len(cancel_retries)} cancellations in " 

309 f"{SHUTDOWN_CANCEL_RETRY_DELAY_SEC}s..." 

310 ) 

311 try: 

312 await asyncio.sleep(SHUTDOWN_CANCEL_RETRY_DELAY_SEC) 

313 except asyncio.CancelledError: 

314 logger.error("Shutdown forced during retry wait.") 

315 return 

316 

317 retry_tasks = [ 

318 _bounded_cancel(retry.symbol, retry.client_order_id) 

319 for retry in cancel_retries 

320 ] 

321 

322 if retry_tasks: 

323 await asyncio.gather(*retry_tasks, return_exceptions=True) 

324 logger.info("Retry attempts finished.") 

325 

326 raise SystemExit(0) 

327 

328 async def _insert_equity_metric( 

329 self, oms_id: str, balance_endpoint: Endpoints 

330 ) -> None: 

331 """ 

332 Read the wallet balance from the exchange and record equity. 

333 

334 The rate-limit guard deliberately covers the EXCHANGE call only. It used 

335 to wrap the aegis write too, which meant a ClickHouse outage -- nothing 

336 to do with the exchange -- was charged to the wallet rate-limit pool. 

337 Because the limiter only adopts the exchange's own quota headers on 

338 success, every such failure decremented the pool without ever learning 

339 the real quota, and once it reached zero the pool refused for the life of 

340 the process. Reported from production as "UID_WALLET: 0/None, never 

341 recovers", with the dashboard balance falling back to the static 

342 initial_balance -- one root cause, two symptoms. 

343 

344 Extracted from on_aegis_update so that guard scope is explicit and 

345 testable, since the scope IS the bug. 

346 

347 The balance is read fresh here on every call by design; see the spec's 

348 Balances section. No streamed balance substitutes for it. 

349 """ 

350 async with self.rate_limiter.guard(endpoint=balance_endpoint): 

351 balance = await self.config.exchange.get_wallet_balance() 

352 

353 logger.debug("Inserting equity to aegis") 

354 await self.metric_builder.create_equity( 

355 oms_id=oms_id, equity=str(balance.margin_balance) 

356 ) 

357 

358 def _supervise_task(self, task: asyncio.Task, label: str) -> None: 

359 """ 

360 Log when a long-lived stream task ends. Logging only, never restarts. 

361 

362 Nothing ever awaits these tasks, so without this "started and healthy" and 

363 "started and dead" are indistinguishable from the OMS. Both endings are 

364 silent today: 

365 

366 - cybotrade's `_stream` catches stream errors, prints and breaks, so 

367 `start()` returns *normally* once a live socket dies. 

368 - `persist_conn_with` retries only *reconnects*; it uses `?` on the initial 

369 connect, and in `_stream` that await sits outside the try, so failing to 

370 connect at all escapes into the task, where nobody retrieves it. 

371 

372 Deliberately does not restart: a restart loop against a hard-failing 

373 connect would hammer the exchange, and the price feed already fails safe 

374 (PriceFeed withholds quotes once liveness lapses, so reads go to REST). 

375 Retrieving the exception here also stops asyncio complaining about it. 

376 """ 

377 

378 def _on_done(finished: asyncio.Task) -> None: 

379 if finished.cancelled(): 

380 logger.info(f"[{label}] Task cancelled") 

381 return 

382 exc = finished.exception() 

383 if exc is not None: 

384 logger.error(f"[{label}] Task failed: {exc!r}", exc_info=exc) 

385 return 

386 logger.error( 

387 f"[{label}] Task ended on its own, so the stream stopped " 

388 f"delivering and this connection will not come back" 

389 ) 

390 

391 task.add_done_callback(_on_done) 

392 

393 async def on_price_feed_event(self, event: Event) -> None: 

394 """ 

395 Translate public market-data events into the feed. 

396 

397 Every event refreshes liveness, including ones that carry no quote: an 

398 unparseable frame still proves the socket is delivering. Subscribed means 

399 the connection just (re)established, so nothing cached survives it. 

400 

401 Async to satisfy the adapter's callback contract — BinancePublicWS does 

402 `await self.on_event(event)`, and a plain def there raises on every frame 

403 — not because the translation itself needs to await anything. 

404 """ 

405 match event.event_type: 

406 case EventType.BookTicker: 

407 book_ticker = event.data 

408 self._note_event_time_divergence(book_ticker) 

409 self.price_feed.apply( 

410 book_ticker.symbol, book_ticker.bid, book_ticker.ask 

411 ) 

412 case EventType.Subscribed: 

413 logger.info(f"[PRICE_FEED] (Re)subscribed to {event.data}, clearing") 

414 self.price_feed.clear() 

415 self.price_feed.note_message() 

416 case _: 

417 self.price_feed.note_message() 

418 

419 def _note_event_time_divergence(self, book_ticker) -> None: 

420 """ 

421 Record a frame whose exchange timestamp trails our synced clock, and 

422 report at most one summary per PRICE_FEED_DIVERGENCE_LOG_INTERVAL_SEC. 

423 

424 Detection is per frame; reporting is not. One stall delays every frame 

425 behind it, so a per-frame line describes a single event once per frame -- 

426 156 of them in a minute, observed in production. Count plus worst value 

427 is both quieter and more informative. 

428 

429 Severity escalates rather than starting loud: an isolated stall (startup, 

430 a GC pause) is not actionable, whereas the same thing two windows running 

431 means the loop is persistently behind and quote age is understating 

432 staleness for real. 

433 

434 The whole body is guarded. "Never let a metric break the feed" only holds 

435 if every step is covered, and a raise here would propagate through 

436 on_price_feed_event into the adapter, where start() swallows it per frame 

437 -- losing the quote and leaving no OMS-level trace. 

438 """ 

439 try: 

440 now = time.monotonic() 

441 exchange_ms = int(book_ticker.event_time.timestamp() * 1000) 

442 divergence_ms = self.rate_limiter.get_synced_time_ms() - exchange_ms 

443 

444 if divergence_ms > PRICE_FEED_EVENT_TIME_DIVERGENCE_WARN_MS: 

445 if self._divergence_window_started_at is None: 

446 self._divergence_window_started_at = now 

447 # A stall long after the last report is a fresh incident, not 

448 # a continuation, so it starts quiet again. 

449 last = self._divergence_last_emit_at 

450 if ( 

451 last is not None 

452 and now - last > 2 * PRICE_FEED_DIVERGENCE_LOG_INTERVAL_SEC 

453 ): 

454 self._divergence_emitted_previous_window = False 

455 self._divergence_count += 1 

456 if divergence_ms > self._divergence_worst_ms: 

457 self._divergence_worst_ms = divergence_ms 

458 self._divergence_worst_symbol = str(book_ticker.symbol) 

459 

460 started = self._divergence_window_started_at 

461 if ( 

462 started is None 

463 or now - started < PRICE_FEED_DIVERGENCE_LOG_INTERVAL_SEC 

464 ): 

465 return 

466 

467 count = self._divergence_count 

468 worst = self._divergence_worst_ms 

469 symbol = self._divergence_worst_symbol 

470 elapsed = now - started 

471 # Reset before emitting, so a failure inside logging cannot wedge the 

472 # window open and suppress every later report. 

473 self._divergence_count = 0 

474 self._divergence_worst_ms = 0 

475 self._divergence_worst_symbol = None 

476 self._divergence_window_started_at = None 

477 self._divergence_last_emit_at = now 

478 

479 message = ( 

480 f"[PRICE_FEED] Event-time divergence: {count} frame(s) in the " 

481 f"last {elapsed:.0f}s, worst {worst}ms ({symbol}). The loop " 

482 f"stalled or the clock has drifted, so quote age understates " 

483 f"real staleness" 

484 ) 

485 if self._divergence_emitted_previous_window: 

486 logger.warning(message) 

487 else: 

488 logger.info(message) 

489 self._divergence_emitted_previous_window = True 

490 except Exception: 

491 return 

492 

493 def _start_price_feed(self) -> None: 

494 """ 

495 Start the public feed for the configured symbols. 

496 

497 Only Binance and Bybit have an adapter today; on any other exchange 

498 the OMS runs exactly as before, reading prices from REST. 

499 """ 

500 if self.price_feed_task is not None and not self.price_feed_task.done(): 

501 # Starting without stopping would orphan the running socket, which 

502 # would keep writing into the shared cache with nobody able to cancel it. 

503 logger.warning( 

504 "[PRICE_FEED] Already running, not starting a second connection" 

505 ) 

506 return 

507 symbols = list(self.config.config.base_asset_to_symbol_table.values()) 

508 self.price_feed_ws = self.config.config.credentials.to_public_exchange_event( 

509 symbols=symbols 

510 ) 

511 if self.price_feed_ws is None: 

512 logger.info( 

513 "[PRICE_FEED] No public feed for " 

514 f"{self.config.config.credentials.exchange}, using REST prices" 

515 ) 

516 return 

517 self.price_feed_ws.on_event = self.on_price_feed_event 

518 self.price_feed_task = asyncio.create_task(self.price_feed_ws.start()) 

519 self._supervise_task(self.price_feed_task, "PRICE_FEED") 

520 # .topics (Bybit) and .streams (Binance) are mutually exclusive today, 

521 # but `or` short-circuits rather than evaluating both eagerly, so a 

522 # future adapter that grows both wouldn't have `topics` silently win. 

523 subscriptions = getattr(self.price_feed_ws, "topics", None) or getattr( 

524 self.price_feed_ws, "streams", None 

525 ) 

526 logger.info( 

527 f"[PRICE_FEED] Started {type(self.price_feed_ws).__name__} for {subscriptions}" 

528 ) 

529 

530 def _stop_price_feed(self) -> None: 

531 if self.price_feed_task is not None: 

532 self.price_feed_task.cancel() 

533 self.price_feed_task = None 

534 # Drop the adapter with the task: after a switch away from any 

535 # exchange with a public feed, this would otherwise still point at 

536 # the dead connection's adapter. 

537 self.price_feed_ws = None 

538 self.price_feed.clear() 

539 

540 async def on_refresh_config(self): 

541 """To refresh and update config if there are any changes made during runtime""" 

542 old_config = copy.deepcopy(self.config.config) 

543 await self.config.refresh() 

544 self.opm.executor.config = self.config.config 

545 

546 # Both branches below can invalidate the feed's connection. Restart it 

547 # once, after they have run, so a refresh that changes both does not open 

548 # a socket only to cancel it. 

549 needs_price_feed_restart = False 

550 

551 # Restart with a new exchange events when credentials have changed 

552 if old_config.credentials != self.config.config.credentials: 

553 logger.info( 

554 "Detected credentials update, refreshing exchange event handler..." 

555 ) 

556 if self.exchange_events_task is not None: 

557 self.exchange_events_task.cancel() 

558 self.exchange_event = self.config.config.credentials.to_exchange_event() 

559 self.exchange_event.on_event = self.opm.on_exchange_event 

560 self.exchange_events_task = asyncio.create_task(self.exchange_event.start()) 

561 self._supervise_task(self.exchange_events_task, "EXCHANGE_EVENTS") 

562 self.opm.executor.exchange = self.config.exchange 

563 self.opm.order_pools.exchange = self.config.exchange 

564 # The exchange and the testnet flag both live in credentials, and both 

565 # decide whether there is a feed at all and which book it reads. A 

566 # feed left running after this would quote the previous exchange. 

567 needs_price_feed_restart = True 

568 

569 if ( 

570 old_config.base_asset_to_symbol_table 

571 != self.config.config.base_asset_to_symbol_table 

572 ): 

573 try: 

574 await self.config.update_symbol_info(self.rate_limiter, force=True) 

575 except Exception as e: 

576 logger.warning(f"update symbol info failed due to {e}") 

577 

578 await self.init() 

579 

580 # The stream list is encoded in the feed's URL, so a changed symbol 

581 # table needs a new connection. 

582 needs_price_feed_restart = True 

583 

584 if needs_price_feed_restart: 

585 # clear() happens in _stop_price_feed: quotes for symbols we may no 

586 # longer trade, or from an exchange we no longer talk to, must not linger. 

587 self._stop_price_feed() 

588 self._start_price_feed() 

589 

590 async def on_portfolio_signal(self, msg: Msg): 

591 """To store the latest signal from portfolio server""" 

592 try: 

593 payload = json.loads(msg.data.decode()) 

594 portfolio_signal = PortfolioSignal( 

595 assets=payload["assets"], 

596 timestamp=int(payload["timestamp"]), 

597 ) 

598 logger.info(f"Portfolio signal retrieved {portfolio_signal}") 

599 

600 if sum(abs(p) for p in portfolio_signal.assets.values()) > Decimal("1"): 

601 logger.warning( 

602 "[ON_PROCESS_LATEST_SIGNAL] Current sum of positions in latest signal is more than 1" 

603 ) 

604 return 

605 

606 unknown_assets = [ 

607 asset 

608 for asset in portfolio_signal.assets 

609 if asset not in self.config.config.base_asset_to_symbol_table 

610 ] 

611 if unknown_assets: 

612 logger.error( 

613 f"[ON_PORTFOLIO_SIGNAL] Rejecting signal with unknown asset(s) " 

614 f"{unknown_assets}, not in base_asset_to_symbol_table" 

615 ) 

616 return 

617 

618 self.latest_signal = portfolio_signal 

619 if self.observer is not None: 

620 self.observer.on_signal_received(portfolio_signal) 

621 except Exception as e: 

622 logger.error(f"Failed to process portfolio signal due to {e}") 

623 

624 async def on_command(self, msg: Msg): 

625 """Control-plane callback: an operator sent a command on this OMS's 

626 command subject. Run the matching operation and, when the caller used 

627 request-reply (msg.reply set), publish the outcome back. 

628 

629 One reply point with a result dict keeps every path — parse error, 

630 unknown command, failure, success — replying exactly once. 

631 """ 

632 try: 

633 payload = json.loads(msg.data.decode()) 

634 command = payload["command"] 

635 except Exception as e: 

636 logger.error(f"[ON_COMMAND] Failed to parse command message: {e}") 

637 await self._reply(msg, {"status": "error", "error": f"bad message: {e}"}) 

638 return 

639 

640 try: 

641 if command == "rebalance": 

642 await self.rebalance() 

643 result = {"status": "ok", "command": command} 

644 else: 

645 logger.warning(f"[ON_COMMAND] Unknown command: {command}") 

646 result = {"status": "error", "error": f"unknown command: {command}"} 

647 except Exception as e: 

648 logger.error(f"[ON_COMMAND] Command '{command}' failed: {e}") 

649 result = {"status": "error", "command": command, "error": str(e)} 

650 

651 await self._reply(msg, result) 

652 

653 async def _reply(self, msg: Msg, result: dict): 

654 """Publish a result to the requester's reply inbox. No-op for 

655 fire-and-forget callers, whose msg.reply is empty.""" 

656 reply_to = getattr(msg, "reply", "") 

657 if not reply_to: 

658 return 

659 try: 

660 await self.metric_stream.publish(reply_to, json.dumps(result).encode()) 

661 except Exception as e: 

662 logger.error(f"[ON_COMMAND] Failed to reply on {reply_to}: {e}") 

663 

664 async def on_process_latest_signal(self): 

665 """ 

666 Periodic: recompute desired from the latest signal, skipping assets 

667 whose weight is unchanged. Use rebalance() to force a recompute at the 

668 current price even when the signal is static. 

669 """ 

670 await self._recompute_desired(force=False) 

671 

672 async def _recompute_desired(self, force: bool = False): 

673 """ 

674 Turn the latest signal into desired positions at the current price. 

675 

676 Desired quantity is price-dependent (balance * leverage * weight / 

677 price), so with force=True the weight-unchanged skip is bypassed and 

678 desired is repriced against a moved market even if the signal has not 

679 changed. 

680 """ 

681 if not self.latest_signal: 

682 logger.warning("[RECOMPUTE_DESIRED] There is no signal to act on") 

683 return 

684 

685 market_quotes: dict[Symbol, Decimal] = {} 

686 PRECISION_4 = Decimal("0.0001") 

687 symbol_table = self.config.config.base_asset_to_symbol_table 

688 

689 # First pass: decide which assets need a recompute, before touching 

690 # any shared state 

691 to_recompute: dict[str, Decimal] = {} 

692 for asset, weightage in self.latest_signal.assets.items(): 

693 symbol = Symbol(symbol_table[asset]) 

694 # .get: a newly added asset has no entry in the previous signal 

695 previous_weightage = ( 

696 self.previous_signal.assets.get(asset) if self.previous_signal else None 

697 ) 

698 if ( 

699 not force 

700 and previous_weightage is not None 

701 and previous_weightage.quantize(exp=PRECISION_4) 

702 == weightage.quantize(exp=PRECISION_4) 

703 ): 

704 logger.info( 

705 f"There is no significant change in signal weights for {symbol}, it is ignored" 

706 ) 

707 if self.observer is not None: 

708 self.observer.on_signal_skipped(symbol, "no_change") 

709 continue 

710 to_recompute[asset] = weightage 

711 

712 if to_recompute: 

713 recompute_symbols = { 

714 str(Symbol(symbol_table[asset])) for asset in to_recompute 

715 } 

716 async with self.opm.order_pools.get_order_backlog() as order_backlog: 

717 order_backlog[:] = [ 

718 b for b in order_backlog if b.symbol not in recompute_symbols 

719 ] 

720 

721 for asset, weightage in to_recompute.items(): 

722 symbol = Symbol(symbol_table[asset]) 

723 if symbol not in market_quotes.keys(): 

724 price = await self.opm.executor.get_current_price(symbol=symbol) 

725 if price is None: 

726 logger.warning( 

727 f"[RECOMPUTE_DESIRED] No price for {symbol}, skipping this cycle" 

728 ) 

729 return 

730 market_quotes[symbol] = price 

731 

732 quantity = self.position.compute_base_quantity( 

733 price=market_quotes[symbol], weightage=weightage 

734 ) 

735 self.position.desired[symbol] = Position( 

736 symbol=symbol, 

737 quantity=quantity, 

738 entry_price=market_quotes[symbol], 

739 updated_time=datetime.now(tz=timezone.utc), 

740 ) 

741 desired_position = self.position.desired[symbol] 

742 logger.debug( 

743 f"[RECOMPUTE_DESIRED] {self.config.config.portfolio_id} wants {desired_position.quantity} of {desired_position.symbol} at the price of {desired_position.entry_price}" 

744 ) 

745 if self.observer is not None: 

746 self.observer.on_desired_updated(symbol, desired_position.quantity) 

747 

748 self.position.desired = self.risk.cap_desired(self.position.desired) 

749 self.previous_signal = self.latest_signal 

750 

751 async def rebalance(self): 

752 """ 

753 Reprice desired against the current market even when the signal is 

754 static (bypassing the weight-unchanged skip). Placement is left to the 

755 scheduled on_order_placement tick, so this never races the cron and 

756 needs no placement lock. 

757 """ 

758 logger.info("[REBALANCE] repricing desired at current market") 

759 await self._recompute_desired(force=True) 

760 

761 async def on_aegis_update(self): 

762 """ 

763 To upsert latest equity value from exchange to aegis`` 

764 """ 

765 logger.info("[ON_AEGIS_UPDATE]") 

766 

767 # The 60s cadence this handler runs on is the liveness anchor for the 

768 # streamed exchange position (see POSITION_ANCHOR_MAX_AGE_SEC) -- it 

769 # must not be contingent on the aegis/metrics path below, which reads 

770 # from a separate DB and exchange call that can fail independently. 

771 # update_exchange() already swallows and logs its own failures, so a 

772 # bad read here degrades to the anchor aging out rather than raising. 

773 try: 

774 await self.position.update_exchange() 

775 except Exception as e: 

776 logger.warning(f"Failed to refresh exchange positions due to {e}") 

777 

778 async def _check_and_sync_trade( 

779 sem, 

780 oms_id: str, 

781 package_id: str, 

782 record: tuple[Symbol, str, Decimal, datetime], 

783 ): 

784 """ 

785 Worker function. 

786 Returns: (package_id, record) IF it is fully filled and synced or 0-filled. 

787 Returns: None IF it is pending or failed. 

788 """ 

789 async with sem: 

790 symbol, client_order_id, start_price, start_time = ( 

791 record[0], 

792 record[1], 

793 record[2], 

794 record[3], 

795 ) 

796 asset = parts[0] if symbol and (parts := symbol.split()) else "" 

797 

798 try: 

799 async with self.rate_limiter.guard( 

800 endpoint=Endpoints.GET_ORDER_DETAILS 

801 ): 

802 result = ( 

803 await self.config.exchange.get_order_details_from_history( 

804 symbol=symbol, 

805 client_order_id=client_order_id, 

806 ) 

807 ) 

808 except Exception as e: 

809 logger.warning(f"Failed to fetch {client_order_id}: {e}") 

810 return None # Failed fetch, keep in pool 

811 

812 if not result: 

813 age = datetime.now(timezone.utc) - start_time.astimezone( 

814 timezone.utc 

815 ) 

816 if age > AEGIS_NOT_FOUND_GIVE_UP: 

817 logger.error( 

818 f"[ON_AEGIS_UPDATE] {client_order_id} still not in exchange history after {age}, dropping record" 

819 ) 

820 return (package_id, record) 

821 logger.warning( 

822 f"[ON_AEGIS_UPDATE] {client_order_id} NOT FOUND IN EXCHANGE HISTORY" 

823 ) 

824 return None # Not found, keep in pool 

825 

826 if result.status in ( 

827 OrderStatus.CREATED, 

828 OrderStatus.PARTIALLY_FILLED, 

829 ): 

830 # Still live; syncing now would freeze a partial fill as 

831 # the final trade and later fills would never reach aegis 

832 return None 

833 

834 if result.filled_size == Decimal("0"): 

835 # Terminal with nothing filled, no trade will ever come 

836 return (package_id, record) 

837 

838 match result.side: 

839 case OrderSide.BUY: 

840 order_side = 1 

841 case OrderSide.SELL: 

842 order_side = -1 

843 case OrderSide.NONE: 

844 logger.error("[ON_AEGIS_UPDATE] Orderside shouldn't be NONE") 

845 order_side = 0 

846 

847 try: 

848 await self.metric_builder.create_trade( 

849 oms_id=oms_id, 

850 package_id=package_id,