dbos 1.11.0a4__tar.gz → 1.11.0a6__tar.gz

{dbos-1.11.0a4 → dbos-1.11.0a6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 1.11.0a4
+Version: 1.11.0a6
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-1.11.0a4 → dbos-1.11.0a6}/dbos/_core.py

@@ -157,7 +157,7 @@ class WorkflowHandlePolling(Generic[R]):
 
 class WorkflowHandleAsyncTask(Generic[R]):
 
-    def __init__(self, workflow_id: str, task: asyncio.Task[R], dbos: "DBOS"):
+    def __init__(self, workflow_id: str, task: asyncio.Future[R], dbos: "DBOS"):
         self.workflow_id = workflow_id
         self.task = task
         self.dbos = dbos
@@ -698,7 +698,8 @@ async def start_workflow_async(
         return WorkflowHandleAsyncPolling(new_wf_id, dbos)
 
     coro = _execute_workflow_async(dbos, status, func, new_wf_ctx, *args, **kwargs)
-    task = asyncio.create_task(coro)
+    # Shield the workflow task from cancellation
+    task = asyncio.shield(asyncio.create_task(coro))
     return WorkflowHandleAsyncTask(new_wf_id, task, dbos)
 
 
@@ -1072,7 +1073,8 @@ def decorate_step(
 
             def on_exception(attempt: int, error: BaseException) -> float:
                 dbos.logger.warning(
-                    f"Step being automatically retried. (attempt {attempt + 1} of {attempts}). {traceback.format_exc()}"
+                    f"Step being automatically retried (attempt {attempt + 1} of {attempts})",
+                    exc_info=error,
                 )
                 ctx = assert_current_dbos_context()
                 span = ctx.get_current_span()

{dbos-1.11.0a4 → dbos-1.11.0a6}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "1.11.0a4"
+version = "1.11.0a6"
 
 [project.license]
 text = "MIT"

{dbos-1.11.0a4 → dbos-1.11.0a6}/tests/test_async.py

@@ -595,3 +595,37 @@ async def test_main_loop(dbos: DBOS, config: DBOSConfig) -> None:
     # Verify the enqueued task is submitted into the main event loop
     handle = await queue.enqueue_async(test_workflow)
     assert await handle.get_result() == id(asyncio.get_running_loop())
+
+
+@pytest.mark.asyncio
+async def test_workflow_with_task_cancellation(dbos: DBOS) -> None:
+    @DBOS.workflow()
+    async def test_workflow(duration: float) -> str:
+        await DBOS.sleep_async(duration)
+        return "completed"
+
+    # Run the workflow in an asyncio task
+    wfid = str(uuid.uuid4())
+    event = asyncio.Event()
+
+    async def run_workflow_task() -> str:
+        with SetWorkflowID(wfid):
+            handle = await DBOS.start_workflow_async(test_workflow, 1.0)
+        event.set()
+        return await handle.get_result()
+
+    task = asyncio.create_task(run_workflow_task())
+
+    # Wait for the workflow to start
+    await event.wait()
+
+    # Cancel the task
+    task.cancel()
+
+    # Verify the task was cancelled
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # Verify the workflow completes despite the task cancellation
+    handle: WorkflowHandleAsync[str] = await DBOS.retrieve_workflow_async(wfid)
+    assert await handle.get_result() == "completed"