diff --git a/CHANGELOG.md b/CHANGELOG.md index 19f761152..118926c62 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,10 @@ to include examples, links to docs, or any other relevant information. - Added GCP Cloud Run serverless-worker OpenTelemetry plugin in `temporalio.contrib.opentelemetry`. - Added new options to ActivityHandle.describe() to retrieve associated payloads, such as activity input and outcome. - New properties and methods in ActivityExecution and ActivityExecutionDescription. +- Added experimental `temporalio.converter.NexusSerializationContext` support for Nexus callers + and handlers. Callers use it for inputs, results, and failures; handlers use it for inputs, + synchronous results, and failures. Asynchronous handler results and detached standalone handles + are not yet supported. Standalone `USE_EXISTING` handles use their start request's context. ### Changed @@ -36,6 +40,8 @@ to include examples, links to docs, or any other relevant information. - System Nexus Signal-with-Start Workflow operations now invoke `WorkflowOutboundInterceptor.start_system_nexus_operation` after their typed interception point. They continue not to invoke `WorkflowOutboundInterceptor.start_nexus_operation`. +- The experimental `GetNexusOperationResultInput` now includes the Nexus endpoint, service, and + operation. ### Deprecated diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index b5f6ab677..848975dc9 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -1549,6 +1549,12 @@ async def start_nexus_operation( self, input: StartNexusOperationInput ) -> NexusOperationHandle[Any]: """Start a nexus operation and return a handle to it.""" + nexus_context = temporalio.converter.NexusSerializationContext( + endpoint=input.endpoint, + service=input.service, + operation=input.operation, + ) + data_converter = self._client.data_converter.with_context(nexus_context) req = temporalio.api.workflowservice.v1.StartNexusOperationExecutionRequest( namespace=self._client.namespace, identity=self._client.identity, @@ -1575,7 +1581,7 @@ async def start_nexus_operation( req.start_to_close_timeout.FromTimedelta(input.start_to_close_timeout) # Set input payload - encoded = await self._client.data_converter.encode([input.arg]) + encoded = await data_converter.encode([input.arg]) if encoded: req.input.CopyFrom(encoded[0]) @@ -1620,6 +1626,7 @@ async def start_nexus_operation( result_type=input.result_type, endpoint=input.endpoint, service=input.service, + operation=input.operation, ) async def describe_nexus_operation( @@ -1637,15 +1644,31 @@ async def describe_nexus_operation( metadata=input.rpc_metadata, timeout=input.rpc_timeout, ) + data_converter = self._client.data_converter.with_context( + temporalio.converter.NexusSerializationContext( + endpoint=resp.info.endpoint, + service=resp.info.service, + operation=resp.info.operation, + ) + ) return await NexusOperationExecutionDescription._from_execution_info( info=resp.info, - data_converter=self._client.data_converter, + data_converter=data_converter, ) async def get_nexus_operation_result( self, input: GetNexusOperationResultInput ) -> Any: """Poll for nexus operation result until it's available.""" + data_converter = self._client.data_converter + if input.endpoint and input.service and input.operation: + data_converter = data_converter.with_context( + temporalio.converter.NexusSerializationContext( + endpoint=input.endpoint, + service=input.service, + operation=input.operation, + ) + ) req = temporalio.api.workflowservice.v1.PollNexusOperationExecutionRequest( namespace=self._client.namespace, operation_id=input.operation_id, @@ -1667,21 +1690,14 @@ async def get_nexus_operation_result( match res.WhichOneof("outcome"): case "result": type_hints = [input.result_type] if input.result_type else None - [result] = await self._client.data_converter.decode( - [res.result], type_hints - ) + [result] = await data_converter.decode([res.result], type_hints) return result - case "failure": raise NexusOperationFailureError( - cause=await self._client.data_converter.decode_failure( - res.failure - ) + cause=await data_converter.decode_failure(res.failure) ) - case None: - # poll again - pass + continue except RPCError as err: match err.status: case RPCStatusCode.DEADLINE_EXCEEDED: diff --git a/temporalio/client/_interceptor.py b/temporalio/client/_interceptor.py index b9d82d6ea..2ca62e10a 100644 --- a/temporalio/client/_interceptor.py +++ b/temporalio/client/_interceptor.py @@ -18,9 +18,7 @@ import temporalio.api.common.v1 import temporalio.api.workflowservice.v1 import temporalio.common -from temporalio.converter import ( - DataConverter, -) +from temporalio.converter import DataConverter if TYPE_CHECKING: from ._activity import ( @@ -657,6 +655,9 @@ class GetNexusOperationResultInput: operation_id: str run_id: str | None + endpoint: str + service: str + operation: str rpc_metadata: Mapping[str, str | bytes] rpc_timeout: timedelta | None result_type: type[Any] | None diff --git a/temporalio/client/_nexus.py b/temporalio/client/_nexus.py index 7eea155a9..1f0a7338e 100644 --- a/temporalio/client/_nexus.py +++ b/temporalio/client/_nexus.py @@ -1066,6 +1066,7 @@ def __init__( result_type: type | None = None, endpoint: str = "", service: str = "", + operation: str = "", ) -> None: """Create nexus operation handle.""" self._client = client @@ -1074,6 +1075,7 @@ def __init__( self._result_type = result_type self._endpoint = endpoint self._service = service + self._operation = operation # the default value is `_arg_unset` because ReturnType could be None self._known_outcome: ReturnType | NexusOperationFailureError | object = ( temporalio.common._arg_unset @@ -1136,9 +1138,12 @@ async def result( GetNexusOperationResultInput( operation_id=self._operation_id, run_id=self._run_id, - result_type=self._result_type, + endpoint=self._endpoint, + service=self._service, + operation=self._operation, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout, + result_type=self._result_type, ) ) ) diff --git a/temporalio/converter/__init__.py b/temporalio/converter/__init__.py index 99e55a775..324b477f2 100644 --- a/temporalio/converter/__init__.py +++ b/temporalio/converter/__init__.py @@ -50,6 +50,7 @@ ) from temporalio.converter._serialization_context import ( ActivitySerializationContext, + NexusSerializationContext, SerializationContext, WithSerializationContext, WorkflowSerializationContext, @@ -82,6 +83,7 @@ "JSONProtoPayloadConverter", "JSONTypeConverter", "JSONTypeConverterUnhandled", + "NexusSerializationContext", "PayloadCodec", "PayloadConverter", "SerializationContext", diff --git a/temporalio/converter/_serialization_context.py b/temporalio/converter/_serialization_context.py index 73a4a7104..8046a814c 100644 --- a/temporalio/converter/_serialization_context.py +++ b/temporalio/converter/_serialization_context.py @@ -28,6 +28,10 @@ class SerializationContext(ABC): context type is :py:class:`ActivitySerializationContext` and the workflow ID is that of the currently-executing workflow. ActivitySerializationContext is also set on data converter operations in the activity context. + + When operating on a Nexus operation payload, the context type is + :py:class:`NexusSerializationContext` and identifies the Nexus endpoint, service, and + resolved operation name. """ pass @@ -94,6 +98,38 @@ class ActivitySerializationContext(SerializationContext): """Whether the activity is a local activity started from a workflow.""" +@dataclass(frozen=True) +class NexusSerializationContext(SerializationContext): + """Serialization context for Nexus operation payloads. + + Callers receive this context when encoding inputs and decoding results or failures. The context + is not propagated to a handler that completes an asynchronous operation. Handlers receive it + when decoding inputs, encoding synchronous results, and encoding failures produced while + handling a Nexus task. + + A standalone operation handle retains the context used to start the operation and uses it to + decode the result, including when the start request returns an existing operation. A handle + created with :py:meth:`temporalio.client.Client.get_nexus_operation_handle` has no endpoint, + service, or operation information and therefore decodes without Nexus context. + + A failure encoded by a handler is later decoded by a caller. Because some operation paths may + lack this context, contextual encodings must be self-describing and decoders must continue to + accept payloads encoded without context. + + .. warning:: + This API is experimental and unstable. + """ + + endpoint: str + """Nexus endpoint name.""" + + service: str + """Nexus service name.""" + + operation: str + """Nexus operation name.""" + + class WithSerializationContext(ABC): """Interface for classes that can use serialization context. diff --git a/temporalio/worker/_command_aware_visitor.py b/temporalio/worker/_command_aware_visitor.py index 500fc4db5..7c03c2cd4 100644 --- a/temporalio/worker/_command_aware_visitor.py +++ b/temporalio/worker/_command_aware_visitor.py @@ -24,6 +24,7 @@ ScheduleNexusOperation, SignalExternalWorkflowExecution, StartChildWorkflowExecution, + WorkflowCommand, ) @@ -115,6 +116,18 @@ async def _visit_coresdk_workflow_commands_ScheduleNexusOperation( with current_command(CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq): await super()._visit_coresdk_workflow_commands_ScheduleNexusOperation(fs, o) + async def _visit_coresdk_workflow_commands_WorkflowCommand( + self, fs: VisitorFunctions, o: WorkflowCommand + ) -> None: + if o.HasField("schedule_nexus_operation"): + with current_command( + CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, + o.schedule_nexus_operation.seq, + ): + await super()._visit_coresdk_workflow_commands_WorkflowCommand(fs, o) + else: + await super()._visit_coresdk_workflow_commands_WorkflowCommand(fs, o) + # Workflow activation jobs with payloads async def _visit_coresdk_workflow_activation_ResolveActivity( self, fs: VisitorFunctions, o: ResolveActivity diff --git a/temporalio/worker/_nexus.py b/temporalio/worker/_nexus.py index 90ba40382..a2f4b8ca7 100644 --- a/temporalio/worker/_nexus.py +++ b/temporalio/worker/_nexus.py @@ -230,18 +230,31 @@ async def _complete_task( await asyncio.shield(self._bridge_worker().complete_nexus_task(completion)) async def _encode_completion( - self, completion: temporalio.bridge.proto.nexus.NexusTaskCompletion + self, + completion: temporalio.bridge.proto.nexus.NexusTaskCompletion, + data_converter: temporalio.converter.DataConverter, ) -> None: """Apply the payload codec then external storage to the completion's payloads.""" - dc = self._data_converter await PayloadVisitor(skip_search_attributes=True, skip_headers=True).visit( - _PayloadTransformVisitor(dc._encode_payload_sequence), completion + _PayloadTransformVisitor(data_converter._encode_payload_sequence), + completion, ) await PayloadVisitor(skip_search_attributes=True).visit( - _PayloadTransformVisitor(dc._external_store_payload_sequence), + _PayloadTransformVisitor(data_converter._external_store_payload_sequence), completion, ) + def _data_converter_for_nexus_task( + self, endpoint: str, service: str, operation: str + ) -> temporalio.converter.DataConverter: + return self._data_converter.with_context( + temporalio.converter.NexusSerializationContext( + endpoint=endpoint, + service=service, + operation=operation, + ) + ) + # TODO(nexus-preview): stack trace pruning. See sdk-typescript NexusHandler.execute # "Any call up to this function and including this one will be trimmed out of stack traces."" @@ -272,6 +285,9 @@ async def _handle_cancel_operation_task( task_cancellation=task_cancellation, request_deadline=request_deadline, ) + data_converter = self._data_converter_for_nexus_task( + endpoint, request.service, request.operation + ) temporalio.nexus._operation_context._TemporalCancelOperationContext( info=lambda: Info( endpoint=endpoint, @@ -293,7 +309,7 @@ async def _handle_cancel_operation_task( ), ) # No-op but keeps the cancel covered if it ever carries a payload. - await self._encode_completion(completion) + await self._encode_completion(completion, data_converter) except asyncio.CancelledError: completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( task_token=task_token, @@ -305,12 +321,12 @@ async def _handle_cancel_operation_task( completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( task_token=task_token, ) - self._data_converter.failure_converter.to_failure( + data_converter.failure_converter.to_failure( handler_error, - self._data_converter.payload_converter, + data_converter.payload_converter, completion.failure, ) - await self._encode_completion(completion) + await self._encode_completion(completion, data_converter) await self._complete_task(completion) except Exception: logger.exception("Failed to send Nexus task completion") @@ -336,6 +352,9 @@ async def _handle_start_operation_task( Attempt to execute the user start_operation method and invoke the data converter on the result. Handle errors and send the task completion. """ + data_converter = self._data_converter_for_nexus_task( + endpoint, start_request.service, start_request.operation + ) try: try: start_response = await self._start_operation( @@ -344,6 +363,7 @@ async def _handle_start_operation_task( task_cancellation, request_deadline, endpoint, + data_converter, ) completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( task_token=task_token, @@ -351,7 +371,7 @@ async def _handle_start_operation_task( start_operation=start_response ), ) - await self._encode_completion(completion) + await self._encode_completion(completion, data_converter) except asyncio.CancelledError: completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( task_token=task_token, @@ -363,15 +383,15 @@ async def _handle_start_operation_task( task_token=task_token, ) handler_error = _exception_to_handler_error(err) - self._data_converter.failure_converter.to_failure( + data_converter.failure_converter.to_failure( handler_error, - self._data_converter.payload_converter, + data_converter.payload_converter, completion.failure, ) if isinstance(err, concurrent.futures.BrokenExecutor): self._fail_worker_exception_queue.put_nowait(err) - await self._encode_completion(completion) + await self._encode_completion(completion, data_converter) await self._complete_task(completion) except Exception: @@ -391,6 +411,7 @@ async def _start_operation( cancellation: nexusrpc.handler.OperationTaskCancellation, request_deadline: datetime | None, endpoint: str, + data_converter: temporalio.converter.DataConverter, ) -> temporalio.api.nexus.v1.StartOperationResponse: """Invoke the Nexus handler's start_operation method and construct the StartOperationResponse. @@ -430,7 +451,7 @@ async def _start_operation( ).set() input = LazyValue( serializer=_NexusPayloadSerializer( - data_converter=self._data_converter, + data_converter=data_converter, payload=start_request.payload, ), headers={}, @@ -450,9 +471,7 @@ async def _start_operation( ) ) elif isinstance(result, nexusrpc.handler.StartOperationResultSync): - [payload] = self._data_converter.payload_converter.to_payloads( - [result.value] - ) + [payload] = data_converter.payload_converter.to_payloads([result.value]) return temporalio.api.nexus.v1.StartOperationResponse( sync_success=temporalio.api.nexus.v1.StartOperationResponse.Sync( payload=payload, @@ -481,9 +500,9 @@ async def _start_operation( ) from err.__cause__ except FailureError as new_err: response = temporalio.api.nexus.v1.StartOperationResponse() - self._data_converter.failure_converter.to_failure( + data_converter.failure_converter.to_failure( new_err, - self._data_converter.payload_converter, + data_converter.payload_converter, response.failure, ) return response diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 80d77f103..89625d64c 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -2194,14 +2194,26 @@ async def operation_handle_fn() -> OutputT: user_payload_converter, user_failure_converter, ) + failure_converter = user_failure_converter else: - payload_converter = self._context_free_payload_converter + serialization_context = temporalio.converter.NexusSerializationContext( + endpoint=input.endpoint, + service=input.service, + operation=input.operation_name, + ) + payload_converter = self._payload_converter_with_context( + serialization_context + ) + failure_converter = self._failure_converter_with_context( + serialization_context + ) handle = _NexusOperationHandle( self, self._next_seq("nexus_operation"), input, operation_handle_fn(), payload_converter, + failure_converter, ) handle._apply_schedule_command() self._pending_nexus_operations[handle._seq] = handle @@ -2454,9 +2466,11 @@ def get_serialization_context( nexus_operation._input.operation_name, nexus_operation._input.input, ) - # Other Nexus operations have no context because the caller workflow context is - # unavailable on the handler side for decryption. - return None + return temporalio.converter.NexusSerializationContext( + endpoint=nexus_operation._input.endpoint, + service=nexus_operation._input.service, + operation=nexus_operation._input.operation_name, + ) else: # Use payload codec with workflow context for all other payloads @@ -3648,6 +3662,7 @@ def __init__( input: StartNexusOperationInput[Any, OutputT], fn: Coroutine[Any, Any, OutputT], payload_converter: temporalio.converter.PayloadConverter, + failure_converter: temporalio.converter.FailureConverter, ): self._instance = instance self._seq = seq @@ -3656,7 +3671,7 @@ def __init__( self._start_fut: asyncio.Future[str | None] = instance.create_future() self._result_fut: asyncio.Future[OutputT | None] = instance.create_future() self._payload_converter = payload_converter - self._failure_converter = self._instance._context_free_failure_converter + self._failure_converter = failure_converter @property def operation_token(self) -> str | None: diff --git a/tests/nexus/test_link_propagation.py b/tests/nexus/test_link_propagation.py index 554620b95..f656f57c1 100644 --- a/tests/nexus/test_link_propagation.py +++ b/tests/nexus/test_link_propagation.py @@ -884,6 +884,9 @@ async def test_sync_response_includes_signal_backlinks() -> None: cancellation=_NexusTaskCancellation(), request_deadline=None, endpoint="endpoint", + data_converter=worker._data_converter_for_nexus_task( + "endpoint", "_BacklinkStashingService", "sync_op" + ), ) assert response.HasField("sync_success") assert len(response.sync_success.links) == 1 @@ -898,6 +901,9 @@ async def test_async_response_includes_signal_backlinks() -> None: cancellation=_NexusTaskCancellation(), request_deadline=None, endpoint="endpoint", + data_converter=worker._data_converter_for_nexus_task( + "endpoint", "_BacklinkStashingService", "async_op" + ), ) assert response.HasField("async_success") assert response.async_success.operation_token == "op-token" diff --git a/tests/nexus/test_standalone_operations.py b/tests/nexus/test_standalone_operations.py index 26a8316b4..f12a7fd85 100644 --- a/tests/nexus/test_standalone_operations.py +++ b/tests/nexus/test_standalone_operations.py @@ -869,7 +869,9 @@ async def get_nexus_operation_result( self, input: GetNexusOperationResultInput ) -> Any: self._parent.result_calls.append(input) - return await super().get_nexus_operation_result(input) + result = await super().get_nexus_operation_result(input) + self._parent.result_outputs.append(result) + return result async def cancel_nexus_operation(self, input: CancelNexusOperationInput) -> None: self._parent.cancel_calls.append(input) @@ -898,6 +900,7 @@ def __init__(self) -> None: self.start_calls: list[StartNexusOperationInput] = [] self.describe_calls: list[DescribeNexusOperationInput] = [] self.result_calls: list[GetNexusOperationResultInput] = [] + self.result_outputs: list[Any] = [] self.cancel_calls: list[CancelNexusOperationInput] = [] self.terminate_calls: list[TerminateNexusOperationInput] = [] self.list_calls: list[ListNexusOperationsInput] = [] @@ -981,6 +984,21 @@ async def test_interceptor_receives_inputs(client: Client, env: WorkflowEnvironm assert isinstance(result_input, GetNexusOperationResultInput) assert result_input.operation_id == op_id assert result_input.result_type == EchoOutput + assert result_input.endpoint == endpoint_name + assert result_input.service == "StandaloneTestService" + assert result_input.operation == "blocking_async" + + # Interceptors receive successfully decoded results. + value = f"interceptor-success-{uuid.uuid4()}" + handle = await nexus_client.start_operation( + StandaloneTestService.echo_sync, + EchoInput(value=value), + id=str(uuid.uuid4()), + schedule_to_close_timeout=timedelta(seconds=30), + ) + result = await handle.result() + assert result == EchoOutput(value=value) + assert interceptor.result_outputs == [EchoOutput(value=value)] # Start another so we can terminate it previous_start_count = len(interceptor.start_calls) diff --git a/tests/test_serialization_context.py b/tests/test_serialization_context.py index 8d65d5f1f..793663ca9 100644 --- a/tests/test_serialization_context.py +++ b/tests/test_serialization_context.py @@ -28,6 +28,10 @@ from temporalio.client import ( AsyncActivityHandle, Client, + GetNexusOperationResultInput, + Interceptor, + NexusOperationFailureError, + OutboundInterceptor, WorkflowFailureError, WorkflowUpdateFailedError, ) @@ -41,17 +45,17 @@ DefaultPayloadConverter, EncodingPayloadConverter, JSONPlainPayloadConverter, + NexusSerializationContext, PayloadCodec, PayloadConverter, SerializationContext, WithSerializationContext, WorkflowSerializationContext, ) -from temporalio.exceptions import ApplicationError +from temporalio.exceptions import ApplicationError, NexusOperationError from temporalio.testing import WorkflowEnvironment -from temporalio.worker import Worker +from temporalio.worker import Replayer, Worker from temporalio.worker._workflow_instance import UnsandboxedWorkflowRunner -from tests.helpers.nexus import make_nexus_endpoint_name @dataclass @@ -1688,25 +1692,88 @@ async def test_decode_context_matches_encode_context( # Test nexus payload codec -class AssertNexusLacksContextPayloadCodec(PayloadCodec, WithSerializationContext): - def __init__(self): - self.context = None +class _NexusResultDecodingInterceptor(Interceptor): + def __init__(self) -> None: + super().__init__() + self.decoded_results: list[Any] = [] + + def intercept_client(self, next: OutboundInterceptor) -> OutboundInterceptor: + return _NexusResultDecodingOutboundInterceptor(next, self) + + +class _NexusResultDecodingOutboundInterceptor(OutboundInterceptor): + def __init__( + self, next: OutboundInterceptor, parent: _NexusResultDecodingInterceptor + ) -> None: + super().__init__(next) + self._parent = parent + + async def get_nexus_operation_result( + self, input: GetNexusOperationResultInput + ) -> Any: + result = await super().get_nexus_operation_result(input) + self._parent.decoded_results.append(result) + return result + + +class NexusContextMarkerPayloadCodec(PayloadCodec, WithSerializationContext): + MARKER_KEY = "nexus-context-marker" + + def __init__( + self, + markers: dict[NexusSerializationContext, bytes], + context: SerializationContext | None = None, + ): + self.markers = markers + self.context = context def with_context( self, context: SerializationContext - ) -> AssertNexusLacksContextPayloadCodec: - codec = AssertNexusLacksContextPayloadCodec() - codec.context = context - return codec + ) -> NexusContextMarkerPayloadCodec: + return NexusContextMarkerPayloadCodec(self.markers, context) + + def _marker(self) -> bytes | None: + if not isinstance(self.context, NexusSerializationContext): + return None + try: + return self.markers[self.context] + except KeyError: + raise AssertionError( + f"No Nexus payload codec configured for {self.context!r}" + ) from None - async def _assert_context_iff_not_nexus( + async def encode( self, payloads: Sequence[temporalio.api.common.v1.Payload] ) -> list[temporalio.api.common.v1.Payload]: - [payload] = payloads - assert bool(self.context) == (payload.data.decode() != '"nexus-data"') - return list(payloads) + marker = self._marker() + if marker is None: + return list(payloads) + encoded = [] + for payload in payloads: + marked_payload = temporalio.api.common.v1.Payload() + marked_payload.CopyFrom(payload) + marked_payload.metadata[self.MARKER_KEY] = marker + encoded.append(marked_payload) + return encoded - encode = decode = _assert_context_iff_not_nexus + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + marker = self._marker() + if marker is None: + return list(payloads) + decoded = [] + for payload in payloads: + actual_marker = payload.metadata.get(self.MARKER_KEY) + if actual_marker is None: + decoded.append(payload) + continue + assert actual_marker == marker + decoded_payload = temporalio.api.common.v1.Payload() + decoded_payload.CopyFrom(payload) + del decoded_payload.metadata[self.MARKER_KEY] + decoded.append(decoded_payload) + return decoded @nexusrpc.handler.service_handler @@ -1717,52 +1784,365 @@ async def operation( ) -> str: return data + @nexusrpc.handler.sync_operation + async def fail(self, _: nexusrpc.handler.StartOperationContext, data: str) -> str: + raise ApplicationError(data, non_retryable=True) + @workflow.defn class NexusOperationTestWorkflow: @workflow.run - async def run(self, _data: str) -> None: + async def run(self, red_endpoint_name: str, blue_endpoint_name: str) -> list[str]: + red_handle, blue_handle = await asyncio.gather( + workflow.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=red_endpoint_name, + ).start_operation( + NexusOperationTestServiceHandler.operation, + input="nexus-data", + summary="nexus-summary", + ), + workflow.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=blue_endpoint_name, + ).start_operation( + NexusOperationTestServiceHandler.operation, + input="nexus-data", + summary="nexus-summary", + ), + ) + return list(await asyncio.gather(red_handle, blue_handle)) + + +@workflow.defn +class NexusOperationFailureTestWorkflow: + @workflow.run + async def run(self, endpoint_name: str) -> None: nexus_client = workflow.create_nexus_client( service=NexusOperationTestServiceHandler, - endpoint=make_nexus_endpoint_name(workflow.info().task_queue), - ) - await nexus_client.start_operation( - NexusOperationTestServiceHandler.operation, input="nexus-data" + endpoint=endpoint_name, ) + try: + await nexus_client.start_operation( + NexusOperationTestServiceHandler.fail, input="nexus-failure" + ) + except NexusOperationError: + return + raise AssertionError("Nexus operation should have failed") + + +nexus_failure_context_traces: list[tuple[str, NexusSerializationContext]] = [] + + +class NexusFailureConverterWithContext( + DefaultFailureConverter, WithSerializationContext +): + def __init__(self, context: SerializationContext | None = None): + super().__init__() + self.context = context + + def with_context( + self, context: SerializationContext + ) -> NexusFailureConverterWithContext: + return NexusFailureConverterWithContext(context) + + def to_failure( + self, + exception: BaseException, + payload_converter: PayloadConverter, + failure: temporalio.api.failure.v1.Failure, + ) -> None: + if isinstance(self.context, NexusSerializationContext): + nexus_failure_context_traces.append(("to_failure", self.context)) + super().to_failure(exception, payload_converter, failure) + + def from_failure( + self, + failure: temporalio.api.failure.v1.Failure, + payload_converter: PayloadConverter, + ) -> BaseException: + if isinstance(self.context, NexusSerializationContext): + nexus_failure_context_traces.append(("from_failure", self.context)) + return super().from_failure(failure, payload_converter) @pytest.mark.requires_local_server -async def test_nexus_payload_codec_operations_lack_context( +async def test_workflow_nexus_payload_codec_receives_context( env: WorkflowEnvironment, ): - """ - encode() and decode() on nexus payloads should not have any context set. - """ + """Nexus payload codecs get context for workflow inputs, summaries, and results.""" if env.supports_time_skipping: pytest.skip("Nexus tests don't work with the Java test server") + task_queue = "workflow-nexus-context-codec-task-queue" + red_endpoint_name = "workflow-red-nexus-endpoint" + blue_endpoint_name = "workflow-blue-nexus-endpoint" + red_context = NexusSerializationContext( + endpoint=red_endpoint_name, + service="NexusOperationTestServiceHandler", + operation="operation", + ) + blue_context = NexusSerializationContext( + endpoint=blue_endpoint_name, + service="NexusOperationTestServiceHandler", + operation="operation", + ) + payload_codec = NexusContextMarkerPayloadCodec( + {red_context: b"red", blue_context: b"blue"} + ) config = env.client.config() config["data_converter"] = dataclasses.replace( DataConverter.default, - payload_codec=AssertNexusLacksContextPayloadCodec(), + payload_codec=payload_codec, ) client = Client(**config) async with Worker( client, - task_queue=str(uuid.uuid4()), + task_queue=task_queue, workflows=[NexusOperationTestWorkflow], nexus_service_handlers=[NexusOperationTestServiceHandler()], ) as worker: - endpoint_name = make_nexus_endpoint_name(worker.task_queue) + await env.create_nexus_endpoint(red_endpoint_name, worker.task_queue) + await env.create_nexus_endpoint(blue_endpoint_name, worker.task_queue) + handle = await client.start_workflow( + NexusOperationTestWorkflow.run, + args=[red_endpoint_name, blue_endpoint_name], + id=str(uuid.uuid4()), + task_queue=worker.task_queue, + ) + assert await handle.result() == ["nexus-data", "nexus-data"] + + history = await handle.fetch_history() + scheduled_endpoints: dict[int, str] = {} + encoded_summaries: dict[str, temporalio.api.common.v1.Payload] = {} + encoded_results: dict[str, temporalio.api.common.v1.Payload] = {} + for event in history.events: + if event.HasField("nexus_operation_scheduled_event_attributes"): + scheduled_attrs = event.nexus_operation_scheduled_event_attributes + assert scheduled_attrs.service == "NexusOperationTestServiceHandler" + assert scheduled_attrs.operation == "operation" + scheduled_endpoints[event.event_id] = scheduled_attrs.endpoint + assert event.HasField("user_metadata") + assert event.user_metadata.HasField("summary") + encoded_summaries[scheduled_attrs.endpoint] = ( + event.user_metadata.summary + ) + elif event.HasField("nexus_operation_completed_event_attributes"): + completed_attrs = event.nexus_operation_completed_event_attributes + endpoint = scheduled_endpoints[completed_attrs.scheduled_event_id] + encoded_results[endpoint] = completed_attrs.result + assert set(scheduled_endpoints.values()) == { + red_endpoint_name, + blue_endpoint_name, + } + assert { + endpoint: payload.metadata[NexusContextMarkerPayloadCodec.MARKER_KEY] + for endpoint, payload in encoded_summaries.items() + } == { + red_endpoint_name: b"red", + blue_endpoint_name: b"blue", + } + assert { + endpoint: payload.metadata[NexusContextMarkerPayloadCodec.MARKER_KEY] + for endpoint, payload in encoded_results.items() + } == { + red_endpoint_name: b"red", + blue_endpoint_name: b"blue", + } + + scheduled_contexts: dict[int, NexusSerializationContext] = {} + for event in history.events: + if event.HasField("nexus_operation_scheduled_event_attributes"): + scheduled_attrs = event.nexus_operation_scheduled_event_attributes + context = NexusSerializationContext( + endpoint=scheduled_attrs.endpoint, + service=scheduled_attrs.service, + operation=scheduled_attrs.operation, + ) + scheduled_contexts[event.event_id] = context + [decoded] = await payload_codec.with_context(context).decode( + [scheduled_attrs.input] + ) + scheduled_attrs.input.CopyFrom(decoded) + [decoded] = await payload_codec.with_context(context).decode( + [event.user_metadata.summary] + ) + event.user_metadata.summary.CopyFrom(decoded) + elif event.HasField("nexus_operation_completed_event_attributes"): + completed_attrs = event.nexus_operation_completed_event_attributes + context = scheduled_contexts[completed_attrs.scheduled_event_id] + [decoded] = await payload_codec.with_context(context).decode( + [completed_attrs.result] + ) + completed_attrs.result.CopyFrom(decoded) + await Replayer( + workflows=[NexusOperationTestWorkflow], + data_converter=config["data_converter"], + ).replay_workflow(history) + + +@pytest.mark.requires_local_server +async def test_standalone_nexus_payload_codec_receives_context( + env: WorkflowEnvironment, +): + """Nexus payload codecs get context for standalone inputs and results.""" + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + + task_queue = "standalone-nexus-context-codec-task-queue" + red_endpoint_name = "standalone-red-nexus-endpoint" + blue_endpoint_name = "standalone-blue-nexus-endpoint" + red_context = NexusSerializationContext( + endpoint=red_endpoint_name, + service="NexusOperationTestServiceHandler", + operation="operation", + ) + blue_context = NexusSerializationContext( + endpoint=blue_endpoint_name, + service="NexusOperationTestServiceHandler", + operation="operation", + ) + config = env.client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, + payload_codec=NexusContextMarkerPayloadCodec( + {red_context: b"red", blue_context: b"blue"} + ), + ) + result_interceptor = _NexusResultDecodingInterceptor() + config["interceptors"] = list(config.get("interceptors") or []) + [ + result_interceptor + ] + client = Client(**config) + + async with Worker( + client, + task_queue=task_queue, + nexus_service_handlers=[NexusOperationTestServiceHandler()], + ) as worker: + await env.create_nexus_endpoint(red_endpoint_name, worker.task_queue) + await env.create_nexus_endpoint(blue_endpoint_name, worker.task_queue) + red_standalone_result, blue_standalone_result = await asyncio.gather( + client.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=red_endpoint_name, + ).execute_operation( + NexusOperationTestServiceHandler.operation, + "standalone-red", + id=str(uuid.uuid4()), + schedule_to_close_timeout=timedelta(seconds=10), + ), + client.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=blue_endpoint_name, + ).execute_operation( + NexusOperationTestServiceHandler.operation, + "standalone-blue", + id=str(uuid.uuid4()), + schedule_to_close_timeout=timedelta(seconds=10), + ), + ) + assert red_standalone_result == "standalone-red" + assert blue_standalone_result == "standalone-blue" + assert set(result_interceptor.decoded_results) == { + "standalone-red", + "standalone-blue", + } + + +@pytest.mark.requires_local_server +async def test_workflow_nexus_failure_converter_has_context( + env: WorkflowEnvironment, +): + """Workflow Nexus callers and handlers use context for failures.""" + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + + nexus_failure_context_traces.clear() + task_queue = "workflow-nexus-failure-context-task-queue" + endpoint_name = "workflow-failure-nexus-endpoint" + expected_context = NexusSerializationContext( + endpoint=endpoint_name, + service="NexusOperationTestServiceHandler", + operation="fail", + ) + config = env.client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, + failure_converter_class=NexusFailureConverterWithContext, + ) + client = Client(**config) + + async with Worker( + client, + task_queue=task_queue, + workflows=[NexusOperationFailureTestWorkflow], + nexus_service_handlers=[NexusOperationTestServiceHandler()], + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: await env.create_nexus_endpoint(endpoint_name, worker.task_queue) await client.execute_workflow( - NexusOperationTestWorkflow.run, - "workflow-data", + NexusOperationFailureTestWorkflow.run, + endpoint_name, id=str(uuid.uuid4()), task_queue=worker.task_queue, ) + assert ("to_failure", expected_context) in nexus_failure_context_traces + assert ("from_failure", expected_context) in nexus_failure_context_traces + + +@pytest.mark.requires_local_server +async def test_standalone_nexus_failure_converter_has_context( + env: WorkflowEnvironment, +): + """Standalone Nexus callers and handlers use context for failures.""" + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + + nexus_failure_context_traces.clear() + task_queue = "standalone-nexus-failure-context-task-queue" + endpoint_name = "standalone-failure-nexus-endpoint" + expected_context = NexusSerializationContext( + endpoint=endpoint_name, + service="NexusOperationTestServiceHandler", + operation="fail", + ) + config = env.client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, + failure_converter_class=NexusFailureConverterWithContext, + ) + client = Client(**config) + + async with Worker( + client, + task_queue=task_queue, + nexus_service_handlers=[NexusOperationTestServiceHandler()], + ) as worker: + await env.create_nexus_endpoint(endpoint_name, worker.task_queue) + nexus_client = client.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=endpoint_name, + ) + operation_handle = await nexus_client.start_operation( + NexusOperationTestServiceHandler.fail, + "nexus-failure", + id=str(uuid.uuid4()), + schedule_to_close_timeout=timedelta(seconds=10), + ) + with pytest.raises(NexusOperationFailureError): + await operation_handle.result() + + assert ("to_failure", expected_context) in nexus_failure_context_traces + assert ("from_failure", expected_context) in nexus_failure_context_traces + + nexus_failure_context_traces.clear() + description = await operation_handle.describe() + assert description.last_attempt_failure is not None + assert ("from_failure", expected_context) in nexus_failure_context_traces + # Test pydantic converter with context