continuum-task-server-sdk 0.1.0__tar.gz → 1.2.0__tar.gz

{continuum_task_server_sdk-0.1.0 → continuum_task_server_sdk-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: continuum-task-server-sdk
-Version: 0.1.0
+Version: 1.2.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-0.1.0 → continuum_task_server_sdk-1.2.0}/continuum_task_server/__init__.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from .client import ContinuumClient
+from .context import TaskExecutionContext
 from .exceptions import (
     AmbiguousCommandError,
     BadRequestError,
@@ -24,6 +25,7 @@ from .models import (
     QueueEvent,
     QueueEventType,
     QueueItem,
+    QueueItemTrigger,
     TaskItem,
     TaskItemVersion,
     TaskStatus,
@@ -54,8 +56,10 @@ __all__ = [
     "QueueEvent",
     "QueueEventType",
     "QueueItem",
+    "QueueItemTrigger",
     "RateLimitError",
     "ServerError",
+    "TaskExecutionContext",
     "TaskItem",
     "TaskItemVersion",
     "TaskServer",

{continuum_task_server_sdk-0.1.0 → continuum_task_server_sdk-1.2.0}/continuum_task_server/client.py

@@ -11,6 +11,7 @@ from uuid import UUID
 import httpx
 
 from ._http import HttpClient, drop_none
+from .exceptions import ContinuumError
 from .models import Content, QueueItem, TaskItem, TaskItemVersion, TaskStatus, TaskType
 
 
@@ -230,6 +231,7 @@ class QueueApi:
         input_data: Any = None,
         priority: int | None = None,
         parent_id: UUID | str | None = None,
+        idempotency_token: str | None = None,
     ) -> QueueItem:
         """Add a queue item. At least one of task_name or task_item_name is required."""
         if task_name is None and task_item_name is None:
@@ -241,6 +243,7 @@ class QueueApi:
                 "taskItemName": task_item_name,
                 "priority": priority,
                 "inputData": _encode_input(input_data),
+                "idempotencyToken": idempotency_token,
             }
         )
         return QueueItem.model_validate(self._http.post_json(self._MGMT, body))
@@ -262,9 +265,15 @@ class QueueApi:
             return None
         return QueueItem.model_validate(result)
 
-    def heartbeat(self, queue_item_id: UUID | str) -> QueueItem:
+    def heartbeat(
+        self,
+        queue_item_id: UUID | str,
+        *,
+        claim_token: UUID | str | None = None,
+    ) -> QueueItem:
+        body = drop_none({"claimToken": str(claim_token) if claim_token else None})
         return QueueItem.model_validate(
-            self._http.post_json(f"{self._QUEUE}/queue-items/{queue_item_id}/heartbeat", {})
+            self._http.post_json(f"{self._QUEUE}/queue-items/{queue_item_id}/heartbeat", body)
         )
 
     def update_status(
@@ -273,16 +282,50 @@ class QueueApi:
         status: TaskStatus,
         *,
         output_data: Any = None,
+        claim_token: UUID | str | None = None,
     ) -> QueueItem:
         # Server expects `status` on this endpoint (not `taskStatus` on queue item JSON).
         body: dict[str, Any] = {"status": status.value}
         encoded = _encode_input(output_data)
         if encoded is not None:
             body["outputData"] = encoded
+        if claim_token is not None:
+            body["claimToken"] = str(claim_token)
         return QueueItem.model_validate(
             self._http.post_json(f"{self._QUEUE}/queue-items/{queue_item_id}/status", body)
         )
 
+    def cancel(self, queue_item_id: UUID | str) -> QueueItem:
+        result = self._http.post_json(f"{self._MGMT}/{queue_item_id}/cancel", {})
+        status = 0 if result is None else int(result.get("status", 0))
+        if result is None or status != 0:
+            message = None if result is None else result.get("message")
+            raise ContinuumError(message or "cancel failed")
+        return QueueItem.model_validate(result["item"])
+
+    def complete_control(
+        self,
+        queue_item_id: UUID | str,
+        *,
+        claim_token: UUID | str | None = None,
+        output_data: Any = None,
+    ) -> QueueItem:
+        body = drop_none(
+            {
+                "claimToken": str(claim_token) if claim_token else None,
+                "outputData": _encode_input(output_data),
+            }
+        )
+        result = self._http.post_json(
+            f"{self._QUEUE}/queue-items/{queue_item_id}/control-complete",
+            body,
+        )
+        status = 0 if result is None else int(result.get("status", 0))
+        if result is None or status != 0:
+            message = None if result is None else result.get("message")
+            raise ContinuumError(message or "control-complete failed")
+        return QueueItem.model_validate(result["item"])
+
 
 class ContentStoreApi:
     _PATH = "/api/management/content-store"

continuum_task_server_sdk-1.2.0/continuum_task_server/context.py

@@ -0,0 +1,79 @@
+"""Per-claim execution context with an idempotent cancellation callback."""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+from uuid import UUID
+
+from .models import QueueItem
+
+_callback_lock = threading.Lock()
+_callback_pool: list[threading.Thread] = []
+
+
+class TaskExecutionContext:
+    """Cancellation token/callback for one claimed queue item."""
+
+    def __init__(self, item: QueueItem) -> None:
+        self._item = item
+        self._cancelled = threading.Event()
+        self._callback: Callable[[], None] | None = None
+        self._callback_started = threading.Event()
+        self._cleanup_done = threading.Event()
+        self._lock = threading.Lock()
+
+    @property
+    def item(self) -> QueueItem:
+        return self._item
+
+    @property
+    def claim_token(self) -> UUID | None:
+        return self._item.claim_token
+
+    def is_cancelled(self) -> bool:
+        return self._cancelled.is_set()
+
+    def on_cancelled(self, handler: Callable[[], None] | None) -> None:
+        if handler is None:
+            return
+        with self._lock:
+            if self._callback is not None:
+                return
+            self._callback = handler
+            already = self._cancelled.is_set()
+        if already:
+            self._start_callback()
+
+    def request_cancel(self) -> None:
+        with self._lock:
+            if self._cancelled.is_set():
+                return
+            self._cancelled.set()
+            has_callback = self._callback is not None
+        if has_callback:
+            self._start_callback()
+
+    def await_cleanup(self, timeout: float) -> None:
+        if not self._cancelled.is_set() or self._callback is None:
+            return
+        self._cleanup_done.wait(timeout=max(0.001, timeout))
+
+    def _start_callback(self) -> None:
+        with self._lock:
+            if self._callback_started.is_set():
+                return
+            handler = self._callback
+            if handler is None:
+                self._cleanup_done.set()
+                return
+            self._callback_started.set()
+
+        def _run() -> None:
+            try:
+                handler()
+            finally:
+                self._cleanup_done.set()
+
+        thread = threading.Thread(target=_run, name="continuum-cancel-cb", daemon=True)
+        thread.start()

{continuum_task_server_sdk-0.1.0 → continuum_task_server_sdk-1.2.0}/continuum_task_server/models.py

@@ -36,6 +36,8 @@ class TaskType(_Base):
     is_global: bool = Field(default=False, alias="isGlobal")
     max_duration_sec: int | None = Field(default=None, alias="maxDurationSec")
     heartbeat_timeout_sec: int | None = Field(default=None, alias="heartbeatTimeoutSec")
+    claim_start_timeout_sec: int | None = Field(default=None, alias="claimStartTimeoutSec")
+    control_grace_sec: int | None = Field(default=None, alias="controlGraceSec")
 
 
 class TaskItem(_Base):
@@ -69,6 +71,18 @@ class TaskItemVersion(_Base):
             return self.item_definition
 
 
+class QueueItemTrigger(_Base):
+    type: str | None = None
+    source: str | None = None
+    schedule_id: UUID | None = Field(default=None, alias="scheduleId")
+    schedule_name: str | None = Field(default=None, alias="scheduleName")
+    occurrence_id: UUID | None = Field(default=None, alias="occurrenceId")
+    scheduled_for: datetime | None = Field(default=None, alias="scheduledFor")
+    enqueued_at: datetime | None = Field(default=None, alias="enqueuedAt")
+    time_zone: str | None = Field(default=None, alias="timeZone")
+    misfire: bool | None = None
+
+
 class QueueItem(_Base):
     id: UUID
     depth: int = 0
@@ -81,6 +95,14 @@ class QueueItem(_Base):
     created_by: UUID | None = Field(default=None, alias="createdBy")
     input_data: str | None = Field(default=None, alias="inputData")
     output_data: str | None = Field(default=None, alias="outputData")
+    claim_token: UUID | None = Field(default=None, alias="claimToken")
+    claim_attempt: int | None = Field(default=None, alias="claimAttempt")
+    control_request: str | None = Field(default=None, alias="controlRequest")
+    control_reason: str | None = Field(default=None, alias="controlReason")
+    heartbeat_dtm: datetime | None = Field(default=None, alias="heartbeatDtm")
+    start_dtm: datetime | None = Field(default=None, alias="startDtm")
+    end_dtm: datetime | None = Field(default=None, alias="endDtm")
+    trigger: QueueItemTrigger | None = None
 
     @property
     def input_data_json(self) -> Any:
@@ -132,6 +154,7 @@ class QueueEventType(str, Enum):
 
     WORK_AVAILABLE = "work.available"
     COMPLETED = "queue.completed"
+    CONTROL_REQUESTED = "queue.control.requested"
 
 
 class WaitTargetStatus(_Base):
@@ -157,6 +180,8 @@ class QueueEvent(_Base):
     correlation: str | None = None
     wait_mode: WaitMode | None = Field(default=None, alias="waitMode")
     targets: list[WaitTargetStatus] | None = None
+    control_request: str | None = Field(default=None, alias="controlRequest")
+    claim_token: UUID | None = Field(default=None, alias="claimToken")
 
 
 class WorkSubscriptionResult(_Base):

{continuum_task_server_sdk-0.1.0 → continuum_task_server_sdk-1.2.0}/continuum_task_server/server.py

@@ -13,7 +13,9 @@ the claim stays alive across process restarts.
 
 from __future__ import annotations
 
+import inspect
 import logging
+import os
 import signal
 import threading
 import time
@@ -27,13 +29,14 @@ from typing import Any
 import httpx
 
 from .client import ContinuumClient
+from .context import TaskExecutionContext
 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]
+Handler = Callable[..., Any]
 
 
 @dataclass
@@ -61,10 +64,10 @@ class TaskServer:
     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.
+    With ``auto_complete=False``, the handler returns without ENDED. Automatic
+    heartbeats stop when the handler returns. Heartbeat from the completing
+    process with ``client.queue.heartbeat`` until ``complete_queue_item`` or
+    ``fail_queue_item``.
     """
 
     def __init__(
@@ -119,6 +122,7 @@ class TaskServer:
         self._claim_loop_lock = threading.Lock()
         self._base_url = base_url
         self._api_key = api_key
+        self._active: dict[uuid.UUID, _ActiveClaim] = {}
 
     @property
     def client(self) -> ContinuumClient:
@@ -179,7 +183,14 @@ class TaskServer:
         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)
+        claim = self._active.get(qid)
+        if claim is not None and claim.control_seen.is_set():
+            claim.acknowledge_control()
+            return
+        token = None if claim is None else claim.claim_token
+        self._safe_update_status_by_id(qid, TaskStatus.ENDED, output=output_data, claim_token=token)
+        if claim is not None:
+            claim.mark_completed()
 
     def fail_queue_item(
         self,
@@ -189,6 +200,10 @@ class TaskServer:
     ) -> None:
         """Mark a queue item KILLED."""
         qid = self._normalize_queue_id(queue_item_id)
+        claim = self._active.get(qid)
+        if claim is not None and claim.control_seen.is_set():
+            claim.acknowledge_control()
+            return
         payload: dict[str, Any]
         if isinstance(error, dict):
             payload = dict(error)
@@ -204,7 +219,10 @@ class TaskServer:
             }
         else:
             payload = {"error": "failed"}
-        self._safe_update_status_by_id(qid, TaskStatus.KILLED, output=payload)
+        token = None if claim is None else claim.claim_token
+        self._safe_update_status_by_id(qid, TaskStatus.KILLED, output=payload, claim_token=token)
+        if claim is not None:
+            claim.mark_completed()
 
     @staticmethod
     def _normalize_queue_id(queue_item_id: uuid.UUID | str) -> uuid.UUID:
@@ -243,6 +261,8 @@ class TaskServer:
         if not self._stop_event.is_set():
             logger.info("Continuum task server shutdown requested")
         self._stop_event.set()
+        for claim in list(self._active.values()):
+            claim.deferred_gate.set()
 
     def _install_signal_handlers(self) -> None:
         def _handler(signum: int, _frame: object) -> None:
@@ -311,6 +331,9 @@ class TaskServer:
             self._handle_queue_event(event)
 
     def _handle_queue_event(self, event: QueueEvent) -> None:
+        if event.event_type is QueueEventType.CONTROL_REQUESTED:
+            self._apply_control_event(event)
+            return
         if event.event_type is not QueueEventType.WORK_AVAILABLE:
             return
         name = event.task_name
@@ -320,6 +343,17 @@ class TaskServer:
             self._pending_hints.add(name)
         self._try_claim_pending()
 
+    def _apply_control_event(self, event: QueueEvent) -> None:
+        if event.queue_item_id is None:
+            return
+        claim = self._active.get(event.queue_item_id)
+        if claim is None:
+            return
+        token = event.claim_token
+        if token is not None and claim.claim_token is not None and token != claim.claim_token:
+            return
+        claim.signal_control()
+
     def _try_claim_pending(self) -> None:
         while not self._stop_event.is_set():
             if not self._claim_loop_lock.acquire(blocking=False):
@@ -394,57 +428,132 @@ class TaskServer:
             self._inflight.discard(future)
 
     def _run_item(self, registration: _Registration, item: QueueItem) -> None:
+        claim = _ActiveClaim(item=item, server=self)
+        self._active[item.id] = claim
         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),
+                args=(claim,),
                 name=f"hb-{item.id}",
                 daemon=True,
             )
             heartbeat_thread.start()
 
             try:
-                self._safe_update_status(item, TaskStatus.STARTED)
+                self._subscribe_control(claim)
+                self._inspect_control(
+                    self._safe_update_status(
+                        item, TaskStatus.STARTED, claim_token=claim.claim_token
+                    )
+                )
                 try:
-                    result = registration.handler(item)
+                    result = self._invoke_handler(registration.handler, item, claim.context)
                 except Exception as e:
+                    if claim.control_seen.is_set() or claim.context.is_cancelled():
+                        claim.acknowledge_control()
+                        return
                     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)
+                    self._safe_update_status(
+                        item, TaskStatus.KILLED, output=error_payload, claim_token=claim.claim_token
+                    )
+                    claim.mark_completed()
+                    return
+
+                if claim.control_seen.is_set() or claim.context.is_cancelled():
+                    claim.acknowledge_control()
                     return
 
                 if registration.auto_complete:
-                    self._safe_update_status(item, TaskStatus.ENDED, output=result)
+                    self._safe_update_status(
+                        item, TaskStatus.ENDED, output=result, claim_token=claim.claim_token
+                    )
                     logger.info("Completed queue item %s", item.id)
+                    claim.mark_completed()
                 else:
                     logger.info(
                         "Handler returned for queue item %s without ENDED "
-                        "(auto_complete=False); caller must heartbeat and complete",
+                        "(auto_complete=False); remaining subscribed until complete",
                         item.id,
                     )
+                    self._stop_heartbeat(claim, heartbeat_thread)
+                    self._wait_for_deferred(claim)
             finally:
-                stop_heartbeat.set()
-                heartbeat_thread.join(timeout=1.0)
+                self._stop_heartbeat(claim, heartbeat_thread)
+                self._unsubscribe_control(claim)
         finally:
+            self._active.pop(item.id, None)
             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):
+    def _wait_for_deferred(self, claim: _ActiveClaim) -> None:
+        while (
+            not self._stop_event.is_set()
+            and not claim.completed.is_set()
+            and not claim.control_seen.is_set()
+        ):
+            if claim.deferred_gate.wait(timeout=0.1):
+                break
+        if claim.control_seen.is_set() or claim.context.is_cancelled():
+            claim.acknowledge_control()
+
+    def _stop_heartbeat(self, claim: _ActiveClaim, heartbeat_thread: threading.Thread) -> None:
+        claim.stop_heartbeat.set()
+        if heartbeat_thread is threading.current_thread():
+            return
+        heartbeat_thread.join(timeout=1.0)
+
+    def _invoke_handler(
+        self, handler: Handler, item: QueueItem, context: TaskExecutionContext
+    ) -> Any:
+        if _accepts_context(handler):
+            return handler(item, context)
+        return handler(item)
+
+    def _heartbeat_loop(self, claim: _ActiveClaim) -> None:
+        while not claim.stop_heartbeat.wait(timeout=self._heartbeat_interval):
             try:
+                latest: QueueItem | None
                 if self._ws is not None:
-                    self._ws.heartbeat(item.id)
+                    latest = self._ws.heartbeat(claim.item.id, claim_token=claim.claim_token)
                 else:
-                    self._client.queue.heartbeat(item.id)
+                    latest = self._client.queue.heartbeat(
+                        claim.item.id, claim_token=claim.claim_token
+                    )
+                self._inspect_control(latest)
             except ContinuumError as e:
-                logger.warning("heartbeat for %s failed: %s", item.id, e)
+                logger.warning("heartbeat for %s failed: %s", claim.item.id, e)
+
+    def _inspect_control(self, latest: QueueItem | None) -> None:
+        if latest is None or latest.control_request in (None, ""):
+            return
+        claim = self._active.get(latest.id)
+        if claim is None:
+            return
+        claim.signal_control()
+
+    def _subscribe_control(self, claim: _ActiveClaim) -> None:
+        if self._ws is None:
+            return
+        try:
+            self._ws.subscribe_control(claim.item.id, claim_token=claim.claim_token)
+            claim.control_subscribed = True
+        except ContinuumError as e:
+            logger.warning("subscribeControl(%s) failed: %s", claim.item.id, e)
+
+    def _unsubscribe_control(self, claim: _ActiveClaim) -> None:
+        if self._ws is None or not claim.control_subscribed:
+            return
+        try:
+            self._ws.unsubscribe_control(claim.item.id, claim_token=claim.claim_token)
+        except ContinuumError as e:
+            logger.debug("unsubscribeControl(%s) failed: %s", claim.item.id, e)
 
     def _safe_update_status(
         self,
@@ -452,8 +561,11 @@ class TaskServer:
         status: TaskStatus,
         *,
         output: Any = None,
-    ) -> None:
-        self._safe_update_status_by_id(item.id, status, output=output)
+        claim_token: uuid.UUID | None = None,
+    ) -> QueueItem | None:
+        return self._safe_update_status_by_id(
+            item.id, status, output=output, claim_token=claim_token
+        )
 
     def _safe_update_status_by_id(
         self,
@@ -461,14 +573,23 @@ class TaskServer:
         status: TaskStatus,
         *,
         output: Any = None,
-    ) -> None:
+        claim_token: uuid.UUID | None = None,
+    ) -> QueueItem | None:
         try:
+            latest: QueueItem | None
             if self._ws is not None:
-                self._ws.update_status(queue_item_id, status, output_data=output)
+                latest = self._ws.update_status(
+                    queue_item_id, status, output_data=output, claim_token=claim_token
+                )
             else:
-                self._client.queue.update_status(queue_item_id, status, output_data=output)
+                latest = self._client.queue.update_status(
+                    queue_item_id, status, output_data=output, claim_token=claim_token
+                )
+            self._inspect_control(latest)
+            return latest
         except ContinuumError as e:
             logger.error("update_status(%s, %s) failed: %s", queue_item_id, status.value, e)
+            return None
 
     def _drain(self) -> None:
         if self._executor is None:
@@ -494,9 +615,78 @@ class TaskServer:
         logger.info("Continuum task server stopped")
 
 
+class _ActiveClaim:
+    def __init__(self, *, item: QueueItem, server: TaskServer) -> None:
+        self.item = item
+        self.server = server
+        self.context = TaskExecutionContext(item)
+        self.stop_heartbeat = threading.Event()
+        self.control_seen = threading.Event()
+        self.completed = threading.Event()
+        self.deferred_gate = threading.Event()
+        self.control_subscribed = False
+        self._control_acked = threading.Event()
+
+    @property
+    def claim_token(self) -> uuid.UUID | None:
+        return self.item.claim_token
+
+    def signal_control(self) -> None:
+        if self.control_seen.is_set():
+            return
+        self.control_seen.set()
+        self.context.request_cancel()
+        self.deferred_gate.set()
+
+    def mark_completed(self) -> None:
+        self.completed.set()
+        self.deferred_gate.set()
+
+    def acknowledge_control(self) -> None:
+        if self._control_acked.is_set():
+            return
+        self._control_acked.set()
+        self.context.request_cancel()
+        self.context.await_cleanup(max(0.001, self.server._shutdown_timeout))
+        try:
+            if self.server._ws is not None:
+                self.server._ws.complete_control(self.item.id, claim_token=self.claim_token)
+            else:
+                self.server._client.queue.complete_control(
+                    self.item.id, claim_token=self.claim_token
+                )
+            logger.info("Acknowledged control for queue item %s", self.item.id)
+        except ContinuumError as e:
+            logger.warning("completeControl(%s) failed: %s", self.item.id, e)
+        self.mark_completed()
+
+
+def _accepts_context(handler: Handler) -> bool:
+    try:
+        signature = inspect.signature(handler)
+    except (TypeError, ValueError):
+        return False
+    positional = [
+        param
+        for param in signature.parameters.values()
+        if param.kind
+        in (
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            inspect.Parameter.VAR_POSITIONAL,
+        )
+    ]
+    if any(param.kind == inspect.Parameter.VAR_POSITIONAL for param in positional):
+        return True
+    return len(positional) >= 2
+
+
 def _coerce_transport(value: TransportMode | str | None) -> TransportMode:
     if value is None:
-        return TransportMode.HTTP
+        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()

{continuum_task_server_sdk-0.1.0 → continuum_task_server_sdk-1.2.0}/continuum_task_server/websocket.py

@@ -127,6 +127,7 @@ class ContinuumWebSocketClient:
         self._events_cv = threading.Condition()
         self._event_callbacks: list[EventCallback] = []
         self._work_subscriptions: set[str] = set()
+        self._control_subscriptions: set[tuple[str, str | None]] = set()
         self._wait_subscribers: set[str] = set()
         self._seen_completed: OrderedDict[str, None] = OrderedDict()
         self._connected = threading.Event()
@@ -225,6 +226,7 @@ class ContinuumWebSocketClient:
         priority: int | None = None,
         parent_id: UUID | str | None = None,
         request_id: str | None = None,
+        idempotency_token: str | None = None,
     ) -> QueueItem:
         if task_name is None and task_item_name is None:
             raise ValueError("at least one of task_name or task_item_name is required")
@@ -235,6 +237,7 @@ class ContinuumWebSocketClient:
                 "taskItemName": task_item_name,
                 "priority": priority,
                 "inputData": _encode_input(input_data),
+                "idempotencyToken": idempotency_token,
             }
         )
         payload = self._command("/app/enqueue", body, request_id=request_id, retry=False)
@@ -255,6 +258,7 @@ class ContinuumWebSocketClient:
         mode: WaitMode | str | None = None,
         correlation: str | None = None,
         request_id: str | None = None,
+        idempotency_token: str | None = None,
     ) -> EnqueueAndSubscribeWaitResult:
         if task_name is None and task_item_name is None:
             raise ValueError("at least one of task_name or task_item_name is required")
@@ -271,6 +275,7 @@ class ContinuumWebSocketClient:
                 "subscriberId": subscriber_id,
                 "mode": mode_value,
                 "correlation": correlation,
+                "idempotencyToken": idempotency_token,
             }
         )
         payload = self._command(
@@ -295,10 +300,21 @@ class ContinuumWebSocketClient:
             return None
         return QueueItem.model_validate(item)
 
-    def heartbeat(self, queue_item_id: UUID | str, *, request_id: str | None = None) -> QueueItem:
+    def heartbeat(
+        self,
+        queue_item_id: UUID | str,
+        *,
+        claim_token: UUID | str | None = None,
+        request_id: str | None = None,
+    ) -> QueueItem:
         payload = self._command(
             "/app/heartbeat",
-            {"queueItemId": str(queue_item_id)},
+            drop_none(
+                {
+                    "queueItemId": str(queue_item_id),
+                    "claimToken": str(claim_token) if claim_token else None,
+                }
+            ),
             request_id=request_id,
             retry=True,
         )
@@ -310,6 +326,7 @@ class ContinuumWebSocketClient:
         status: TaskStatus,
         *,
         output_data: Any = None,
+        claim_token: UUID | str | None = None,
         request_id: str | None = None,
     ) -> QueueItem:
         body: dict[str, Any] = {
@@ -319,6 +336,8 @@ class ContinuumWebSocketClient:
         encoded = _encode_input(output_data)
         if encoded is not None:
             body["outputData"] = encoded
+        if claim_token is not None:
+            body["claimToken"] = str(claim_token)
         payload = self._command("/app/status", body, request_id=request_id, retry=True)
         return QueueItem.model_validate(payload["item"])
 
@@ -348,6 +367,74 @@ class ContinuumWebSocketClient:
             self._work_subscriptions.discard(task_name)
         return WorkSubscriptionResult.model_validate(payload)
 
+    def subscribe_control(
+        self,
+        queue_item_id: UUID | str,
+        *,
+        claim_token: UUID | str | None = None,
+        request_id: str | None = None,
+    ) -> WorkSubscriptionResult:
+        payload = self._command(
+            "/app/subscribeControl",
+            drop_none(
+                {
+                    "queueItemId": str(queue_item_id),
+                    "claimToken": str(claim_token) if claim_token else None,
+                }
+            ),
+            request_id=request_id,
+            retry=True,
+        )
+        with self._state_lock:
+            self._control_subscriptions.add(
+                (str(queue_item_id), str(claim_token) if claim_token else None)
+            )
+        return WorkSubscriptionResult.model_validate(payload)
+
+    def unsubscribe_control(
+        self,
+        queue_item_id: UUID | str,
+        *,
+        claim_token: UUID | str | None = None,
+        request_id: str | None = None,
+    ) -> WorkSubscriptionResult:
+        payload = self._command(
+            "/app/unsubscribeControl",
+            drop_none(
+                {
+                    "queueItemId": str(queue_item_id),
+                    "claimToken": str(claim_token) if claim_token else None,
+                }
+            ),
+            request_id=request_id,
+            retry=True,
+        )
+        with self._state_lock:
+            self._control_subscriptions.discard(
+                (str(queue_item_id), str(claim_token) if claim_token else None)
+            )
+        return WorkSubscriptionResult.model_validate(payload)
+
+    def complete_control(
+        self,
+        queue_item_id: UUID | str,
+        *,
+        claim_token: UUID | str | None = None,
+        request_id: str | None = None,
+    ) -> QueueItem:
+        payload = self._command(
+            "/app/completeControl",
+            drop_none(
+                {
+                    "queueItemId": str(queue_item_id),
+                    "claimToken": str(claim_token) if claim_token else None,
+                }
+            ),
+            request_id=request_id,
+            retry=True,
+        )
+        return QueueItem.model_validate(payload["item"])
+
     def subscribe_wait(
         self,
         subscriber_id: str,
@@ -593,6 +680,7 @@ class ContinuumWebSocketClient:
         with self._state_lock:
             work = list(self._work_subscriptions)
             waits = list(self._wait_subscribers)
+            controls = list(self._control_subscriptions)
         for task_name in work:
             try:
                 self._command_inline("/app/subscribeWork", {"taskName": task_name})
@@ -603,6 +691,21 @@ class ContinuumWebSocketClient:
                 self._command_inline("/app/subscribeWait", {"subscriberId": subscriber_id})
             except ContinuumError:
                 logger.error("Failed to restore wait subscriber %s", subscriber_id, exc_info=True)
+        for queue_item_id, claim_token in controls:
+            try:
+                self._command_inline(
+                    "/app/subscribeControl",
+                    drop_none(
+                        {
+                            "queueItemId": queue_item_id,
+                            "claimToken": claim_token,
+                        }
+                    ),
+                )
+            except ContinuumError:
+                logger.error(
+                    "Failed to restore control subscription %s", queue_item_id, exc_info=True
+                )
 
     def _command_inline(self, destination: str, body: dict[str, Any]) -> dict[str, Any]:
         rid = str(uuid.uuid4())

{continuum_task_server_sdk-0.1.0 → continuum_task_server_sdk-1.2.0}/pyproject.toml

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