From 72e6f1efc5907a0f5fa18a897f2b78d5bbc42ccc Mon Sep 17 00:00:00 2001 From: Joshua Frenchwood Date: Wed, 9 Sep 2026 15:01:14 -0500 Subject: [PATCH 1/4] NexusSerializationContext for data/failure converters --- CHANGELOG.md | 4 + temporalio/client/_impl.py | 30 +- temporalio/client/_interceptor.py | 2 + temporalio/client/_nexus.py | 29 +- temporalio/converter/__init__.py | 2 + .../converter/_serialization_context.py | 33 ++ temporalio/worker/_nexus.py | 64 ++- temporalio/worker/_workflow_instance.py | 32 +- tests/test_serialization_context.py | 419 ++++++++++++++++-- 9 files changed, 544 insertions(+), 71 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 19f761152..063543a86 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 diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index b5f6ab677..75abe021e 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, + _nexus_serialization_context=nexus_context, ) async def describe_nexus_operation( @@ -1637,15 +1644,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, @@ -1667,16 +1687,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 b9d82d6ea..cf2179f10 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: @@ -660,6 +661,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 80d77f103..7fb3148aa 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -2194,14 +2194,29 @@ async def operation_handle_fn() -> OutputT: user_payload_converter, user_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 @@ -2454,9 +2469,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 +3665,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 @@ -3656,7 +3675,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: @@ -3718,7 +3738,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 From 4e8e392010f6358ad82c8d8788bdabfdb909c5b3 Mon Sep 17 00:00:00 2001 From: Joshua Frenchwood Date: Thu, 10 Sep 2026 16:20:14 -0500 Subject: [PATCH 2/4] Updating the nexus context for summary --- temporalio/worker/_command_aware_visitor.py | 13 +++++++++++ temporalio/worker/_nexus.py | 6 ----- temporalio/worker/_workflow_instance.py | 8 ++----- tests/test_serialization_context.py | 25 ++++++++++++++++++++- 4 files changed, 39 insertions(+), 13 deletions(-) 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 7614ccaa5..23be79bd6 100644 --- a/temporalio/worker/_nexus.py +++ b/temporalio/worker/_nexus.py @@ -247,12 +247,6 @@ async def _encode_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, diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 7fb3148aa..3e078334b 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -2175,6 +2175,7 @@ async def operation_handle_fn() -> OutputT: ), ) + # TODO: Extend system endpoint converter handling for worker callbacks. if temporalio.nexus.system.is_system_endpoint(input.endpoint): serialization_context = temporalio.nexus.system._get_serialization_context( input.service, @@ -2194,7 +2195,6 @@ async def operation_handle_fn() -> OutputT: user_payload_converter, user_failure_converter, ) - summary_payload_converter = payload_converter failure_converter = self._context_free_failure_converter else: serialization_context = temporalio.converter.NexusSerializationContext( @@ -2205,7 +2205,6 @@ async def operation_handle_fn() -> OutputT: 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 ) @@ -2215,7 +2214,6 @@ async def operation_handle_fn() -> OutputT: input, operation_handle_fn(), payload_converter, - summary_payload_converter, failure_converter, ) handle._apply_schedule_command() @@ -3665,7 +3663,6 @@ 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 @@ -3675,7 +3672,6 @@ 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._summary_payload_converter = summary_payload_converter self._failure_converter = failure_converter @property @@ -3738,7 +3734,7 @@ def _apply_schedule_command(self) -> None: if self._input.summary: command.user_metadata.summary.CopyFrom( - self._summary_payload_converter.to_payload(self._input.summary) + self._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 370f2cf13..52170b773 100644 --- a/tests/test_serialization_context.py +++ b/tests/test_serialization_context.py @@ -1804,6 +1804,7 @@ async def run(self, hmac_endpoint_name: str, zlib_endpoint_name: str) -> list[st ).start_operation( NexusOperationTestServiceHandler.operation, input="nexus-data", + summary="nexus-summary", ), workflow.create_nexus_client( service=NexusOperationTestServiceHandler, @@ -1811,6 +1812,7 @@ async def run(self, hmac_endpoint_name: str, zlib_endpoint_name: str) -> list[st ).start_operation( NexusOperationTestServiceHandler.operation, input="nexus-data", + summary="nexus-summary", ), ) return list(await asyncio.gather(hmac_handle, zlib_handle)) @@ -1872,7 +1874,7 @@ def from_failure( async def test_workflow_nexus_payload_codec_selects_codec_from_context( env: WorkflowEnvironment, ): - """Nexus context selects codecs for workflow inputs and results.""" + """Nexus context selects codecs for workflow inputs, summaries, and results.""" if env.supports_time_skipping: pytest.skip("Nexus tests don't work with the Java test server") @@ -1917,6 +1919,7 @@ async def test_workflow_nexus_payload_codec_selects_codec_from_context( 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"): @@ -1924,6 +1927,11 @@ async def test_workflow_nexus_payload_codec_selects_codec_from_context( 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] @@ -1932,6 +1940,17 @@ async def test_workflow_nexus_payload_codec_selects_codec_from_context( hmac_endpoint_name, zlib_endpoint_name, } + assert { + endpoint: payload.metadata["encoding"] + for endpoint, payload in encoded_summaries.items() + } == { + hmac_endpoint_name: NexusContextPayloadCodecSelector.HMAC_ENCODING, + zlib_endpoint_name: NexusContextPayloadCodecSelector.ZLIB_ENCODING, + } + assert ( + encoded_summaries[hmac_endpoint_name].data + != encoded_summaries[zlib_endpoint_name].data + ) assert { endpoint: payload.metadata["encoding"] for endpoint, payload in encoded_results.items() @@ -1958,6 +1977,10 @@ async def test_workflow_nexus_payload_codec_selects_codec_from_context( [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] From 117a7fd41cb0eb81653226c01070630ff6e9cc25 Mon Sep 17 00:00:00 2001 From: Joshua Frenchwood Date: Thu, 10 Sep 2026 16:49:56 -0500 Subject: [PATCH 3/4] Moving decoding logic to the nexus operation handle --- temporalio/client/_impl.py | 24 ++---------- temporalio/client/_interceptor.py | 11 ++---- temporalio/client/_nexus.py | 45 ++++++++++++++++------- tests/nexus/test_standalone_operations.py | 4 +- 4 files changed, 40 insertions(+), 44 deletions(-) diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index 75abe021e..fde2061e4 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -114,7 +114,6 @@ NexusOperationExecutionAsyncIterator, NexusOperationExecutionCount, NexusOperationExecutionDescription, - NexusOperationFailureError, NexusOperationHandle, ) from ._schedule import ( @@ -1659,13 +1658,8 @@ async def describe_nexus_operation( async def get_nexus_operation_result( self, input: GetNexusOperationResultInput - ) -> Any: + ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse: """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, @@ -1684,20 +1678,8 @@ async def get_nexus_operation_result( timeout=input.rpc_timeout, ) ) - match res.WhichOneof("outcome"): - case "result": - type_hints = [input.result_type] if input.result_type else None - [result] = await data_converter.decode([res.result], type_hints) - return result - - case "failure": - raise NexusOperationFailureError( - cause=await data_converter.decode_failure(res.failure) - ) - - case None: - # poll again - pass + if res.WhichOneof("outcome") is not None: + return res except RPCError as err: match err.status: case RPCStatusCode.DEADLINE_EXCEEDED: diff --git a/temporalio/client/_interceptor.py b/temporalio/client/_interceptor.py index cf2179f10..e58f8cede 100644 --- a/temporalio/client/_interceptor.py +++ b/temporalio/client/_interceptor.py @@ -18,10 +18,7 @@ import temporalio.api.common.v1 import temporalio.api.workflowservice.v1 import temporalio.common -from temporalio.converter import ( - DataConverter, - NexusSerializationContext, -) +from temporalio.converter import DataConverter if TYPE_CHECKING: from ._activity import ( @@ -660,8 +657,6 @@ class GetNexusOperationResultInput: run_id: str | None rpc_metadata: Mapping[str, str | bytes] rpc_timeout: timedelta | None - result_type: type[Any] | None - _nexus_serialization_context: NexusSerializationContext | None = None @dataclass @@ -1005,9 +1000,11 @@ async def describe_nexus_operation( async def get_nexus_operation_result( self, input: GetNexusOperationResultInput - ) -> Any: + ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse: """Called for every :py:meth:`NexusOperationHandle.result` call. + The raw response is decoded by the handle after interception. + .. warning:: This API is experimental and unstable. """ diff --git a/temporalio/client/_nexus.py b/temporalio/client/_nexus.py index 5d3a291fb..5eb055e2b 100644 --- a/temporalio/client/_nexus.py +++ b/temporalio/client/_nexus.py @@ -1138,21 +1138,38 @@ async def result( RPCError: Operation result could not be fetched for some reason. """ 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, - _nexus_serialization_context=self._nexus_serialization_context, - ) + response = await self._client._impl.get_nexus_operation_result( + GetNexusOperationResultInput( + operation_id=self._operation_id, + run_id=self._run_id, + rpc_metadata=rpc_metadata, + rpc_timeout=rpc_timeout, + ) + ) + + data_converter = self._client.data_converter + if self._nexus_serialization_context is not None: + data_converter = data_converter.with_context( + self._nexus_serialization_context ) - return cast(ReturnType, self._known_outcome) - except NexusOperationFailureError as failure: - self._known_outcome = failure - raise + match response.WhichOneof("outcome"): + case "result": + type_hints = [self._result_type] if self._result_type else None + [result] = await data_converter.decode( + [response.result], type_hints + ) + self._known_outcome = result + return cast(ReturnType, result) + case "failure": + operation_failure = NexusOperationFailureError( + cause=await data_converter.decode_failure(response.failure) + ) + self._known_outcome = operation_failure + raise operation_failure + case None: + raise RuntimeError( + "Nexus operation result response did not contain an outcome" + ) elif isinstance(self._known_outcome, NexusOperationFailureError): raise self._known_outcome else: diff --git a/tests/nexus/test_standalone_operations.py b/tests/nexus/test_standalone_operations.py index 26a8316b4..a211473a3 100644 --- a/tests/nexus/test_standalone_operations.py +++ b/tests/nexus/test_standalone_operations.py @@ -22,6 +22,7 @@ ) import temporalio.api.enums.v1 +import temporalio.api.workflowservice.v1 from temporalio import nexus, workflow from temporalio.client import ( CancelNexusOperationInput, @@ -867,7 +868,7 @@ async def describe_nexus_operation( async def get_nexus_operation_result( self, input: GetNexusOperationResultInput - ) -> Any: + ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse: self._parent.result_calls.append(input) return await super().get_nexus_operation_result(input) @@ -980,7 +981,6 @@ async def test_interceptor_receives_inputs(client: Client, env: WorkflowEnvironm result_input = interceptor.result_calls[0] assert isinstance(result_input, GetNexusOperationResultInput) assert result_input.operation_id == op_id - assert result_input.result_type == EchoOutput # Start another so we can terminate it previous_start_count = len(interceptor.start_calls) From d7560489aaaa141029201dd14045b9ac319cf28b Mon Sep 17 00:00:00 2001 From: Joshua Frenchwood Date: Fri, 11 Sep 2026 15:35:49 -0500 Subject: [PATCH 4/4] Adding GetNexusOperationResultOutput for the interceptor --- CHANGELOG.md | 3 ++ temporalio/client/__init__.py | 2 + temporalio/client/_impl.py | 20 ++++++++-- temporalio/client/_interceptor.py | 28 ++++++++++++-- temporalio/client/_nexus.py | 46 ++++++++++++----------- tests/nexus/test_standalone_operations.py | 27 +++++++++++-- tests/test_serialization_context.py | 41 ++++++++++++++++++++ 7 files changed, 136 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 063543a86..cd18abd14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,9 @@ to include examples, links to docs, or any other relevant information. ### :boom: Breaking Changes +- The experimental `OutboundInterceptor.get_nexus_operation_result` method now returns + `GetNexusOperationResultOutput`, containing the raw result or failure and the data + converter used to decode it. - Experimental external storage: `ExternalStorage.driver_selector` is now called with a `StorageDriverSelectContext` instead of a `StorageDriverStoreContext`. Update the annotation; the new type carries the same `target` field. Since selectors are plain callables, a stale diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index 2eef41a39..2e947dfd1 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -112,6 +112,7 @@ FailAsyncActivityInput, FetchWorkflowHistoryEventsInput, GetNexusOperationResultInput, + GetNexusOperationResultOutput, GetWorkerBuildIdCompatibilityInput, GetWorkerTaskReachabilityInput, HeartbeatAsyncActivityInput, @@ -311,6 +312,7 @@ "StartNexusOperationInput", "DescribeNexusOperationInput", "GetNexusOperationResultInput", + "GetNexusOperationResultOutput", "CancelNexusOperationInput", "TerminateNexusOperationInput", "ListNexusOperationsInput", diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index fde2061e4..add692791 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -80,6 +80,7 @@ FailAsyncActivityInput, FetchWorkflowHistoryEventsInput, GetNexusOperationResultInput, + GetNexusOperationResultOutput, GetWorkerBuildIdCompatibilityInput, GetWorkerTaskReachabilityInput, HeartbeatAsyncActivityInput, @@ -1658,7 +1659,7 @@ async def describe_nexus_operation( async def get_nexus_operation_result( self, input: GetNexusOperationResultInput - ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse: + ) -> GetNexusOperationResultOutput: """Poll for nexus operation result until it's available.""" req = temporalio.api.workflowservice.v1.PollNexusOperationExecutionRequest( namespace=self._client.namespace, @@ -1678,8 +1679,21 @@ async def get_nexus_operation_result( timeout=input.rpc_timeout, ) ) - if res.WhichOneof("outcome") is not None: - return res + match res.WhichOneof("outcome"): + case "result": + return GetNexusOperationResultOutput( + raw_result=res.result, + raw_failure=None, + data_converter=input._data_converter, + ) + case "failure": + return GetNexusOperationResultOutput( + raw_result=None, + raw_failure=res.failure, + data_converter=input._data_converter, + ) + case None: + 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 e58f8cede..c90db27b5 100644 --- a/temporalio/client/_interceptor.py +++ b/temporalio/client/_interceptor.py @@ -8,7 +8,7 @@ Mapping, Sequence, ) -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import timedelta from typing import ( TYPE_CHECKING, @@ -16,6 +16,7 @@ ) import temporalio.api.common.v1 +import temporalio.api.failure.v1 import temporalio.api.workflowservice.v1 import temporalio.common from temporalio.converter import DataConverter @@ -657,6 +658,26 @@ class GetNexusOperationResultInput: run_id: str | None rpc_metadata: Mapping[str, str | bytes] rpc_timeout: timedelta | None + result_type: type[Any] | None + _data_converter: DataConverter = field(repr=False, compare=False) + + +@dataclass +class GetNexusOperationResultOutput: + """Output for :py:meth:`OutboundInterceptor.get_nexus_operation_result`. + + .. warning:: + This API is experimental and unstable. + """ + + raw_result: temporalio.api.common.v1.Payload | None + """Raw result payload if the operation succeeded.""" + + raw_failure: temporalio.api.failure.v1.Failure | None + """Raw failure if the operation failed.""" + + data_converter: DataConverter = field(repr=False, compare=False) + """Data converter for decoding the result or failure.""" @dataclass @@ -1000,10 +1021,11 @@ async def describe_nexus_operation( async def get_nexus_operation_result( self, input: GetNexusOperationResultInput - ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse: + ) -> GetNexusOperationResultOutput: """Called for every :py:meth:`NexusOperationHandle.result` call. - The raw response is decoded by the handle after interception. + The returned result or failure can be decoded with the output's + :py:attr:`GetNexusOperationResultOutput.data_converter`. .. warning:: This API is experimental and unstable. diff --git a/temporalio/client/_nexus.py b/temporalio/client/_nexus.py index 5eb055e2b..9d157d147 100644 --- a/temporalio/client/_nexus.py +++ b/temporalio/client/_nexus.py @@ -1138,38 +1138,40 @@ async def result( RPCError: Operation result could not be fetched for some reason. """ if self._known_outcome is temporalio.common._arg_unset: + # TODO: Use server-provided Nexus serialization context once available + # so reconstructed handles can also decode with context. + data_converter = self._client.data_converter + if self._nexus_serialization_context is not None: + data_converter = data_converter.with_context( + self._nexus_serialization_context + ) response = await self._client._impl.get_nexus_operation_result( GetNexusOperationResultInput( operation_id=self._operation_id, run_id=self._run_id, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout, + result_type=self._result_type, + _data_converter=data_converter, ) ) - - data_converter = self._client.data_converter - if self._nexus_serialization_context is not None: - data_converter = data_converter.with_context( - self._nexus_serialization_context + if response.raw_result is not None: + type_hints = [self._result_type] if self._result_type else None + [result] = await response.data_converter.decode( + [response.raw_result], type_hints ) - match response.WhichOneof("outcome"): - case "result": - type_hints = [self._result_type] if self._result_type else None - [result] = await data_converter.decode( - [response.result], type_hints - ) - self._known_outcome = result - return cast(ReturnType, result) - case "failure": - operation_failure = NexusOperationFailureError( - cause=await data_converter.decode_failure(response.failure) - ) - self._known_outcome = operation_failure - raise operation_failure - case None: - raise RuntimeError( - "Nexus operation result response did not contain an outcome" + self._known_outcome = result + return cast(ReturnType, result) + elif response.raw_failure is not None: + operation_failure = NexusOperationFailureError( + cause=await response.data_converter.decode_failure( + response.raw_failure ) + ) + self._known_outcome = operation_failure + raise operation_failure + else: + raise RuntimeError("Nexus operation result did not contain an outcome") elif isinstance(self._known_outcome, NexusOperationFailureError): raise self._known_outcome else: diff --git a/tests/nexus/test_standalone_operations.py b/tests/nexus/test_standalone_operations.py index a211473a3..635c49459 100644 --- a/tests/nexus/test_standalone_operations.py +++ b/tests/nexus/test_standalone_operations.py @@ -22,7 +22,6 @@ ) import temporalio.api.enums.v1 -import temporalio.api.workflowservice.v1 from temporalio import nexus, workflow from temporalio.client import ( CancelNexusOperationInput, @@ -30,6 +29,7 @@ CountNexusOperationsInput, DescribeNexusOperationInput, GetNexusOperationResultInput, + GetNexusOperationResultOutput, Interceptor, ListNexusOperationsInput, NexusOperationExecutionDescription, @@ -868,9 +868,11 @@ async def describe_nexus_operation( async def get_nexus_operation_result( self, input: GetNexusOperationResultInput - ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse: + ) -> GetNexusOperationResultOutput: self._parent.result_calls.append(input) - return await super().get_nexus_operation_result(input) + output = await super().get_nexus_operation_result(input) + self._parent.result_outputs.append(output) + return output async def cancel_nexus_operation(self, input: CancelNexusOperationInput) -> None: self._parent.cancel_calls.append(input) @@ -899,6 +901,7 @@ def __init__(self) -> None: self.start_calls: list[StartNexusOperationInput] = [] self.describe_calls: list[DescribeNexusOperationInput] = [] self.result_calls: list[GetNexusOperationResultInput] = [] + self.result_outputs: list[GetNexusOperationResultOutput] = [] self.cancel_calls: list[CancelNexusOperationInput] = [] self.terminate_calls: list[TerminateNexusOperationInput] = [] self.list_calls: list[ListNexusOperationsInput] = [] @@ -981,6 +984,24 @@ async def test_interceptor_receives_inputs(client: Client, env: WorkflowEnvironm result_input = interceptor.result_calls[0] assert isinstance(result_input, GetNexusOperationResultInput) assert result_input.operation_id == op_id + assert result_input.result_type == EchoOutput + assert len(interceptor.result_outputs) == 1 + assert interceptor.result_outputs[0].raw_result is None + assert interceptor.result_outputs[0].raw_failure is not None + + # Successful raw results and their data converter are also available. + 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 len(interceptor.result_outputs) == 2 + assert interceptor.result_outputs[1].raw_result is not None + assert interceptor.result_outputs[1].raw_failure is None # 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 52170b773..6b5894445 100644 --- a/tests/test_serialization_context.py +++ b/tests/test_serialization_context.py @@ -31,7 +31,11 @@ from temporalio.client import ( AsyncActivityHandle, Client, + GetNexusOperationResultInput, + GetNexusOperationResultOutput, + Interceptor, NexusOperationFailureError, + OutboundInterceptor, WorkflowFailureError, WorkflowUpdateFailedError, ) @@ -1692,6 +1696,35 @@ async def test_decode_context_matches_encode_context( # Test nexus payload codec +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 + ) -> GetNexusOperationResultOutput: + output = await super().get_nexus_operation_result(input) + if output.raw_result is not None: + type_hints = [input.result_type] if input.result_type else None + [result] = await output.data_converter.decode( + [output.raw_result], type_hints + ) + self._parent.decoded_results.append(result) + return output + + class NexusContextPayloadCodecSelector(PayloadCodec, WithSerializationContext): HMAC_ENCODING = b"binary/nexus-context-hmac" ZLIB_ENCODING = b"binary/nexus-context-zlib" @@ -2022,6 +2055,10 @@ async def test_standalone_nexus_payload_codec_selects_codec_from_context( {hmac_context: "hmac", zlib_context: "zlib"} ), ) + result_interceptor = _NexusResultDecodingInterceptor() + config["interceptors"] = list(config.get("interceptors") or []) + [ + result_interceptor + ] client = Client(**config) async with Worker( @@ -2053,6 +2090,10 @@ async def test_standalone_nexus_payload_codec_selects_codec_from_context( ) assert hmac_standalone_result == "standalone-hmac" assert zlib_standalone_result == "standalone-zlib" + assert set(result_interceptor.decoded_results) == { + "standalone-hmac", + "standalone-zlib", + } @pytest.mark.requires_local_server