jon conc

routers/gemini.py

@@ -372,8 +372,8 @@ async def get_job_status(
         response["output"] = job.output_data
         
         # For video jobs, add download URL
-        if job.job_type == "video" and job.output_data and (job.output_data.get("filename") or job.output_data.get("video_url")):
-            response["download_url"] = f"/gemini/download/{job.job_id}"
+        # if job.job_type == "video" and job.output_data and (job.output_data.get("filename") or job.output_data.get("video_url")):
+        #     response["download_url"] = f"/gemini/download/{job.job_id}"
     
     if job.status == "failed":
         response["error"] = job.error_message

services/priority_worker_pool.py

@@ -73,6 +73,7 @@ class WorkerConfig:
     medium_interval: int = 10   # seconds
     slow_interval: int = 15     # seconds
     max_retries: int = 60       # Max retry attempts before failing
+    job_per_api_key: int = 1    # Max concurrent jobs per API key
     
     @classmethod
     def from_env(cls) -> 'WorkerConfig':
@@ -84,6 +85,7 @@ class WorkerConfig:
             fast_interval=int(os.getenv("FAST_INTERVAL", "5")),
             medium_interval=int(os.getenv("MEDIUM_INTERVAL", "30")),
             slow_interval=int(os.getenv("SLOW_INTERVAL", "60")),
+            job_per_api_key=int(os.getenv("JOB_PER_API_KEY", "1")),
         )
 
 
@@ -150,7 +152,8 @@ class PriorityWorker(Generic[JobType]):
         job_model: type,
         job_processor: JobProcessor[JobType],
         max_retries: int = 60,
-        wake_event: Optional[asyncio.Event] = None
+        wake_event: Optional[asyncio.Event] = None,
+        config: Optional[WorkerConfig] = None
     ):
         self.worker_id = worker_id
         self.priority = priority
@@ -162,6 +165,7 @@ class PriorityWorker(Generic[JobType]):
         self._running = False
         self._current_job_id: Optional[str] = None
         self._wake_event = wake_event  # Event to wake worker immediately when new jobs arrive
+        self._config = config or WorkerConfig.from_env()
     
     async def start(self):
         """Start the worker polling loop."""
@@ -210,12 +214,33 @@ class PriorityWorker(Generic[JobType]):
     async def _process_one_job(self) -> bool:
         """Find and process one job.
         
+        Enforces constraints:
+        1. Only one job per user can be in processing state at a time
+        2. Total processing jobs limited to JOB_PER_API_KEY * number of API keys
+        
         Returns:
             True if a job was found and processed, False if no jobs available
         """
         async with self.session_maker() as session:
+            from sqlalchemy import func
+            
             now = datetime.utcnow()
             
+            # Get number of API keys for capacity calculation
+            try:
+                from services.api_key_manager import get_key_count
+                num_api_keys = get_key_count()
+                max_processing = self._config.job_per_api_key * num_api_keys
+            except ImportError:
+                max_processing = 10  # Default fallback
+            
+            # Check if we're at max processing capacity (only for new jobs being picked up)
+            count_query = select(func.count()).where(
+                self.job_model.status == "processing"
+            )
+            count_result = await session.execute(count_query)
+            current_processing = count_result.scalar() or 0
+            
             # Query for jobs matching this priority tier
             query = select(self.job_model).where(
                 and_(
@@ -234,6 +259,27 @@ class PriorityWorker(Generic[JobType]):
             if not job:
                 return False
             
+            # For queued jobs, apply the constraints
+            if job.status == "queued":
+                # Constraint 1: Check if this user already has a job in processing
+                user_processing_query = select(func.count()).where(
+                    and_(
+                        self.job_model.user_id == job.user_id,
+                        self.job_model.status == "processing"
+                    )
+                )
+                user_result = await session.execute(user_processing_query)
+                user_processing_count = user_result.scalar() or 0
+                
+                if user_processing_count > 0:
+                    logger.debug(f"Worker {self.worker_id}: User {job.user_id} already has a job processing, skipping")
+                    return False
+                
+                # Constraint 2: Check if we're at max total processing capacity
+                if current_processing >= max_processing:
+                    logger.debug(f"Worker {self.worker_id}: At max capacity ({current_processing}/{max_processing}), skipping new job")
+                    return False
+            
             self._current_job_id = job.job_id
             
             try:
@@ -400,7 +446,8 @@ class PriorityWorkerPool(Generic[JobType]):
                 job_model=self.job_model,
                 job_processor=self.job_processor,
                 max_retries=self.config.max_retries,
-                wake_event=self._wake_events["fast"]
+                wake_event=self._wake_events["fast"],
+                config=self.config
             )
             self.workers.append(worker)
             await worker.start()
@@ -416,7 +463,8 @@ class PriorityWorkerPool(Generic[JobType]):
                 job_model=self.job_model,
                 job_processor=self.job_processor,
                 max_retries=self.config.max_retries,
-                wake_event=self._wake_events["medium"]
+                wake_event=self._wake_events["medium"],
+                config=self.config
             )
             self.workers.append(worker)
             await worker.start()
@@ -432,7 +480,8 @@ class PriorityWorkerPool(Generic[JobType]):
                 job_model=self.job_model,
                 job_processor=self.job_processor,
                 max_retries=self.config.max_retries,
-                wake_event=self._wake_events["slow"]
+                wake_event=self._wake_events["slow"],
+                config=self.config
             )
             self.workers.append(worker)
             await worker.start()