continuum-task-server-sdk 1.1.0__py3-none-any.whl

continuum_task_server/server.py

@@ -0,0 +1,511 @@
+"""TaskServer - worker-loop abstraction for handling Continuum queue items.
+
+Register handlers per task-type name with ``@server.task("name")`` and call
+``server.run()``. The server polls for OPEN items, claims them, heartbeats while
+the handler runs, then reports ENDED (return value) or KILLED (exception).
+
+Use ``@server.task(..., auto_complete=False)`` when the handler records work
+elsewhere and you will call ``complete_queue_item`` / ``fail_queue_item`` later.
+The SDK does **not** heartbeat after the handler returns; you must call
+``client.queue.heartbeat`` yourself (for example on each pass of your DB poll) so
+the claim stays alive across process restarts.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import signal
+import threading
+import time
+import traceback
+import uuid
+from collections.abc import Callable
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass, field
+from typing import Any
+
+import httpx
+
+from .client import ContinuumClient
+from .exceptions import ContinuumError
+from .models import QueueEvent, QueueEventType, QueueItem, TaskStatus, TransportMode
+from .websocket import ContinuumWebSocketClient, WebSocketOptions
+
+logger = logging.getLogger("continuum_task_server")
+
+Handler = Callable[[QueueItem], Any]
+
+
+@dataclass
+class _Registration:
+    name: str
+    handler: Handler
+    concurrency: int
+    auto_complete: bool
+    semaphore: threading.Semaphore = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.semaphore = threading.Semaphore(self.concurrency)
+
+
+class TaskServer:
+    """Worker-first task server.
+
+    Example:
+        >>> server = TaskServer(base_url="http://localhost:8080", api_key="...")
+        >>> @server.task("echo")
+        ... def echo(item):
+        ...     return {"echoed": item.input_data_json}
+        >>> server.run()  # blocks until SIGINT/SIGTERM
+
+    With ``auto_complete=True`` (default), handler return values become ``outputData``
+    and the task is marked ENDED. Raising any exception marks the task KILLED.
+
+    With ``auto_complete=False``, the handler returns without ENDED; you are
+    responsible for heartbeats until you call ``complete_queue_item`` or
+    ``fail_queue_item`` (see README). Heartbeats only run while the handler is
+    executing, not after it returns.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str,
+        *,
+        max_workers: int = 4,
+        poll_interval: float = 1.0,
+        max_poll_interval: float = 5.0,
+        heartbeat_interval: float = 15.0,
+        shutdown_timeout: float = 30.0,
+        client: ContinuumClient | None = None,
+        connect_timeout: float = 10.0,
+        request_timeout: float = 30.0,
+        transport: TransportMode | str | httpx.BaseTransport | None = None,
+        websocket_options: WebSocketOptions | None = None,
+        websocket_client: Any = None,
+    ) -> None:
+        http_transport: httpx.BaseTransport | None = None
+        if isinstance(transport, httpx.BaseTransport):
+            http_transport = transport
+            self._mode = TransportMode.HTTP
+        else:
+            self._mode = _coerce_transport(transport)
+        self._client = client or ContinuumClient(
+            base_url=base_url,
+            api_key=api_key,
+            connect_timeout=connect_timeout,
+            request_timeout=request_timeout,
+            transport=http_transport,
+        )
+        self._owns_client = client is None
+        self._max_workers = max_workers
+        self._poll_interval = poll_interval
+        self._max_poll_interval = max_poll_interval
+        self._heartbeat_interval = heartbeat_interval
+        self._shutdown_timeout = shutdown_timeout
+        self._handlers: dict[str, _Registration] = {}
+        self._stop_event = threading.Event()
+        self._executor: ThreadPoolExecutor | None = None
+        self._inflight: set[Future[Any]] = set()
+        self._inflight_lock = threading.Lock()
+        self._ws = websocket_client
+        self._ws_options = websocket_options or WebSocketOptions(
+            connect_timeout=connect_timeout,
+            request_timeout=request_timeout,
+        )
+        self._owns_ws = websocket_client is None
+        self._hint_lock = threading.Lock()
+        self._pending_hints: set[str] = set()
+        self._claim_loop_lock = threading.Lock()
+        self._base_url = base_url
+        self._api_key = api_key
+
+    @property
+    def client(self) -> ContinuumClient:
+        """Underlying ContinuumClient. Useful for management operations from handlers."""
+        return self._client
+
+    def task(
+        self,
+        name: str,
+        *,
+        concurrency: int = 1,
+        auto_complete: bool = True,
+    ) -> Callable[[Handler], Handler]:
+        """Decorator registering ``handler`` for queue items of task-type ``name``.
+
+        ``concurrency`` is the max number of in-flight items for this task name.
+
+        If ``auto_complete`` is False, the handler returns without the server sending
+        ENDED; call ``complete_queue_item`` or ``fail_queue_item`` when done. You must
+        heartbeat the queue item yourself until then (see README).
+        """
+        if concurrency < 1:
+            raise ValueError("concurrency must be >= 1")
+
+        def decorator(handler: Handler) -> Handler:
+            if name in self._handlers:
+                raise ValueError(f"a handler for task {name!r} is already registered")
+            self._handlers[name] = _Registration(
+                name=name,
+                handler=handler,
+                concurrency=concurrency,
+                auto_complete=auto_complete,
+            )
+            return handler
+
+        return decorator
+
+    def register(
+        self,
+        name: str,
+        handler: Handler,
+        *,
+        concurrency: int = 1,
+        auto_complete: bool = True,
+    ) -> None:
+        """Imperative alternative to ``@task``."""
+        self.task(name, concurrency=concurrency, auto_complete=auto_complete)(handler)
+
+    def complete_queue_item(
+        self,
+        queue_item_id: uuid.UUID | str,
+        *,
+        output_data: Any = None,
+    ) -> None:
+        """Mark a queue item ENDED.
+
+        Use after ``auto_complete=False`` when work finished successfully. Uses the
+        same ``ContinuumClient`` (and API key) as this ``TaskServer``.
+        """
+        qid = self._normalize_queue_id(queue_item_id)
+        self._safe_update_status_by_id(qid, TaskStatus.ENDED, output=output_data)
+
+    def fail_queue_item(
+        self,
+        queue_item_id: uuid.UUID | str,
+        *,
+        error: BaseException | str | dict[str, Any] | None = None,
+    ) -> None:
+        """Mark a queue item KILLED."""
+        qid = self._normalize_queue_id(queue_item_id)
+        payload: dict[str, Any]
+        if isinstance(error, dict):
+            payload = dict(error)
+        elif isinstance(error, str):
+            payload = {"error": error}
+        elif isinstance(error, BaseException):
+            payload = {
+                "error": str(error),
+                "type": type(error).__name__,
+                "traceback": "".join(
+                    traceback.format_exception(type(error), error, error.__traceback__)
+                ),
+            }
+        else:
+            payload = {"error": "failed"}
+        self._safe_update_status_by_id(qid, TaskStatus.KILLED, output=payload)
+
+    @staticmethod
+    def _normalize_queue_id(queue_item_id: uuid.UUID | str) -> uuid.UUID:
+        if isinstance(queue_item_id, uuid.UUID):
+            return queue_item_id
+        return uuid.UUID(str(queue_item_id))
+
+    def run(self, *, install_signal_handlers: bool = True) -> None:
+        """Block and poll until ``stop()`` is called or a shutdown signal arrives."""
+        if not self._handlers:
+            raise RuntimeError("no task handlers registered; use @server.task(...) first")
+
+        if install_signal_handlers:
+            self._install_signal_handlers()
+
+        self._stop_event.clear()
+        self._executor = ThreadPoolExecutor(max_workers=self._max_workers)
+
+        logger.info(
+            "Continuum task server started: tasks=%s, max_workers=%d, transport=%s",
+            list(self._handlers),
+            self._max_workers,
+            self._mode.value,
+        )
+
+        try:
+            if self._mode is TransportMode.WEBSOCKET:
+                self._websocket_loop()
+            else:
+                self._poll_loop()
+        finally:
+            self._drain()
+
+    def stop(self) -> None:
+        """Signal the server to stop polling and drain in-flight handlers."""
+        if not self._stop_event.is_set():
+            logger.info("Continuum task server shutdown requested")
+        self._stop_event.set()
+
+    def _install_signal_handlers(self) -> None:
+        def _handler(signum: int, _frame: object) -> None:
+            logger.info("Received signal %d, stopping", signum)
+            self.stop()
+
+        try:
+            signal.signal(signal.SIGINT, _handler)
+            signal.signal(signal.SIGTERM, _handler)
+        except ValueError:
+            logger.debug("signal handlers not installed (not on main thread)")
+
+    def _poll_loop(self) -> None:
+        names = list(self._handlers)
+        index = 0
+        backoff = self._poll_interval
+
+        while not self._stop_event.is_set():
+            registration = self._handlers[names[index % len(names)]]
+            index += 1
+
+            if not registration.semaphore.acquire(blocking=False):
+                if index % len(names) == 0:
+                    self._stop_event.wait(timeout=self._poll_interval)
+                continue
+
+            try:
+                item = self._client.queue.claim(registration.name)
+            except ContinuumError as e:
+                registration.semaphore.release()
+                logger.error("claim(%s) failed: %s", registration.name, e)
+                self._stop_event.wait(timeout=min(backoff, self._max_poll_interval))
+                backoff = min(backoff * 2, self._max_poll_interval)
+                continue
+
+            if item is None:
+                registration.semaphore.release()
+                if index % len(names) == 0:
+                    self._stop_event.wait(timeout=backoff)
+                    backoff = min(backoff * 2, self._max_poll_interval)
+                continue
+
+            backoff = self._poll_interval
+            self._dispatch(registration, item)
+
+    def _ensure_websocket(self) -> ContinuumWebSocketClient:
+        if self._ws is None:
+            self._ws = ContinuumWebSocketClient(
+                self._base_url,
+                self._api_key,
+                options=self._ws_options,
+            )
+            self._owns_ws = True
+        return self._ws
+
+    def _websocket_loop(self) -> None:
+        ws = self._ensure_websocket()
+        ws.connect()
+        for name in self._handlers:
+            ws.subscribe_work(name)
+        logger.info("WebSocket subscribed to work: %s", list(self._handlers))
+        while not self._stop_event.is_set():
+            event = ws.next_event(timeout=0.25)
+            if event is None:
+                continue
+            self._handle_queue_event(event)
+
+    def _handle_queue_event(self, event: QueueEvent) -> None:
+        if event.event_type is not QueueEventType.WORK_AVAILABLE:
+            return
+        name = event.task_name
+        if name is None or name not in self._handlers:
+            return
+        with self._hint_lock:
+            self._pending_hints.add(name)
+        self._try_claim_pending()
+
+    def _try_claim_pending(self) -> None:
+        while not self._stop_event.is_set():
+            if not self._claim_loop_lock.acquire(blocking=False):
+                return
+            try:
+                self._drain_claims()
+            finally:
+                self._claim_loop_lock.release()
+            if not self._has_pending_with_capacity():
+                return
+
+    def _has_pending_with_capacity(self) -> bool:
+        with self._hint_lock:
+            names = list(self._pending_hints)
+        for name in names:
+            registration = self._handlers.get(name)
+            if registration is None:
+                continue
+            if registration.semaphore.acquire(blocking=False):
+                registration.semaphore.release()
+                return True
+        return False
+
+    def _drain_claims(self) -> None:
+        ws = self._ws
+        if ws is None:
+            return
+        while not self._stop_event.is_set():
+            registration = self._pop_pending_with_capacity()
+            if registration is None:
+                return
+            try:
+                item = ws.claim(registration.name)
+            except ContinuumError as e:
+                registration.semaphore.release()
+                logger.error("claim(%s) failed: %s", registration.name, e)
+                with self._hint_lock:
+                    self._pending_hints.add(registration.name)
+                return
+            if item is None:
+                registration.semaphore.release()
+                continue
+            with self._hint_lock:
+                self._pending_hints.add(registration.name)
+            self._dispatch(registration, item)
+
+    def _pop_pending_with_capacity(self) -> _Registration | None:
+        with self._hint_lock:
+            names = list(self._pending_hints)
+        for name in names:
+            registration = self._handlers.get(name)
+            if registration is None:
+                with self._hint_lock:
+                    self._pending_hints.discard(name)
+                continue
+            if not registration.semaphore.acquire(blocking=False):
+                continue
+            with self._hint_lock:
+                self._pending_hints.discard(name)
+            return registration
+        return None
+
+    def _dispatch(self, registration: _Registration, item: QueueItem) -> None:
+        assert self._executor is not None
+        future = self._executor.submit(self._run_item, registration, item)
+        with self._inflight_lock:
+            self._inflight.add(future)
+        future.add_done_callback(self._on_done)
+
+    def _on_done(self, future: Future[Any]) -> None:
+        with self._inflight_lock:
+            self._inflight.discard(future)
+
+    def _run_item(self, registration: _Registration, item: QueueItem) -> None:
+        try:
+            logger.info("Claimed queue item %s (task=%s)", item.id, registration.name)
+            stop_heartbeat = threading.Event()
+            heartbeat_thread = threading.Thread(
+                target=self._heartbeat_loop,
+                args=(item, stop_heartbeat),
+                name=f"hb-{item.id}",
+                daemon=True,
+            )
+            heartbeat_thread.start()
+
+            try:
+                self._safe_update_status(item, TaskStatus.STARTED)
+                try:
+                    result = registration.handler(item)
+                except Exception as e:
+                    logger.exception("Handler raised for queue item %s", item.id)
+                    error_payload = {
+                        "error": str(e),
+                        "type": type(e).__name__,
+                        "traceback": traceback.format_exc(),
+                    }
+                    self._safe_update_status(item, TaskStatus.KILLED, output=error_payload)
+                    return
+
+                if registration.auto_complete:
+                    self._safe_update_status(item, TaskStatus.ENDED, output=result)
+                    logger.info("Completed queue item %s", item.id)
+                else:
+                    logger.info(
+                        "Handler returned for queue item %s without ENDED "
+                        "(auto_complete=False); caller must heartbeat and complete",
+                        item.id,
+                    )
+            finally:
+                stop_heartbeat.set()
+                heartbeat_thread.join(timeout=1.0)
+        finally:
+            registration.semaphore.release()
+            if self._mode is TransportMode.WEBSOCKET:
+                self._try_claim_pending()
+
+    def _heartbeat_loop(self, item: QueueItem, stop_event: threading.Event) -> None:
+        while not stop_event.wait(timeout=self._heartbeat_interval):
+            try:
+                if self._ws is not None:
+                    self._ws.heartbeat(item.id)
+                else:
+                    self._client.queue.heartbeat(item.id)
+            except ContinuumError as e:
+                logger.warning("heartbeat for %s failed: %s", item.id, e)
+
+    def _safe_update_status(
+        self,
+        item: QueueItem,
+        status: TaskStatus,
+        *,
+        output: Any = None,
+    ) -> None:
+        self._safe_update_status_by_id(item.id, status, output=output)
+
+    def _safe_update_status_by_id(
+        self,
+        queue_item_id: uuid.UUID,
+        status: TaskStatus,
+        *,
+        output: Any = None,
+    ) -> None:
+        try:
+            if self._ws is not None:
+                self._ws.update_status(queue_item_id, status, output_data=output)
+            else:
+                self._client.queue.update_status(queue_item_id, status, output_data=output)
+        except ContinuumError as e:
+            logger.error("update_status(%s, %s) failed: %s", queue_item_id, status.value, e)
+
+    def _drain(self) -> None:
+        if self._executor is None:
+            return
+        logger.info("Draining in-flight tasks (timeout=%.1fs)", self._shutdown_timeout)
+        deadline = time.monotonic() + self._shutdown_timeout
+        with self._inflight_lock:
+            futures = list(self._inflight)
+        for future in futures:
+            remaining = max(0.0, deadline - time.monotonic())
+            try:
+                future.result(timeout=remaining)
+            except Exception:
+                logger.debug("drain: handler future raised", exc_info=True)
+        self._executor.shutdown(wait=False, cancel_futures=True)
+        self._executor = None
+        if self._ws is not None:
+            self._ws.close()
+            if self._owns_ws:
+                self._ws = None
+        if self._owns_client:
+            self._client.close()
+        logger.info("Continuum task server stopped")
+
+
+def _coerce_transport(value: TransportMode | str | None) -> TransportMode:
+    if value is None:
+        env = os.environ.get("CONTINUUM_TRANSPORT")
+        if env is None or not str(env).strip():
+            return TransportMode.HTTP
+        value = env
+    if isinstance(value, TransportMode):
+        return value
+    normalized = str(value).strip().lower()
+    if normalized in {"http", "https", "poll"}:
+        return TransportMode.HTTP
+    if normalized in {"websocket", "ws", "wss", "stomp"}:
+        return TransportMode.WEBSOCKET
+    raise ValueError(f"unknown transport {value!r}; expected 'http' or 'websocket'")