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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: continuum-task-server-sdk
-Version: 1.1.0
+Version: 1.4.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
@@ -112,6 +112,14 @@ for row in db.outstanding_rows():
 
 Standalone process (no `TaskServer`): build a `ContinuumClient` with the worker key and call `client.queue.heartbeat` / `client.queue.update_status` the same way.
 
+By default the deferred claim still occupies the per-task ``concurrency`` slot until ``complete_queue_item``. Set ``occupy_slot_until_complete=False`` so the handler is still gated (you will not insert 10k rows at once) but the slot is freed when the function returns:
+
+```python
+@server.task("discord", auto_complete=False, occupy_slot_until_complete=False)
+def discord(item):
+    db.insert_await(queue_item_id=str(item.id), payload=item.input_data_json)
+```
+
 ### TaskServer options
 
 ```python

{continuum_task_server_sdk-1.1.0 → continuum_task_server_sdk-1.4.0}/README.md

@@ -83,6 +83,14 @@ for row in db.outstanding_rows():
 
 Standalone process (no `TaskServer`): build a `ContinuumClient` with the worker key and call `client.queue.heartbeat` / `client.queue.update_status` the same way.
 
+By default the deferred claim still occupies the per-task ``concurrency`` slot until ``complete_queue_item``. Set ``occupy_slot_until_complete=False`` so the handler is still gated (you will not insert 10k rows at once) but the slot is freed when the function returns:
+
+```python
+@server.task("discord", auto_complete=False, occupy_slot_until_complete=False)
+def discord(item):
+    db.insert_await(queue_item_id=str(item.id), payload=item.input_data_json)
+```
+
 ### TaskServer options
 
 ```python

{continuum_task_server_sdk-1.1.0 → continuum_task_server_sdk-1.4.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,
@@ -21,9 +22,11 @@ from .models import (
     Content,
     EnqueueAndSubscribeWaitResult,
     EventAckResult,
+    QueueDescendantsCancelResult,
     QueueEvent,
     QueueEventType,
     QueueItem,
+    QueueItemTrigger,
     TaskItem,
     TaskItemVersion,
     TaskStatus,
@@ -51,11 +54,14 @@ __all__ = [
     "ForbiddenError",
     "NotFoundError",
     "ProtocolError",
+    "QueueDescendantsCancelResult",
     "QueueEvent",
     "QueueEventType",
     "QueueItem",
+    "QueueItemTrigger",
     "RateLimitError",
     "ServerError",
+    "TaskExecutionContext",
     "TaskItem",
     "TaskItemVersion",
     "TaskServer",
@@ -71,4 +77,4 @@ __all__ = [
     "WorkSubscriptionResult",
 ]
 
-__version__ = "0.1.0"
+__version__ = "1.4.0"

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

@@ -11,7 +11,16 @@ from uuid import UUID
 import httpx
 
 from ._http import HttpClient, drop_none
-from .models import Content, QueueItem, TaskItem, TaskItemVersion, TaskStatus, TaskType
+from .exceptions import ContinuumError
+from .models import (
+    Content,
+    QueueDescendantsCancelResult,
+    QueueItem,
+    TaskItem,
+    TaskItemVersion,
+    TaskStatus,
+    TaskType,
+)
 
 
 def _encode_input(data: Any) -> str | None:
@@ -264,9 +273,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(
@@ -275,16 +290,68 @@ 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"])
+
+    def cancel_descendants(
+        self,
+        queue_item_id: UUID | str,
+        *,
+        claim_token: UUID | str | None = None,
+    ) -> QueueDescendantsCancelResult:
+        """Cancel descendants of a claimed item without cancelling the root."""
+        body = drop_none({"claimToken": str(claim_token) if claim_token else None})
+        result = self._http.post_json(
+            f"{self._QUEUE}/queue-items/{queue_item_id}/cancel-descendants",
+            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 "cancel-descendants failed")
+        return QueueDescendantsCancelResult.model_validate(result)
+
 
 class ContentStoreApi:
     _PATH = "/api/management/content-store"

continuum_task_server_sdk-1.4.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-1.1.0 → continuum_task_server_sdk-1.4.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:
@@ -103,6 +125,14 @@ class QueueItem(_Base):
             return self.output_data
 
 
+class QueueDescendantsCancelResult(_Base):
+    """Result of ``POST /api/queue/queue-items/{id}/cancel-descendants``."""
+
+    status: int = 0
+    message: str | None = None
+    cancelled_count: int | None = Field(default=None, alias="cancelledCount")
+
+
 class Content(_Base):
     """Raw content returned from a content-store or queue-item content endpoint."""
 
@@ -132,6 +162,7 @@ class QueueEventType(str, Enum):
 
     WORK_AVAILABLE = "work.available"
     COMPLETED = "queue.completed"
+    CONTROL_REQUESTED = "queue.control.requested"
 
 
 class WaitTargetStatus(_Base):
@@ -157,6 +188,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-1.1.0 → continuum_task_server_sdk-1.4.0}/continuum_task_server/server.py

@@ -9,10 +9,15 @@ 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.
+
+By default a deferred claim still occupies the per-task ``concurrency`` slot
+until complete. Pass ``occupy_slot_until_complete=False`` to free the slot when
+the handler returns so another item can be claimed while the first stays open.
 """
 
 from __future__ import annotations
 
+import inspect
 import logging
 import os
 import signal
@@ -28,13 +33,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
@@ -43,6 +49,7 @@ class _Registration:
     handler: Handler
     concurrency: int
     auto_complete: bool
+    occupy_slot_until_complete: bool
     semaphore: threading.Semaphore = field(init=False)
 
     def __post_init__(self) -> None:
@@ -62,10 +69,11 @@ 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``. The concurrency slot stays occupied until complete
+    unless ``occupy_slot_until_complete=False``.
     """
 
     def __init__(
@@ -120,6 +128,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:
@@ -132,6 +141,7 @@ class TaskServer:
         *,
         concurrency: int = 1,
         auto_complete: bool = True,
+        occupy_slot_until_complete: bool = True,
     ) -> Callable[[Handler], Handler]:
         """Decorator registering ``handler`` for queue items of task-type ``name``.
 
@@ -140,6 +150,10 @@ class TaskServer:
         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).
+
+        ``occupy_slot_until_complete`` (default True) keeps the concurrency slot until
+        that complete/fail. Set False so the slot is freed when the handler returns;
+        the claim stays in ``_active`` for complete/cancel.
         """
         if concurrency < 1:
             raise ValueError("concurrency must be >= 1")
@@ -152,6 +166,7 @@ class TaskServer:
                 handler=handler,
                 concurrency=concurrency,
                 auto_complete=auto_complete,
+                occupy_slot_until_complete=occupy_slot_until_complete,
             )
             return handler
 
@@ -164,9 +179,15 @@ class TaskServer:
         *,
         concurrency: int = 1,
         auto_complete: bool = True,
+        occupy_slot_until_complete: bool = True,
     ) -> None:
         """Imperative alternative to ``@task``."""
-        self.task(name, concurrency=concurrency, auto_complete=auto_complete)(handler)
+        self.task(
+            name,
+            concurrency=concurrency,
+            auto_complete=auto_complete,
+            occupy_slot_until_complete=occupy_slot_until_complete,
+        )(handler)
 
     def complete_queue_item(
         self,
@@ -180,7 +201,15 @@ 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()
+            self._finish_claim(claim)
 
     def fail_queue_item(
         self,
@@ -190,6 +219,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)
@@ -205,7 +238,11 @@ 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()
+            self._finish_claim(claim)
 
     @staticmethod
     def _normalize_queue_id(queue_item_id: uuid.UUID | str) -> uuid.UUID:
@@ -244,6 +281,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:
@@ -312,6 +351,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
@@ -321,6 +363,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):
@@ -395,57 +448,151 @@ class TaskServer:
             self._inflight.discard(future)
 
     def _run_item(self, registration: _Registration, item: QueueItem) -> None:
+        claim = _ActiveClaim(
+            item=item,
+            server=self,
+            occupy_slot_until_complete=registration.occupy_slot_until_complete,
+        )
+        self._active[item.id] = claim
+        slot_held = True
         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)
+                    claim.handler_returned.set()
+                    if registration.occupy_slot_until_complete:
+                        self._wait_for_deferred(claim)
+                    else:
+                        slot_held = False
+                        self._release_slot(registration)
             finally:
-                stop_heartbeat.set()
-                heartbeat_thread.join(timeout=1.0)
+                self._stop_heartbeat(claim, heartbeat_thread)
+                if slot_held:
+                    self._unsubscribe_control(claim)
         finally:
-            registration.semaphore.release()
-            if self._mode is TransportMode.WEBSOCKET:
-                self._try_claim_pending()
+            if slot_held:
+                self._active.pop(item.id, None)
+                self._release_slot(registration)
+
+    def _release_slot(self, registration: _Registration) -> None:
+        registration.semaphore.release()
+        if self._mode is TransportMode.WEBSOCKET:
+            self._try_claim_pending()
+
+    def _finish_claim(self, claim: _ActiveClaim) -> None:
+        self._unsubscribe_control(claim)
+        self._active.pop(claim.item.id, None)
+
+    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, item: QueueItem, stop_event: threading.Event) -> None:
-        while not stop_event.wait(timeout=self._heartbeat_interval):
+    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,
@@ -453,8 +600,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,
@@ -462,14 +612,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:
@@ -492,9 +651,92 @@ class TaskServer:
                 self._ws = None
         if self._owns_client:
             self._client.close()
+        for leftover in list(self._active.values()):
+            self._finish_claim(leftover)
         logger.info("Continuum task server stopped")
 
 
+class _ActiveClaim:
+    def __init__(
+        self,
+        *,
+        item: QueueItem,
+        server: TaskServer,
+        occupy_slot_until_complete: bool = True,
+    ) -> None:
+        self.item = item
+        self.server = server
+        self.occupy_slot_until_complete = occupy_slot_until_complete
+        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()
+        self.handler_returned = 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()
+        if self.handler_returned.is_set() and not self.occupy_slot_until_complete:
+            threading.Thread(
+                target=self.acknowledge_control,
+                name=f"control-ack-{self.item.id}",
+                daemon=True,
+            ).start()
+
+    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()
+        self.server._finish_claim(self)
+
+
+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:
         env = os.environ.get("CONTINUUM_TRANSPORT")

{continuum_task_server_sdk-1.1.0 → continuum_task_server_sdk-1.4.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()
@@ -299,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,
         )
@@ -314,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] = {
@@ -323,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"])
 
@@ -352,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,
@@ -597,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})
@@ -607,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-1.1.0 → continuum_task_server_sdk-1.4.0}/pyproject.toml

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