task_queue.py

"""Durable work queue pattern using NATS JetStream.

Distributes tasks to worker agents with acknowledgment, retry,
and exponential backoff. Built for AI agent task delegation.
"""

from __future__ import annotations

import json
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable, Awaitable

from nats.aio.client import Client as NATSClient
from nats.js.api import (
    ConsumerConfig,
    DeliverPolicy,
    AckPolicy,
    StreamConfig,
    RetentionPolicy,
)


@dataclass
class Task:
    """A unit of work to be processed by a worker agent."""

    type: str
    payload: dict[str, Any]
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    priority: int = 0
    created_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )
    retry_count: int = 0

    def serialize(self) -> bytes:
        return json.dumps(asdict(self)).encode()

    @classmethod
    def deserialize(cls, data: bytes) -> Task:
        d = json.loads(data.decode())
        return cls(**d)


class TaskQueue:
    """Durable task queue backed by NATS JetStream.

    Creates a stream for task distribution with pull-based consumers,
    acknowledgment, and automatic redelivery on failure.

    Usage:
        queue = TaskQueue(nc, stream_name="agent-tasks")
        await queue.setup()

        # Producer
        await queue.enqueue(Task(type="analyze", payload={"url": "..."}))

        # Consumer
        await queue.worker(my_handler, durable_name="worker-1")
    """

    MAX_RETRIES = 3
    BACKOFF_BASE_SECONDS = 2.0

    def __init__(
        self,
        nc: NATSClient,
        stream_name: str = "TASKS",
        org: str = "default",
    ) -> None:
        self._nc = nc
        self._stream_name = stream_name
        self._org = org
        self._js = None
        self._subject = f"{org}.tasks.>"

    @property
    def subject_prefix(self) -> str:
        return f"{self._org}.tasks"

    async def setup(self) -> None:
        """Initialize the JetStream context and create the stream."""
        self._js = self._nc.jetstream()
        await self._js.add_stream(
            StreamConfig(
                name=self._stream_name,
                subjects=[self._subject],
                retention=RetentionPolicy.WORK_QUEUE,
            )
        )

    async def enqueue(self, task: Task, task_subject: str | None = None) -> None:
        """Publish a task to the queue.

        Args:
            task: The Task to enqueue.
            task_subject: Optional subject suffix (defaults to task.type).
        """
        if self._js is None:
            raise RuntimeError("TaskQueue not set up. Call setup() first.")
        subject = task_subject or f"{self.subject_prefix}.{task.type}"
        await self._js.publish(subject, task.serialize())

    async def worker(
        self,
        handler: Callable[[Task], Awaitable[None]],
        durable_name: str,
        task_type: str = "*",
    ):
        """Start consuming tasks from the queue.

        Tasks are acknowledged on success. On failure, the message is
        nak'd with a backoff delay so JetStream redelivers it.

        Args:
            handler: Async function that processes a Task. Raise to signal failure.
            durable_name: Durable consumer name (survives restarts).
            task_type: Filter to specific task type, or '*' for all.

        Returns:
            The pull subscription (for lifecycle management).
        """
        if self._js is None:
            raise RuntimeError("TaskQueue not set up. Call setup() first.")

        filter_subject = f"{self.subject_prefix}.{task_type}"
        sub = await self._js.pull_subscribe(
            filter_subject,
            durable=durable_name,
            config=ConsumerConfig(
                durable_name=durable_name,
                ack_policy=AckPolicy.EXPLICIT,
                deliver_policy=DeliverPolicy.ALL,
                max_deliver=self.MAX_RETRIES + 1,
                ack_wait=30,
            ),
        )
        return _WorkerSubscription(sub, handler, max_retries=self.MAX_RETRIES)


class _WorkerSubscription:
    """Wraps a pull subscription with processing logic."""

    def __init__(self, sub, handler, max_retries: int = 3) -> None:
        self._sub = sub
        self._handler = handler
        self._max_retries = max_retries
        self._running = False

    async def process_batch(self, batch_size: int = 10, timeout: float = 5.0) -> int:
        """Fetch and process a batch of tasks.

        Args:
            batch_size: Max number of tasks to fetch at once.
            timeout: How long to wait for messages.

        Returns:
            Number of tasks successfully processed in this batch.
        """
        try:
            messages = await self._sub.fetch(batch=batch_size, timeout=timeout)
        except Exception:
            return 0

        processed = 0
        for msg in messages:
            task = Task.deserialize(msg.data)
            try:
                await self._handler(task)
                await msg.ack()
                processed += 1
            except Exception:
                retry = msg.metadata.num_delivered if msg.metadata else task.retry_count
                if retry >= self._max_retries:
                    await msg.term()
                else:
                    backoff = TaskQueue.BACKOFF_BASE_SECONDS * (2 ** retry)
                    await msg.nak(delay=backoff)
        return processed

    async def run(self, batch_size: int = 10, timeout: float = 5.0) -> None:
        """Continuously process tasks until stopped."""
        self._running = True
        while self._running:
            await self.process_batch(batch_size, timeout)

    def stop(self) -> None:
        """Signal the worker to stop after the current batch."""
        self._running = False