From 66aa62f66ae1bda844f043cc042f3c45a47f77a7 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 13 Jul 2026 17:35:55 -0700 Subject: [PATCH] fix(ptu): coerce naive datetimes on reservation requests to UTC Pydantic accepted naive datetimes on effective_from and effective_to; comparing those against Prisma's UTC-aware row values raised TypeError, turning /ptu_reservation/close into a 500 on inputs like '2026-08-15T00:00:00'. Adds a field validator on both the request types and the domain model that stamps missing tzinfo as UTC and converts tz-aware values into UTC. Three regression tests pin the new contract (create with naive input, close with naive input, close with naive input before effective_from). --- litellm/models/ptu_reservation.py | 13 +++++- .../management_endpoints/ptu_reservation.py | 22 ++++++++- .../test_ptu_reservation_endpoints.py | 45 +++++++++++++++++++ 3 files changed, 76 insertions(+), 4 deletions(-) diff --git a/litellm/models/ptu_reservation.py b/litellm/models/ptu_reservation.py index 7513e248a36..85790653c37 100644 --- a/litellm/models/ptu_reservation.py +++ b/litellm/models/ptu_reservation.py @@ -4,10 +4,10 @@ PTU Reservation table model. Canonical definition for ``litellm_ptureservation``. """ -from datetime import datetime +from datetime import datetime, timezone from typing import Literal -from pydantic import ConfigDict, model_validator +from pydantic import ConfigDict, field_validator, model_validator from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -32,6 +32,15 @@ class LiteLLM_PTUReservation(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) + @field_validator("effective_from", "effective_to", mode="after") + @classmethod + def _coerce_utc(cls, value: datetime | None) -> datetime | None: + if value is None: + return None + if value.tzinfo is None: + return value.replace(tzinfo=timezone.utc) + return value.astimezone(timezone.utc) + @model_validator(mode="after") def _enforce_cost_source_fields(self) -> "LiteLLM_PTUReservation": if self.cost_source == "manual": diff --git a/litellm/types/proxy/management_endpoints/ptu_reservation.py b/litellm/types/proxy/management_endpoints/ptu_reservation.py index 39db7a282e1..8967b2522ae 100644 --- a/litellm/types/proxy/management_endpoints/ptu_reservation.py +++ b/litellm/types/proxy/management_endpoints/ptu_reservation.py @@ -1,13 +1,21 @@ """Request/response types for /ptu_reservation/* endpoints.""" -from datetime import datetime +from datetime import datetime, timezone from typing import Literal -from pydantic import ConfigDict, Field +from pydantic import ConfigDict, Field, field_validator from litellm.types.llms.base import LiteLLMPydanticObjectBase +def _to_utc(value: datetime | None) -> datetime | None: + if value is None: + return None + if value.tzinfo is None: + return value.replace(tzinfo=timezone.utc) + return value.astimezone(timezone.utc) + + class PTUReservationNewRequest(LiteLLMPydanticObjectBase): """Body accepted by POST /ptu_reservation/new.""" @@ -37,6 +45,11 @@ class PTUReservationNewRequest(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) + @field_validator("effective_from", "effective_to", mode="after") + @classmethod + def _coerce_utc(cls, value: datetime | None) -> datetime | None: + return _to_utc(value) + class PTUReservationListRequest(LiteLLMPydanticObjectBase): """Query params for GET /ptu_reservation/list.""" @@ -54,3 +67,8 @@ class PTUReservationCloseRequest(LiteLLMPydanticObjectBase): default=None, description="Timestamp to set as effective_to. Defaults to now (UTC) if omitted.", ) + + @field_validator("effective_to", mode="after") + @classmethod + def _coerce_utc(cls, value: datetime | None) -> datetime | None: + return _to_utc(value) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py index fe7402558e2..0fcba1074dd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py @@ -413,3 +413,48 @@ async def test_create_after_close_of_same_team_model(client_and_mocks): ) assert resp.status_code == 200, resp.text mock_table.create.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_new_reservation_accepts_naive_datetime_and_coerces_to_utc(client_and_mocks): + client, _, mock_table = client_and_mocks + + resp = client.post( + "/ptu_reservation/new", + json=_valid_new_payload(effective_from="2026-08-01T00:00:00"), + ) + assert resp.status_code == 200, resp.text + mock_table.create.assert_awaited_once() + created = mock_table.create.await_args.kwargs["data"]["effective_from"] + assert created.tzinfo is not None + assert created.utcoffset() == timedelta(0) + + +@pytest.mark.asyncio +async def test_close_with_naive_datetime_does_not_500(client_and_mocks): + client, _, mock_table = client_and_mocks + existing = _row(id="res_1", effective_from=_dt(), effective_to=None) + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.post( + "/ptu_reservation/close", + json={"id": "res_1", "effective_to": "2026-07-16T00:00:00"}, + ) + assert resp.status_code == 200, resp.text + stored = mock_table.update.await_args.kwargs["data"]["effective_to"] + assert stored.tzinfo is not None + assert stored.utcoffset() == timedelta(0) + + +@pytest.mark.asyncio +async def test_close_rejects_naive_effective_to_before_from(client_and_mocks): + client, _, mock_table = client_and_mocks + existing = _row(id="res_1", effective_from=_dt(days=10), effective_to=None) + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.post( + "/ptu_reservation/close", + json={"id": "res_1", "effective_to": "2026-07-05T00:00:00"}, + ) + assert resp.status_code == 400, resp.text + mock_table.update.assert_not_awaited()