diff --git a/scripts/qualification/qualify_paper.py b/scripts/qualification/qualify_paper.py index bcc9ae9..9b68f58 100644 --- a/scripts/qualification/qualify_paper.py +++ b/scripts/qualification/qualify_paper.py @@ -43,6 +43,8 @@ SOAK_SNAPSHOT_INTERVAL_SECONDS = 5 * 60 SOAK_RSS_GROWTH_LIMIT_BYTES = 25 * 1024 * 1024 SOAK_SHUTDOWN_LIMIT_SECONDS = 5.0 +SOAK_RECONNECT_ATTEMPTS = 3 +SOAK_RECONNECT_RETRY_SECONDS = 2.0 EXERCISE_STEP_SEQUENCE = ( "installed_candidate", "connect", @@ -742,6 +744,18 @@ async def _sleep_until(deadline: float) -> None: await asyncio.sleep(remaining) +async def _connect_for_controlled_reconnect(provider: str, broker: Any) -> None: + attempts = SOAK_RECONNECT_ATTEMPTS if provider == "ib" else 1 + for attempt in range(1, attempts + 1): + try: + await broker.connect() + return + except RuntimeError: + if attempt == attempts: + raise + await asyncio.sleep(SOAK_RECONNECT_RETRY_SECONDS) + + async def run_provider_soak( *, provider: str, candidate: dict[str, Any], checkout_root: Path ) -> dict[str, Any]: @@ -797,7 +811,7 @@ async def run_provider_soak( maximum_shutdown_seconds = max( maximum_shutdown_seconds, time.monotonic() - shutdown_started ) - await broker.connect() + await _connect_for_controlled_reconnect(provider, broker) broker.assert_paper_trading() reconnect_count += 1 reconnected = True diff --git a/src/ml4t/live/brokers/ib.py b/src/ml4t/live/brokers/ib.py index 4a3bf95..bb37bd3 100644 --- a/src/ml4t/live/brokers/ib.py +++ b/src/ml4t/live/brokers/ib.py @@ -186,10 +186,14 @@ async def connect(self) -> None: timeout=20, # Outer timeout wrapper ) except (TimeoutError, ConnectionRefusedError) as e: + self.ib.disconnect() + self._connected = False detail = str(redact_sensitive(str(e))) logger.error("IBBroker: Connection failed: %s", detail) raise RuntimeError(f"IB connection failed: {detail}") from None except Exception as e: + self.ib.disconnect() + self._connected = False detail = str(redact_sensitive(str(e))) logger.error("IBBroker: Unexpected connect error: %s", detail) raise RuntimeError(f"IB connection failed: {detail}") from None diff --git a/tests/unit/test_ib_broker.py b/tests/unit/test_ib_broker.py index 7edcbc3..08eafe4 100644 --- a/tests/unit/test_ib_broker.py +++ b/tests/unit/test_ib_broker.py @@ -124,6 +124,22 @@ async def test_connect_with_account(self, mock_ib_class): assert broker._connected assert broker._account == "DU67890" # Pre-specified account + @pytest.mark.asyncio + @patch("ml4t.live.brokers.ib.IB") + async def test_failed_connect_cleans_up_partial_vendor_state(self, mock_ib_class): + mock_ib = MagicMock() + mock_ib_class.return_value = mock_ib + mock_ib.connectAsync = AsyncMock(side_effect=TimeoutError) + + broker = IBBroker() + broker.ib = mock_ib + + with pytest.raises(RuntimeError, match="IB connection failed"): + await broker.connect() + + mock_ib.disconnect.assert_called_once() + assert broker._connected is False + @pytest.mark.asyncio @patch("ml4t.live.brokers.ib.IB") async def test_disconnect(self, mock_ib_class): diff --git a/tests/unit/test_paper_qualification.py b/tests/unit/test_paper_qualification.py index 7dde814..40063b9 100644 --- a/tests/unit/test_paper_qualification.py +++ b/tests/unit/test_paper_qualification.py @@ -471,6 +471,37 @@ async def fake_snapshot(provider: str, captured_broker: object) -> dict: assert broker.disconnect_count == 2 +@pytest.mark.asyncio +async def test_controlled_ib_reconnect_recovers_from_transient_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + broker = MagicMock() + broker.connect = AsyncMock(side_effect=[RuntimeError("transient timeout"), None]) + sleep = AsyncMock() + monkeypatch.setattr(paper_qualification.asyncio, "sleep", sleep) + + await paper_qualification._connect_for_controlled_reconnect("ib", broker) + + assert broker.connect.await_count == 2 + sleep.assert_awaited_once_with(paper_qualification.SOAK_RECONNECT_RETRY_SECONDS) + + +@pytest.mark.asyncio +async def test_controlled_ib_reconnect_fails_after_bounded_attempts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + broker = MagicMock() + broker.connect = AsyncMock(side_effect=RuntimeError("persistent timeout")) + sleep = AsyncMock() + monkeypatch.setattr(paper_qualification.asyncio, "sleep", sleep) + + with pytest.raises(RuntimeError, match="persistent timeout"): + await paper_qualification._connect_for_controlled_reconnect("ib", broker) + + assert broker.connect.await_count == paper_qualification.SOAK_RECONNECT_ATTEMPTS + assert sleep.await_count == paper_qualification.SOAK_RECONNECT_ATTEMPTS - 1 + + @pytest.mark.asyncio async def test_capability_and_policy_rejections_never_reach_provider(tmp_path: Path) -> None: class FakeIB: