diff --git a/CHANGELOG.md b/CHANGELOG.md index 156d13b0a..d9a7356bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,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 diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index e4ff9f3f5..af9473e75 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -1453,6 +1453,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, @@ -1479,7 +1485,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]) @@ -1524,6 +1530,7 @@ async def start_nexus_operation( result_type=input.result_type, endpoint=input.endpoint, service=input.service, + _nexus_serialization_context=nexus_context, ) async def describe_nexus_operation( @@ -1541,15 +1548,28 @@ async def describe_nexus_operation( metadata=input.rpc_metadata, timeout=input.rpc_timeout, ) + nexus_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, + failure_data_converter=self._client.data_converter.with_context( + nexus_context + ), ) 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._nexus_serialization_context is not None: + data_converter = data_converter.with_context( + input._nexus_serialization_context + ) req = temporalio.api.workflowservice.v1.PollNexusOperationExecutionRequest( namespace=self._client.namespace, operation_id=input.operation_id, @@ -1571,16 +1591,12 @@ 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: diff --git a/temporalio/client/_interceptor.py b/temporalio/client/_interceptor.py index 449d0004c..015264b11 100644 --- a/temporalio/client/_interceptor.py +++ b/temporalio/client/_interceptor.py @@ -20,6 +20,7 @@ import temporalio.common from temporalio.converter import ( DataConverter, + NexusSerializationContext, ) if TYPE_CHECKING: @@ -611,6 +612,7 @@ class GetNexusOperationResultInput: rpc_metadata: Mapping[str, str | bytes] rpc_timeout: timedelta | None result_type: type[Any] | None + _nexus_serialization_context: NexusSerializationContext | None = None @dataclass diff --git a/temporalio/client/_nexus.py b/temporalio/client/_nexus.py index 7eea155a9..5d3a291fb 100644 --- a/temporalio/client/_nexus.py +++ b/temporalio/client/_nexus.py @@ -291,8 +291,10 @@ async def _from_execution_info( cls, info: temporalio.api.nexus.v1.NexusOperationExecutionInfo, data_converter: temporalio.converter.DataConverter, + failure_data_converter: temporalio.converter.DataConverter | None = None, ) -> Self: """Create from raw proto nexus operation execution info.""" + failure_data_converter = failure_data_converter or data_converter return cls( _data_converter=data_converter, operation_id=info.operation_id, @@ -360,7 +362,9 @@ async def _from_execution_info( last_attempt_failure=( cast( BaseException | None, - await data_converter.decode_failure(info.last_attempt_failure), + await failure_data_converter.decode_failure( + info.last_attempt_failure + ), ) if info.HasField("last_attempt_failure") else None @@ -376,7 +380,7 @@ async def _from_execution_info( identity=info.identity, cancellation_info=( await NexusOperationExecutionCancellationInfo._from_cancellation_info( - info.cancellation_info, data_converter + info.cancellation_info, failure_data_converter ) if info.HasField("cancellation_info") else None @@ -1066,6 +1070,9 @@ def __init__( result_type: type | None = None, endpoint: str = "", service: str = "", + _nexus_serialization_context: ( + temporalio.converter.NexusSerializationContext | None + ) = None, ) -> None: """Create nexus operation handle.""" self._client = client @@ -1074,6 +1081,7 @@ def __init__( self._result_type = result_type self._endpoint = endpoint self._service = service + self._nexus_serialization_context = _nexus_serialization_context # the default value is `_arg_unset` because ReturnType could be None self._known_outcome: ReturnType | NexusOperationFailureError | object = ( temporalio.common._arg_unset @@ -1131,15 +1139,14 @@ async def result( """ if self._known_outcome is temporalio.common._arg_unset: try: - self._known_outcome = ( - await self._client._impl.get_nexus_operation_result( - GetNexusOperationResultInput( - operation_id=self._operation_id, - run_id=self._run_id, - result_type=self._result_type, - rpc_metadata=rpc_metadata, - rpc_timeout=rpc_timeout, - ) + self._known_outcome = await self._client._impl.get_nexus_operation_result( + GetNexusOperationResultInput( + operation_id=self._operation_id, + run_id=self._run_id, + result_type=self._result_type, + rpc_metadata=rpc_metadata, + rpc_timeout=rpc_timeout, + _nexus_serialization_context=self._nexus_serialization_context, ) ) return cast(ReturnType, self._known_outcome) 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..da80f1af3 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,35 @@ 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. Handlers + receive it when decoding inputs, encoding synchronous results, and encoding failures produced + while handling a Nexus task. + + The context is not propagated to the eventual result of an asynchronous operation. Standalone + operation handles use the context of their start request, including when an existing operation + is returned, while handles created without starting an operation do not receive it. + + Callers and handlers receive this context on opposite sides of failure conversion. Contextual + encodings should therefore be self-describing and support legacy payloads 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/_nexus.py b/temporalio/worker/_nexus.py index 90ba40382..7614ccaa5 100644 --- a/temporalio/worker/_nexus.py +++ b/temporalio/worker/_nexus.py @@ -230,18 +230,37 @@ 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: + service_handler = self._handler.service_handlers.get(service) + if ( + service_handler is None + or operation not in service_handler.service.operation_definitions + ): + return self._data_converter + 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 +291,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 +315,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 +327,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 +358,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 +369,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 +377,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 +389,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 +417,7 @@ async def _start_operation( cancellation: nexusrpc.handler.OperationTaskCancellation, request_deadline: datetime | None, endpoint: str, + data_converter: temporalio.converter.DataConverter | None = None, ) -> temporalio.api.nexus.v1.StartOperationResponse: """Invoke the Nexus handler's start_operation method and construct the StartOperationResponse. @@ -398,6 +425,9 @@ async def _start_operation( All other exceptions are handled by a caller of this function. """ + data_converter = data_converter or self._data_converter_for_nexus_task( + endpoint, start_request.service, start_request.operation + ) # Create the worker shutdown event if not created if not self._worker_shutdown_event: self._worker_shutdown_event = temporalio.common._CompositeEvent( @@ -430,7 +460,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 +480,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 +509,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 f7cd66cfe..8829a4a7d 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -2180,14 +2180,29 @@ async def operation_handle_fn() -> OutputT: self._workflow_context_payload_converter, self._workflow_context_failure_converter, ) + summary_payload_converter = payload_converter + failure_converter = self._context_free_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 + ) + summary_payload_converter = self._context_free_payload_converter + failure_converter = self._failure_converter_with_context( + serialization_context + ) handle = _NexusOperationHandle( self, self._next_seq("nexus_operation"), input, operation_handle_fn(), payload_converter, + summary_payload_converter, + failure_converter, ) handle._apply_schedule_command() self._pending_nexus_operations[handle._seq] = handle @@ -2440,9 +2455,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 @@ -3634,6 +3651,8 @@ def __init__( input: StartNexusOperationInput[Any, OutputT], fn: Coroutine[Any, Any, OutputT], payload_converter: temporalio.converter.PayloadConverter, + summary_payload_converter: temporalio.converter.PayloadConverter, + failure_converter: temporalio.converter.FailureConverter, ): self._instance = instance self._seq = seq @@ -3642,7 +3661,8 @@ 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._summary_payload_converter = summary_payload_converter + self._failure_converter = failure_converter @property def operation_token(self) -> str | None: @@ -3704,7 +3724,7 @@ def _apply_schedule_command(self) -> None: if self._input.summary: command.user_metadata.summary.CopyFrom( - self._payload_converter.to_payload(self._input.summary) + self._summary_payload_converter.to_payload(self._input.summary) ) def _apply_cancel_command( diff --git a/tests/test_serialization_context.py b/tests/test_serialization_context.py index 8d65d5f1f..370f2cf13 100644 --- a/tests/test_serialization_context.py +++ b/tests/test_serialization_context.py @@ -9,8 +9,11 @@ import asyncio import dataclasses +import hashlib +import hmac import json import uuid +import zlib from collections import defaultdict from collections.abc import Sequence from dataclasses import dataclass, field @@ -28,6 +31,7 @@ from temporalio.client import ( AsyncActivityHandle, Client, + NexusOperationFailureError, 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,92 @@ async def test_decode_context_matches_encode_context( # Test nexus payload codec -class AssertNexusLacksContextPayloadCodec(PayloadCodec, WithSerializationContext): - def __init__(self): - self.context = None +class NexusContextPayloadCodecSelector(PayloadCodec, WithSerializationContext): + HMAC_ENCODING = b"binary/nexus-context-hmac" + ZLIB_ENCODING = b"binary/nexus-context-zlib" + HMAC_KEY = b"nexus-context-test-key" + + def __init__( + self, + codecs: dict[NexusSerializationContext, Literal["hmac", "zlib"]], + context: SerializationContext | None = None, + ): + self.codecs = codecs + self.context = context def with_context( self, context: SerializationContext - ) -> AssertNexusLacksContextPayloadCodec: - codec = AssertNexusLacksContextPayloadCodec() - codec.context = context - return codec + ) -> NexusContextPayloadCodecSelector: + return NexusContextPayloadCodecSelector(self.codecs, context) - async def _assert_context_iff_not_nexus( + def _codec(self) -> Literal["hmac", "zlib"] | None: + if not isinstance(self.context, NexusSerializationContext): + return None + try: + return self.codecs[self.context] + except KeyError: + raise AssertionError( + f"No Nexus payload codec configured for {self.context!r}" + ) from None + + 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) + codec = self._codec() + if codec is None: + return list(payloads) + encoded = [] + for payload in payloads: + serialized = payload.SerializeToString(deterministic=True) + if codec == "hmac": + signature = hmac.new(self.HMAC_KEY, serialized, hashlib.sha256).digest() + encoded.append( + temporalio.api.common.v1.Payload( + metadata={"encoding": self.HMAC_ENCODING}, + data=signature + serialized, + ) + ) + else: + encoded.append( + temporalio.api.common.v1.Payload( + metadata={"encoding": self.ZLIB_ENCODING}, + data=zlib.compress(serialized), + ) + ) + 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]: + codec = self._codec() + if codec is None: + return list(payloads) + decoded = [] + for payload in payloads: + encoding = payload.metadata.get("encoding") + if encoding not in (self.HMAC_ENCODING, self.ZLIB_ENCODING): + decoded.append(payload) + continue + expected_encoding = ( + self.HMAC_ENCODING if codec == "hmac" else self.ZLIB_ENCODING + ) + assert encoding == expected_encoding + if codec == "hmac": + digest_size = hashlib.sha256().digest_size + signature, serialized = ( + payload.data[:digest_size], + payload.data[digest_size:], + ) + assert hmac.compare_digest( + signature, + hmac.new(self.HMAC_KEY, serialized, hashlib.sha256).digest(), + ) + else: + serialized = zlib.decompress(payload.data) + decoded_payload = temporalio.api.common.v1.Payload() + decoded_payload.ParseFromString(serialized) + decoded.append(decoded_payload) + return decoded @nexusrpc.handler.service_handler @@ -1717,52 +1788,342 @@ 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, hmac_endpoint_name: str, zlib_endpoint_name: str) -> list[str]: + hmac_handle, zlib_handle = await asyncio.gather( + workflow.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=hmac_endpoint_name, + ).start_operation( + NexusOperationTestServiceHandler.operation, + input="nexus-data", + ), + workflow.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=zlib_endpoint_name, + ).start_operation( + NexusOperationTestServiceHandler.operation, + input="nexus-data", + ), + ) + return list(await asyncio.gather(hmac_handle, zlib_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_selects_codec_from_context( env: WorkflowEnvironment, ): - """ - encode() and decode() on nexus payloads should not have any context set. - """ + """Nexus context selects codecs for workflow inputs 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" + hmac_endpoint_name = "workflow-hmac-nexus-endpoint" + zlib_endpoint_name = "workflow-zlib-nexus-endpoint" + hmac_context = NexusSerializationContext( + endpoint=hmac_endpoint_name, + service="NexusOperationTestServiceHandler", + operation="operation", + ) + zlib_context = NexusSerializationContext( + endpoint=zlib_endpoint_name, + service="NexusOperationTestServiceHandler", + operation="operation", + ) + payload_codec = NexusContextPayloadCodecSelector( + {hmac_context: "hmac", zlib_context: "zlib"} + ) 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(hmac_endpoint_name, worker.task_queue) + await env.create_nexus_endpoint(zlib_endpoint_name, worker.task_queue) + handle = await client.start_workflow( + NexusOperationTestWorkflow.run, + args=[hmac_endpoint_name, zlib_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_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 + 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()) == { + hmac_endpoint_name, + zlib_endpoint_name, + } + assert { + endpoint: payload.metadata["encoding"] + for endpoint, payload in encoded_results.items() + } == { + hmac_endpoint_name: NexusContextPayloadCodecSelector.HMAC_ENCODING, + zlib_endpoint_name: NexusContextPayloadCodecSelector.ZLIB_ENCODING, + } + assert ( + encoded_results[hmac_endpoint_name].data + != encoded_results[zlib_endpoint_name].data + ) + + 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) + 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_selects_codec_from_context( + env: WorkflowEnvironment, +): + """Nexus context selects codecs 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" + hmac_endpoint_name = "standalone-hmac-nexus-endpoint" + zlib_endpoint_name = "standalone-zlib-nexus-endpoint" + hmac_context = NexusSerializationContext( + endpoint=hmac_endpoint_name, + service="NexusOperationTestServiceHandler", + operation="operation", + ) + zlib_context = NexusSerializationContext( + endpoint=zlib_endpoint_name, + service="NexusOperationTestServiceHandler", + operation="operation", + ) + config = env.client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, + payload_codec=NexusContextPayloadCodecSelector( + {hmac_context: "hmac", zlib_context: "zlib"} + ), + ) + client = Client(**config) + + async with Worker( + client, + task_queue=task_queue, + nexus_service_handlers=[NexusOperationTestServiceHandler()], + ) as worker: + await env.create_nexus_endpoint(hmac_endpoint_name, worker.task_queue) + await env.create_nexus_endpoint(zlib_endpoint_name, worker.task_queue) + hmac_standalone_result, zlib_standalone_result = await asyncio.gather( + client.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=hmac_endpoint_name, + ).execute_operation( + NexusOperationTestServiceHandler.operation, + "standalone-hmac", + id=str(uuid.uuid4()), + schedule_to_close_timeout=timedelta(seconds=10), + ), + client.create_nexus_client( + service=NexusOperationTestServiceHandler, + endpoint=zlib_endpoint_name, + ).execute_operation( + NexusOperationTestServiceHandler.operation, + "standalone-zlib", + id=str(uuid.uuid4()), + schedule_to_close_timeout=timedelta(seconds=10), + ), + ) + assert hmac_standalone_result == "standalone-hmac" + assert zlib_standalone_result == "standalone-zlib" + + +@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