diff --git a/app/modules/quota_planner/repository.py b/app/modules/quota_planner/repository.py index a666eb74d..50258f1d1 100644 --- a/app/modules/quota_planner/repository.py +++ b/app/modules/quota_planner/repository.py @@ -7,7 +7,7 @@ from sqlalchemy import Integer, and_, cast, func, literal, select, update from sqlalchemy.ext.asyncio import AsyncSession -from app.core.utils.time import utcnow +from app.core.utils.time import to_utc_naive, utcnow from app.db.models import ( Account, QuotaPlannerDecision, @@ -100,8 +100,8 @@ async def log_decision( mode=mode, action=action, account_id=account_id, - scheduled_at=scheduled_at, - executed_at=executed_at, + scheduled_at=_quota_planner_db_timestamp(scheduled_at), + executed_at=_quota_planner_db_timestamp(executed_at), score=score, reason=reason, forecast_snapshot_hash=forecast_snapshot_hash, @@ -138,7 +138,7 @@ async def update_decision_status( if reason is not None: values["reason"] = reason if executed_at is not None: - values["executed_at"] = executed_at + values["executed_at"] = _quota_planner_db_timestamp(executed_at) if state_after_json is not None: values["state_after_json"] = state_after_json stmt = update(QuotaPlannerDecision).where(QuotaPlannerDecision.id == decision_id).values(**values) @@ -160,6 +160,7 @@ async def update_decision_status( return row async def count_executed_warmups_since(self, since: datetime) -> int: + since = to_utc_naive(since) stmt = select(func.count(QuotaPlannerDecision.id)).where( and_( QuotaPlannerDecision.action == "warmup", @@ -170,6 +171,7 @@ async def count_executed_warmups_since(self, since: datetime) -> int: return int(await self._session.scalar(stmt) or 0) async def warmup_cost_since(self, since: datetime) -> float: + since = to_utc_naive(since) stmt = select(func.coalesce(func.sum(RequestLog.cost_usd), 0.0)).where( and_( RequestLog.request_kind == "warmup", @@ -215,7 +217,7 @@ async def add_window_observation( ) -> QuotaWindowObservation: row = QuotaWindowObservation( account_id=account_id, - observed_at=observed_at or utcnow(), + observed_at=_quota_planner_db_timestamp(observed_at) or utcnow(), model=model, primary_remaining_percent=primary_remaining_percent, primary_reset_at=primary_reset_at, @@ -236,7 +238,7 @@ async def aggregate_demand_bins( since: datetime | None = None, bucket_seconds: int = 900, ) -> list[DemandBin]: - since = since or (utcnow() - timedelta(days=28)) + since = to_utc_naive(since) if since is not None else (utcnow() - timedelta(days=28)) bind = self._session.get_bind() dialect = bind.dialect.name if bind else "sqlite" if dialect == "postgresql": @@ -296,6 +298,12 @@ async def aggregate_demand_bins( ] +def _quota_planner_db_timestamp(value: datetime | None) -> datetime | None: + if value is None: + return None + return to_utc_naive(value) + + def _settings_from_row(row: QuotaPlannerSettings) -> PlannerSettings: return PlannerSettings( mode=row.mode, diff --git a/tests/unit/test_quota_planner.py b/tests/unit/test_quota_planner.py index 036305066..fc352e779 100644 --- a/tests/unit/test_quota_planner.py +++ b/tests/unit/test_quota_planner.py @@ -5,7 +5,8 @@ import pytest from app.core.balancer import AccountState -from app.db.models import AccountStatus +from app.db.models import Account, AccountStatus, RequestLog +from app.db.session import SessionLocal from app.modules.quota_planner.logic import ( DemandForecastSlot, PlannerAction, @@ -18,7 +19,7 @@ plan_shadow_actions, simulate_pool, ) -from app.modules.quota_planner.repository import DemandBin +from app.modules.quota_planner.repository import DemandBin, QuotaPlannerRepository, _quota_planner_db_timestamp pytestmark = pytest.mark.unit @@ -92,6 +93,94 @@ def test_planner_settings_default_to_nonblocking_shadow_mode() -> None: assert settings.dry_run is True +def test_quota_planner_db_timestamp_normalizes_aware_datetimes_to_naive_utc() -> None: + scheduled_at = datetime(2026, 6, 18, 8, 30, tzinfo=timezone(timedelta(hours=3))) + + assert _quota_planner_db_timestamp(scheduled_at) == datetime(2026, 6, 18, 5, 30) + assert _quota_planner_db_timestamp(None) is None + + +@pytest.mark.asyncio +async def test_quota_planner_repository_normalizes_aware_datetimes_at_session_boundary(db_setup) -> None: + del db_setup + tz_plus_3 = timezone(timedelta(hours=3)) + scheduled_at = datetime(2026, 6, 18, 8, 30, tzinfo=tz_plus_3) + executed_at = datetime(2026, 6, 18, 9, 0, tzinfo=tz_plus_3) + observed_at = datetime(2026, 6, 18, 9, 15, tzinfo=tz_plus_3) + aware_since = datetime(2026, 6, 18, 8, 30, tzinfo=tz_plus_3) + + async with SessionLocal() as session: + session.add( + Account( + id="quota-planner-aware-account", + email="quota-planner-aware@example.com", + plan_type="plus", + access_token_encrypted=b"access", + refresh_token_encrypted=b"refresh", + id_token_encrypted=b"id", + last_refresh=datetime(2026, 6, 18, 5, 0), + status=AccountStatus.ACTIVE, + ) + ) + await session.commit() + + repo = QuotaPlannerRepository(session) + decision = await repo.log_decision( + mode="shadow", + action="warmup", + idempotency_key="quota-planner-aware-boundary", + account_id="quota-planner-aware-account", + scheduled_at=scheduled_at, + status="planned", + ) + + assert decision.scheduled_at == datetime(2026, 6, 18, 5, 30) + + updated = await repo.update_decision_status( + decision.id, + status="executed", + executed_at=executed_at, + expected_status="planned", + ) + + assert updated is not None + assert updated.executed_at == datetime(2026, 6, 18, 6, 0) + + observation = await repo.add_window_observation( + account_id="quota-planner-aware-account", + source="warmup_probe", + observed_at=observed_at, + model="gpt-5.4-mini", + confidence="observed", + ) + + assert observation.observed_at == datetime(2026, 6, 18, 6, 15) + + session.add( + RequestLog( + request_id="quota-planner-aware-warmup", + request_kind="warmup", + requested_at=datetime(2026, 6, 18, 6, 30), + model="gpt-5.4-mini", + status="ok", + input_tokens=12, + cached_input_tokens=3, + output_tokens=4, + cost_usd=0.25, + ) + ) + await session.commit() + + assert await repo.count_executed_warmups_since(aware_since) == 1 + assert await repo.warmup_cost_since(aware_since) == pytest.approx(0.25) + + bins = await repo.aggregate_demand_bins(since=aware_since) + + assert len(bins) == 1 + assert bins[0].request_count == 1 + assert bins[0].cost_usd == pytest.approx(0.25) + + def test_candidate_start_times_do_not_floor_now_into_the_past() -> None: settings = PlannerSettings( mode="shadow",