| 1 | from helpers.api import ApiHandler, Input, Output, Request |
| 2 | from helpers.task_scheduler import ( |
| 3 | TaskScheduler, ScheduledTask, AdHocTask, PlannedTask, TaskState, |
| 4 | serialize_task, parse_task_schedule, parse_task_plan |
| 5 | ) |
| 6 | from helpers.localization import Localization |
| 7 | |
| 8 | |
| 9 | class SchedulerTaskUpdate(ApiHandler): |
| 10 | async def process(self, input: Input, request: Request) -> Output: |
| 11 | """ |
| 12 | Update an existing task in the scheduler |
| 13 | """ |
| 14 | # Get timezone from input (do not set if not provided, we then rely on poll() to set it) |
| 15 | if timezone := input.get("timezone", None): |
| 16 | Localization.get().set_timezone(timezone) |
| 17 | |
| 18 | scheduler = TaskScheduler.get() |
| 19 | await scheduler.reload() |
| 20 | |
| 21 | # Get task ID from input |
| 22 | task_id: str = input.get("task_id", "") |
| 23 | |
| 24 | if not task_id: |
| 25 | return {"error": "Missing required field: task_id"} |
| 26 | |
| 27 | # Get the task to update |
| 28 | task = scheduler.get_task_by_uuid(task_id) |
| 29 | |
| 30 | if not task: |
| 31 | return {"error": f"Task with ID {task_id} not found"} |
| 32 | |
| 33 | # Update fields if provided using the task's update method |
| 34 | update_params = {} |
| 35 | |
| 36 | if "name" in input: |
| 37 | update_params["name"] = input.get("name", "") |
| 38 | |
| 39 | if "state" in input: |
| 40 | update_params["state"] = TaskState(input.get("state", TaskState.IDLE)) |
| 41 | |
| 42 | if "system_prompt" in input: |
| 43 | update_params["system_prompt"] = input.get("system_prompt", "") |
| 44 | |
| 45 | if "prompt" in input: |
| 46 | update_params["prompt"] = input.get("prompt", "") |
| 47 | |
| 48 | if "attachments" in input: |
| 49 | update_params["attachments"] = input.get("attachments", []) |
| 50 | |
| 51 | if "project_name" in input or "project_color" in input: |
| 52 | return {"error": "Project changes are not allowed"} |
| 53 | |
| 54 | # Update schedule if this is a scheduled task and schedule is provided |
| 55 | if isinstance(task, ScheduledTask) and "schedule" in input: |
| 56 | schedule_data = input.get("schedule", {}) |
| 57 | try: |
| 58 | # Parse the schedule with timezone handling |
| 59 | task_schedule = parse_task_schedule(schedule_data) |
| 60 | |
| 61 | # Set the timezone from the request if not already in schedule_data |
| 62 | if not schedule_data.get('timezone', None) and timezone: |
| 63 | task_schedule.timezone = timezone |
| 64 | |
| 65 | update_params["schedule"] = task_schedule |
| 66 | except ValueError as e: |
| 67 | return {"error": f"Invalid schedule format: {str(e)}"} |
| 68 | elif isinstance(task, AdHocTask) and "token" in input: |
| 69 | token_value = input.get("token", "") |
| 70 | if token_value: # Only update if non-empty |
| 71 | update_params["token"] = token_value |
| 72 | elif isinstance(task, PlannedTask) and "plan" in input: |
| 73 | plan_data = input.get("plan", {}) |
| 74 | try: |
| 75 | # Parse the plan data |
| 76 | task_plan = parse_task_plan(plan_data) |
| 77 | update_params["plan"] = task_plan |
| 78 | except ValueError as e: |
| 79 | return {"error": f"Invalid plan format: {str(e)}"} |
| 80 | |
| 81 | # Use atomic update method to apply changes |
| 82 | updated_task = await scheduler.update_task(task_id, **update_params) |
| 83 | |
| 84 | if not updated_task: |
| 85 | return {"error": f"Task with ID {task_id} not found or could not be updated"} |
| 86 | |
| 87 | # Return the updated task using our standardized serialization function |
| 88 | task_dict = serialize_task(updated_task) |
| 89 | |
| 90 | return { |
| 91 | "ok": True, |
| 92 | "task": task_dict |
| 93 | } |