422 lines
14 KiB
Python
422 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Scheduled Task Manager CLI Tool
|
|
Add, delete, modify, and query user scheduled tasks. Data is stored in tasks.yaml.
|
|
|
|
Environment variables:
|
|
ASSISTANT_ID: Current bot ID
|
|
USER_IDENTIFIER: Current user identifier
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import yaml
|
|
import json
|
|
import string
|
|
import random
|
|
from datetime import datetime, timezone, timedelta
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from croniter import croniter
|
|
except ImportError:
|
|
print("Error: croniter is required. Install with: pip install croniter", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
# Language to timezone mapping
|
|
LANGUAGE_TIMEZONE_MAP = {
|
|
'zh': 'Asia/Shanghai',
|
|
'ja': 'Asia/Tokyo',
|
|
'jp': 'Asia/Tokyo',
|
|
'en': 'UTC',
|
|
}
|
|
|
|
# Timezone to UTC offset (hours)
|
|
TIMEZONE_OFFSET_MAP = {
|
|
'Asia/Shanghai': 8,
|
|
'Asia/Tokyo': 9,
|
|
'UTC': 0,
|
|
'America/New_York': -5,
|
|
'America/Los_Angeles': -8,
|
|
'Europe/London': 0,
|
|
'Europe/Berlin': 1,
|
|
}
|
|
|
|
|
|
def get_tasks_dir(bot_id: str, user_id: str) -> Path:
|
|
"""Get user task directory path"""
|
|
return Path("users") / user_id
|
|
|
|
|
|
def get_tasks_file(bot_id: str, user_id: str) -> Path:
|
|
"""Get tasks.yaml file path"""
|
|
return get_tasks_dir(bot_id, user_id) / "tasks.yaml"
|
|
|
|
|
|
def get_logs_dir(bot_id: str, user_id: str) -> Path:
|
|
"""Get task logs directory"""
|
|
return get_tasks_dir(bot_id, user_id) / "task_logs"
|
|
|
|
|
|
def load_tasks(bot_id: str, user_id: str) -> dict:
|
|
"""Load tasks.yaml"""
|
|
tasks_file = get_tasks_file(bot_id, user_id)
|
|
if not tasks_file.exists():
|
|
return {"tasks": []}
|
|
with open(tasks_file, 'r', encoding='utf-8') as f:
|
|
data = yaml.safe_load(f)
|
|
return data if data and "tasks" in data else {"tasks": []}
|
|
|
|
|
|
def save_tasks(bot_id: str, user_id: str, data: dict):
|
|
"""Save tasks.yaml"""
|
|
tasks_file = get_tasks_file(bot_id, user_id)
|
|
tasks_file.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(tasks_file, 'w', encoding='utf-8') as f:
|
|
yaml.dump(data, f, allow_unicode=True, default_flow_style=False, sort_keys=False)
|
|
|
|
|
|
def generate_task_id() -> str:
|
|
"""Generate unique task ID"""
|
|
ts = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
|
|
rand = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
|
|
return f"task_{ts}_{rand}"
|
|
|
|
|
|
def parse_timezone_offset(tz: str) -> int:
|
|
"""Get UTC offset hours for a timezone"""
|
|
return TIMEZONE_OFFSET_MAP.get(tz, 0)
|
|
|
|
|
|
def compute_next_run_cron(schedule: str, tz: str, after: datetime = None) -> str:
|
|
"""
|
|
Calculate next execution UTC time based on cron expression and timezone.
|
|
|
|
The cron expression is based on user's local time, so we first calculate
|
|
the next trigger in local time, then convert to UTC.
|
|
"""
|
|
offset_hours = parse_timezone_offset(tz)
|
|
offset = timedelta(hours=offset_hours)
|
|
|
|
# Current UTC time
|
|
now_utc = after or datetime.now(timezone.utc)
|
|
# Convert to user's local time (naive)
|
|
now_local = (now_utc + offset).replace(tzinfo=None)
|
|
|
|
# Calculate next cron trigger in local time
|
|
cron = croniter(schedule, now_local)
|
|
next_local = cron.get_next(datetime)
|
|
|
|
# Convert back to UTC
|
|
next_utc = next_local - offset
|
|
return next_utc.replace(tzinfo=timezone.utc).isoformat()
|
|
|
|
|
|
def parse_scheduled_at(scheduled_at_str: str) -> str:
|
|
"""Parse one-time task time string, return UTC ISO format"""
|
|
# Try to parse ISO format with timezone offset
|
|
try:
|
|
dt = datetime.fromisoformat(scheduled_at_str)
|
|
if dt.tzinfo is None:
|
|
# No timezone info, assume UTC
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return dt.astimezone(timezone.utc).isoformat()
|
|
except ValueError:
|
|
pass
|
|
|
|
# Try to parse common formats
|
|
for fmt in ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M"]:
|
|
try:
|
|
dt = datetime.strptime(scheduled_at_str, fmt).replace(tzinfo=timezone.utc)
|
|
return dt.isoformat()
|
|
except ValueError:
|
|
continue
|
|
|
|
raise ValueError(f"Cannot parse time format: {scheduled_at_str}")
|
|
|
|
|
|
def cmd_list(args, bot_id: str, user_id: str):
|
|
"""List all tasks"""
|
|
data = load_tasks(bot_id, user_id)
|
|
tasks = data.get("tasks", [])
|
|
|
|
if not tasks:
|
|
print("No scheduled tasks.")
|
|
return
|
|
|
|
if args.format == "brief":
|
|
# Brief format for PrePrompt hook
|
|
print(f"Scheduled tasks ({len(tasks)}):")
|
|
for t in tasks:
|
|
status_icon = {"active": "✅", "paused": "⏸️", "done": "✔️", "expired": "⏰"}.get(t["status"], "❓")
|
|
if t["type"] == "cron":
|
|
print(f" {status_icon} [{t['id']}] {t['name']} | cron: {t['schedule']} ({t.get('timezone', 'UTC')}) | executed {t.get('execution_count', 0)} times")
|
|
else:
|
|
print(f" {status_icon} [{t['id']}] {t['name']} | once: {t.get('scheduled_at', 'N/A')}")
|
|
else:
|
|
# Detail format
|
|
for t in tasks:
|
|
print(f"--- Task: {t['name']} ---")
|
|
print(f" ID: {t['id']}")
|
|
print(f" Type: {t['type']}")
|
|
print(f" Status: {t['status']}")
|
|
if t["type"] == "cron":
|
|
print(f" Cron: {t['schedule']}")
|
|
print(f" Timezone: {t.get('timezone', 'UTC')}")
|
|
else:
|
|
print(f" Scheduled at: {t.get('scheduled_at', 'N/A')}")
|
|
print(f" Message: {t['message']}")
|
|
print(f" Next run: {t.get('next_run_at', 'N/A')}")
|
|
print(f" Last run: {t.get('last_executed_at', 'N/A')}")
|
|
print(f" Executions: {t.get('execution_count', 0)}")
|
|
print(f" Created at: {t.get('created_at', 'N/A')}")
|
|
print()
|
|
|
|
|
|
def cmd_add(args, bot_id: str, user_id: str):
|
|
"""Add a task"""
|
|
data = load_tasks(bot_id, user_id)
|
|
|
|
task_id = generate_task_id()
|
|
now_utc = datetime.now(timezone.utc).isoformat()
|
|
|
|
task = {
|
|
"id": task_id,
|
|
"name": args.name,
|
|
"type": args.type,
|
|
"schedule": None,
|
|
"scheduled_at": None,
|
|
"timezone": args.timezone or "UTC",
|
|
"message": args.message,
|
|
"status": "active",
|
|
"created_at": now_utc,
|
|
"last_executed_at": None,
|
|
"next_run_at": None,
|
|
"execution_count": 0,
|
|
}
|
|
|
|
if args.type == "cron":
|
|
if not args.schedule:
|
|
print("Error: --schedule is required for cron type tasks", file=sys.stderr)
|
|
sys.exit(1)
|
|
# Validate cron expression
|
|
try:
|
|
croniter(args.schedule)
|
|
except (ValueError, KeyError) as e:
|
|
print(f"Error: Invalid cron expression '{args.schedule}': {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
task["schedule"] = args.schedule
|
|
task["next_run_at"] = compute_next_run_cron(args.schedule, task["timezone"])
|
|
|
|
elif args.type == "once":
|
|
if not args.scheduled_at:
|
|
print("Error: --scheduled-at is required for once type tasks", file=sys.stderr)
|
|
sys.exit(1)
|
|
try:
|
|
utc_time = parse_scheduled_at(args.scheduled_at)
|
|
except ValueError as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
task["scheduled_at"] = utc_time
|
|
task["next_run_at"] = utc_time
|
|
|
|
data["tasks"].append(task)
|
|
save_tasks(bot_id, user_id, data)
|
|
print(f"Task created: {task_id}")
|
|
print(f" Name: {args.name}")
|
|
print(f" Type: {args.type}")
|
|
print(f" Next run (UTC): {task['next_run_at']}")
|
|
|
|
|
|
def cmd_edit(args, bot_id: str, user_id: str):
|
|
"""Edit a task"""
|
|
data = load_tasks(bot_id, user_id)
|
|
|
|
task = None
|
|
for t in data["tasks"]:
|
|
if t["id"] == args.task_id:
|
|
task = t
|
|
break
|
|
|
|
if not task:
|
|
print(f"Error: Task {args.task_id} not found", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if args.name:
|
|
task["name"] = args.name
|
|
if args.message:
|
|
task["message"] = args.message
|
|
if args.schedule:
|
|
try:
|
|
croniter(args.schedule)
|
|
except (ValueError, KeyError) as e:
|
|
print(f"Error: Invalid cron expression '{args.schedule}': {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
task["schedule"] = args.schedule
|
|
if args.timezone:
|
|
task["timezone"] = args.timezone
|
|
if args.scheduled_at:
|
|
try:
|
|
task["scheduled_at"] = parse_scheduled_at(args.scheduled_at)
|
|
task["next_run_at"] = task["scheduled_at"]
|
|
except ValueError as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Recalculate next_run_at (if cron and schedule or timezone changed)
|
|
if task["type"] == "cron" and task.get("schedule"):
|
|
task["next_run_at"] = compute_next_run_cron(task["schedule"], task.get("timezone", "UTC"))
|
|
|
|
save_tasks(bot_id, user_id, data)
|
|
print(f"Task updated: {args.task_id}")
|
|
|
|
|
|
def cmd_delete(args, bot_id: str, user_id: str):
|
|
"""Delete a task"""
|
|
data = load_tasks(bot_id, user_id)
|
|
original_count = len(data["tasks"])
|
|
data["tasks"] = [t for t in data["tasks"] if t["id"] != args.task_id]
|
|
|
|
if len(data["tasks"]) == original_count:
|
|
print(f"Error: Task {args.task_id} not found", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
save_tasks(bot_id, user_id, data)
|
|
print(f"Task deleted: {args.task_id}")
|
|
|
|
|
|
def cmd_toggle(args, bot_id: str, user_id: str):
|
|
"""Pause/resume task"""
|
|
data = load_tasks(bot_id, user_id)
|
|
|
|
for t in data["tasks"]:
|
|
if t["id"] == args.task_id:
|
|
if t["status"] == "active":
|
|
t["status"] = "paused"
|
|
print(f"Task paused: {args.task_id}")
|
|
elif t["status"] == "paused":
|
|
t["status"] = "active"
|
|
# Recalculate next_run_at when resuming
|
|
if t["type"] == "cron" and t.get("schedule"):
|
|
t["next_run_at"] = compute_next_run_cron(t["schedule"], t.get("timezone", "UTC"))
|
|
print(f"Task resumed: {args.task_id}")
|
|
else:
|
|
print(f"Task status is {t['status']}, cannot toggle", file=sys.stderr)
|
|
sys.exit(1)
|
|
save_tasks(bot_id, user_id, data)
|
|
return
|
|
|
|
print(f"Error: Task {args.task_id} not found", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
def cmd_logs(args, bot_id: str, user_id: str):
|
|
"""View execution logs"""
|
|
logs_dir = get_logs_dir(bot_id, user_id)
|
|
log_file = logs_dir / "execution.log"
|
|
|
|
if not log_file.exists():
|
|
print("No execution logs yet.")
|
|
return
|
|
|
|
with open(log_file, 'r', encoding='utf-8') as f:
|
|
logs = yaml.safe_load(f)
|
|
|
|
if not logs:
|
|
print("No execution logs yet.")
|
|
return
|
|
|
|
# Filter by task ID
|
|
if args.task_id:
|
|
logs = [l for l in logs if l.get("task_id") == args.task_id]
|
|
|
|
# Limit count
|
|
limit = args.limit or 10
|
|
logs = logs[-limit:]
|
|
|
|
if not logs:
|
|
print("No matching logs found.")
|
|
return
|
|
|
|
for log in logs:
|
|
status_icon = "✅" if log.get("status") == "success" else "❌"
|
|
print(f"{status_icon} [{log.get('executed_at', 'N/A')}] {log.get('task_name', 'N/A')}")
|
|
response = log.get("response", "")
|
|
# Truncate long responses
|
|
if len(response) > 200:
|
|
response = response[:200] + "..."
|
|
print(f" {response}")
|
|
if log.get("duration_ms"):
|
|
print(f" Duration: {log['duration_ms']}ms")
|
|
print()
|
|
|
|
|
|
def main():
|
|
bot_id = os.getenv("ASSISTANT_ID", "")
|
|
user_id = os.getenv("USER_IDENTIFIER", "")
|
|
|
|
if not bot_id or not user_id:
|
|
print("Error: ASSISTANT_ID and USER_IDENTIFIER environment variables must be set", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
parser = argparse.ArgumentParser(description="Scheduled task manager")
|
|
subparsers = parser.add_subparsers(dest="command", help="Available commands")
|
|
|
|
# list
|
|
p_list = subparsers.add_parser("list", help="List all tasks")
|
|
p_list.add_argument("--format", choices=["brief", "detail"], default="detail", help="Output format")
|
|
|
|
# add
|
|
p_add = subparsers.add_parser("add", help="Add a task")
|
|
p_add.add_argument("--name", required=True, help="Task name")
|
|
p_add.add_argument("--type", required=True, choices=["cron", "once"], help="Task type")
|
|
p_add.add_argument("--schedule", help="Cron expression (required for cron type)")
|
|
p_add.add_argument("--scheduled-at", help="Execution time in ISO 8601 format (required for once type)")
|
|
p_add.add_argument("--timezone", help="Timezone (e.g. Asia/Tokyo), default UTC")
|
|
p_add.add_argument("--message", required=True, help="Message content to send to AI")
|
|
|
|
# edit
|
|
p_edit = subparsers.add_parser("edit", help="Edit a task")
|
|
p_edit.add_argument("task_id", help="Task ID")
|
|
p_edit.add_argument("--name", help="New task name")
|
|
p_edit.add_argument("--schedule", help="New cron expression")
|
|
p_edit.add_argument("--scheduled-at", help="New execution time")
|
|
p_edit.add_argument("--timezone", help="New timezone")
|
|
p_edit.add_argument("--message", help="New message content")
|
|
|
|
# delete
|
|
p_delete = subparsers.add_parser("delete", help="Delete a task")
|
|
p_delete.add_argument("task_id", help="Task ID")
|
|
|
|
# toggle
|
|
p_toggle = subparsers.add_parser("toggle", help="Pause/resume a task")
|
|
p_toggle.add_argument("task_id", help="Task ID")
|
|
|
|
# logs
|
|
p_logs = subparsers.add_parser("logs", help="View execution logs")
|
|
p_logs.add_argument("--task-id", help="Filter by task ID")
|
|
p_logs.add_argument("--limit", type=int, default=10, help="Number of entries to show")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not args.command:
|
|
parser.print_help()
|
|
sys.exit(1)
|
|
|
|
commands = {
|
|
"list": cmd_list,
|
|
"add": cmd_add,
|
|
"edit": cmd_edit,
|
|
"delete": cmd_delete,
|
|
"toggle": cmd_toggle,
|
|
"logs": cmd_logs,
|
|
}
|
|
|
|
commands[args.command](args, bot_id, user_id)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|