credit system

core/models.py

@@ -125,6 +125,10 @@ class GeminiJob(Base):
     created_at = Column(DateTime(timezone=True), server_default=func.now())
     started_at = Column(DateTime(timezone=True), nullable=True)
     completed_at = Column(DateTime(timezone=True), nullable=True)
+    
+    # Credit tracking for reservation pattern
+    credits_reserved = Column(Integer, default=0)  # Credits reserved for this job
+    credits_refunded = Column(Boolean, default=False)  # Whether credits were refunded
 
     def __repr__(self):
         return f"<GeminiJob(job_id={self.job_id}, type={self.job_type}, status={self.status}, priority={self.priority})>"

routers/gemini.py

@@ -86,7 +86,8 @@ async def create_job(
         job_type=job_type,
         status="queued",
         priority=priority,
-        input_data=input_data
+        input_data=input_data,
+        credits_reserved=1  # Track that 1 credit is reserved for this job
     )
     db.add(job)
     await db.commit()

services/credit_service.py

@@ -0,0 +1,257 @@
+"""
+Credit Service - Manages credit reservation, confirmation, and refunding.
+
+Implements the Credit Reservation Pattern:
+1. Reserve credits when job is created (deduct from user, track in job)
+2. Confirm credits only on successful completion
+3. Refund credits on refundable errors (server-side issues)
+4. Keep credits on non-refundable errors (user-caused issues)
+"""
+import logging
+from typing import Optional
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Error Categories for Refund Decisions
+# =============================================================================
+
+# Refundable errors - User gets credits back (server/API issues)
+REFUNDABLE_ERROR_PATTERNS = [
+    "API_KEY_INVALID",
+    "QUOTA_EXCEEDED",
+    "INTERNAL_ERROR", 
+    "CONNECTION_FAILED",
+    "SERVER_SHUTDOWN",
+    "TIMEOUT",
+    "Server Authentication Error",
+    "Network error",
+    "Connection refused",
+    "Connection reset",
+    "Service unavailable",
+    "503",
+    "500",
+    "429",  # Rate limit (our quota, not user's fault)
+]
+
+# Non-refundable error patterns - User's input/content issue
+NON_REFUNDABLE_ERROR_PATTERNS = [
+    "safety",
+    "blocked",
+    "SAFETY_FILTER",
+    "INVALID_INPUT",
+    "Invalid image",
+    "Bad request",
+    "400",
+    "cancelled",
+    "User cancelled",
+]
+
+
+def is_refundable_error(error_message: Optional[str]) -> bool:
+    """
+    Determine if an error should result in a credit refund.
+    
+    Args:
+        error_message: The error message from the failed job
+        
+    Returns:
+        True if the error is refundable (server/API issue)
+        False if non-refundable (user's fault) or no error message
+    """
+    if not error_message:
+        return False
+    
+    error_lower = error_message.lower()
+    
+    # Check for REFUNDABLE patterns FIRST (specific server errors take precedence)
+    # This ensures API_KEY_INVALID is caught before generic "400" matcher
+    for pattern in REFUNDABLE_ERROR_PATTERNS:
+        if pattern.lower() in error_lower:
+            logger.debug(f"Error matched refundable pattern '{pattern}': {error_message[:100]}")
+            return True
+    
+    # Check for non-refundable patterns (user-caused issues)
+    for pattern in NON_REFUNDABLE_ERROR_PATTERNS:
+        if pattern.lower() in error_lower:
+            logger.debug(f"Error matched non-refundable pattern '{pattern}': {error_message[:100]}")
+            return False
+    
+    # Default: Max retries exceeded is refundable (we consumed API resources trying)
+    if "max retries" in error_lower:
+        return True
+    
+    # Default: Unknown errors are NOT refundable to prevent abuse
+    # If it's an unknown error, it's more likely user-caused
+    logger.debug(f"Unknown error (not refundable): {error_message[:100]}")
+    return False
+
+
+async def reserve_credit(session: AsyncSession, user, amount: int = 1) -> bool:
+    """
+    Reserve credits for a job (deduct from user's balance).
+    
+    The credits are deducted but tracked in the job's credits_reserved field.
+    If the job fails with a refundable error, they can be restored.
+    
+    Args:
+        session: Database session
+        user: User model instance
+        amount: Number of credits to reserve (default: 1)
+        
+    Returns:
+        True if credits were successfully reserved
+        False if user has insufficient credits
+    """
+    if user.credits < amount:
+        logger.warning(f"User {user.user_id} has insufficient credits ({user.credits}) to reserve {amount}")
+        return False
+    
+    user.credits -= amount
+    logger.info(f"Reserved {amount} credit(s) for user {user.user_id}. Remaining: {user.credits}")
+    # Note: Don't commit here - let caller handle transaction
+    return True
+
+
+async def confirm_credit(session: AsyncSession, job) -> None:
+    """
+    Confirm that credits were legitimately used for a completed job.
+    
+    This is called when a job completes successfully. The credits stay
+    deducted (they were already deducted during reservation).
+    
+    Args:
+        session: Database session
+        job: GeminiJob model instance
+    """
+    if job.credits_reserved > 0:
+        # Credits were used - clear the reservation tracking
+        credits_used = job.credits_reserved
+        job.credits_reserved = 0
+        logger.info(f"Confirmed {credits_used} credit(s) used for job {job.job_id}")
+    # Note: Don't commit here - let caller handle transaction
+
+
+async def refund_credit(session: AsyncSession, job, reason: str) -> bool:
+    """
+    Refund reserved credits back to the user.
+    
+    Called when a job fails due to a refundable error (server-side issue).
+    
+    Args:
+        session: Database session  
+        job: GeminiJob model instance
+        reason: Reason for the refund (for logging)
+        
+    Returns:
+        True if credits were refunded
+        False if no credits to refund or already refunded
+    """
+    if job.credits_reserved <= 0:
+        logger.debug(f"Job {job.job_id} has no credits to refund")
+        return False
+    
+    if job.credits_refunded:
+        logger.warning(f"Job {job.job_id} was already refunded")
+        return False
+    
+    # Get the user to restore credits
+    from core.models import User
+    
+    result = await session.execute(
+        select(User).where(User.user_id == job.user_id)
+    )
+    user = result.scalar_one_or_none()
+    
+    if not user:
+        logger.error(f"Cannot refund job {job.job_id}: User {job.user_id} not found")
+        return False
+    
+    # Restore credits
+    credits_to_refund = job.credits_reserved
+    user.credits += credits_to_refund
+    job.credits_reserved = 0
+    job.credits_refunded = True
+    
+    logger.info(
+        f"Refunded {credits_to_refund} credit(s) to user {user.user_id} for job {job.job_id}. "
+        f"Reason: {reason[:100]}. New balance: {user.credits}"
+    )
+    
+    # Note: Don't commit here - let caller handle transaction
+    return True
+
+
+async def handle_job_completion(session: AsyncSession, job) -> None:
+    """
+    Handle credit finalization when a job completes or fails.
+    
+    This is the main entry point called by the job worker.
+    
+    Args:
+        session: Database session
+        job: GeminiJob model instance with final status
+    """
+    if job.status == "completed":
+        # Success - confirm credits were used
+        await confirm_credit(session, job)
+        
+    elif job.status == "failed":
+        # Failure - check if refundable
+        if is_refundable_error(job.error_message):
+            await refund_credit(session, job, job.error_message or "Unknown error")
+        else:
+            # Non-refundable - confirm credits were used (user's fault)
+            await confirm_credit(session, job)
+            logger.info(f"Job {job.job_id} failed with non-refundable error, credits kept")
+            
+    elif job.status == "cancelled":
+        # Cancelled jobs get refunds only if they were never started
+        if job.started_at is None:
+            await refund_credit(session, job, "Job cancelled before processing")
+        else:
+            # Was processing - keep credits (API may have been consumed)
+            await confirm_credit(session, job)
+            logger.info(f"Job {job.job_id} cancelled during processing, credits kept")
+
+
+async def refund_orphaned_jobs(session: AsyncSession) -> int:
+    """
+    Refund credits for jobs that were abandoned due to server shutdown.
+    
+    Called during graceful shutdown to ensure no credits are lost.
+    
+    Args:
+        session: Database session
+        
+    Returns:
+        Number of jobs that were refunded
+    """
+    from core.models import GeminiJob
+    
+    # Find jobs that are still processing with reserved credits
+    result = await session.execute(
+        select(GeminiJob).where(
+            GeminiJob.status == "processing",
+            GeminiJob.credits_reserved > 0,
+            GeminiJob.credits_refunded == False
+        )
+    )
+    orphaned_jobs = result.scalars().all()
+    
+    refund_count = 0
+    for job in orphaned_jobs:
+        if await refund_credit(session, job, "SERVER_SHUTDOWN: Job orphaned during server shutdown"):
+            # Mark job as failed
+            job.status = "failed"
+            job.error_message = "Server shutdown during processing. Credits refunded."
+            refund_count += 1
+    
+    if refund_count > 0:
+        await session.commit()
+        logger.info(f"Refunded {refund_count} orphaned job(s) during shutdown")
+    
+    return refund_count

services/priority_worker_pool.py

@@ -319,7 +319,26 @@ class PriorityWorker(Generic[JobType]):
             job.error_message = f"Max retries ({self.max_retries}) exceeded"
             job.completed_at = datetime.utcnow()
         
+        # Handle credit finalization for jobs with reserved credits
+        if job.status in ("completed", "failed", "cancelled"):
+            await self._handle_job_credits(session, job)
+        
         await session.commit()
+    
+    async def _handle_job_credits(self, session: AsyncSession, job: JobType):
+        """Handle credit finalization when job reaches terminal state."""
+        # Check if job has credits_reserved attribute (credit-enabled jobs)
+        if not hasattr(job, 'credits_reserved') or job.credits_reserved <= 0:
+            return
+        
+        try:
+            from services.credit_service import handle_job_completion
+            await handle_job_completion(session, job)
+        except ImportError:
+            # Credit service not available - skip
+            logger.debug(f"Credit service not available for job {job.job_id}")
+        except Exception as e:
+            logger.error(f"Error handling credits for job {job.job_id}: {e}")
 
 
 class PriorityWorkerPool(Generic[JobType]):
@@ -438,11 +457,28 @@ class PriorityWorkerPool(Generic[JobType]):
             logger.debug(f"Notified {priority} workers of new job")
     
     async def stop(self):
-        """Stop all workers."""
+        """Stop all workers and refund orphaned jobs."""
         self._running = False
+        
+        # Refund credits for any jobs that were processing when server stopped
+        await self._refund_orphaned_jobs()
+        
         for worker in self.workers:
             await worker.stop()
         logger.info("PriorityWorkerPool stopped")
+    
+    async def _refund_orphaned_jobs(self):
+        """Refund credits for jobs abandoned during shutdown."""
+        try:
+            from services.credit_service import refund_orphaned_jobs
+            async with self.session_maker() as session:
+                refund_count = await refund_orphaned_jobs(session)
+                if refund_count > 0:
+                    logger.info(f"Shutdown: Refunded {refund_count} orphaned job(s)")
+        except ImportError:
+            logger.debug("Credit service not available for orphaned job refunds")
+        except Exception as e:
+            logger.error(f"Error refunding orphaned jobs during shutdown: {e}")
 
 
 # Convenience functions for priority mapping