fix: Address all points raised in PR review

Author: BhaveshKukreja29

docker/init-db.sh

@@ -9,7 +9,6 @@ psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-E
     CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
 
     -- You can add other initialization SQL here
-    CREATE DATABASE mxtoaitest;
 
     GRANT ALL PRIVILEGES ON DATABASE $POSTGRES_DB TO $POSTGRES_USER;
 EOSQL

mxgo/api.py

@@ -895,10 +895,12 @@ def _build_newsletter_instructions(request: CreateNewsletterRequest) -> str:
 
 async def _validate_newsletter_limits(user_email: str, cron_expressions: list[str]):
     """Validates the user's plan limits for newsletters."""
+    # Get user plan and corresponding limits from config
     user_plan = await user.get_user_plan(user_email)
     plan_limits = NEWSLETTER_LIMITS_BY_PLAN.get(user_plan, NEWSLETTER_LIMITS_BY_PLAN[UserPlan.BETA])
     min_interval = timedelta(days=plan_limits["min_interval_days"])
 
+    # Check total task count against the plan's max tasks
     db_connection = init_db_connection()
     with db_connection.get_session() as session:
         active_task_count = crud.count_active_tasks_for_user(session, user_email)
@@ -911,7 +913,9 @@ async def _validate_newsletter_limits(user_email: str, cron_expressions: list[st
             f"(max: {plan_limits['max_tasks']}).",
         )
 
+    # Loop through each cron expression to validate its frequency
     for cron_expr in cron_expressions:
+        # One-time tasks don't have a recurring interval, so we skip them
         if not is_one_time_task(cron_expr):
             interval = calculate_cron_interval(cron_expr)
             if interval < min_interval:
@@ -922,14 +926,14 @@ async def _validate_newsletter_limits(user_email: str, cron_expressions: list[st
                 )
 
 
-def _create_and_schedule_task(user_email: str, cron_expr: str, distilled_instructions: str, prompt: str) -> str:
+def _create_and_schedule_task(user_email: str, cron_expr: str, distilled_instructions: str, prompt: str) -> str:  # noqa: ARG001
     """Creates a single newsletter task and schedules it."""
     task_id = str(uuid.uuid4())
     scheduler_job_id = f"task_{task_id}"
     email_for_task = EmailRequest(
         from_email=user_email,
         to="[email protected]",
-        subject=f"Newsletter: {prompt[:50]}...",
+        subject="Generate Newsletter as per following Instructions",
         distilled_processing_instructions=distilled_instructions,
         distilled_alias=HandlerAlias.ASK,
         messageId=f"<newsletter-{task_id}-{datetime.now(timezone.utc).isoformat()}@mxgo.ai>",
@@ -959,15 +963,20 @@ def _create_and_schedule_task(user_email: str, cron_expr: str, distilled_instruc
 
 
 async def _handle_post_creation_action(
-    user_email: str, *, is_whitelisted: bool, first_task_id: str, distilled_instructions: str, prompt: str
+    user_email: str,
+    *,
+    is_whitelisted: bool,
+    first_task_id: str,
+    distilled_instructions: str,
+    prompt: str,  # noqa: ARG001
 ):
     """Sends a sample email if the user is whitelisted, otherwise triggers verification."""
     if is_whitelisted:
         logger.info(f"User {user_email} is whitelisted. Sending sample newsletter.")
         sample_email_request = EmailRequest(
             from_email=user_email,
             to="[email protected]",
-            subject=f"[SAMPLE] Newsletter: {prompt[:40]}...",
+            subject="[SAMPLE] Generate Newsletter as per following Instructions",
             distilled_processing_instructions=distilled_instructions,
             distilled_alias=HandlerAlias.ASK,
             messageId=f"<newsletter-sample-{first_task_id}-{datetime.now(timezone.utc).isoformat()}@mxgo.ai>",
@@ -1001,6 +1010,27 @@ async def create_newsletter(
     user_email = current_user.email
     logger.info(f"Received newsletter creation request for user: {user_email}")
 
+    if validators.redis_client:
+        redis_key = f"newsletter_request:{request.request_id}"
+        existing_task_ids_json = await validators.redis_client.get(redis_key)
+        if existing_task_ids_json:
+            logger.info(
+                f"Duplicate request_id {request.request_id} detected for {user_email}, returning existing tasks from Redis"
+            )
+            existing_task_ids = json.loads(existing_task_ids_json)
+            return Response(
+                content=json.dumps(
+                    {
+                        "message": "This request has already been processed.",
+                        "status": "duplicate",
+                        "request_id": request.request_id,
+                        "scheduled_task_ids": existing_task_ids,
+                    }
+                ),
+                status_code=status.HTTP_409_CONFLICT,
+                media_type="application/json",
+            )
+
     distilled_instructions = _build_newsletter_instructions(request)
 
     try:
@@ -1022,10 +1052,18 @@ async def create_newsletter(
         logger.error(f"Failed to schedule one or more newsletter tasks for {user_email}: {e}")
 
         # Rollback created tasks if any failed
+        scheduler = Scheduler()
         db_connection = init_db_connection()
         with db_connection.get_session() as session:
             for tid in created_task_ids:
                 crud.delete_task(session, tid)
+
+                try:
+                    scheduler.remove_job(f"task_{tid}")
+                    logger.info(f"Removed scheduler job for rolled-back task {tid}")
+                except Exception as scheduler_e:
+                    logger.error(f"Failed to remove scheduler job for task {tid}: {scheduler_e}")
+
         raise HTTPException(status_code=500, detail="Failed to schedule one or more newsletter tasks.") from e
 
     sample_email_sent = False
@@ -1038,6 +1076,11 @@ async def create_newsletter(
             prompt=request.prompt,
         )
 
+    if validators.redis_client and created_task_ids:
+        redis_key = f"newsletter_request:{request.request_id}"
+        # Store for 24 hours
+        await validators.redis_client.setex(redis_key, 86400, json.dumps(created_task_ids))
+
     return CreateNewsletterResponse(
         is_scheduled=bool(created_task_ids),
         is_whitelisted=is_whitelisted,

mxgo/prompts/template_prompts.py

@@ -1177,6 +1177,14 @@
 - **distilled_future_task_instructions**: Clear, detailed instructions about how the task should be processed when executed in the future. This should include the processing approach, any specific requirements, and what the expected outcome should be. **CRITICAL: If the original email contains attachments, you MUST include detailed context about the attachments in these instructions since attachments will not be available during scheduled execution. Include attachment names, types, sizes, and any relevant content or context from the attachments.**
 - **start_time**: (Optional) Start time for the task in ISO 8601 format - task will not execute before this time (e.g., "2024-09-01T00:00:00Z")
 - **end_time**: (Optional) End time for the task in ISO 8601 format - task will not execute after this time (e.g., "2024-12-31T23:59:59Z")
+- **future_handle_alias**: The specific email handle to use for the future task. Select one of the following based on the user's instructions:
+    - **`ask`**: (Default) For general questions, custom tasks, or when no other handle fits.
+    - **`news`**: For fetching the latest news on a specific topic (e.g., "send me the news on AI every Friday").
+    - **`research`**: For conducting in-depth research on a topic at a later time.
+    - **`summarize`**: For summarizing the content of the email at a future date.
+    - **`fact-check`**: For verifying claims in the email content later.
+    - **`simplify`**: For explaining the content of the email in simple terms at a scheduled time.
+    - **`translate`**: For translating the email content at a future time.
 
 **Response (Successful Scheduling):**
 1. Confirmation message with:

mxgo/schemas.py

@@ -503,18 +503,21 @@ class ScheduleType(str, Enum):
 
 
 class RecurringWeekly(BaseModel):
-    days: list[str] = Field(..., description="List of weekdays, e.g., ['monday', 'wednesday', 'friday']")
+    days: list[int] = Field(..., description="List of weekdays as integers (0=Sunday, 1=Monday, ..., 6=Saturday)")
     time: str = Field(..., description="Time in HH:MM format, e.g., '09:30'")
 
     @field_validator("days")
     @classmethod
     def validate_days(cls, v):
-        valid_days = {"monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"}
+        if len(v) > 7:  # noqa: PLR2004
+            msg = "Cannot specify more than 7 days"
+            raise ValueError(msg)
+
         for day in v:
-            if day.lower() not in valid_days:
-                msg = f"Invalid day: {day}. Must be one of {valid_days}"
+            if not 0 <= day <= 6:  # noqa: PLR2004
+                msg = f"Invalid day: {day}. Must be an integer between 0 and 6."
                 raise ValueError(msg)
-        return [day.lower() for day in v]
+        return v
 
     @field_validator("time")
     @classmethod
@@ -527,15 +530,16 @@ def validate_time(cls, v):
 
 class ScheduleOptions(BaseModel):
     type: ScheduleType
-    specific_dates: list[str] | None = Field(
-        None, description="List of ISO 8601 datetime strings for specific, non-recurring schedules."
+    specific_datetime: str | None = Field(
+        None, description="List of ISO 8601 datetime string for a specific, non-recurring schedules."
     )
-    recurring_weekly: RecurringWeekly | None = Field(None, description="Configuration for a recurring weekly schedule.")
+    weekly_schedule: RecurringWeekly | None = Field(None, description="Configuration for a recurring weekly schedule.")
 
 
 class CreateNewsletterRequest(BaseModel):
     """Request model for creating a newsletter based on the new UI mock."""
 
+    request_id: str = Field(..., description="Unique request identifier for idempotency")
     prompt: str = Field(..., description="The main instructions for the newsletter content.")
     estimated_read_time: int | None = Field(None, description="Estimated read time in minutes for the newsletter.")
     sources: list[str] | None = Field(None, description="A list of source names or URLs to use.")

mxgo/tools/scheduled_tasks_tool.py

@@ -260,6 +260,9 @@ def forward(
                     logger.info(f"Resolved future handle alias '{future_handle_alias}' to '{resolved_handle_alias}'")
                 except Exception:
                     logger.warning(f"Could not resolve handle alias '{future_handle_alias}'. Validation might fail.")
+                    resolved_handle_alias = "ask"
+            else:
+                resolved_handle_alias = "ask"
 
             # Validate input using Pydantic
             input_data = ScheduledTaskInput(

mxgo/utils.py

@@ -1,5 +1,7 @@
 from datetime import datetime, timedelta, timezone
 
+from croniter import croniter
+
 from mxgo.schemas import ScheduleOptions, ScheduleType
 
 
@@ -107,11 +109,28 @@ def calculate_cron_interval(cron_expression: str) -> timedelta:  # noqa: PLR0912
         # Daily pattern (specific hour and minute, every day)
         elif day == "*" and month == "*" and (weekday in {"*", "?"}):
             interval = timedelta(days=1)
-        # Weekly pattern (specific weekday)
+        # Weekly pattern with multiple days
         elif day == "*" and month == "*" and weekday not in {"*", "?"}:
-            # If multiple days are specified (e.g., "1-5" or "1,3,5"), the minimum interval is 1 day.
-            # Otherwise, it's a single day of the week, so the interval is 7 days.
-            interval = timedelta(days=1) if "," in weekday or "-" in weekday else timedelta(weeks=1)
+            # This handles cases like "1,3,5" (Mon, Wed, Fri)
+            if "," in weekday:
+                days_of_week = sorted([int(d) for d in weekday.split(",")])
+                if len(days_of_week) <= 1:
+                    return timedelta(weeks=1)
+
+                min_diff = 7
+                # Calculate interval between consecutive days
+                for i in range(1, len(days_of_week)):
+                    diff = days_of_week[i] - days_of_week[i - 1]
+                    min_diff = min(min_diff, diff)
+
+                # Calculate wrap-around interval (e.g., from Friday to Monday)
+                wrap_around_diff = (days_of_week[0] + 7) - days_of_week[-1]
+                min_diff = min(min_diff, wrap_around_diff)
+
+                return timedelta(days=min_diff)
+
+            # Single day of the week
+            return timedelta(weeks=1)
 
         # Monthly pattern (specific day of month)
         elif day != "*" and month == "*":
@@ -131,34 +150,67 @@ def calculate_cron_interval(cron_expression: str) -> timedelta:  # noqa: PLR0912
 
 
 def convert_schedule_to_cron_list(schedule: ScheduleOptions) -> list[str]:
-    """Converts schedule options from the newsletter request into a list of cron expressions."""
+    """
+    Converts schedule options from the newsletter request into a list of valid cron expressions.
+
+    Raises:
+        ValueError: If the schedule configuration is invalid or results in an invalid cron expression.
+
+    """
+    cron_expressions = []
+
     if schedule.type == ScheduleType.IMMEDIATE:
-        # Schedule for 1 minute in the future to be executed ASAP
+        # Schedule for 1 minute in the future to be executed as soon as possible.
         now = datetime.now(timezone.utc) + timedelta(minutes=1)
-        return [f"{now.minute} {now.hour} {now.day} {now.month} *"]
+        cron_str = f"{now.minute} {now.hour} {now.day} {now.month} *"
+        cron_expressions.append(cron_str)
 
-    if schedule.type == ScheduleType.SPECIFIC_DATES:
-        cron_list = []
-        if not schedule.specific_dates:
-            msg = "specific_dates must be provided for SPECIFIC_DATES schedule type."
+    elif schedule.type == ScheduleType.SPECIFIC_DATES:
+        if not schedule.specific_datetime:
+            msg = "specific_datetime must be provided for SPECIFIC_DATES schedule type."
             raise ValueError(msg)
-        for dt_str in schedule.specific_dates:
-            dt = datetime.fromisoformat(dt_str)
-            if dt.tzinfo is None:
-                dt = dt.replace(tzinfo=timezone.utc)  # Assume UTC if naive
-            cron_list.append(f"{dt.minute} {dt.hour} {dt.day} {dt.month} *")
-        return cron_list
-
-    if schedule.type == ScheduleType.RECURRING_WEEKLY:
-        if not schedule.recurring_weekly or not schedule.recurring_weekly.days:
-            msg = "recurring_weekly with at least one day must be provided for RECURRING_WEEKLY schedule type."
+
+        dt = datetime.fromisoformat(schedule.specific_datetime)
+        if dt.tzinfo is None:
+            dt = dt.replace(tzinfo=timezone.utc)  # Assume UTC if timezone is not specified
+
+        # Ensure the scheduled time is not in the past
+        if dt <= datetime.now(timezone.utc):
+            msg = "Specific datetime for a one-time schedule must be in the future."
+            raise ValueError(msg)
+
+        cron_str = f"{dt.minute} {dt.hour} {dt.day} {dt.month} *"
+        cron_expressions.append(cron_str)
+
+    elif schedule.type == ScheduleType.RECURRING_WEEKLY:
+        if not schedule.weekly_schedule or not schedule.weekly_schedule.days:
+            msg = "weekly_schedule with at least one day must be provided for RECURRING_WEEKLY type."
             raise ValueError(msg)
 
-        day_map = {"monday": 1, "tuesday": 2, "wednesday": 3, "thursday": 4, "friday": 5, "saturday": 6, "sunday": 0}
-        days_of_week = ",".join(str(day_map[day]) for day in schedule.recurring_weekly.days)
-        hour, minute = schedule.recurring_weekly.time.split(":")
+        # Days are now expected to be integers (0=Sunday, 1=Monday, ..., 6=Saturday)
+        days_of_week = ",".join(map(str, sorted(schedule.weekly_schedule.days)))
 
-        return [f"{minute} {hour} * * {days_of_week}"]
+        try:
+            hour, minute = schedule.weekly_schedule.time.split(":")
+            # Basic validation for time format
+            if not (0 <= int(hour) <= 23 and 0 <= int(minute) <= 59):  # noqa: PLR2004
+                msg = "Invalid time format."
+                raise ValueError(msg)
+        except ValueError as e:
+            msg = f"Invalid time format in weekly_schedule: {schedule.weekly_schedule.time}"
+            raise ValueError(msg) from e
+
+        cron_str = f"{minute} {hour} * * {days_of_week}"
+        cron_expressions.append(cron_str)
+
+    else:
+        msg = f"Unsupported schedule type: {schedule.type}"
+        raise ValueError(msg)
+
+    # Validate all generated cron expressions before returning
+    for cron_expr in cron_expressions:
+        if not croniter.is_valid(cron_expr):
+            msg = f"Generated an invalid cron expression: '{cron_expr}'"
+            raise ValueError(msg)
 
-    msg = f"Unsupported schedule type: {schedule.type}"
-    raise ValueError(msg)
+    return cron_expressions