128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
|
|
from django.core.management.base import BaseCommand
|
|
|
|
from core.models import ExternalSyncEvent, TaskProviderConfig
|
|
from core.tasks.providers import get_provider
|
|
from core.util import logs
|
|
|
|
log = logs.get_logger("codex_worker")
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Process queued external sync events for worker-backed providers (codex_cli)."
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument("--once", action="store_true", default=False)
|
|
parser.add_argument("--sleep-seconds", type=float, default=2.0)
|
|
parser.add_argument("--batch-size", type=int, default=20)
|
|
parser.add_argument("--provider", default="codex_cli")
|
|
|
|
def _claim_batch(self, provider: str, batch_size: int) -> list[str]:
|
|
ids: list[str] = []
|
|
rows = list(
|
|
ExternalSyncEvent.objects.filter(
|
|
provider=provider,
|
|
status__in=["pending", "retrying"],
|
|
)
|
|
.order_by("updated_at")[: max(1, batch_size)]
|
|
.values_list("id", flat=True)
|
|
)
|
|
for row_id in rows:
|
|
updated = ExternalSyncEvent.objects.filter(
|
|
id=row_id,
|
|
provider=provider,
|
|
status__in=["pending", "retrying"],
|
|
).update(status="retrying")
|
|
if updated:
|
|
ids.append(str(row_id))
|
|
return ids
|
|
|
|
def _run_event(self, event: ExternalSyncEvent) -> None:
|
|
provider = get_provider(event.provider)
|
|
if not bool(getattr(provider, "run_in_worker", False)):
|
|
return
|
|
|
|
cfg = (
|
|
TaskProviderConfig.objects.filter(
|
|
user=event.user,
|
|
provider=event.provider,
|
|
enabled=True,
|
|
)
|
|
.order_by("-updated_at")
|
|
.first()
|
|
)
|
|
if cfg is None:
|
|
event.status = "failed"
|
|
event.error = "provider_disabled_or_missing"
|
|
event.save(update_fields=["status", "error", "updated_at"])
|
|
return
|
|
|
|
payload = dict(event.payload or {})
|
|
action = str(payload.get("action") or "append_update").strip().lower()
|
|
provider_payload = dict(payload.get("provider_payload") or payload)
|
|
|
|
if action == "create":
|
|
result = provider.create_task(dict(cfg.settings or {}), provider_payload)
|
|
elif action == "complete":
|
|
result = provider.mark_complete(dict(cfg.settings or {}), provider_payload)
|
|
elif action == "link_task":
|
|
result = provider.link_task(dict(cfg.settings or {}), provider_payload)
|
|
else:
|
|
result = provider.append_update(dict(cfg.settings or {}), provider_payload)
|
|
|
|
event.status = "ok" if result.ok else "failed"
|
|
event.error = str(result.error or "")
|
|
event.payload = dict(
|
|
payload,
|
|
worker_processed=True,
|
|
result=dict(result.payload or {}),
|
|
)
|
|
event.save(update_fields=["status", "error", "payload", "updated_at"])
|
|
|
|
if result.ok and result.external_key and event.task_id and not str(event.task.external_key or "").strip():
|
|
event.task.external_key = str(result.external_key)
|
|
event.task.save(update_fields=["external_key"])
|
|
|
|
def handle(self, *args, **options):
|
|
once = bool(options.get("once"))
|
|
sleep_seconds = max(0.2, float(options.get("sleep_seconds") or 2.0))
|
|
batch_size = max(1, int(options.get("batch_size") or 20))
|
|
provider_name = str(options.get("provider") or "codex_cli").strip().lower()
|
|
|
|
log.info(
|
|
"codex_worker started provider=%s once=%s sleep=%s batch_size=%s",
|
|
provider_name,
|
|
once,
|
|
sleep_seconds,
|
|
batch_size,
|
|
)
|
|
|
|
while True:
|
|
claimed_ids = self._claim_batch(provider_name, batch_size)
|
|
if not claimed_ids:
|
|
if once:
|
|
log.info("codex_worker exiting: no pending events")
|
|
return
|
|
time.sleep(sleep_seconds)
|
|
continue
|
|
|
|
for row_id in claimed_ids:
|
|
event = ExternalSyncEvent.objects.filter(id=row_id).select_related("task", "user").first()
|
|
if event is None:
|
|
continue
|
|
try:
|
|
self._run_event(event)
|
|
except Exception as exc:
|
|
log.exception("codex_worker failed processing id=%s", row_id)
|
|
ExternalSyncEvent.objects.filter(id=row_id).update(
|
|
status="failed",
|
|
error=f"worker_exception:{exc}",
|
|
)
|
|
|
|
if once:
|
|
log.info("codex_worker processed %s event(s)", len(claimed_ids))
|
|
return
|