continuum-task-server-sdk 1.4.0__tar.gz → 1.5.0__tar.gz

{continuum_task_server_sdk-1.4.0 → continuum_task_server_sdk-1.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: continuum-task-server-sdk
-Version: 1.4.0
+Version: 1.5.0
 Summary: Python SDK for the Continuum Task Server
 Project-URL: Homepage, https://github.com/ContinuumWorkflow/continuum-task-server-sdk-python
 Project-URL: Issues, https://github.com/ContinuumWorkflow/continuum-task-server-sdk-python/issues

{continuum_task_server_sdk-1.4.0 → continuum_task_server_sdk-1.5.0}/continuum_task_server/_http.py

@@ -46,6 +46,21 @@ class HttpClient:
     def close(self) -> None:
         self._client.close()
 
+    def set_task_server_token(self, token: str | None) -> None:
+        if token:
+            self._client.headers["Task-Server-Token"] = token
+        else:
+            self._client.headers.pop("Task-Server-Token", None)
+
+    def delete(self, path: str) -> None:
+        try:
+            response = self._client.delete(path)
+        except httpx.HTTPError as e:
+            raise ContinuumError(f"Request failed: {e}") from e
+        if response.status_code == 204:
+            return
+        self._raise_for_status(response)
+
     def __enter__(self) -> HttpClient:
         return self
 

{continuum_task_server_sdk-1.4.0 → continuum_task_server_sdk-1.5.0}/continuum_task_server/client.py

@@ -408,6 +408,16 @@ class ContinuumClient:
     def base_url(self) -> str:
         return self._http.base_url
 
+    def set_task_server_token(self, token: str | None) -> None:
+        """Attach the secret from task-server registration. Blank clears it."""
+        self._http.set_task_server_token(token)
+
+    def post_raw(self, path: str, body: Any) -> Any:
+        return self._http.post_json(path, body)
+
+    def delete(self, path: str) -> None:
+        self._http.delete(path)
+
     def close(self) -> None:
         self._http.close()
 

continuum_task_server_sdk-1.5.0/continuum_task_server/presence.py

@@ -0,0 +1,175 @@
+"""Registers this process with the queue API and refreshes the lease.
+
+A missing or failing registration never stops the task server from claiming work.
+"""
+
+from __future__ import annotations
+
+import logging
+import socket
+import threading
+import uuid
+from collections.abc import Callable, Iterable
+from importlib.metadata import PackageNotFoundError, version
+
+from .client import ContinuumClient
+from .exceptions import ContinuumError, NotFoundError
+from .websocket import ContinuumWebSocketClient
+
+logger = logging.getLogger("continuum_task_server")
+
+_PATH = "/api/queue/task-servers"
+
+
+def sdk_version() -> str:
+    try:
+        return version("continuum-task-server-sdk")
+    except PackageNotFoundError:
+        return "unknown"
+
+
+class TaskServerPresence:
+    def __init__(
+        self,
+        client: ContinuumClient,
+        *,
+        name: str | None,
+        instance_id: str | None,
+        labels: dict[str, str] | None,
+        task_names: Callable[[], Iterable[str]],
+    ) -> None:
+        self._client = client
+        self._name = name or None
+        self._instance_id = instance_id or str(uuid.uuid4())
+        self._labels = _labels(labels)
+        self._task_names = task_names
+        self._enabled = True
+        self._stopped = threading.Event()
+        self._token: str | None = None
+        self._heartbeat_seconds = 10.0
+        self._websocket: ContinuumWebSocketClient | None = None
+        self._thread: threading.Thread | None = None
+
+    def attach(self, websocket: ContinuumWebSocketClient | None) -> None:
+        self._websocket = websocket
+        if websocket is not None and self._token:
+            websocket.set_task_server_token(self._token)
+
+    def start(self) -> None:
+        if self._thread is not None:
+            return
+        self._thread = threading.Thread(target=self._loop, name="continuum-presence", daemon=True)
+        self._thread.start()
+
+    def close(self) -> None:
+        self._stopped.set()
+        if self._thread is not None:
+            self._thread.join(timeout=1)
+        current = self._token
+        self._token = None
+        if not current:
+            self._client.set_task_server_token(None)
+            return
+        self._client.set_task_server_token(current)
+        try:
+            self._client.delete(_PATH)
+        except ContinuumError as exc:
+            logger.debug("task server deregister failed: %s", exc)
+        finally:
+            self._client.set_task_server_token(None)
+            if self._websocket is not None:
+                self._websocket.set_task_server_token(None)
+
+    def reregister(self) -> bool:
+        if not self._enabled or self._stopped.is_set():
+            return False
+        self._token = None
+        self._client.set_task_server_token(None)
+        return self._register()
+
+    def _loop(self) -> None:
+        backoff = 1.0
+        while not self._stopped.is_set() and self._enabled:
+            if self._token is None:
+                if not self._register():
+                    if not self._enabled or self._stopped.is_set():
+                        return
+                    if self._stopped.wait(backoff):
+                        return
+                    backoff = min(backoff * 2, 30.0)
+                    continue
+                backoff = 1.0
+            if self._stopped.wait(self._heartbeat_seconds) or self._token is None:
+                continue
+            if not self._heartbeat():
+                self._token = None
+                self._client.set_task_server_token(None)
+
+    def _register(self) -> bool:
+        body = {
+            "instanceId": self._instance_id,
+            "name": self._name,
+            "taskNames": list(self._task_names()),
+            "labels": self._labels,
+        }
+        try:
+            payload = self._client.post_raw(_PATH, body) or {}
+            if payload.get("status") != 0:
+                logger.warning("task server register rejected: %s", payload.get("message"))
+                return False
+            issued = payload.get("token")
+            if not isinstance(issued, str) or not issued:
+                logger.warning("task server register returned no token")
+                return False
+            seconds = payload.get("heartbeatSeconds")
+            self._heartbeat_seconds = (
+                float(seconds) if isinstance(seconds, (int, float)) and seconds > 0 else 10.0
+            )
+            self._token = issued
+            self._client.set_task_server_token(issued)
+            if self._websocket is not None:
+                self._websocket.set_task_server_token(issued)
+            server = (
+                payload.get("taskServer") if isinstance(payload.get("taskServer"), dict) else {}
+            )
+            logger.info("Registered task server %s", server.get("id"))
+            return True
+        except NotFoundError:
+            self._enabled = False
+            self._token = None
+            self._client.set_task_server_token(None)
+            logger.info("Queue API does not support task server presence; claiming without it")
+            return False
+        except Exception as exc:
+            logger.warning("task server register failed: %s", exc)
+            return False
+
+    def _heartbeat(self) -> bool:
+        try:
+            self._client.post_raw(_PATH + "/heartbeat", {})
+            return True
+        except NotFoundError:
+            self._enabled = False
+            logger.info("Queue API does not support task server presence; claiming without it")
+            return False
+        except Exception as exc:
+            logger.warning("task server heartbeat failed: %s", exc)
+            return False
+
+
+def _labels(user: dict[str, str] | None) -> dict[str, str]:
+    merged: dict[str, str] = {}
+    if user:
+        merged.update(user)
+    merged["continuum.sdk"] = "python"
+    merged["continuum.sdkVersion"] = sdk_version()
+    merged["continuum.hostname"] = _hostname()
+    return merged
+
+
+def _hostname() -> str:
+    try:
+        host = socket.gethostname()
+    except OSError:
+        return "unknown"
+    return host or "unknown"

{continuum_task_server_sdk-1.4.0 → continuum_task_server_sdk-1.5.0}/continuum_task_server/server.py

@@ -34,8 +34,9 @@ import httpx
 
 from .client import ContinuumClient
 from .context import TaskExecutionContext
-from .exceptions import ContinuumError
+from .exceptions import ConflictError, ContinuumError
 from .models import QueueEvent, QueueEventType, QueueItem, TaskStatus, TransportMode
+from .presence import TaskServerPresence
 from .websocket import ContinuumWebSocketClient, WebSocketOptions
 
 logger = logging.getLogger("continuum_task_server")
@@ -92,6 +93,10 @@ class TaskServer:
         transport: TransportMode | str | httpx.BaseTransport | None = None,
         websocket_options: WebSocketOptions | None = None,
         websocket_client: Any = None,
+        name: str | None = None,
+        labels: dict[str, str] | None = None,
+        instance_id: str | None = None,
+        presence: bool = True,
     ) -> None:
         http_transport: httpx.BaseTransport | None = None
         if isinstance(transport, httpx.BaseTransport):
@@ -129,6 +134,19 @@ class TaskServer:
         self._base_url = base_url
         self._api_key = api_key
         self._active: dict[uuid.UUID, _ActiveClaim] = {}
+        self._presence = (
+            TaskServerPresence(
+                self._client,
+                name=name,
+                instance_id=instance_id,
+                labels=labels,
+                task_names=self._handlers.keys,
+            )
+            if presence
+            else None
+        )
+        if self._presence is not None and self._ws is not None:
+            self._presence.attach(self._ws)
 
     @property
     def client(self) -> ContinuumClient:
@@ -269,6 +287,8 @@ class TaskServer:
         )
 
         try:
+            if self._presence is not None:
+                self._presence.start()
             if self._mode is TransportMode.WEBSOCKET:
                 self._websocket_loop()
             else:
@@ -310,7 +330,7 @@ class TaskServer:
                 continue
 
             try:
-                item = self._client.queue.claim(registration.name)
+                item = self._claim_http(registration.name)
             except ContinuumError as e:
                 registration.semaphore.release()
                 logger.error("claim(%s) failed: %s", registration.name, e)
@@ -336,6 +356,8 @@ class TaskServer:
                 options=self._ws_options,
             )
             self._owns_ws = True
+        if self._presence is not None:
+            self._presence.attach(self._ws)
         return self._ws
 
     def _websocket_loop(self) -> None:
@@ -406,7 +428,7 @@ class TaskServer:
             if registration is None:
                 return
             try:
-                item = ws.claim(registration.name)
+                item = self._claim_socket(ws, registration.name)
             except ContinuumError as e:
                 registration.semaphore.release()
                 logger.error("claim(%s) failed: %s", registration.name, e)
@@ -630,6 +652,22 @@ class TaskServer:
             logger.error("update_status(%s, %s) failed: %s", queue_item_id, status.value, e)
             return None
 
+    def _claim_http(self, task_name: str) -> QueueItem | None:
+        try:
+            return self._client.queue.claim(task_name)
+        except ConflictError:
+            if self._presence is None or not self._presence.reregister():
+                raise
+            return self._client.queue.claim(task_name)
+
+    def _claim_socket(self, ws: ContinuumWebSocketClient, task_name: str) -> QueueItem | None:
+        try:
+            return ws.claim(task_name)
+        except ConflictError:
+            if self._presence is None or not self._presence.reregister():
+                raise
+            return ws.claim(task_name)
+
     def _drain(self) -> None:
         if self._executor is None:
             return
@@ -645,6 +683,8 @@ class TaskServer:
                 logger.debug("drain: handler future raised", exc_info=True)
         self._executor.shutdown(wait=False, cancel_futures=True)
         self._executor = None
+        if self._presence is not None:
+            self._presence.close()
         if self._ws is not None:
             self._ws.close()
             if self._owns_ws:

{continuum_task_server_sdk-1.4.0 → continuum_task_server_sdk-1.5.0}/continuum_task_server/websocket.py

@@ -137,6 +137,7 @@ class ContinuumWebSocketClient:
         self._heartbeat_stop = threading.Event()
         self._decoder = _stomp.StompDecoder()
         self._reconnect_attempt = 0
+        self._task_server_token: str | None = None
         self._client_heartbeat_ms = self._options.heartbeat_ms
         self._generation = 0
 
@@ -603,24 +604,35 @@ class ContinuumWebSocketClient:
                 return
             backoff = min(backoff * 2, self._options.max_reconnect_backoff)
 
+    def set_task_server_token(self, token: str | None) -> None:
+        """Secret from registration. The next CONNECT sends it, including after reconnect."""
+        self._task_server_token = token or None
+
     def _open_and_handshake(self) -> None:
         url = derive_websocket_url(self._base_url, self._options.endpoint)
+        headers = {"Api-Key": self._api_key}
+        token = self._task_server_token
+        if token:
+            headers["Task-Server-Token"] = token
         socket = self._connector(
             url,
-            {"Api-Key": self._api_key},
+            headers,
             self._options.connect_timeout,
         )
         self._socket = socket
         self._decoder = _stomp.StompDecoder()
         self._generation += 1
         cx = max(0, int(self._options.heartbeat_ms))
+        connect_headers = {
+            "accept-version": "1.2,1.1,1.0",
+            "heart-beat": f"{cx},{cx}",
+            "host": urlparse(url).hostname or "localhost",
+        }
+        if token:
+            connect_headers["Task-Server-Token"] = token
         connect = _stomp.StompFrame(
             command="CONNECT",
-            headers={
-                "accept-version": "1.2,1.1,1.0",
-                "heart-beat": f"{cx},{cx}",
-                "host": urlparse(url).hostname or "localhost",
-            },
+            headers=connect_headers,
         )
         self._send_raw(_stomp.encode_frame(connect))
         connected = self._recv_until_command("CONNECTED", timeout=self._options.connect_timeout)

{continuum_task_server_sdk-1.4.0 → continuum_task_server_sdk-1.5.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "continuum-task-server-sdk"
-version = "1.4.0"
+version = "1.5.0"
 description = "Python SDK for the Continuum Task Server"
 readme = "README.md"
 requires-python = ">=3.10"