Expose Nexus Endpoint in a Nexus Operation Handler (#1437)

Expose Nexus Endpoint in a Nexus Operation Handler

Author: Quinn-With-Two-Ns

temporalio/bridge/proto/nexus/nexus_pb2.py

@@ -240,6 +240,7 @@ class NexusTask(google.protobuf.message.Message):
     TASK_FIELD_NUMBER: builtins.int
     CANCEL_TASK_FIELD_NUMBER: builtins.int
     REQUEST_DEADLINE_FIELD_NUMBER: builtins.int
+    ENDPOINT_FIELD_NUMBER: builtins.int
     @property
     def task(
         self,
@@ -265,13 +266,18 @@ class NexusTask(google.protobuf.message.Message):
         Only set when variant is `task` and the header was present with a valid value.
         Represented as an absolute timestamp.
         """
+    endpoint: builtins.str
+    """The endpoint this request was addressed to. Extracted from the request for convenient access.
+    Only set when variant is `task`.
+    """
     def __init__(
         self,
         *,
         task: temporalio.api.workflowservice.v1.request_response_pb2.PollNexusTaskQueueResponse
         | None = ...,
         cancel_task: global___CancelNexusTask | None = ...,
         request_deadline: google.protobuf.timestamp_pb2.Timestamp | None = ...,
+        endpoint: builtins.str = ...,
     ) -> None: ...
     def HasField(
         self,
@@ -291,6 +297,8 @@ class NexusTask(google.protobuf.message.Message):
         field_name: typing_extensions.Literal[
             "cancel_task",
             b"cancel_task",
+            "endpoint",
+            b"endpoint",
             "request_deadline",
             b"request_deadline",
             "task",

temporalio/bridge/proto/nexus/nexus_pb2.pyi

@@ -79,6 +79,9 @@ class Info:
     Retrieved inside a Nexus operation handler via :py:func:`info`.
     """
 
+    endpoint: str
+    """The endpoint this Nexus request was addressed to."""
+
     namespace: str
     """The namespace of the worker handling this Nexus operation."""
 

temporalio/bridge/sdk-core

@@ -140,6 +140,7 @@ async def raise_from_exception_queue() -> NoReturn:
                                 headers=dict(task.request.header),
                                 task_cancellation=task_cancellation,
                                 request_deadline=request_deadline,
+                                endpoint=nexus_task.endpoint,
                             )
                         )
                         self._running_tasks[task.task_token] = _RunningNexusTask(
@@ -154,6 +155,7 @@ async def raise_from_exception_queue() -> NoReturn:
                                 headers=dict(task.request.header),
                                 task_cancellation=task_cancellation,
                                 request_deadline=request_deadline,
+                                endpoint=nexus_task.endpoint,
                             )
                         )
                         self._running_tasks[task.task_token] = _RunningNexusTask(
@@ -224,6 +226,7 @@ async def _handle_cancel_operation_task(
         headers: Mapping[str, str],
         task_cancellation: nexusrpc.handler.OperationTaskCancellation,
         request_deadline: datetime | None,
+        endpoint: str,
     ) -> None:
         """Handle a cancel operation task.
 
@@ -244,7 +247,11 @@ async def _handle_cancel_operation_task(
             request_deadline=request_deadline,
         )
         temporalio.nexus._operation_context._TemporalCancelOperationContext(
-            info=lambda: Info(namespace=self._namespace, task_queue=self._task_queue),
+            info=lambda: Info(
+                endpoint=endpoint,
+                namespace=self._namespace,
+                task_queue=self._task_queue,
+            ),
             nexus_context=ctx,
             client=self._client,
             _runtime_metric_meter=self._metric_meter,
@@ -293,6 +300,7 @@ async def _handle_start_operation_task(
         headers: Mapping[str, str],
         task_cancellation: nexusrpc.handler.OperationTaskCancellation,
         request_deadline: datetime | None,
+        endpoint: str,
     ) -> None:
         """Handle a start operation task.
 
@@ -302,7 +310,11 @@ async def _handle_start_operation_task(
         try:
             try:
                 start_response = await self._start_operation(
-                    start_request, headers, task_cancellation, request_deadline
+                    start_request,
+                    headers,
+                    task_cancellation,
+                    request_deadline,
+                    endpoint,
                 )
             except asyncio.CancelledError:
                 completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
@@ -346,6 +358,7 @@ async def _start_operation(
         headers: Mapping[str, str],
         cancellation: nexusrpc.handler.OperationTaskCancellation,
         request_deadline: datetime | None,
+        endpoint: str,
     ) -> temporalio.api.nexus.v1.StartOperationResponse:
         """Invoke the Nexus handler's start_operation method and construct the StartOperationResponse.
 
@@ -375,7 +388,11 @@ async def _start_operation(
         temporalio.nexus._operation_context._TemporalStartOperationContext(
             nexus_context=ctx,
             client=self._client,
-            info=lambda: Info(namespace=self._namespace, task_queue=self._task_queue),
+            info=lambda: Info(
+                endpoint=endpoint,
+                namespace=self._namespace,
+                task_queue=self._task_queue,
+            ),
             _runtime_metric_meter=self._metric_meter,
             _worker_shutdown_event=self._worker_shutdown_event,
         ).set()

temporalio/nexus/_operation_context.py

@@ -703,7 +703,11 @@ async def get_info(
         self, _ctx: StartOperationContext, _input: None
     ) -> dict[str, str]:
         info = nexus.info()
-        return {"namespace": info.namespace, "task_queue": info.task_queue}
+        return {
+            "endpoint": info.endpoint,
+            "namespace": info.namespace,
+            "task_queue": info.task_queue,
+        }
 
 
 @workflow.defn
@@ -733,6 +737,9 @@ async def test_nexus_info_includes_namespace(client: Client, env: WorkflowEnviro
             id=str(uuid.uuid4()),
             task_queue=task_queue,
         )
+        if not env.supports_time_skipping:
+            # Time-skipping server doesn't send the endpoint yet.
+            assert result["endpoint"] == endpoint_name
         assert result["namespace"] == client.namespace
         assert result["task_queue"] == task_queue