fix: harden post-stop reentry lockdown

app/services/trade_policy.py

@@ -2,13 +2,16 @@
 
 from __future__ import annotations
 
-from datetime import date, datetime, time, timezone
+from collections import defaultdict
+from datetime import date, datetime, timezone
 
-from sqlalchemy import select
+from sqlalchemy import func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from app.models.benchmark_price import BenchmarkPrice
 from app.models.ohlcv import OHLCVRecord
 from app.models.paper_trade import PaperTrade
+from app.services.benchmark_service import BENCHMARK_SYMBOL
 
 # A ticker stopped at its initial stop may qualify again immediately, but the
 # July 2026 event study showed that waiting five market sessions materially
@@ -17,52 +20,101 @@ from app.models.paper_trade import PaperTrade
 REENTRY_LOCKDOWN_SESSIONS = 5
 
 
+async def get_reentry_lockdowns(
+    db: AsyncSession,
+    *,
+    as_of: date | None = None,
+    sessions: int = REENTRY_LOCKDOWN_SESSIONS,
+) -> dict[int, int]:
+    """Return ``{ticker_id: remaining_sessions}`` for active lockdowns.
+
+    SPY is the canonical calendar for the platform's US-equity universe. When
+    the stored benchmark history does not reach an older stop, only that
+    ticker's own OHLCV dates are used as a conservative fallback. Unrelated
+    ticker dates can therefore never shorten a lockdown.
+    """
+    sessions = max(0, int(sessions))
+    if sessions == 0:
+        return {}
+
+    session_cutoff = as_of or datetime.now(timezone.utc).date()
+    stop_result = await db.execute(
+        select(
+            PaperTrade.ticker_id,
+            func.max(PaperTrade.closed_at).label("last_stop_at"),
+        )
+        .where(
+            PaperTrade.status == "closed",
+            PaperTrade.close_reason == "stop",
+            PaperTrade.closed_at.is_not(None),
+        )
+        .group_by(PaperTrade.ticker_id)
+    )
+    stop_dates = {
+        ticker_id: stopped_at.date()
+        for ticker_id, stopped_at in stop_result.all()
+        if stopped_at is not None and stopped_at.date() <= session_cutoff
+    }
+    if not stop_dates:
+        return {}
+
+    benchmark_result = await db.execute(
+        select(BenchmarkPrice.date)
+        .where(
+            BenchmarkPrice.symbol == BENCHMARK_SYMBOL,
+            BenchmarkPrice.date <= session_cutoff,
+        )
+        .order_by(BenchmarkPrice.date.asc())
+    )
+    benchmark_dates = [row[0] for row in benchmark_result.all()]
+
+    lockdowns: dict[int, int] = {}
+    fallback_stops: dict[int, date] = {}
+    first_benchmark_date = benchmark_dates[0] if benchmark_dates else None
+    for ticker_id, stop_date in stop_dates.items():
+        completed = sum(day > stop_date for day in benchmark_dates)
+        if completed >= sessions:
+            continue
+        if first_benchmark_date is not None and first_benchmark_date <= stop_date:
+            lockdowns[ticker_id] = sessions - completed
+        else:
+            # The benchmark table starts after this stop (or is empty), so it
+            # cannot prove how many sessions elapsed. Resolve only this ticker
+            # against its own bars instead of using universe-wide dates.
+            fallback_stops[ticker_id] = stop_date
+
+    if fallback_stops:
+        own_session_result = await db.execute(
+            select(OHLCVRecord.ticker_id, OHLCVRecord.date)
+            .where(
+                OHLCVRecord.ticker_id.in_(fallback_stops),
+                OHLCVRecord.date > min(fallback_stops.values()),
+                OHLCVRecord.date <= session_cutoff,
+            )
+            .distinct()
+        )
+        own_dates: dict[int, set[date]] = defaultdict(set)
+        for ticker_id, market_date in own_session_result.all():
+            own_dates[ticker_id].add(market_date)
+        for ticker_id, stop_date in fallback_stops.items():
+            completed = sum(day > stop_date for day in own_dates[ticker_id])
+            if completed < sessions:
+                lockdowns[ticker_id] = sessions - completed
+
+    return lockdowns
+
+
 async def get_reentry_lockdown_ticker_ids(
     db: AsyncSession,
     *,
     as_of: date | None = None,
     sessions: int = REENTRY_LOCKDOWN_SESSIONS,
 ) -> set[int]:
-    """Ticker ids still inside the post-initial-stop market-session lockdown.
-
-    The market calendar is derived from stored OHLCV dates, not calendar days.
-    A stop on session D is released once five later stored sessions exist. Only
-    an initial-stop close (``close_reason == "stop"``) starts the lockdown;
-    trailing, target, time, and manual exits do not.
-    """
-    sessions = max(0, int(sessions))
-    if sessions == 0:
-        return set()
-
-    session_cutoff = as_of or datetime.now(timezone.utc).date()
-    session_result = await db.execute(
-        select(OHLCVRecord.date)
-        .where(OHLCVRecord.date <= session_cutoff)
-        .distinct()
-        .order_by(OHLCVRecord.date.desc())
-        .limit(sessions)
-    )
-    recent_sessions = [row[0] for row in session_result.all()]
-    if not recent_sessions:
-        return set()
-
-    # Stops on or after the oldest of the latest N sessions have fewer than N
-    # later completed sessions. Once that oldest session rolls forward, the
-    # corresponding stop automatically leaves the result set.
-    lockdown_threshold = min(recent_sessions)
-    threshold_start = datetime.combine(
-        lockdown_threshold,
-        time.min,
-        tzinfo=timezone.utc,
-    )
-    result = await db.execute(
-        select(PaperTrade.ticker_id)
-        .where(
-            PaperTrade.status == "closed",
-            PaperTrade.close_reason == "stop",
-            PaperTrade.closed_at.is_not(None),
-            PaperTrade.closed_at >= threshold_start,
+    """Compatibility wrapper for callers that only need blocked ticker ids."""
+    return set(
+        await get_reentry_lockdowns(
+            db,
+            as_of=as_of,
+            sessions=sessions,
         )
-        .distinct()
     )
-    return {ticker_id for ticker_id, in result.all()}