Research Wizard and CopilotKit mitigation review

backend/services/subscription/limit_validation.py

@@ -390,17 +390,44 @@ class LimitValidator:
             
             logger.info(f"[Pre-flight Check] 📅 Billing Period: {current_period} (for user {user_id})")
             
+            # Ensure schema columns exist before querying
+            try:
+                from services.subscription.schema_utils import ensure_usage_summaries_columns
+                ensure_usage_summaries_columns(self.db)
+            except Exception as schema_err:
+                logger.warning(f"Schema check failed, will retry on query error: {schema_err}")
+            
             # Explicitly expire any cached objects and refresh from DB to ensure fresh data
             self.db.expire_all()
             
-            usage = self.db.query(UsageSummary).filter(
-                UsageSummary.user_id == user_id,
-                UsageSummary.billing_period == current_period
-            ).first()
+            try:
+                usage = self.db.query(UsageSummary).filter(
+                    UsageSummary.user_id == user_id,
+                    UsageSummary.billing_period == current_period
+                ).first()
             
-            # CRITICAL: Explicitly refresh from database to get latest values (clears SQLAlchemy cache)
-            if usage:
-                self.db.refresh(usage)
+                # CRITICAL: Explicitly refresh from database to get latest values (clears SQLAlchemy cache)
+                if usage:
+                    self.db.refresh(usage)
+            except Exception as query_err:
+                error_str = str(query_err).lower()
+                if 'no such column' in error_str and 'exa_calls' in error_str:
+                    logger.warning("Missing column detected in usage query, fixing schema and retrying...")
+                    import sqlite3
+                    import services.subscription.schema_utils as schema_utils
+                    schema_utils._checked_usage_summaries_columns = False
+                    from services.subscription.schema_utils import ensure_usage_summaries_columns
+                    ensure_usage_summaries_columns(self.db)
+                    self.db.expire_all()
+                    # Retry the query
+                    usage = self.db.query(UsageSummary).filter(
+                        UsageSummary.user_id == user_id,
+                        UsageSummary.billing_period == current_period
+                    ).first()
+                    if usage:
+                        self.db.refresh(usage)
+                else:
+                    raise
             
             # Log what we actually read from database
             if usage:
@@ -718,8 +745,40 @@ class LimitValidator:
             
         except Exception as e:
             error_type = type(e).__name__
-            error_message = str(e)
-            logger.error(f"[Pre-flight Check] ❌ Error during comprehensive limit check: {error_type}: {error_message}", exc_info=True)
+            error_message = str(e).lower()
+            
+            # Handle missing column errors with schema fix and retry
+            if 'operationalerror' in error_type.lower() or 'operationalerror' in error_message:
+                if 'no such column' in error_message and 'exa_calls' in error_message:
+                    logger.warning("Missing column detected in limit check, attempting schema fix...")
+                    try:
+                        import sqlite3
+                        import services.subscription.schema_utils as schema_utils
+                        schema_utils._checked_usage_summaries_columns = False
+                        from services.subscription.schema_utils import ensure_usage_summaries_columns
+                        ensure_usage_summaries_columns(self.db)
+                        self.db.expire_all()
+                        
+                        # Retry the query
+                        usage = self.db.query(UsageSummary).filter(
+                            UsageSummary.user_id == user_id,
+                            UsageSummary.billing_period == current_period
+                        ).first()
+                        
+                        if usage:
+                            self.db.refresh(usage)
+                        
+                        # Continue with the rest of the validation using the retried usage
+                        # (The rest of the function logic continues from here)
+                        # For now, we'll let it fall through to return the error since we'd need to duplicate the entire validation logic
+                        # Instead, we'll just log and return, but the next call should succeed
+                        logger.info(f"[Pre-flight Check] Schema fixed, but need to retry validation on next call")
+                        return False, f"Schema updated, please retry: Database schema was updated. Please try again.", {'error_type': 'schema_update', 'retry': True}
+                    except Exception as retry_err:
+                        logger.error(f"Schema fix and retry failed: {retry_err}")
+                        return False, f"Failed to validate limits: {error_type}: {str(e)}", {}
+            
+            logger.error(f"[Pre-flight Check] ❌ Error during comprehensive limit check: {error_type}: {str(e)}", exc_info=True)
             logger.error(f"[Pre-flight Check] ❌ User: {user_id}, Operations count: {len(operations) if operations else 0}")
-            return False, f"Failed to validate limits: {error_type}: {error_message}", {}
+            return False, f"Failed to validate limits: {error_type}: {str(e)}", {}
 

backend/services/subscription/schema_utils.py

@@ -1,8 +1,11 @@
 from typing import Set
 from sqlalchemy.orm import Session
+from sqlalchemy import text
+from loguru import logger
 
 
 _checked_subscription_plan_columns: bool = False
+_checked_usage_summaries_columns: bool = False
 
 
 def ensure_subscription_plan_columns(db: Session) -> None:
@@ -17,9 +20,11 @@ def ensure_subscription_plan_columns(db: Session) -> None:
         return
 
     try:
-        # Discover existing columns
-        result = db.execute("PRAGMA table_info(subscription_plans)")
+        # Discover existing columns using PRAGMA
+        result = db.execute(text("PRAGMA table_info(subscription_plans)"))
         cols: Set[str] = {row[1] for row in result}
+        
+        logger.debug(f"Schema check: Found {len(cols)} columns in subscription_plans table")
 
         # Columns we may reference in models but might be missing in older DBs
         required_columns = {
@@ -28,12 +33,81 @@ def ensure_subscription_plan_columns(db: Session) -> None:
 
         for col_name, ddl in required_columns.items():
             if col_name not in cols:
-                db.execute(f"ALTER TABLE subscription_plans ADD COLUMN {col_name} {ddl}")
-        db.commit()
-    except Exception:
-        # Do not block app if pragma/alter fails; let normal errors surface
-        db.rollback()
-    finally:
+                logger.info(f"Adding missing column {col_name} to subscription_plans table")
+                try:
+                    db.execute(text(f"ALTER TABLE subscription_plans ADD COLUMN {col_name} {ddl}"))
+                    db.commit()
+                    logger.info(f"Successfully added column {col_name}")
+                except Exception as alter_err:
+                    logger.error(f"Failed to add column {col_name}: {alter_err}")
+                    db.rollback()
+                    # Don't set flag on error - allow retry
+                    raise
+            else:
+                logger.debug(f"Column {col_name} already exists")
+        
+        # Only set flag if we successfully completed the check
         _checked_subscription_plan_columns = True
+    except Exception as e:
+        logger.error(f"Error ensuring subscription_plan columns: {e}", exc_info=True)
+        db.rollback()
+        # Don't set the flag if there was an error, so we retry next time
+        _checked_subscription_plan_columns = False
+        raise
+
+
+def ensure_usage_summaries_columns(db: Session) -> None:
+    """Ensure required columns exist on usage_summaries for runtime safety.
+
+    This is a defensive guard for environments where migrations have not yet
+    been applied. If columns are missing (e.g., exa_calls, exa_cost), we add them
+    with a safe default so ORM queries do not fail.
+    """
+    global _checked_usage_summaries_columns
+    if _checked_usage_summaries_columns:
+        return
+
+    try:
+        # Discover existing columns using PRAGMA
+        result = db.execute(text("PRAGMA table_info(usage_summaries)"))
+        cols: Set[str] = {row[1] for row in result}
+        
+        logger.debug(f"Schema check: Found {len(cols)} columns in usage_summaries table")
+
+        # Columns we may reference in models but might be missing in older DBs
+        required_columns = {
+            "exa_calls": "INTEGER DEFAULT 0",
+            "exa_cost": "REAL DEFAULT 0.0",
+        }
+
+        for col_name, ddl in required_columns.items():
+            if col_name not in cols:
+                logger.info(f"Adding missing column {col_name} to usage_summaries table")
+                try:
+                    db.execute(text(f"ALTER TABLE usage_summaries ADD COLUMN {col_name} {ddl}"))
+                    db.commit()
+                    logger.info(f"Successfully added column {col_name}")
+                except Exception as alter_err:
+                    logger.error(f"Failed to add column {col_name}: {alter_err}")
+                    db.rollback()
+                    # Don't set flag on error - allow retry
+                    raise
+            else:
+                logger.debug(f"Column {col_name} already exists")
+        
+        # Only set flag if we successfully completed the check
+        _checked_usage_summaries_columns = True
+    except Exception as e:
+        logger.error(f"Error ensuring usage_summaries columns: {e}", exc_info=True)
+        db.rollback()
+        # Don't set the flag if there was an error, so we retry next time
+        _checked_usage_summaries_columns = False
+        raise
+
+
+def ensure_all_schema_columns(db: Session) -> None:
+    """Ensure all required columns exist in subscription-related tables."""
+    ensure_subscription_plan_columns(db)
+    ensure_usage_summaries_columns(db)