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

{continuum_task_server_sdk-1.2.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.2.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.2.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.2.0 → continuum_task_server_sdk-1.4.0}/continuum_task_server/__init__.py

@@ -22,6 +22,7 @@ from .models import (
     Content,
     EnqueueAndSubscribeWaitResult,
     EventAckResult,
+    QueueDescendantsCancelResult,
     QueueEvent,
     QueueEventType,
     QueueItem,
@@ -53,6 +54,7 @@ __all__ = [
     "ForbiddenError",
     "NotFoundError",
     "ProtocolError",
+    "QueueDescendantsCancelResult",
     "QueueEvent",
     "QueueEventType",
     "QueueItem",
@@ -75,4 +77,4 @@ __all__ = [
     "WorkSubscriptionResult",
 ]
 
-__version__ = "0.1.0"
+__version__ = "1.4.0"

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

@@ -12,7 +12,15 @@ import httpx
 
 from ._http import HttpClient, drop_none
 from .exceptions import ContinuumError
-from .models import Content, QueueItem, TaskItem, TaskItemVersion, TaskStatus, TaskType
+from .models import (
+    Content,
+    QueueDescendantsCancelResult,
+    QueueItem,
+    TaskItem,
+    TaskItemVersion,
+    TaskStatus,
+    TaskType,
+)
 
 
 def _encode_input(data: Any) -> str | None:
@@ -326,6 +334,24 @@ class QueueApi:
             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.2.0 → continuum_task_server_sdk-1.4.0}/continuum_task_server/models.py

@@ -125,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."""
 

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

@@ -9,6 +9,10 @@ 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
@@ -45,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:
@@ -67,7 +72,8 @@ class TaskServer:
     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``.
+    ``fail_queue_item``. The concurrency slot stays occupied until complete
+    unless ``occupy_slot_until_complete=False``.
     """
 
     def __init__(
@@ -135,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``.
 
@@ -143,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")
@@ -155,6 +166,7 @@ class TaskServer:
                 handler=handler,
                 concurrency=concurrency,
                 auto_complete=auto_complete,
+                occupy_slot_until_complete=occupy_slot_until_complete,
             )
             return handler
 
@@ -167,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,
@@ -191,6 +209,7 @@ class TaskServer:
         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,
@@ -223,6 +242,7 @@ class TaskServer:
         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:
@@ -428,8 +448,13 @@ class TaskServer:
             self._inflight.discard(future)
 
     def _run_item(self, registration: _Registration, item: QueueItem) -> None:
-        claim = _ActiveClaim(item=item, server=self)
+        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)
             heartbeat_thread = threading.Thread(
@@ -482,15 +507,29 @@ class TaskServer:
                         item.id,
                     )
                     self._stop_heartbeat(claim, heartbeat_thread)
-                    self._wait_for_deferred(claim)
+                    claim.handler_returned.set()
+                    if registration.occupy_slot_until_complete:
+                        self._wait_for_deferred(claim)
+                    else:
+                        slot_held = False
+                        self._release_slot(registration)
             finally:
                 self._stop_heartbeat(claim, heartbeat_thread)
-                self._unsubscribe_control(claim)
+                if slot_held:
+                    self._unsubscribe_control(claim)
         finally:
-            self._active.pop(item.id, None)
-            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 (
@@ -612,13 +651,22 @@ 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) -> None:
+    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()
@@ -626,6 +674,7 @@ class _ActiveClaim:
         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:
@@ -637,6 +686,12 @@ class _ActiveClaim:
         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()
@@ -659,6 +714,7 @@ class _ActiveClaim:
         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:

{continuum_task_server_sdk-1.2.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.2.0"
+version = "1.4.0"
 description = "Python SDK for the Continuum Task Server"
 readme = "README.md"
 requires-python = ">=3.10"