From 30fe7b4a12d186372f49c7a30a1d19010a3cb659 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 13 Jul 2026 16:04:13 -0700 Subject: [PATCH] feat(ptu): add PTU reservation table and admin CRUD endpoints Adds a new LiteLLM_PTUReservation table storing admin-registered PTU reservations for a (team, model) pair over a time window, plus admin CRUD endpoints (new, list, info, close). Feature-gated by enable_ptu_cost_attribution in general_settings; default off. Stage 1 storage + endpoints only. No spend behavior change: no daily rollup job, no writes to LiteLLM_DailyTeamSpend, no impact on the per-request cost tracking pipeline. - schema.prisma + migration for LiteLLM_PTUReservation - Pydantic domain models in litellm/models/ptu_reservation.py - Repository in litellm/repositories/ptu_reservation_repository.py - Endpoints in litellm/proxy/management_endpoints/ptu_reservation_endpoints.py - Request/response types under litellm/types/proxy/management_endpoints/ - 28 unit tests covering validation, feature-flag gate, admin-only auth, overlap detection, close/create semantics - Routes listed under management_routes in litellm/proxy/_types.py - Router mounted in proxy_server.py alongside budget_management_router --- .../migration.sql | 27 ++ litellm/models/ptu_reservation.py | 60 +++ litellm/proxy/_types.py | 5 + .../ptu_reservation_endpoints.py | 242 ++++++++++ litellm/proxy/proxy_server.py | 4 + .../ptu_reservation_repository.py | 66 +++ .../management_endpoints/ptu_reservation.py | 56 +++ schema.prisma | 23 +- .../test_ptu_reservation_endpoints.py | 415 ++++++++++++++++++ 9 files changed, 897 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260713000000_add_litellm_ptu_reservation_table/migration.sql create mode 100644 litellm/models/ptu_reservation.py create mode 100644 litellm/proxy/management_endpoints/ptu_reservation_endpoints.py create mode 100644 litellm/repositories/ptu_reservation_repository.py create mode 100644 litellm/types/proxy/management_endpoints/ptu_reservation.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713000000_add_litellm_ptu_reservation_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713000000_add_litellm_ptu_reservation_table/migration.sql new file mode 100644 index 00000000000..4ece18e8857 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713000000_add_litellm_ptu_reservation_table/migration.sql @@ -0,0 +1,27 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_PTUReservation" ( + "id" TEXT NOT NULL, + "team_id" TEXT NOT NULL, + "model" TEXT NOT NULL, + "cost_source" TEXT NOT NULL DEFAULT 'manual', + "ptu_count" INTEGER, + "cost_per_ptu" DOUBLE PRECISION, + "azure_resource_id" TEXT, + "effective_from" TIMESTAMP(3) NOT NULL, + "effective_to" TIMESTAMP(3), + "created_by" TEXT NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_PTUReservation_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_PTUReservation_team_id_model_effective_from_idx" ON "LiteLLM_PTUReservation"("team_id", "model", "effective_from"); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_PTUReservation_cost_source_azure_resource_id_idx" ON "LiteLLM_PTUReservation"("cost_source", "azure_resource_id"); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_PTUReservation_effective_from_effective_to_idx" ON "LiteLLM_PTUReservation"("effective_from", "effective_to"); diff --git a/litellm/models/ptu_reservation.py b/litellm/models/ptu_reservation.py new file mode 100644 index 00000000000..7513e248a36 --- /dev/null +++ b/litellm/models/ptu_reservation.py @@ -0,0 +1,60 @@ +""" +PTU Reservation table model. + +Canonical definition for ``litellm_ptureservation``. +""" + +from datetime import datetime +from typing import Literal + +from pydantic import ConfigDict, model_validator + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + +CostSource = Literal["manual", "azure_billing"] + + +class LiteLLM_PTUReservation(LiteLLMPydanticObjectBase): + """Represents user-controllable params for a LiteLLM_PTUReservation record.""" + + id: str | None = None + team_id: str + model: str + cost_source: CostSource = "manual" + + ptu_count: int | None = None + cost_per_ptu: float | None = None + + azure_resource_id: str | None = None + + effective_from: datetime + effective_to: datetime | None = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def _enforce_cost_source_fields(self) -> "LiteLLM_PTUReservation": + if self.cost_source == "manual": + if self.ptu_count is None or self.cost_per_ptu is None: + raise ValueError("manual reservations require both ptu_count and cost_per_ptu") + if self.ptu_count <= 0: + raise ValueError("ptu_count must be positive") + if self.cost_per_ptu < 0: + raise ValueError("cost_per_ptu must be non-negative") + elif self.cost_source == "azure_billing": + if self.azure_resource_id is None: + raise ValueError("azure_billing reservations require azure_resource_id") + if self.ptu_count is not None or self.cost_per_ptu is not None: + raise ValueError("azure_billing reservations must not set ptu_count or cost_per_ptu") + if self.effective_to is not None and self.effective_to <= self.effective_from: + raise ValueError("effective_to must be strictly after effective_from") + return self + + +class LiteLLM_PTUReservationFull(LiteLLM_PTUReservation): + """LiteLLM_PTUReservation + server-managed fields returned on API responses.""" + + created_by: str + created_at: datetime + updated_by: str + updated_at: datetime diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5e3ea4b7dcb..25023e6b70e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -611,6 +611,11 @@ class LiteLLMRoutes(enum.Enum): "/jwt/key/mapping/delete", "/jwt/key/mapping/list", "/jwt/key/mapping/info", + # ptu reservations + "/ptu_reservation/new", + "/ptu_reservation/list", + "/ptu_reservation/info", + "/ptu_reservation/close", ] + key_management_routes + mcp_management_routes diff --git a/litellm/proxy/management_endpoints/ptu_reservation_endpoints.py b/litellm/proxy/management_endpoints/ptu_reservation_endpoints.py new file mode 100644 index 00000000000..c40289a13f2 --- /dev/null +++ b/litellm/proxy/management_endpoints/ptu_reservation_endpoints.py @@ -0,0 +1,242 @@ +""" +PTU RESERVATION MANAGEMENT + +All /ptu_reservation management endpoints. + +/ptu_reservation/new +/ptu_reservation/list +/ptu_reservation/info +/ptu_reservation/close +""" + +from datetime import datetime, timezone +from typing import Annotated + +from fastapi import APIRouter, Depends, HTTPException + +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.ptu_reservation_repository import PTUReservationRepository +from litellm.types.proxy.management_endpoints.ptu_reservation import ( + PTUReservationCloseRequest, + PTUReservationNewRequest, +) + +router = APIRouter() + +CurrentUser = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)] + + +def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={ + "error": "{}, your role={}".format( + CommonProxyErrors.not_allowed_access.value, + user_api_key_dict.user_role, + ) + }, + ) + + +def _require_feature_enabled() -> None: + from litellm.proxy.proxy_server import general_settings + + if not general_settings.get("enable_ptu_cost_attribution", False): + raise HTTPException( + status_code=403, + detail={ + "error": ( + "PTU cost attribution is not enabled. Set 'enable_ptu_cost_attribution: true' in general_settings." + ) + }, + ) + + +def _require_db() -> "object": + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + return prisma_client + + +@router.post( + "/ptu_reservation/new", + tags=["ptu reservation management"], + dependencies=[Depends(user_api_key_auth)], +) +async def new_ptu_reservation( + body: PTUReservationNewRequest, + user_api_key_dict: CurrentUser, +): + """Create a new PTU reservation for a (team, model) pair. + + Parameters: + - team_id (str, required) + - model (str, required) + - cost_source (str): "manual" (default). "azure_billing" is reserved. + - ptu_count (int): required for cost_source="manual", positive + - cost_per_ptu (float): required for cost_source="manual", non-negative USD/month + - azure_resource_id (str): reserved; must be null for manual + - effective_from (datetime, required): inclusive UTC start + - effective_to (datetime, optional): exclusive UTC end; null = still active + """ + _require_feature_enabled() + _require_proxy_admin(user_api_key_dict) + prisma_client = _require_db() + + if body.cost_source == "azure_billing": + raise HTTPException( + status_code=400, + detail={"error": "cost_source='azure_billing' is not supported in this release"}, + ) + + try: + validated = body.model_dump(exclude_none=False) + validated_reservation = _validated_domain_model(validated) + except ValueError as e: + raise HTTPException(status_code=400, detail={"error": str(e)}) + + repo = PTUReservationRepository(prisma_client) + overlapping = await repo.find_overlapping( + team_id=validated_reservation["team_id"], + model=validated_reservation["model"], + effective_from=validated_reservation["effective_from"], + effective_to=validated_reservation["effective_to"], + ) + if overlapping: + raise HTTPException( + status_code=409, + detail={ + "error": "reservation overlaps existing active reservation(s) for the same (team, model)", + "overlapping_ids": [r.id for r in overlapping], + }, + ) + + actor = user_api_key_dict.user_id or "admin" + create_data = { + **{k: v for k, v in validated_reservation.items() if k != "id" and v is not None}, + "created_by": actor, + "updated_by": actor, + } + return await repo.table.create(data=create_data) + + +@router.get( + "/ptu_reservation/list", + tags=["ptu reservation management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_ptu_reservations( + user_api_key_dict: CurrentUser, + team_id: str | None = None, + model: str | None = None, + active_only: bool = False, +): + """List PTU reservations. + + Query parameters: + - team_id (optional): filter by team + - model (optional): filter by model + - active_only (optional, default false): only reservations live right now + """ + _require_feature_enabled() + _require_proxy_admin(user_api_key_dict) + prisma_client = _require_db() + + repo = PTUReservationRepository(prisma_client) + if active_only: + return await repo.find_active(as_of=datetime.now(timezone.utc), team_id=team_id, model=model) + + where: dict = {} + if team_id is not None: + where["team_id"] = team_id + if model is not None: + where["model"] = model + return await repo.table.find_many(where=where) + + +@router.get( + "/ptu_reservation/info", + tags=["ptu reservation management"], + dependencies=[Depends(user_api_key_auth)], +) +async def info_ptu_reservation( + id: str, + user_api_key_dict: CurrentUser, +): + """Get a single reservation by id. + + Query parameter: + - id (str, required) + """ + _require_feature_enabled() + _require_proxy_admin(user_api_key_dict) + prisma_client = _require_db() + + repo = PTUReservationRepository(prisma_client) + row = await repo.table.find_unique(where={"id": id}) + if row is None: + raise HTTPException(status_code=404, detail={"error": f"reservation '{id}' not found"}) + return row + + +@router.post( + "/ptu_reservation/close", + tags=["ptu reservation management"], + dependencies=[Depends(user_api_key_auth)], +) +async def close_ptu_reservation( + body: PTUReservationCloseRequest, + user_api_key_dict: CurrentUser, +): + """Close a reservation by setting effective_to. + + Parameters: + - id (str, required) + - effective_to (datetime, optional): defaults to now UTC + """ + _require_feature_enabled() + _require_proxy_admin(user_api_key_dict) + prisma_client = _require_db() + + repo = PTUReservationRepository(prisma_client) + row = await repo.table.find_unique(where={"id": body.id}) + if row is None: + raise HTTPException(status_code=404, detail={"error": f"reservation '{body.id}' not found"}) + if row.effective_to is not None: + raise HTTPException( + status_code=400, + detail={"error": f"reservation '{body.id}' is already closed at {row.effective_to.isoformat()}"}, + ) + + close_at = body.effective_to or datetime.now(timezone.utc) + if close_at <= row.effective_from: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"effective_to ({close_at.isoformat()}) must be strictly after " + f"effective_from ({row.effective_from.isoformat()})" + ) + }, + ) + + actor = user_api_key_dict.user_id or "admin" + return await repo.table.update( + where={"id": body.id}, + data={"effective_to": close_at, "updated_by": actor}, + ) + + +def _validated_domain_model(payload: dict) -> dict: + """Validate a reservation payload against the domain model and return its dict form.""" + from litellm.models.ptu_reservation import LiteLLM_PTUReservation + + reservation = LiteLLM_PTUReservation(**payload) + return reservation.model_dump(exclude_none=False) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6661474d215..a4cdca1f869 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -357,6 +357,9 @@ from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( from litellm.proxy.management_endpoints.budget_management_endpoints import ( router as budget_management_router, ) +from litellm.proxy.management_endpoints.ptu_reservation_endpoints import ( + router as ptu_reservation_router, +) from litellm.proxy.management_endpoints.cache_settings_endpoints import ( router as cache_settings_router, ) @@ -16079,6 +16082,7 @@ app.include_router(rust_control_plane_router) app.include_router(ui_crud_endpoints_router) app.include_router(team_callback_router) app.include_router(budget_management_router) +app.include_router(ptu_reservation_router) app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) diff --git a/litellm/repositories/ptu_reservation_repository.py b/litellm/repositories/ptu_reservation_repository.py new file mode 100644 index 00000000000..2af503b4050 --- /dev/null +++ b/litellm/repositories/ptu_reservation_repository.py @@ -0,0 +1,66 @@ +""" +PTU Reservation repository for database operations on LiteLLM_PTUReservation. +""" + +from datetime import datetime +from typing import Any + +from litellm.models.ptu_reservation import LiteLLM_PTUReservationFull +from litellm.repositories.base_repository import BaseRepository + + +class PTUReservationRepository(BaseRepository[LiteLLM_PTUReservationFull]): + """Repository for PTU reservation database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_ptureservation + + @property + def model_class(self) -> type[LiteLLM_PTUReservationFull]: + return LiteLLM_PTUReservationFull + + async def find_active( + self, + as_of: datetime, + team_id: str | None = None, + model: str | None = None, + ) -> list[Any]: + """Return reservations live at ``as_of``.""" + where: dict = { + "effective_from": {"lte": as_of}, + "OR": [ + {"effective_to": None}, + {"effective_to": {"gt": as_of}}, + ], + } + if team_id is not None: + where["team_id"] = team_id + if model is not None: + where["model"] = model + return await self.table.find_many(where=where) + + async def find_overlapping( + self, + team_id: str, + model: str, + effective_from: datetime, + effective_to: datetime | None, + ) -> list[Any]: + """Return existing reservations overlapping the proposed window for the same (team, model).""" + where: dict = { + "team_id": team_id, + "model": model, + } + if effective_to is None: + where["OR"] = [ + {"effective_to": None}, + {"effective_to": {"gt": effective_from}}, + ] + else: + where["effective_from"] = {"lt": effective_to} + where["OR"] = [ + {"effective_to": None}, + {"effective_to": {"gt": effective_from}}, + ] + return await self.table.find_many(where=where) diff --git a/litellm/types/proxy/management_endpoints/ptu_reservation.py b/litellm/types/proxy/management_endpoints/ptu_reservation.py new file mode 100644 index 00000000000..39db7a282e1 --- /dev/null +++ b/litellm/types/proxy/management_endpoints/ptu_reservation.py @@ -0,0 +1,56 @@ +"""Request/response types for /ptu_reservation/* endpoints.""" + +from datetime import datetime +from typing import Literal + +from pydantic import ConfigDict, Field + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class PTUReservationNewRequest(LiteLLMPydanticObjectBase): + """Body accepted by POST /ptu_reservation/new.""" + + team_id: str = Field(description="Team that owns the reservation.") + model: str = Field(description="Model name the reservation covers.") + cost_source: Literal["manual", "azure_billing"] = Field( + default="manual", + description="Source of the cost figures. Only 'manual' is supported in this release.", + ) + ptu_count: int | None = Field( + default=None, + description="Number of provisioned throughput units. Required for manual.", + ) + cost_per_ptu: float | None = Field( + default=None, + description="Monthly cost per PTU in USD. Required for manual.", + ) + azure_resource_id: str | None = Field( + default=None, + description="Azure resource ARM id, for cost_source='azure_billing'.", + ) + effective_from: datetime = Field(description="Inclusive UTC start of the reservation.") + effective_to: datetime | None = Field( + default=None, + description="Exclusive UTC end of the reservation. Null = still active.", + ) + + model_config = ConfigDict(protected_namespaces=()) + + +class PTUReservationListRequest(LiteLLMPydanticObjectBase): + """Query params for GET /ptu_reservation/list.""" + + team_id: str | None = None + model: str | None = None + active_only: bool = False + + +class PTUReservationCloseRequest(LiteLLMPydanticObjectBase): + """Body accepted by POST /ptu_reservation/close.""" + + id: str = Field(description="Reservation id to close.") + effective_to: datetime | None = Field( + default=None, + description="Timestamp to set as effective_to. Defaults to now (UTC) if omitted.", + ) diff --git a/schema.prisma b/schema.prisma index a23cecc3911..acedf8029b6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -30,7 +30,28 @@ model LiteLLM_BudgetTable { end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team - organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization +} + +// Admin-registered PTU reservations: flat prepaid cost for a (team, model) window +model LiteLLM_PTUReservation { + id String @id @default(uuid()) + team_id String + model String + cost_source String @default("manual") + ptu_count Int? + cost_per_ptu Float? + azure_resource_id String? + effective_from DateTime + effective_to DateTime? + created_by String + created_at DateTime @default(now()) @map("created_at") + updated_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + + @@index([team_id, model, effective_from]) + @@index([cost_source, azure_resource_id]) + @@index([effective_from, effective_to]) } // Models on proxy 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 new file mode 100644 index 00000000000..fe7402558e2 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py @@ -0,0 +1,415 @@ +import os +import sys +import types +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +import litellm.proxy.proxy_server as ps +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.proxy_server import app + +sys.path.insert(0, os.path.abspath("../../../")) + + +def _dt(days: int = 0, hours: int = 0) -> datetime: + return datetime(2026, 7, 1, tzinfo=timezone.utc) + timedelta(days=days, hours=hours) + + +def _iso(days: int = 0, hours: int = 0) -> str: + return _dt(days=days, hours=hours).isoformat() + + +def _row( + *, + id: str = "res_1", + team_id: str = "team_x", + model: str = "gpt-4", + cost_source: str = "manual", + ptu_count: int | None = 1, + cost_per_ptu: float | None = 200.0, + azure_resource_id: str | None = None, + effective_from: datetime | None = None, + effective_to: datetime | None = None, + created_by: str = "admin", + updated_by: str = "admin", +): + row = MagicMock() + row.id = id + row.team_id = team_id + row.model = model + row.cost_source = cost_source + row.ptu_count = ptu_count + row.cost_per_ptu = cost_per_ptu + row.azure_resource_id = azure_resource_id + row.effective_from = effective_from or _dt() + row.effective_to = effective_to + row.created_by = created_by + row.updated_by = updated_by + return row + + +@pytest.fixture +def client_and_mocks(monkeypatch): + mock_prisma = MagicMock() + mock_table = MagicMock() + mock_table.create = AsyncMock(side_effect=lambda *, data: _row(**{k: v for k, v in data.items() if k != "created_at" and k != "updated_at"})) + mock_table.update = AsyncMock(side_effect=lambda *, where, data: _row(id=where["id"], **{k: v for k, v in data.items() if k != "updated_at"})) + mock_table.find_many = AsyncMock(return_value=[]) + mock_table.find_unique = AsyncMock(return_value=None) + + mock_prisma.db = types.SimpleNamespace(litellm_ptureservation=mock_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + monkeypatch.setitem(ps.general_settings, "enable_ptu_cost_attribution", True) + + admin_user = UserAPIKeyAuth(user_id="test_admin", user_role=LitellmUserRoles.PROXY_ADMIN) + app.dependency_overrides[ps.user_api_key_auth] = lambda: admin_user + + client = TestClient(app) + yield client, mock_prisma, mock_table + + app.dependency_overrides.clear() + + +def _switch_to_internal_user() -> None: + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="not_admin", user_role=LitellmUserRoles.INTERNAL_USER + ) + + +def _valid_new_payload(**overrides): + payload = { + "team_id": "team_x", + "model": "gpt-4", + "ptu_count": 1, + "cost_per_ptu": 200.0, + "effective_from": _iso(), + } + payload.update(overrides) + return payload + + +@pytest.mark.asyncio +async def test_new_reservation_success(client_and_mocks): + client, _, mock_table = client_and_mocks + + resp = client.post("/ptu_reservation/new", json=_valid_new_payload()) + assert resp.status_code == 200, resp.text + + mock_table.find_many.assert_awaited_once() + mock_table.create.assert_awaited_once() + body = mock_table.create.await_args.kwargs["data"] + assert body["team_id"] == "team_x" + assert body["model"] == "gpt-4" + assert body["ptu_count"] == 1 + assert body["cost_per_ptu"] == 200.0 + assert body["cost_source"] == "manual" + assert body["created_by"] == "test_admin" + assert body["updated_by"] == "test_admin" + assert "id" not in body + + +@pytest.mark.asyncio +async def test_new_reservation_rejects_manual_without_ptu_count(client_and_mocks): + client, _, _ = client_and_mocks + + payload = _valid_new_payload() + del payload["ptu_count"] + resp = client.post("/ptu_reservation/new", json=payload) + assert resp.status_code == 400, resp.text + assert "ptu_count" in resp.json()["detail"]["error"] + + +@pytest.mark.asyncio +async def test_new_reservation_rejects_manual_without_cost_per_ptu(client_and_mocks): + client, _, _ = client_and_mocks + + payload = _valid_new_payload() + del payload["cost_per_ptu"] + resp = client.post("/ptu_reservation/new", json=payload) + assert resp.status_code == 400, resp.text + assert "cost_per_ptu" in resp.json()["detail"]["error"] + + +@pytest.mark.asyncio +async def test_new_reservation_rejects_non_positive_ptu_count(client_and_mocks): + client, _, _ = client_and_mocks + resp = client.post("/ptu_reservation/new", json=_valid_new_payload(ptu_count=0)) + assert resp.status_code == 400, resp.text + + +@pytest.mark.asyncio +async def test_new_reservation_rejects_negative_cost_per_ptu(client_and_mocks): + client, _, _ = client_and_mocks + resp = client.post("/ptu_reservation/new", json=_valid_new_payload(cost_per_ptu=-1.0)) + assert resp.status_code == 400, resp.text + + +@pytest.mark.asyncio +async def test_new_reservation_rejects_effective_to_before_from(client_and_mocks): + client, _, _ = client_and_mocks + resp = client.post( + "/ptu_reservation/new", + json=_valid_new_payload(effective_from=_iso(days=10), effective_to=_iso(days=1)), + ) + assert resp.status_code == 400, resp.text + assert "effective_to" in resp.json()["detail"]["error"] + + +@pytest.mark.asyncio +async def test_new_reservation_rejects_effective_to_equal_from(client_and_mocks): + client, _, _ = client_and_mocks + t = _iso() + resp = client.post( + "/ptu_reservation/new", + json=_valid_new_payload(effective_from=t, effective_to=t), + ) + assert resp.status_code == 400, resp.text + + +@pytest.mark.asyncio +async def test_new_reservation_rejects_azure_billing_mode(client_and_mocks): + client, _, mock_table = client_and_mocks + resp = client.post( + "/ptu_reservation/new", + json={ + "team_id": "team_x", + "model": "gpt-4", + "cost_source": "azure_billing", + "azure_resource_id": "/subscriptions/x/deployments/gpt-4-ptu", + "effective_from": _iso(), + }, + ) + assert resp.status_code == 400, resp.text + mock_table.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_reservation_rejects_overlap(client_and_mocks): + client, _, mock_table = client_and_mocks + mock_table.find_many = AsyncMock(return_value=[_row(id="existing_1")]) + + resp = client.post("/ptu_reservation/new", json=_valid_new_payload()) + assert resp.status_code == 409, resp.text + body = resp.json() + assert body["detail"]["overlapping_ids"] == ["existing_1"] + mock_table.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_reservation_allows_non_overlapping_same_team_model(client_and_mocks): + client, _, mock_table = client_and_mocks + mock_table.find_many = AsyncMock(return_value=[]) + + resp = client.post( + "/ptu_reservation/new", + json=_valid_new_payload(effective_from=_iso(days=40)), + ) + assert resp.status_code == 200, resp.text + mock_table.create.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_feature_flag_off_rejects_new(client_and_mocks, monkeypatch): + client, _, mock_table = client_and_mocks + monkeypatch.setitem(ps.general_settings, "enable_ptu_cost_attribution", False) + + resp = client.post("/ptu_reservation/new", json=_valid_new_payload()) + assert resp.status_code == 403, resp.text + mock_table.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_feature_flag_off_rejects_list(client_and_mocks, monkeypatch): + client, _, _ = client_and_mocks + monkeypatch.setitem(ps.general_settings, "enable_ptu_cost_attribution", False) + + resp = client.get("/ptu_reservation/list") + assert resp.status_code == 403, resp.text + + +@pytest.mark.asyncio +async def test_feature_flag_off_rejects_info(client_and_mocks, monkeypatch): + client, _, _ = client_and_mocks + monkeypatch.setitem(ps.general_settings, "enable_ptu_cost_attribution", False) + + resp = client.get("/ptu_reservation/info?id=res_1") + assert resp.status_code == 403, resp.text + + +@pytest.mark.asyncio +async def test_feature_flag_off_rejects_close(client_and_mocks, monkeypatch): + client, _, _ = client_and_mocks + monkeypatch.setitem(ps.general_settings, "enable_ptu_cost_attribution", False) + + resp = client.post("/ptu_reservation/close", json={"id": "res_1"}) + assert resp.status_code == 403, resp.text + + +@pytest.mark.asyncio +async def test_non_admin_forbidden_from_new(client_and_mocks): + client, _, mock_table = client_and_mocks + _switch_to_internal_user() + + resp = client.post("/ptu_reservation/new", json=_valid_new_payload()) + assert resp.status_code == 403, resp.text + mock_table.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_non_admin_forbidden_from_list(client_and_mocks): + client, _, _ = client_and_mocks + _switch_to_internal_user() + + resp = client.get("/ptu_reservation/list") + assert resp.status_code == 403, resp.text + + +@pytest.mark.asyncio +async def test_non_admin_forbidden_from_close(client_and_mocks): + client, _, _ = client_and_mocks + _switch_to_internal_user() + + resp = client.post("/ptu_reservation/close", json={"id": "res_1"}) + assert resp.status_code == 403, resp.text + + +@pytest.mark.asyncio +async def test_list_no_filters(client_and_mocks): + client, _, mock_table = client_and_mocks + mock_table.find_many = AsyncMock(return_value=[_row(id="a"), _row(id="b")]) + + resp = client.get("/ptu_reservation/list") + assert resp.status_code == 200, resp.text + mock_table.find_many.assert_awaited_once_with(where={}) + + +@pytest.mark.asyncio +async def test_list_with_team_and_model_filters(client_and_mocks): + client, _, mock_table = client_and_mocks + mock_table.find_many = AsyncMock(return_value=[]) + + resp = client.get("/ptu_reservation/list?team_id=team_x&model=gpt-4") + assert resp.status_code == 200, resp.text + mock_table.find_many.assert_awaited_once_with(where={"team_id": "team_x", "model": "gpt-4"}) + + +@pytest.mark.asyncio +async def test_list_active_only(client_and_mocks): + client, _, mock_table = client_and_mocks + mock_table.find_many = AsyncMock(return_value=[]) + + resp = client.get("/ptu_reservation/list?active_only=true&team_id=team_x") + assert resp.status_code == 200, resp.text + where = mock_table.find_many.await_args.kwargs["where"] + assert where["team_id"] == "team_x" + assert "effective_from" in where + assert "OR" in where + + +@pytest.mark.asyncio +async def test_info_not_found(client_and_mocks): + client, _, mock_table = client_and_mocks + mock_table.find_unique = AsyncMock(return_value=None) + + resp = client.get("/ptu_reservation/info?id=missing") + assert resp.status_code == 404, resp.text + + +@pytest.mark.asyncio +async def test_info_returns_row(client_and_mocks): + client, _, mock_table = client_and_mocks + mock_table.find_unique = AsyncMock(return_value=_row(id="res_9")) + + resp = client.get("/ptu_reservation/info?id=res_9") + assert resp.status_code == 200, resp.text + + +@pytest.mark.asyncio +async def test_close_sets_effective_to(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"}) + assert resp.status_code == 200, resp.text + + mock_table.update.assert_awaited_once() + update_kwargs = mock_table.update.await_args.kwargs + assert update_kwargs["where"] == {"id": "res_1"} + assert "effective_to" in update_kwargs["data"] + assert update_kwargs["data"]["updated_by"] == "test_admin" + + +@pytest.mark.asyncio +async def test_close_with_explicit_effective_to(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) + + when = _iso(days=15) + resp = client.post("/ptu_reservation/close", json={"id": "res_1", "effective_to": when}) + assert resp.status_code == 200, resp.text + + update_kwargs = mock_table.update.await_args.kwargs + assert update_kwargs["data"]["effective_to"].isoformat() == when + + +@pytest.mark.asyncio +async def test_close_refuses_already_closed(client_and_mocks): + client, _, mock_table = client_and_mocks + existing = _row(id="res_1", effective_from=_dt(), effective_to=_dt(days=5)) + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.post("/ptu_reservation/close", json={"id": "res_1"}) + assert resp.status_code == 400, resp.text + mock_table.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_close_refuses_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": _iso(days=1)}, + ) + assert resp.status_code == 400, resp.text + mock_table.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_close_not_found(client_and_mocks): + client, _, mock_table = client_and_mocks + mock_table.find_unique = AsyncMock(return_value=None) + + resp = client.post("/ptu_reservation/close", json={"id": "missing"}) + assert resp.status_code == 404, resp.text + + +@pytest.mark.asyncio +async def test_create_after_close_of_same_team_model(client_and_mocks): + client, _, mock_table = client_and_mocks + + closed = _row(id="res_old", effective_from=_dt(), effective_to=_dt(days=15)) + mock_table.find_unique = AsyncMock(return_value=closed) + resp = client.post("/ptu_reservation/close", json={"id": "res_old", "effective_to": _iso(days=14)}) + assert resp.status_code == 400, resp.text + + mock_table.find_unique = AsyncMock(return_value=_row(id="res_active", effective_from=_dt(), effective_to=None)) + close_time = _iso(days=15) + resp = client.post("/ptu_reservation/close", json={"id": "res_active", "effective_to": close_time}) + assert resp.status_code == 200 + + mock_table.find_many = AsyncMock(return_value=[]) + resp = client.post( + "/ptu_reservation/new", + json=_valid_new_payload(effective_from=close_time, ptu_count=100), + ) + assert resp.status_code == 200, resp.text + mock_table.create.assert_awaited_once()