Referral Program Reward: 5 Referrals = 6 Months Free

alembic/versions/33ae457b2ddf_add_referral_columns.py

@@ -5,6 +5,7 @@
 Create Date: 2025-12-26 10:37:46.325765
 
 """
+
 from typing import Sequence, Union
 
 from alembic import op
@@ -13,26 +14,28 @@
 from sqlalchemy.ext.declarative import declarative_base
 
 # revision identifiers, used by Alembic.
-revision: str = '33ae457b2ddf'
-down_revision: Union[str, Sequence[str], None] = '8b9c2e1f4c1c'
+revision: str = "33ae457b2ddf"
+down_revision: Union[str, Sequence[str], None] = "8b9c2e1f4c1c"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
 # Define a minimal model for data migration
 Base = declarative_base()
 
+
 class Profile(Base):
-    __tablename__ = 'profiles'
+    __tablename__ = "profiles"
     user_id = sa.Column(sa.UUID, primary_key=True)
     referral_code = sa.Column(sa.String)
     referral_count = sa.Column(sa.Integer)
 
+
 def upgrade() -> None:
     """Upgrade schema."""
     # 1. Add columns as nullable first
-    op.add_column('profiles', sa.Column('referral_code', sa.String(), nullable=True))
-    op.add_column('profiles', sa.Column('referrer_id', sa.UUID(), nullable=True))
-    op.add_column('profiles', sa.Column('referral_count', sa.Integer(), nullable=True))
+    op.add_column("profiles", sa.Column("referral_code", sa.String(), nullable=True))
+    op.add_column("profiles", sa.Column("referrer_id", sa.UUID(), nullable=True))
+    op.add_column("profiles", sa.Column("referral_count", sa.Integer(), nullable=True))
 
     # 2. Backfill existing rows with 0 count
     bind = op.get_bind()
@@ -45,10 +48,12 @@ def upgrade() -> None:
     # 3. Alter columns
     # referral_code stays nullable=True
     # referral_count becomes nullable=False
-    op.alter_column('profiles', 'referral_count', nullable=False)
+    op.alter_column("profiles", "referral_count", nullable=False)
 
     # 4. Create unique constraint and index
-    op.create_unique_constraint("uq_profiles_referral_code", "profiles", ["referral_code"])
+    op.create_unique_constraint(
+        "uq_profiles_referral_code", "profiles", ["referral_code"]
+    )
     op.create_index("ix_profiles_referral_code", "profiles", ["referral_code"])
 
     # Add foreign key for referrer_id
@@ -62,6 +67,6 @@ def downgrade() -> None:
     op.drop_constraint("fk_profiles_referrer_id", "profiles", type_="foreignkey")
     op.drop_index("ix_profiles_referral_code", table_name="profiles")
     op.drop_constraint("uq_profiles_referral_code", "profiles", type_="unique")
-    op.drop_column('profiles', 'referral_count')
-    op.drop_column('profiles', 'referrer_id')
-    op.drop_column('profiles', 'referral_code')
+    op.drop_column("profiles", "referral_count")
+    op.drop_column("profiles", "referrer_id")
+    op.drop_column("profiles", "referral_code")

common/config_models.py

@@ -128,13 +128,21 @@ class PaymentRetryConfig(BaseModel):
     max_attempts: int
 
 
+class ReferralConfig(BaseModel):
+    """Referral program configuration."""
+
+    referrals_required: int
+    reward_months: int
+
+
 class SubscriptionConfig(BaseModel):
     """Subscription configuration."""
 
     stripe: SubscriptionStripeConfig
     metered: MeteredConfig
     trial_period_days: int
     payment_retry: PaymentRetryConfig
+    referral: ReferralConfig
 
 
 class StripeWebhookConfig(BaseModel):
@@ -161,9 +169,3 @@ class TelegramConfig(BaseModel):
     """Telegram configuration."""
 
     chat_ids: TelegramChatIdsConfig
-
-
-class ServerConfig(BaseModel):
-    """Server configuration."""
-
-    allowed_origins: list[str]

common/global_config.py

@@ -25,7 +25,6 @@
     SubscriptionConfig,
     StripeConfig,
     TelegramConfig,
-    ServerConfig,
 )
 from common.db_uri_resolver import resolve_db_uri
 
@@ -152,7 +151,6 @@ class Config(BaseSettings):
     subscription: SubscriptionConfig
     stripe: StripeConfig
     telegram: TelegramConfig
-    server: ServerConfig
 
     # Environment variables (required)
     DEV_ENV: str

common/global_config.yaml

@@ -92,6 +92,10 @@ subscription:
   trial_period_days: 7
   payment_retry:
     max_attempts: 3
+  # Referral program configuration
+  referral:
+    referrals_required: 5
+    reward_months: 6
 
 ########################################################
 # Stripe
@@ -108,10 +112,3 @@ telegram:
   chat_ids:
     admin_alerts: "1560836485"
     test: "1560836485"
-
-########################################################
-# Server
-########################################################
-server:
-  allowed_origins:
-    - "http://localhost:8080"

docs/REFERRALS.md

@@ -0,0 +1,33 @@
+# Referral Program
+
+The application includes a referral system that rewards users for inviting others to the platform.
+
+## Incentive: Refer 5, Get 6 Months Free (Default)
+
+When a user successfully refers a specific number of new users (default: **5**), they are automatically rewarded with a period of the Plus Tier subscription for free (default: **6 months**).
+
+### Configuration
+
+The referral program parameters are configurable in `common/global_config.yaml` under the `subscription.referral` section:
+
+```yaml
+subscription:
+  referral:
+    referrals_required: 5
+    reward_months: 6
+```
+
+### How it works
+
+1.  **Referral Code**: Each user has a unique referral code.
+2.  **Invitation**: Users share their code with potential new users.
+3.  **Redemption**: When a new user signs up (or enters the code in their settings), they apply the referral code.
+4.  **Tracking**: The system tracks the number of successful referrals for each referrer.
+5.  **Reward Trigger**:
+    *   Once the referrer's count reaches the configured `referrals_required`, the system automatically grants the reward.
+    *   **New Subscription**: If the referrer is on the Free tier, they are upgraded to the Plus tier for `reward_months`.
+    *   **Existing Subscription**: If the referrer already has a Plus tier subscription, their subscription end date is extended by `reward_months`.
+
+### Technical Implementation
+
+The logic is handled in `src/api/services/referral_service.py` within the `apply_referral` method. When the referral count increments to the configured threshold, the `grant_referral_reward` method is called to update the `UserSubscriptions` table.

src/api/auth/workos_auth.py

@@ -135,18 +135,8 @@ async def get_current_workos_user(request: Request) -> WorkOSUser:
         token = auth_header.split(" ", 1)[1]
 
         # Check if we're in test mode (skip signature verification for tests)
-        # Detect test mode by checking if pytest is running or if DEV_ENV is explicitly set to "test"
-        # We also check for 'test' in sys.argv[0] ONLY if we are NOT in production, to avoid security risks
-        # where a script named "test_something.py" could bypass auth in prod.
-        is_pytest = "pytest" in sys.modules
-        is_dev_env_test = global_config.DEV_ENV.lower() == "test"
-
-        # Only check sys.argv if we are definitely not in prod
-        is_script_test = False
-        if global_config.DEV_ENV.lower() != "prod":
-            is_script_test = "test" in sys.argv[0].lower()
-
-        is_test_mode = is_pytest or is_dev_env_test or is_script_test
+        # Detect test mode by checking if pytest is running
+        is_test_mode = "pytest" in sys.modules or "test" in sys.argv[0].lower()
 
         # Determine whether the token declares an audience so we can decide
         # whether to enforce audience verification (access tokens currently omit aud).

src/api/routes/agent/agent.py

@@ -69,17 +69,6 @@ class ConversationPayload(BaseModel):
     conversation: list[ConversationMessage]
 
 
-class AgentLimitResponse(BaseModel):
-    """Response model for agent limit status."""
-
-    tier: str
-    limit_name: str
-    limit_value: int
-    used_today: int
-    remaining: int
-    reset_at: datetime
-
-
 class AgentResponse(BaseModel):
     """Response model for agent endpoint."""
 
@@ -314,33 +303,6 @@ def record_agent_message(
     return message
 
 
-@router.get("/agent/limits", response_model=AgentLimitResponse)
-async def get_agent_limits(
-    request: Request,
-    db: Session = Depends(get_db_session),
-) -> AgentLimitResponse:
-    """
-    Get the current user's agent limit status.
-
-    Returns usage statistics for the daily agent chat limit, including
-    current tier, usage count, remaining quota, and reset time.
-    """
-    auth_user = await get_authenticated_user(request, db)
-    user_id = auth_user.id
-    user_uuid = user_uuid_from_str(user_id)
-
-    limit_status = ensure_daily_limit(db=db, user_uuid=user_uuid, enforce=False)
-
-    return AgentLimitResponse(
-        tier=limit_status.tier,
-        limit_name=limit_status.limit_name,
-        limit_value=limit_status.limit_value,
-        used_today=limit_status.used_today,
-        remaining=limit_status.remaining,
-        reset_at=limit_status.reset_at,
-    )
-
-
 @router.post("/agent", response_model=AgentResponse)  # noqa
 @observe()
 async def agent_endpoint(

src/api/routes/agent/tools/alert_admin.py

@@ -85,17 +85,8 @@ def alert_admin(
         telegram = Telegram()
         # Use test chat during testing to avoid spamming production alerts
         import sys
-        from common import global_config
 
-        is_pytest = "pytest" in sys.modules
-        is_dev_env_test = global_config.DEV_ENV.lower() == "test"
-
-        # Only check sys.argv if we are definitely not in prod
-        is_script_test = False
-        if global_config.DEV_ENV.lower() != "prod":
-            is_script_test = "test" in sys.argv[0].lower()
-
-        is_testing = is_pytest or is_dev_env_test or is_script_test
+        is_testing = "pytest" in sys.modules or "test" in sys.argv[0].lower()
         chat_name = "test" if is_testing else "admin_alerts"
 
         message_id = telegram.send_message_to_chat(

src/api/services/referral_service.py

@@ -2,6 +2,13 @@
 from sqlalchemy.exc import IntegrityError
 from src.db.models.public.profiles import Profiles, generate_referral_code
 from src.db.utils.db_transaction import db_transaction
+from src.db.models.stripe.user_subscriptions import UserSubscriptions
+from src.db.models.stripe.subscription_types import SubscriptionTier
+from common.global_config import global_config
+from datetime import datetime, timedelta, timezone
+from loguru import logger
+from typing import cast
+import uuid
 
 
 class ReferralService:
@@ -18,6 +25,52 @@ def validate_referral_code(
             db.query(Profiles).filter(Profiles.referral_code == referral_code).first()
         )
 
+    @staticmethod
+    def grant_referral_reward(db: Session, user_id: uuid.UUID):
+        """
+        Grant Plus Tier to the user based on configured reward duration.
+        """
+        now = datetime.now(timezone.utc)
+        reward_months = global_config.subscription.referral.reward_months
+        reward_duration = timedelta(days=30 * reward_months)
+
+        subscription = (
+            db.query(UserSubscriptions)
+            .filter(UserSubscriptions.user_id == user_id)
+            .first()
+        )
+
+        if subscription:
+            subscription.subscription_tier = SubscriptionTier.PLUS.value
+            subscription.is_active = True
+
+            # If current subscription is valid and ends in the future, extend it
+            # Otherwise start from now
+            current_end = subscription.subscription_end_date
+            if current_end and current_end.tzinfo is None:
+                current_end = current_end.replace(tzinfo=timezone.utc)
+
+            if current_end and current_end > now:
+                subscription.subscription_end_date = current_end + reward_duration
+            else:
+                subscription.subscription_end_date = now + reward_duration
+
+            logger.info(
+                f"Updated subscription for user {user_id} via referral reward ({reward_months} months)"
+            )
+        else:
+            new_subscription = UserSubscriptions(
+                user_id=user_id,
+                subscription_tier=SubscriptionTier.PLUS.value,
+                is_active=True,
+                subscription_start_date=now,
+                subscription_end_date=now + reward_duration,
+            )
+            db.add(new_subscription)
+            logger.info(
+                f"Created subscription for user {user_id} via referral reward ({reward_months} months)"
+            )
+
     @staticmethod
     def apply_referral(db: Session, user_profile: Profiles, referral_code: str) -> bool:
         """
@@ -46,6 +99,15 @@ def apply_referral(db: Session, user_profile: Profiles, referral_code: str) -> b
 
             db.add(user_profile)
 
+            # Refresh referrer to get updated count and trigger reward if applicable
+            db.refresh(referrer)
+
+            required_referrals = global_config.subscription.referral.referrals_required
+            if referrer.referral_count == required_referrals:
+                # Cast user_id to uuid.UUID to satisfy ty
+                user_id = cast(uuid.UUID, referrer.user_id)
+                ReferralService.grant_referral_reward(db, user_id)
+
         db.refresh(user_profile)
         return True
 

src/db/utils/users.py

@@ -4,13 +4,14 @@
 import uuid
 from loguru import logger
 
+
 def ensure_profile_exists(
     db: Session,
     user_uuid: uuid.UUID,
     email: str | None = None,
     username: str | None = None,
     avatar_url: str | None = None,
-    is_approved: bool = False
+    is_approved: bool = False,
 ) -> Profiles:
     """
     Ensure a profile exists for the given user UUID.
@@ -27,7 +28,7 @@ def ensure_profile_exists(
                 email=email,
                 username=username,
                 avatar_url=avatar_url,
-                is_approved=is_approved
+                is_approved=is_approved,
             )
             db.add(profile)
         # No need for explicit commit/refresh as db_transaction handles commit,

src/server.py

@@ -16,7 +16,9 @@
 # Add CORS middleware with specific allowed origins
 app.add_middleware(  # type: ignore[call-overload]
     CORSMiddleware,  # type: ignore[arg-type]
-    allow_origins=global_config.server.allowed_origins,
+    allow_origins=[
+        "http://localhost:8080",
+    ],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -52,5 +54,5 @@ def include_all_routers():
         host="0.0.0.0",
         port=int(os.getenv("PORT", 8080)),
         log_config=None,  # Disable uvicorn's logging config
-        access_log=True,  # Enable access logs
+        access_log=False,  # Disable access logs
     )