Update product application logic to include start date

- Modified the `apply_product_for_team` function to accept a `start_date` parameter, allowing for accurate tracking of the last payment date.
- Updated the `handle_stripe_event_background` function to pass the correct start date when applying products for new subscriptions and renewals.
- Adjusted related tests to ensure compatibility with the new function signature and validate the correct application of products with the specified start date.

Author: PhiRho

app/core/resource_limits.py

@@ -1,5 +1,5 @@
 from sqlalchemy.orm import Session
-from sqlalchemy import func, and_, or_
+from sqlalchemy import func, or_
 from app.db.models import DBTeam, DBUser, DBPrivateAIKey, DBTeamProduct, DBProduct
 from fastapi import HTTPException, status
 from typing import Optional

app/core/worker.py

@@ -5,7 +5,16 @@
 import logging
 from collections import defaultdict
 from app.core.resource_limits import get_token_restrictions
-from app.services.stripe import get_product_id_from_session, get_product_id_from_subscription, known_events, subscription_success_events, session_failure_events, subscription_failure_events, invoice_failure_events, invoice_success_events
+from app.services.stripe import (
+    get_product_id_from_session,
+    get_product_id_from_subscription,
+    known_events,
+    subscription_success_events,
+    session_failure_events,
+    subscription_failure_events,
+    invoice_failure_events,
+    invoice_success_events
+)
 
 logger = logging.getLogger(__name__)
 
@@ -28,12 +37,14 @@ async def handle_stripe_event_background(event, db: Session):
         if event_type in subscription_success_events:
             # new subscription
             product_id = await get_product_id_from_subscription(event_object.id)
-            await apply_product_for_team(db, customer_id, product_id)
+            start_date = datetime.fromtimestamp(event_object.start_date, tz=UTC)
+            await apply_product_for_team(db, customer_id, product_id, start_date)
         elif event_type in invoice_success_events:
             # subscription renewed
             subscription = event_object.parent.subscription_details.subscription
             product_id = await get_product_id_from_subscription(subscription)
-            await apply_product_for_team(db, customer_id, product_id)
+            start_date = datetime.fromtimestamp(event_object.period_start, tz=UTC)
+            await apply_product_for_team(db, customer_id, product_id, start_date)
         # Failure Events
         elif event_type in session_failure_events:
             product_id = await get_product_id_from_session(event_object.id)
@@ -50,7 +61,7 @@ async def handle_stripe_event_background(event, db: Session):
         logger.error(f"Error in background event handler: {str(e)}")
 
 
-async def apply_product_for_team(db: Session, customer_id: str, product_id: str):
+async def apply_product_for_team(db: Session, customer_id: str, product_id: str, start_date: datetime):
     """
     Apply a product to a team and update their last payment date.
     Also extends all team keys and sets their max budgets via LiteLLM service.
@@ -77,7 +88,7 @@ async def apply_product_for_team(db: Session, customer_id: str, product_id: str)
             return
 
         # Update the last payment date
-        team.last_payment = datetime.now(UTC)
+        team.last_payment = start_date
 
         # Check if the product is already active for the team
         existing_association = db.query(DBTeamProduct).filter(

app/services/stripe.py

@@ -1,11 +1,9 @@
-from typing import Optional, Union
 import stripe
 import os
 import logging
 from urllib.parse import urljoin
 from fastapi import HTTPException, status
 from sqlalchemy.orm import Session
-from stripe._list_object import ListObject
 from app.db.models import DBTeam, DBSystemSecret
 
 # Configure logger

tests/test_resource_limits.py

@@ -2,7 +2,15 @@
 from app.db.models import DBUser, DBProduct, DBTeamProduct, DBPrivateAIKey
 from datetime import datetime, UTC, timedelta
 from fastapi import HTTPException
-from app.core.resource_limits import check_key_limits, check_team_user_limit, check_vector_db_limits, get_token_restrictions, DEFAULT_KEY_DURATION, DEFAULT_MAX_SPEND, DEFAULT_RPM_PER_KEY
+from app.core.resource_limits import (
+    check_key_limits,
+    check_team_user_limit,
+    check_vector_db_limits,
+    get_token_restrictions,
+    DEFAULT_KEY_DURATION,
+    DEFAULT_MAX_SPEND,
+    DEFAULT_RPM_PER_KEY
+)
 
 def test_add_user_within_product_limit(db, test_team, test_product):
     """Test adding a user when within product user limit"""

tests/test_worker.py

@@ -180,7 +180,7 @@ async def test_apply_product_success(db, test_team, test_product):
     db.commit()
 
     # Apply product to team
-    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id)
+    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id, datetime.now(UTC))
 
     # Refresh team from database
     db.refresh(test_team)
@@ -200,7 +200,7 @@ async def test_apply_product_team_not_found(db, test_product):
     THEN: The operation completes without error
     """
     # Try to apply product to non-existent team
-    await apply_product_for_team(db, "cus_nonexistent", test_product.id)
+    await apply_product_for_team(db, "cus_nonexistent", test_product.id, datetime.now(UTC))
     # No assertions needed as function should complete without error
 
 @pytest.mark.asyncio
@@ -217,7 +217,7 @@ async def test_apply_product_product_not_found(db, test_team):
     db.commit()
 
     # Try to apply non-existent product
-    await apply_product_for_team(db, test_team.stripe_customer_id, "prod_nonexistent")
+    await apply_product_for_team(db, test_team.stripe_customer_id, "prod_nonexistent", datetime.now(UTC))
     # No assertions needed as function should complete without error
 
 @pytest.mark.asyncio
@@ -257,7 +257,7 @@ async def test_apply_product_multiple_products(db, test_team, test_product):
 
     # Apply each product to the team
     for product in products:
-        await apply_product_for_team(db, test_team.stripe_customer_id, product.id)
+        await apply_product_for_team(db, test_team.stripe_customer_id, product.id, datetime.now(UTC))
 
     # Refresh team from database
     db.refresh(test_team)
@@ -283,14 +283,14 @@ async def test_apply_product_already_active(db, test_team, test_product):
     db.refresh(test_team)  # Refresh to ensure we have the latest data
 
     # First apply the product
-    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id)
+    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id, datetime.now(UTC))
 
     # Get the initial last payment date
     db.refresh(test_team)
     initial_last_payment = test_team.last_payment
 
     # Apply the same product again
-    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id)
+    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id, datetime.now(UTC))
 
     # Refresh team from database
     db.refresh(test_team)
@@ -357,7 +357,7 @@ async def test_apply_product_extends_keys_and_sets_budget(mock_litellm, db, test
     mock_instance.set_key_restrictions = AsyncMock()
 
     # Apply product to team
-    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id)
+    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id, datetime.now(UTC))
 
     # Verify LiteLLM service was initialized with correct region settings
     mock_litellm.assert_called_once_with(
@@ -400,7 +400,7 @@ async def test_remove_product_success(db, test_team, test_product):
     db.commit()
 
     # First apply the product to ensure it exists
-    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id)
+    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id, datetime.now(UTC))
 
     # Remove the product
     await remove_product_from_team(db, test_team.stripe_customer_id, test_product.id)
@@ -491,8 +491,8 @@ async def test_remove_product_multiple_products(db, test_team, test_product):
     db.commit()
 
     # Apply both products to the team
-    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id)
-    await apply_product_for_team(db, test_team.stripe_customer_id, second_product.id)
+    await apply_product_for_team(db, test_team.stripe_customer_id, test_product.id, datetime.now(UTC))
+    await apply_product_for_team(db, test_team.stripe_customer_id, second_product.id, datetime.now(UTC))
 
     # Remove only the first product
     await remove_product_from_team(db, test_team.stripe_customer_id, test_product.id)