dbos 0.26.0a21__py3-none-any.whl → 0.26.0a22__py3-none-any.whl

dbos/_client.py

@@ -1,5 +1,6 @@
 import asyncio
 import sys
+import time
 import uuid
 from typing import Any, Generic, List, Optional, TypedDict, TypeVar
 
@@ -39,6 +40,7 @@ class EnqueueOptions(TypedDict):
     queue_name: str
     workflow_id: NotRequired[str]
     app_version: NotRequired[str]
+    workflow_timeout: NotRequired[float]
 
 
 class WorkflowHandleClientPolling(Generic[R]):
@@ -107,6 +109,7 @@ class DBOSClient:
         workflow_id = options.get("workflow_id")
         if workflow_id is None:
             workflow_id = str(uuid.uuid4())
+        workflow_timeout = options.get("workflow_timeout", None)
 
         status: WorkflowStatusInternal = {
             "workflow_uuid": workflow_id,
@@ -127,6 +130,10 @@ class DBOSClient:
             "executor_id": None,
             "recovery_attempts": None,
             "app_id": None,
+            "workflow_timeout_ms": (
+                int(workflow_timeout * 1000) if workflow_timeout is not None else None
+            ),
+            "workflow_deadline_epoch_ms": None,
         }
 
         inputs: WorkflowInputs = {
@@ -190,6 +197,8 @@ class DBOSClient:
             "recovery_attempts": None,
             "app_id": None,
             "app_version": None,
+            "workflow_timeout_ms": None,
+            "workflow_deadline_epoch_ms": None,
         }
         with self._sys_db.engine.begin() as conn:
             self._sys_db.insert_workflow_status(

dbos/_context.py

@@ -93,6 +93,11 @@ class DBOSContext:
         self.assumed_role: Optional[str] = None
         self.step_status: Optional[StepStatus] = None
 
+        # A user-specified workflow timeout. Takes priority over a propagated deadline.
+        self.workflow_timeout_ms: Optional[int] = None
+        # A propagated workflow deadline.
+        self.workflow_deadline_epoch_ms: Optional[int] = None
+
     def create_child(self) -> DBOSContext:
         rv = DBOSContext()
         rv.logger = self.logger
@@ -360,11 +365,60 @@ class SetWorkflowID:
         return False  # Did not handle
 
 
+class SetWorkflowTimeout:
+    """
+    Set the workflow timeout (in seconds) to be used for the enclosed workflow invocations.
+
+    Typical Usage
+        ```
+        with SetWorkflowTimeout(<timeout in seconds>):
+            result = workflow_function(...)
+        ```
+    """
+
+    def __init__(self, workflow_timeout_sec: Optional[float]) -> None:
+        if workflow_timeout_sec and not workflow_timeout_sec > 0:
+            raise Exception(
+                f"Invalid workflow timeout {workflow_timeout_sec}. Timeouts must be positive."
+            )
+        self.created_ctx = False
+        self.workflow_timeout_ms = (
+            int(workflow_timeout_sec * 1000)
+            if workflow_timeout_sec is not None
+            else None
+        )
+        self.saved_workflow_timeout: Optional[int] = None
+
+    def __enter__(self) -> SetWorkflowTimeout:
+        # Code to create a basic context
+        ctx = get_local_dbos_context()
+        if ctx is None:
+            self.created_ctx = True
+            _set_local_dbos_context(DBOSContext())
+        ctx = assert_current_dbos_context()
+        self.saved_workflow_timeout = ctx.workflow_timeout_ms
+        ctx.workflow_timeout_ms = self.workflow_timeout_ms
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Literal[False]:
+        assert_current_dbos_context().workflow_timeout_ms = self.saved_workflow_timeout
+        # Code to clean up the basic context if we created it
+        if self.created_ctx:
+            _clear_local_dbos_context()
+        return False  # Did not handle
+
+
 class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
     def __init__(self, attributes: TracedAttributes) -> None:
         self.created_ctx = False
         self.attributes = attributes
         self.is_temp_workflow = attributes["name"] == "temp_wf"
+        self.saved_workflow_timeout: Optional[int] = None
 
     def __enter__(self) -> DBOSContext:
         # Code to create a basic context
@@ -374,6 +428,10 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
             ctx = DBOSContext()
             _set_local_dbos_context(ctx)
         assert not ctx.is_within_workflow()
+        # Unset the workflow_timeout_ms context var so it is not applied to this
+        # workflow's children (instead we propagate the deadline)
+        self.saved_workflow_timeout = ctx.workflow_timeout_ms
+        ctx.workflow_timeout_ms = None
         ctx.start_workflow(
             None, self.attributes, self.is_temp_workflow
         )  # Will get from the context's next workflow ID
@@ -388,6 +446,10 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
         ctx = assert_current_dbos_context()
         assert ctx.is_within_workflow()
         ctx.end_workflow(exc_value, self.is_temp_workflow)
+        # Restore the saved workflow timeout
+        ctx.workflow_timeout_ms = self.saved_workflow_timeout
+        # Clear any propagating timeout
+        ctx.workflow_deadline_epoch_ms = None
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()

dbos/_core.py

@@ -3,6 +3,7 @@ import functools
 import inspect
 import json
 import sys
+import threading
 import time
 import traceback
 from concurrent.futures import Future
@@ -14,11 +15,9 @@ from typing import (
     Coroutine,
     Generic,
     Optional,
-    Tuple,
     TypeVar,
     Union,
     cast,
-    overload,
 )
 
 from dbos._outcome import Immediate, NoResult, Outcome, Pending
@@ -59,7 +58,6 @@ from ._error import (
 )
 from ._registrations import (
     DEFAULT_MAX_RECOVERY_ATTEMPTS,
-    DBOSFuncInfo,
     get_config_name,
     get_dbos_class_name,
     get_dbos_func_name,
@@ -227,12 +225,14 @@ class WorkflowHandleAsyncPolling(Generic[R]):
 def _init_workflow(
     dbos: "DBOS",
     ctx: DBOSContext,
+    *,
     inputs: WorkflowInputs,
     wf_name: str,
     class_name: Optional[str],
     config_name: Optional[str],
-    temp_wf_type: Optional[str],
     queue: Optional[str],
+    workflow_timeout_ms: Optional[int],
+    workflow_deadline_epoch_ms: Optional[int],
     max_recovery_attempts: Optional[int],
 ) -> WorkflowStatusInternal:
     wfid = (
@@ -240,6 +240,15 @@ def _init_workflow(
         if len(ctx.workflow_id) > 0
         else ctx.id_assigned_for_next_workflow
     )
+
+    # In debug mode, just return the existing status
+    if dbos.debug_mode:
+        get_status_result = dbos._sys_db.get_workflow_status(wfid)
+        if get_status_result is None:
+            raise DBOSNonExistentWorkflowError(wfid)
+        return get_status_result
+
+    # Initialize a workflow status object from the context
     status: WorkflowStatusInternal = {
         "workflow_uuid": wfid,
         "status": (
@@ -267,25 +276,47 @@ def _init_workflow(
         "queue_name": queue,
         "created_at": None,
         "updated_at": None,
+        "workflow_timeout_ms": workflow_timeout_ms,
+        "workflow_deadline_epoch_ms": workflow_deadline_epoch_ms,
     }
 
     # If we have a class name, the first arg is the instance and do not serialize
     if class_name is not None:
         inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
 
-    wf_status = status["status"]
-    if dbos.debug_mode:
-        get_status_result = dbos._sys_db.get_workflow_status(wfid)
-        if get_status_result is None:
-            raise DBOSNonExistentWorkflowError(wfid)
-        wf_status = get_status_result["status"]
-    else:
-        wf_status = dbos._sys_db.init_workflow(
-            status,
-            _serialization.serialize_args(inputs),
-            max_recovery_attempts=max_recovery_attempts,
-        )
+    # Synchronously record the status and inputs for workflows
+    wf_status, workflow_deadline_epoch_ms = dbos._sys_db.init_workflow(
+        status,
+        _serialization.serialize_args(inputs),
+        max_recovery_attempts=max_recovery_attempts,
+    )
 
+    if workflow_deadline_epoch_ms is not None:
+        evt = threading.Event()
+        dbos.stop_events.append(evt)
+
+        def timeout_func() -> None:
+            try:
+                assert workflow_deadline_epoch_ms is not None
+                time_to_wait_sec = (
+                    workflow_deadline_epoch_ms - (time.time() * 1000)
+                ) / 1000
+                if time_to_wait_sec > 0:
+                    was_stopped = evt.wait(time_to_wait_sec)
+                    if was_stopped:
+                        return
+                dbos._sys_db.cancel_workflow(wfid)
+            except Exception as e:
+                dbos.logger.warning(
+                    f"Exception in timeout thread for workflow {wfid}: {e}"
+                )
+
+        timeout_thread = threading.Thread(target=timeout_func, daemon=True)
+        timeout_thread.start()
+        dbos._background_threads.append(timeout_thread)
+
+    ctx.workflow_deadline_epoch_ms = workflow_deadline_epoch_ms
+    status["workflow_deadline_epoch_ms"] = workflow_deadline_epoch_ms
     status["status"] = wf_status
     return status
 
@@ -501,6 +532,13 @@ def start_workflow(
         "kwargs": kwargs,
     }
 
+    local_ctx = get_local_dbos_context()
+    workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
+        local_ctx, queue_name
+    )
+    workflow_timeout_ms = (
+        local_ctx.workflow_timeout_ms if local_ctx is not None else None
+    )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
     ctx = new_wf_ctx
@@ -519,8 +557,9 @@ def start_workflow(
         wf_name=get_dbos_func_name(func),
         class_name=get_dbos_class_name(fi, func, args),
         config_name=get_config_name(fi, func, args),
-        temp_wf_type=get_temp_workflow_type(func),
         queue=queue_name,
+        workflow_timeout_ms=workflow_timeout_ms,
+        workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
         max_recovery_attempts=fi.max_recovery_attempts,
     )
 
@@ -583,6 +622,10 @@ async def start_workflow_async(
         "kwargs": kwargs,
     }
 
+    local_ctx = get_local_dbos_context()
+    workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
+        local_ctx, queue_name
+    )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
     ctx = new_wf_ctx
@@ -604,8 +647,9 @@ async def start_workflow_async(
         wf_name=get_dbos_func_name(func),
         class_name=get_dbos_class_name(fi, func, args),
         config_name=get_config_name(fi, func, args),
-        temp_wf_type=get_temp_workflow_type(func),
         queue=queue_name,
+        workflow_timeout_ms=workflow_timeout_ms,
+        workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
         max_recovery_attempts=fi.max_recovery_attempts,
     )
 
@@ -680,6 +724,9 @@ def workflow_wrapper(
             "kwargs": kwargs,
         }
         ctx = get_local_dbos_context()
+        workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
+            ctx, queue=None
+        )
         enterWorkflowCtxMgr = (
             EnterDBOSChildWorkflow if ctx and ctx.is_workflow() else EnterDBOSWorkflow
         )
@@ -717,8 +764,9 @@ def workflow_wrapper(
                 wf_name=get_dbos_func_name(func),
                 class_name=get_dbos_class_name(fi, func, args),
                 config_name=get_config_name(fi, func, args),
-                temp_wf_type=get_temp_workflow_type(func),
                 queue=None,
+                workflow_timeout_ms=workflow_timeout_ms,
+                workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
                 max_recovery_attempts=max_recovery_attempts,
             )
 
@@ -1212,3 +1260,24 @@ def get_event(
     else:
         # Directly call it outside of a workflow
         return dbos._sys_db.get_event(workflow_id, key, timeout_seconds)
+
+
+def _get_timeout_deadline(
+    ctx: Optional[DBOSContext], queue: Optional[str]
+) -> tuple[Optional[int], Optional[int]]:
+    if ctx is None:
+        return None, None
+    # If a timeout is explicitly specified, use it over any propagated deadline
+    if ctx.workflow_timeout_ms:
+        if queue:
+            # Queued workflows are assigned a deadline on dequeue
+            return ctx.workflow_timeout_ms, None
+        else:
+            # Otherwise, compute the deadline immediately
+            return (
+                ctx.workflow_timeout_ms,
+                int(time.time() * 1000) + ctx.workflow_timeout_ms,
+            )
+    # Otherwise, return the propagated deadline, if any
+    else:
+        return None, ctx.workflow_deadline_epoch_ms

dbos/_migrations/versions/83f3732ae8e7_workflow_timeout.py

@@ -0,0 +1,44 @@
+"""workflow_timeout
+
+Revision ID: 83f3732ae8e7
+Revises: f4b9b32ba814
+Create Date: 2025-04-16 17:05:36.642395
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "83f3732ae8e7"
+down_revision: Union[str, None] = "f4b9b32ba814"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "workflow_status",
+        sa.Column(
+            "workflow_timeout_ms",
+            sa.BigInteger(),
+            nullable=True,
+        ),
+        schema="dbos",
+    )
+    op.add_column(
+        "workflow_status",
+        sa.Column(
+            "workflow_deadline_epoch_ms",
+            sa.BigInteger(),
+            nullable=True,
+        ),
+        schema="dbos",
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("workflow_status", "workflow_deadline_epoch_ms", schema="dbos")
+    op.drop_column("workflow_status", "workflow_timeout_ms", schema="dbos")

dbos/_schemas/system_database.py

@@ -54,7 +54,9 @@ class SystemSchema:
             nullable=True,
             server_default=text("'0'::bigint"),
         ),
-        Column("queue_name", Text),
+        Column("queue_name", Text, nullable=True),
+        Column("workflow_timeout_ms", BigInteger, nullable=True),
+        Column("workflow_deadline_epoch_ms", BigInteger, nullable=True),
         Index("workflow_status_created_at_index", "created_at"),
         Index("workflow_status_executor_id_index", "executor_id"),
     )

dbos/_sys_db.py

@@ -128,6 +128,11 @@ class WorkflowStatusInternal(TypedDict):
     app_version: Optional[str]
     app_id: Optional[str]
     recovery_attempts: Optional[int]
+    # The start-to-close timeout of the workflow in ms
+    workflow_timeout_ms: Optional[int]
+    # The deadline of a workflow, computed by adding its timeout to its start time.
+    # Deadlines propagate to children. When the deadline is reached, the workflow is cancelled.
+    workflow_deadline_epoch_ms: Optional[int]
 
 
 class RecordedResult(TypedDict):
@@ -328,10 +333,11 @@ class SystemDatabase:
         conn: sa.Connection,
         *,
         max_recovery_attempts: Optional[int],
-    ) -> WorkflowStatuses:
+    ) -> tuple[WorkflowStatuses, Optional[int]]:
         if self._debug_mode:
             raise Exception("called insert_workflow_status in debug mode")
         wf_status: WorkflowStatuses = status["status"]
+        workflow_deadline_epoch_ms: Optional[int] = status["workflow_deadline_epoch_ms"]
 
         cmd = (
             pg.insert(SystemSchema.workflow_status)
@@ -354,6 +360,8 @@ class SystemDatabase:
                 recovery_attempts=(
                     1 if wf_status != WorkflowStatusString.ENQUEUED.value else 0
                 ),
+                workflow_timeout_ms=status["workflow_timeout_ms"],
+                workflow_deadline_epoch_ms=status["workflow_deadline_epoch_ms"],
             )
             .on_conflict_do_update(
                 index_elements=["workflow_uuid"],
@@ -367,7 +375,7 @@ class SystemDatabase:
             )
         )
 
-        cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name)  # type: ignore
+        cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.workflow_deadline_epoch_ms, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name)  # type: ignore
 
         results = conn.execute(cmd)
 
@@ -377,17 +385,18 @@ class SystemDatabase:
             # A mismatch indicates a workflow starting with the same UUID but different functions, which would throw an exception.
             recovery_attempts: int = row[0]
             wf_status = row[1]
+            workflow_deadline_epoch_ms = row[2]
             err_msg: Optional[str] = None
-            if row[2] != status["name"]:
-                err_msg = f"Workflow already exists with a different function name: {row[2]}, but the provided function name is: {status['name']}"
-            elif row[3] != status["class_name"]:
-                err_msg = f"Workflow already exists with a different class name: {row[3]}, but the provided class name is: {status['class_name']}"
-            elif row[4] != status["config_name"]:
-                err_msg = f"Workflow already exists with a different config name: {row[4]}, but the provided config name is: {status['config_name']}"
-            elif row[5] != status["queue_name"]:
+            if row[3] != status["name"]:
+                err_msg = f"Workflow already exists with a different function name: {row[3]}, but the provided function name is: {status['name']}"
+            elif row[4] != status["class_name"]:
+                err_msg = f"Workflow already exists with a different class name: {row[4]}, but the provided class name is: {status['class_name']}"
+            elif row[5] != status["config_name"]:
+                err_msg = f"Workflow already exists with a different config name: {row[5]}, but the provided config name is: {status['config_name']}"
+            elif row[6] != status["queue_name"]:
                 # This is a warning because a different queue name is not necessarily an error.
                 dbos_logger.warning(
-                    f"Workflow already exists in queue: {row[5]}, but the provided queue name is: {status['queue_name']}. The queue is not updated."
+                    f"Workflow already exists in queue: {row[6]}, but the provided queue name is: {status['queue_name']}. The queue is not updated."
                 )
             if err_msg is not None:
                 raise DBOSConflictingWorkflowError(status["workflow_uuid"], err_msg)
@@ -427,7 +436,7 @@ class SystemDatabase:
                     status["workflow_uuid"], max_recovery_attempts
                 )
 
-        return wf_status
+        return wf_status, workflow_deadline_epoch_ms
 
     def update_workflow_status(
         self,
@@ -485,6 +494,18 @@ class SystemDatabase:
         if self._debug_mode:
             raise Exception("called cancel_workflow in debug mode")
         with self.engine.begin() as c:
+            # Check the status of the workflow. If it is complete, do nothing.
+            row = c.execute(
+                sa.select(
+                    SystemSchema.workflow_status.c.status,
+                ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+            ).fetchone()
+            if (
+                row is None
+                or row[0] == WorkflowStatusString.SUCCESS.value
+                or row[0] == WorkflowStatusString.ERROR.value
+            ):
+                return
             # Remove the workflow from the queues table so it does not block the table
             c.execute(
                 sa.delete(SystemSchema.workflow_queue).where(
@@ -531,11 +552,15 @@ class SystemDatabase:
                     queue_name=INTERNAL_QUEUE_NAME,
                 )
             )
-            # Set the workflow's status to ENQUEUED and clear its recovery attempts.
+            # Set the workflow's status to ENQUEUED and clear its recovery attempts and deadline.
             c.execute(
                 sa.update(SystemSchema.workflow_status)
                 .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
-                .values(status=WorkflowStatusString.ENQUEUED.value, recovery_attempts=0)
+                .values(
+                    status=WorkflowStatusString.ENQUEUED.value,
+                    recovery_attempts=0,
+                    workflow_deadline_epoch_ms=None,
+                )
             )
 
     def get_max_function_id(self, workflow_uuid: str) -> Optional[int]:
@@ -648,6 +673,8 @@ class SystemDatabase:
                     SystemSchema.workflow_status.c.updated_at,
                     SystemSchema.workflow_status.c.application_version,
                     SystemSchema.workflow_status.c.application_id,
+                    SystemSchema.workflow_status.c.workflow_deadline_epoch_ms,
+                    SystemSchema.workflow_status.c.workflow_timeout_ms,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -671,6 +698,8 @@ class SystemDatabase:
                 "updated_at": row[12],
                 "app_version": row[13],
                 "app_id": row[14],
+                "workflow_deadline_epoch_ms": row[15],
+                "workflow_timeout_ms": row[16],
             }
             return status
 
@@ -1100,50 +1129,56 @@ class SystemDatabase:
         *,
         conn: Optional[sa.Connection] = None,
     ) -> Optional[RecordedResult]:
-        # Retrieve the status of the workflow. Additionally, if this step
-        # has run before, retrieve its name, output, and error.
-        sql = (
-            sa.select(
-                SystemSchema.workflow_status.c.status,
-                SystemSchema.operation_outputs.c.output,
-                SystemSchema.operation_outputs.c.error,
-                SystemSchema.operation_outputs.c.function_name,
-            )
-            .select_from(
-                SystemSchema.workflow_status.outerjoin(
-                    SystemSchema.operation_outputs,
-                    (
-                        SystemSchema.workflow_status.c.workflow_uuid
-                        == SystemSchema.operation_outputs.c.workflow_uuid
-                    )
-                    & (SystemSchema.operation_outputs.c.function_id == function_id),
-                )
-            )
-            .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+        # First query: Retrieve the workflow status
+        workflow_status_sql = sa.select(
+            SystemSchema.workflow_status.c.status,
+        ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+
+        # Second query: Retrieve operation outputs if they exist
+        operation_output_sql = sa.select(
+            SystemSchema.operation_outputs.c.output,
+            SystemSchema.operation_outputs.c.error,
+            SystemSchema.operation_outputs.c.function_name,
+        ).where(
+            (SystemSchema.operation_outputs.c.workflow_uuid == workflow_id)
+            & (SystemSchema.operation_outputs.c.function_id == function_id)
         )
-        # If in a transaction, use the provided connection
-        rows: Sequence[Any]
+
+        # Execute both queries
         if conn is not None:
-            rows = conn.execute(sql).all()
+            workflow_status_rows = conn.execute(workflow_status_sql).all()
+            operation_output_rows = conn.execute(operation_output_sql).all()
         else:
             with self.engine.begin() as c:
-                rows = c.execute(sql).all()
-        assert len(rows) > 0, f"Error: Workflow {workflow_id} does not exist"
-        workflow_status, output, error, recorded_function_name = (
-            rows[0][0],
-            rows[0][1],
-            rows[0][2],
-            rows[0][3],
-        )
+                workflow_status_rows = c.execute(workflow_status_sql).all()
+                operation_output_rows = c.execute(operation_output_sql).all()
+
+        # Check if the workflow exists
+        assert (
+            len(workflow_status_rows) > 0
+        ), f"Error: Workflow {workflow_id} does not exist"
+
+        # Get workflow status
+        workflow_status = workflow_status_rows[0][0]
+
         # If the workflow is cancelled, raise the exception
         if workflow_status == WorkflowStatusString.CANCELLED.value:
             raise DBOSWorkflowCancelledError(
                 f"Workflow {workflow_id} is cancelled. Aborting function."
             )
-        # If there is no row for the function, return None
-        if recorded_function_name is None:
+
+        # If there are no operation outputs, return None
+        if not operation_output_rows:
             return None
-        # If the provided and recorded function name are different, throw an exception.
+
+        # Extract operation output data
+        output, error, recorded_function_name = (
+            operation_output_rows[0][0],
+            operation_output_rows[0][1],
+            operation_output_rows[0][2],
+        )
+
+        # If the provided and recorded function name are different, throw an exception
         if function_name != recorded_function_name:
             raise DBOSUnexpectedStepError(
                 workflow_id=workflow_id,
@@ -1151,6 +1186,7 @@ class SystemDatabase:
                 expected_name=function_name,
                 recorded_name=recorded_function_name,
             )
+
         result: RecordedResult = {
             "output": output,
             "error": error,
@@ -1699,6 +1735,17 @@ class SystemDatabase:
                         status=WorkflowStatusString.PENDING.value,
                         application_version=app_version,
                         executor_id=executor_id,
+                        # If a timeout is set, set the deadline on dequeue
+                        workflow_deadline_epoch_ms=sa.case(
+                            (
+                                SystemSchema.workflow_status.c.workflow_timeout_ms.isnot(
+                                    None
+                                ),
+                                sa.func.extract("epoch", sa.func.now()) * 1000
+                                + SystemSchema.workflow_status.c.workflow_timeout_ms,
+                            ),
+                            else_=SystemSchema.workflow_status.c.workflow_deadline_epoch_ms,
+                        ),
                     )
                 )
                 if res.rowcount > 0:
@@ -1821,12 +1868,12 @@ class SystemDatabase:
         inputs: str,
         *,
         max_recovery_attempts: Optional[int],
-    ) -> WorkflowStatuses:
+    ) -> tuple[WorkflowStatuses, Optional[int]]:
         """
         Synchronously record the status and inputs for workflows in a single transaction
         """
         with self.engine.begin() as conn:
-            wf_status = self.insert_workflow_status(
+            wf_status, workflow_deadline_epoch_ms = self.insert_workflow_status(
                 status, conn, max_recovery_attempts=max_recovery_attempts
             )
             # TODO: Modify the inputs if they were changed by `update_workflow_inputs`
@@ -1837,7 +1884,7 @@ class SystemDatabase:
                 and wf_status == WorkflowStatusString.ENQUEUED.value
             ):
                 self.enqueue(status["workflow_uuid"], status["queue_name"], conn)
-        return wf_status
+        return wf_status, workflow_deadline_epoch_ms
 
 
 def reset_system_database(config: ConfigFile) -> None:

{dbos-0.26.0a21.dist-info → dbos-0.26.0a22.dist-info}/METADATA

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

{dbos-0.26.0a21.dist-info → dbos-0.26.0a22.dist-info}/RECORD

@@ -1,17 +1,17 @@
-dbos-0.26.0a21.dist-info/METADATA,sha256=6JPLTUn5uCaKpHI_sEis2zJWOrNqsRUIQ2_4h7ZzWfw,5554
-dbos-0.26.0a21.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-dbos-0.26.0a21.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-0.26.0a21.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.26.0a22.dist-info/METADATA,sha256=Tc4eDxO8J5OnULO6LrEY6mzhVwYZfmFN_okn3ZnSRLc,5554
+dbos-0.26.0a22.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
+dbos-0.26.0a22.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-0.26.0a22.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=VoGS7H9GVtNAnD2S4zseIEioS1dNIJXRovQ4oHlg8og,842
 dbos/__main__.py,sha256=G7Exn-MhGrVJVDbgNlpzhfh8WMX_72t3_oJaFT9Lmt8,653
 dbos/_admin_server.py,sha256=RrbABfR1D3p9c_QLrCSrgFuYce6FKi0fjMRIYLjO_Y8,9038
 dbos/_app_db.py,sha256=obNlgC9IZ20y8tqQeA1q4TjceG3jBFalxz70ieDOWCA,11332
 dbos/_classproperty.py,sha256=f0X-_BySzn3yFDRKB2JpCbLYQ9tLwt1XftfshvY7CBs,626
-dbos/_client.py,sha256=S3tejQ7xAJ9wjo1PhQ0P3UYuloDOdZqXsQwE4YjAQ8s,12124
+dbos/_client.py,sha256=f1n5bbtVO-Mf5dDvI3sNlozxHSUfstWtgPirSqv1kpE,12518
 dbos/_conductor/conductor.py,sha256=HYzVL29IMMrs2Mnms_7cHJynCnmmEN5SDQOMjzn3UoU,16840
 dbos/_conductor/protocol.py,sha256=zEKIuOQdIaSduNqfZKpo8PSD9_1oNpKIPnBNCu3RUyE,6681
-dbos/_context.py,sha256=I8sLkdKTTkZEz7wG-MjynaQB6XEF2bLXuwNksiauP7w,19430
-dbos/_core.py,sha256=SecObOKLjNinNAXDcYVMVUURHcoaPe0js-axLMMNwqY,45098
+dbos/_context.py,sha256=aHzJxO7LLAz9w3G2dkZnOcFW_GG-Yaxd02AaoLu4Et8,21861
+dbos/_core.py,sha256=ylTVSv02h2M5SmDgYEJAZmNiKX35zPq0z-9WA-f4byY,47900
 dbos/_croniter.py,sha256=XHAyUyibs_59sJQfSNWkP7rqQY6_XrlfuuCxk4jYqek,47559
 dbos/_dbos.py,sha256=bbio_FjBfU__Zk1BFegfS16IrPPejFxOKm5rUg5nW1o,47185
 dbos/_dbos_config.py,sha256=m05IFjM0jSwZBsnFMF_4qP2JkjVFc0gqyM2tnotXq20,20636
@@ -29,6 +29,7 @@ dbos/_migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy6HUS9NFvDgl93dM
 dbos/_migrations/versions/04ca4f231047_workflow_queues_executor_id.py,sha256=ICLPl8CN9tQXMsLDsAj8z1TsL831-Z3F8jSBvrR-wyw,736
 dbos/_migrations/versions/50f3227f0b4b_fix_job_queue.py,sha256=ZBYrtTdxy64HxIAlOes89fVIk2P1gNaJack7wuC_epg,873
 dbos/_migrations/versions/5c361fc04708_added_system_tables.py,sha256=Xr9hBDJjkAtymlauOmAy00yUHj0VVUaEz7kNwEM9IwE,6403
+dbos/_migrations/versions/83f3732ae8e7_workflow_timeout.py,sha256=Q_R35pb8AfVI3sg5mzKwyoPfYB88Ychcc8gwxpM9R7A,1035
 dbos/_migrations/versions/a3b18ad34abe_added_triggers.py,sha256=Rv0ZsZYZ_WdgGEULYsPfnp4YzaO5L198gDTgYY39AVA,2022
 dbos/_migrations/versions/d76646551a6b_job_queue_limiter.py,sha256=8PyFi8rd6CN-mUro43wGhsg5wcQWKZPRHD6jw8R5pVc,986
 dbos/_migrations/versions/d76646551a6c_workflow_queue.py,sha256=G942nophZ2uC2vc4hGBC02Ptng1715roTjY3xiyzZU4,729
@@ -43,9 +44,9 @@ dbos/_roles.py,sha256=iOsgmIAf1XVzxs3gYWdGRe1B880YfOw5fpU7Jwx8_A8,2271
 dbos/_scheduler.py,sha256=SR1oRZRcVzYsj-JauV2LA8JtwTkt8mru7qf6H1AzQ1U,2027
 dbos/_schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_schemas/application_database.py,sha256=SypAS9l9EsaBHFn9FR8jmnqt01M74d9AF1AMa4m2hhI,1040
-dbos/_schemas/system_database.py,sha256=W9eSpL7SZzQkxcEZ4W07BOcwkkDr35b9oCjUOgfHWek,5336
+dbos/_schemas/system_database.py,sha256=aChSK7uLECD-v-7BZeOfuZFbtWayllaS3PaowaKDHwY,5490
 dbos/_serialization.py,sha256=YCYv0qKAwAZ1djZisBC7khvKqG-5OcIv9t9EC5PFIog,1743
-dbos/_sys_db.py,sha256=M3BVJVhG0YXkLhw5axSrKjBN1AOS3KmvgWEYn2l94pw,78203
+dbos/_sys_db.py,sha256=SjYTleSEPtZVrPRimgXKeIvTjY8VN9G9jlgbcPT8ghg,80631
 dbos/_templates/dbos-db-starter/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/_templates/dbos-db-starter/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_templates/dbos-db-starter/__package/main.py,sha256=nJMN3ZD2lmwg4Dcgmiwqc-tQGuCJuJal2Xl85iA277U,2453
@@ -65,4 +66,4 @@ dbos/cli/cli.py,sha256=1qCTs__A9LOEfU44XZ6TufwmRwe68ZEwbWEPli3vnVM,17873
 dbos/dbos-config.schema.json,sha256=i7jcxXqByKq0Jzv3nAUavONtj03vTwj6vWP4ylmBr8o,5694
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.26.0a21.dist-info/RECORD,,
+dbos-0.26.0a22.dist-info/RECORD,,