Populate Request Deadline in Nexus Operation Contexts (#1277)

* Nexus error ser.

* Update to nexus-rpc 1.4.0

* Populate Nexus request deadline into operation contexts when present on the task received from Core.

* Add test to confirm request deadline is present in cancel operation contexts

* Update request deadline tests for workflow_run_operation to reflect how users will invoke rather than using a private api

* refactor request deadline out of if branches in nexus worker

* Fix bad rebase in pyproject.toml

* gen protos

* fix some rebase mistakes. Update test to not use the removed http client

* Add request deadline tests for StartOperationContext and CancelOperationContext

Adds tests verifying request_deadline is accessible and timezone-aware (UTC)
in both StartOperationContext and CancelOperationContext. Also strengthens the
existing WorkflowRunOperationContext test to validate the deadline is a proper
UTC datetime rather than just truthy.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Move request_deadline assertion from handler to test body in workflow_run_operation test

Captures the deadline into a list on the handler instance and asserts in
the test body, consistent with the start/cancel operation deadline tests.
This gives clear pytest failure messages instead of opaque handler errors.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Fold start and cancel tests into a single test. Replace sleep with an asyncio.Event to avoid potential flakes

---------

Co-authored-by: Quinn Klassen <klassenq@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 3 people

temporalio/worker/_nexus.py

@@ -8,6 +8,7 @@
 import threading
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
+from datetime import datetime, timezone
 from functools import reduce
 from typing import (
     Any,
@@ -119,14 +120,22 @@ async def raise_from_exception_queue() -> NoReturn:
 
                 if nexus_task.HasField("task"):
                     task = nexus_task.task
+                    request_deadline = (
+                        nexus_task.request_deadline.ToDatetime().replace(
+                            tzinfo=timezone.utc
+                        )
+                        if nexus_task.HasField("request_deadline")
+                        else None
+                    )
                     if task.request.HasField("start_operation"):
                         task_cancellation = _NexusTaskCancellation()
                         start_op_task = asyncio.create_task(
                             self._handle_start_operation_task(
-                                task.task_token,
-                                task.request.start_operation,
-                                dict(task.request.header),
-                                task_cancellation,
+                                task_token=task.task_token,
+                                start_request=task.request.start_operation,
+                                headers=dict(task.request.header),
+                                task_cancellation=task_cancellation,
+                                request_deadline=request_deadline,
                             )
                         )
                         self._running_tasks[task.task_token] = _RunningNexusTask(
@@ -136,10 +145,11 @@ async def raise_from_exception_queue() -> NoReturn:
                         task_cancellation = _NexusTaskCancellation()
                         cancel_op_task = asyncio.create_task(
                             self._handle_cancel_operation_task(
-                                task.task_token,
-                                task.request.cancel_operation,
-                                dict(task.request.header),
-                                task_cancellation,
+                                task_token=task.task_token,
+                                request=task.request.cancel_operation,
+                                headers=dict(task.request.header),
+                                task_cancellation=task_cancellation,
+                                request_deadline=request_deadline,
                             )
                         )
                         self._running_tasks[task.task_token] = _RunningNexusTask(
@@ -209,6 +219,7 @@ async def _handle_cancel_operation_task(
         request: temporalio.api.nexus.v1.CancelOperationRequest,
         headers: Mapping[str, str],
         task_cancellation: nexusrpc.handler.OperationTaskCancellation,
+        request_deadline: datetime | None,
     ) -> None:
         """Handle a cancel operation task.
 
@@ -226,6 +237,7 @@ async def _handle_cancel_operation_task(
             operation=request.operation,
             headers=headers,
             task_cancellation=task_cancellation,
+            request_deadline=request_deadline,
         )
         temporalio.nexus._operation_context._TemporalCancelOperationContext(
             info=lambda: Info(task_queue=self._task_queue),
@@ -276,6 +288,7 @@ async def _handle_start_operation_task(
         start_request: temporalio.api.nexus.v1.StartOperationRequest,
         headers: Mapping[str, str],
         task_cancellation: nexusrpc.handler.OperationTaskCancellation,
+        request_deadline: datetime | None,
     ) -> None:
         """Handle a start operation task.
 
@@ -285,7 +298,7 @@ async def _handle_start_operation_task(
         try:
             try:
                 start_response = await self._start_operation(
-                    start_request, headers, task_cancellation
+                    start_request, headers, task_cancellation, request_deadline
                 )
             except asyncio.CancelledError:
                 completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
@@ -328,6 +341,7 @@ async def _start_operation(
         start_request: temporalio.api.nexus.v1.StartOperationRequest,
         headers: Mapping[str, str],
         cancellation: nexusrpc.handler.OperationTaskCancellation,
+        request_deadline: datetime | None,
     ) -> temporalio.api.nexus.v1.StartOperationResponse:
         """Invoke the Nexus handler's start_operation method and construct the StartOperationResponse.
 
@@ -352,6 +366,7 @@ async def _start_operation(
             ],
             callback_headers=dict(start_request.callback_header),
             task_cancellation=cancellation,
+            request_deadline=request_deadline,
         )
         temporalio.nexus._operation_context._TemporalStartOperationContext(
             nexus_context=ctx,

tests/nexus/test_workflow_caller.py

@@ -7,6 +7,7 @@
 import uuid
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
+from datetime import datetime, timezone
 from enum import IntEnum
 from typing import Any
 from urllib.request import urlopen
@@ -173,6 +174,11 @@ class HeaderTestService:
     cancellable_operation: nexusrpc.Operation[None, str]
 
 
+@nexusrpc.service
+class RequestDeadlineService:
+    cancellable_op: nexusrpc.Operation[None, str]
+
+
 # -----------------------------------------------------------------------------
 # Service implementation
 #
@@ -335,6 +341,46 @@ def cancellable_operation(self) -> OperationHandler[None, str]:
         return CancellableOperationHandler(self.cancel_headers_received)
 
 
+class CancellableDeadlineOperationHandler(OperationHandler[None, str]):
+    """Operation handler that captures request_deadline from start and cancel contexts."""
+
+    def __init__(
+        self,
+        start_deadlines_received: list[datetime | None],
+        cancel_deadlines_received: list[datetime | None],
+        cancel_received: asyncio.Event,
+    ) -> None:
+        self._start_deadlines_received = start_deadlines_received
+        self._cancel_deadlines_received = cancel_deadlines_received
+        self._cancel_received = cancel_received
+
+    async def start(
+        self, ctx: StartOperationContext, input: None
+    ) -> StartOperationResultAsync:
+        self._start_deadlines_received.append(ctx.request_deadline)
+        return StartOperationResultAsync("test-token")
+
+    async def cancel(self, ctx: CancelOperationContext, token: str) -> None:
+        self._cancel_deadlines_received.append(ctx.request_deadline)
+        self._cancel_received.set()
+
+
+@service_handler(service=RequestDeadlineService)
+class RequestDeadlineServiceImpl:
+    def __init__(self) -> None:
+        self.start_deadlines_received: list[datetime | None] = []
+        self.cancel_deadlines_received: list[datetime | None] = []
+        self.cancel_received = asyncio.Event()
+
+    @operation_handler
+    def cancellable_op(self) -> OperationHandler[None, str]:
+        return CancellableDeadlineOperationHandler(
+            self.start_deadlines_received,
+            self.cancel_deadlines_received,
+            self.cancel_received,
+        )
+
+
 # -----------------------------------------------------------------------------
 # Caller workflow
 #
@@ -570,6 +616,30 @@ async def run(self, input: CancelHeaderTestCallerWfInput) -> None:
         await asyncio.sleep(0.1)
 
 
+@workflow.defn
+class CancelDeadlineCallerWorkflow:
+    """Workflow that starts a cancellable operation and then cancels it, for deadline testing."""
+
+    @workflow.run
+    async def run(self, task_queue: str) -> None:
+        nexus_client = workflow.create_nexus_client(
+            service=RequestDeadlineService,
+            endpoint=make_nexus_endpoint_name(task_queue),
+        )
+        op_handle = await nexus_client.start_operation(
+            RequestDeadlineService.cancellable_op,
+            None,
+            cancellation_type=workflow.NexusOperationCancellationType.WAIT_REQUESTED,
+        )
+        # Request cancellation - this sends a cancel operation to the handler
+        op_handle.cancel()
+
+        try:
+            await op_handle
+        except NexusOperationError:
+            pass
+
+
 @workflow.defn
 class WorkflowRunHeaderTestCallerWorkflow:
     """Workflow that calls a workflow_run_operation and verifies headers."""
@@ -2172,3 +2242,47 @@ async def test_task_executor_operation_cancel_method(
             # Verify the workflow completed successfully
             result = await caller_wf_handle.result()
             assert result == "cancelled_successfully"
+
+
+async def test_request_deadline_is_accessible_in_operation(
+    client: Client,
+    env: WorkflowEnvironment,
+):
+    """Test that request_deadline is accessible in StartOperationContext."""
+    if env.supports_time_skipping:
+        pytest.skip("Nexus tests don't work with time-skipping server")
+
+    task_queue = str(uuid.uuid4())
+    service_handler = RequestDeadlineServiceImpl()
+
+    async with Worker(
+        client,
+        nexus_service_handlers=[service_handler],
+        workflows=[CancelDeadlineCallerWorkflow],
+        task_queue=task_queue,
+    ):
+        endpoint_name = make_nexus_endpoint_name(task_queue)
+        await env.create_nexus_endpoint(endpoint_name, task_queue)
+
+        await client.execute_workflow(
+            CancelDeadlineCallerWorkflow.run,
+            task_queue,
+            id=str(uuid.uuid4()),
+            task_queue=task_queue,
+        )
+
+        assert len(service_handler.start_deadlines_received) == 1
+        deadline = service_handler.start_deadlines_received[0]
+        assert (
+            deadline is not None
+        ), "request_deadline should be set in StartOperationContext"
+        assert deadline.tzinfo is timezone.utc, "request_deadline should be in utc"
+
+        await asyncio.wait_for(service_handler.cancel_received.wait(), 1)
+
+        assert len(service_handler.cancel_deadlines_received) == 1
+        deadline = service_handler.cancel_deadlines_received[0]
+        assert (
+            deadline is not None
+        ), "request_deadline should be set in CancelOperationContext"
+        assert deadline.tzinfo is timezone.utc, "request_deadline should be in utc"

tests/nexus/test_workflow_run_operation.py

@@ -1,5 +1,6 @@
 import uuid
 from dataclasses import dataclass
+from datetime import datetime, timezone
 from typing import Any
 
 import nexusrpc
@@ -13,9 +14,9 @@
 )
 from nexusrpc.handler._decorators import operation_handler
 
-from temporalio import workflow
+from temporalio import nexus, workflow
 from temporalio.client import Client
-from temporalio.nexus import WorkflowRunOperationContext
+from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation
 from temporalio.nexus._operation_handlers import WorkflowRunOperationHandler
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
@@ -59,6 +60,42 @@ def op(self) -> OperationHandler[Input, str]:
         return MyOperation()
 
 
+@service
+class RequestDeadlineService:
+    op: Operation[Input, str]
+
+
+@service_handler(service=RequestDeadlineService)
+class RequestDeadlineHandler:
+    def __init__(self) -> None:
+        self.start_deadlines_received: list[datetime | None] = []
+
+    @workflow_run_operation
+    async def op(
+        self, ctx: WorkflowRunOperationContext, input: Input
+    ) -> nexus.WorkflowHandle[str]:
+        self.start_deadlines_received.append(ctx.request_deadline)
+        return await ctx.start_workflow(
+            EchoWorkflow.run,
+            input.value,
+            id=str(uuid.uuid4()),
+        )
+
+
+@workflow.defn
+class RequestDeadlineWorkflow:
+    @workflow.run
+    async def run(self, input: Input, task_queue: str) -> str:
+        client = workflow.create_nexus_client(
+            service=RequestDeadlineService,
+            endpoint=make_nexus_endpoint_name(task_queue),
+        )
+        return await client.execute_operation(
+            RequestDeadlineService.op,
+            input,
+        )
+
+
 @service
 class Service:
     op: Operation[Input, str]
@@ -116,3 +153,36 @@ async def test_workflow_run_operation(
             task_queue=task_queue,
         )
         assert result == "test"
+
+
+async def test_request_deadline_is_accessible_in_workflow_run_operation(
+    client: Client,
+    env: WorkflowEnvironment,
+):
+    """Test that request_deadline is accessible in WorkflowRunOperationContext."""
+    if env.supports_time_skipping:
+        pytest.skip("Nexus tests don't work with time-skipping server")
+
+    task_queue = str(uuid.uuid4())
+    endpoint_name = make_nexus_endpoint_name(task_queue)
+    await env.create_nexus_endpoint(endpoint_name, task_queue)
+    service_handler = RequestDeadlineHandler()
+    async with Worker(
+        env.client,
+        task_queue=task_queue,
+        nexus_service_handlers=[service_handler],
+        workflows=[RequestDeadlineWorkflow, EchoWorkflow],
+    ):
+        await client.execute_workflow(
+            RequestDeadlineWorkflow.run,
+            args=[Input(value="test"), task_queue],
+            task_queue=task_queue,
+            id=str(uuid.uuid4()),
+        )
+
+        assert len(service_handler.start_deadlines_received) == 1
+        deadline = service_handler.start_deadlines_received[0]
+        assert (
+            deadline is not None
+        ), "request_deadline should be set in WorkflowRunOperationContext"
+        assert deadline.tzinfo is timezone.utc, "request_deadline should be in utc"