Add test that confirms token present in callback headers for activities

Author: VegetarianOrc

temporalio/client/_activity.py

@@ -309,13 +309,17 @@ class ActivityExecutionDescription(ActivityExecution):
     long_poll_token: bytes | None
     """Token for follow-on long-poll requests. None if the activity is complete."""
 
+    raw_callbacks: Sequence[temporalio.api.activity.v1.CallbackInfo]
+    """Underlying protobuf callbacks"""
+
     @classmethod
     async def _from_execution_info(
         cls,
         info: temporalio.api.activity.v1.ActivityExecutionInfo,
         long_poll_token: bytes | None,
         namespace: str,
         data_converter: temporalio.converter.DataConverter,
+        callbacks: Sequence[temporalio.api.activity.v1.CallbackInfo],
     ) -> Self:
         """Create from raw proto activity execution info."""
         # Decode heartbeat details if present
@@ -409,6 +413,7 @@ async def _from_execution_info(
             typed_search_attributes=temporalio.converter.decode_typed_search_attributes(
                 info.search_attributes
             ),
+            raw_callbacks=callbacks,
         )
 
 

temporalio/client/_impl.py

@@ -723,6 +723,7 @@ async def describe_activity(
                     workflow_id=input.activity_id,  # Using activity_id as workflow_id for activities not started by a workflow
                 )
             ),
+            callbacks=resp.callbacks,
         )
 
     def list_activities(

tests/nexus/test_temporal_operation.py

@@ -1132,7 +1132,7 @@ async def test_temporal_operation_overloads(
         )
 
 
-async def test_temporal_operation_includes_token_in_callback(
+async def test_temporal_operation_includes_workflow_token_in_callback(
     client: Client, env: WorkflowEnvironment
 ):
     task_queue = str(uuid.uuid4())
@@ -1168,3 +1168,59 @@ async def test_temporal_operation_includes_token_in_callback(
         ).encode()
 
         assert token == expected_token
+
+
+async def test_temporal_operation_includes_activity_token_in_callback(
+    client: Client, env: WorkflowEnvironment
+):
+    task_queue = str(uuid.uuid4())
+    endpoint_name = make_nexus_endpoint_name(task_queue)
+    await env.create_nexus_endpoint(endpoint_name, task_queue)
+
+    @service_handler
+    class ActivityTokenHandler:
+        @nexus.temporal_operation
+        async def echo_activity(
+            self,
+            _ctx: nexus.TemporalStartOperationContext,
+            client: nexus.TemporalNexusClient,
+            input: Input,
+        ) -> nexus.TemporalOperationResult[str]:
+            return await client.start_activity(
+                echo_activity,
+                input,
+                id=input.value,
+                start_to_close_timeout=timedelta(seconds=10),
+            )
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        nexus_service_handlers=[ActivityTokenHandler()],
+        activities=[echo_activity],
+    ):
+        input_value = f"test-{uuid.uuid4()}"
+
+        nexus_client = client.create_nexus_client(ActivityTokenHandler, endpoint_name)
+
+        result = await nexus_client.execute_operation(
+            ActivityTokenHandler.echo_activity,
+            Input(value=input_value, task_queue=task_queue),
+            id=str(uuid.uuid4()),
+        )
+        assert result == input_value
+
+        activity_handle = client.get_activity_handle(input_value)
+
+        desc = await activity_handle.describe()
+        token = desc.raw_callbacks[0].info.callback.nexus.header[
+            "nexus-operation-token"
+        ]
+
+        expected_token = OperationToken(
+            type=OperationTokenType.ACTIVITY,
+            namespace=client.namespace,
+            activity_id=input_value,
+        ).encode()
+
+        assert token == expected_token