diff --git a/temporalio/client/_activity.py b/temporalio/client/_activity.py index 99c9ede31..138d09dc7 100644 --- a/temporalio/client/_activity.py +++ b/temporalio/client/_activity.py @@ -309,6 +309,9 @@ class ActivityExecutionDescription(ActivityExecution): long_poll_token: bytes | None """Token for follow-on long-poll requests. None if the activity is complete.""" + raw_callbacks: Sequence[temporalio.api.activity.v1.CallbackInfo] + """Underlying protobuf callbacks""" + @classmethod async def _from_execution_info( cls, @@ -316,6 +319,7 @@ async def _from_execution_info( long_poll_token: bytes | None, namespace: str, data_converter: temporalio.converter.DataConverter, + callbacks: Sequence[temporalio.api.activity.v1.CallbackInfo], ) -> Self: """Create from raw proto activity execution info.""" # Decode heartbeat details if present @@ -409,6 +413,7 @@ async def _from_execution_info( typed_search_attributes=temporalio.converter.decode_typed_search_attributes( info.search_attributes ), + raw_callbacks=callbacks, ) diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index 7ba54746f..3e142f2b1 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -32,6 +32,7 @@ import temporalio.exceptions import temporalio.nexus import temporalio.nexus._operation_context +import temporalio.nexus._token from temporalio.activity import ActivityCancellationDetails from temporalio.converter import ( ActivitySerializationContext, @@ -246,7 +247,7 @@ async def _build_start_workflow_execution_request( # inside a Nexus operation handler must forward the inbound Nexus task links # explicitly so the started callee's WorkflowExecutionStarted event links back to # the caller. - if not temporalio.nexus._operation_context._in_nexus_backing_workflow_start_context(): + if not temporalio.nexus._operation_context._in_nexus_backing_start_context(): req.links.extend(nexus_ctx._get_request_links()) return req @@ -277,7 +278,7 @@ async def _build_signal_with_start_workflow_execution_request( # If this signal-with-start is issued from inside a Nexus operation handler (but not the # nexus-backing workflow), forward the inbound Nexus task links so both the callee's # WorkflowExecutionStarted and WorkflowExecutionSignaled events link back to the caller. - if not temporalio.nexus._operation_context._in_nexus_backing_workflow_start_context(): + if not temporalio.nexus._operation_context._in_nexus_backing_start_context(): nexus_ctx = ( temporalio.nexus._operation_context._try_start_operation_context() ) @@ -587,6 +588,11 @@ async def start_activity(self, input: StartActivityInput) -> ActivityHandle[Any] input.id, input.activity_type, run_id=details.run_id ) raise + + nexus_ctx = temporalio.nexus._operation_context._try_start_operation_context() + if nexus_ctx is not None: + nexus_ctx._add_response_link(resp.link) + return ActivityHandle( self._client, input.id, @@ -670,6 +676,44 @@ async def _build_start_activity_execution_request( # Set priority req.priority.CopyFrom(input.priority._to_proto()) + nexus_ctx = temporalio.nexus._operation_context._try_start_operation_context() + + request_links: list[temporalio.api.common.v1.Link] = [] + if nexus_ctx is not None: + req.on_conflict_options.attach_request_id = True + req.on_conflict_options.attach_completion_callbacks = True + req.on_conflict_options.attach_links = True + + # Add all Nexus links if we're in a Nexus context, backing or otherwise + request_links = nexus_ctx._get_request_links() + + # Add request ID and callbacks only if we're in a backing Nexus context + if ( + nexus_ctx is not None + and temporalio.nexus._operation_context._in_nexus_backing_start_context() + ): + req.request_id = nexus_ctx.nexus_context.request_id + callbacks = nexus_ctx._get_callbacks( + temporalio.nexus._token.OperationToken( + type=temporalio.nexus._token.OperationTokenType.ACTIVITY, + namespace=self._client.namespace, + activity_id=input.id, + ).encode() + ) + req.completion_callbacks.extend( + temporalio.api.common.v1.Callback( + nexus=temporalio.api.common.v1.Callback.Nexus( + url=callback.url, + header=callback.headers, + ), + links=request_links, + ) + for callback in callbacks + ) + + # Links are duplicated on request for compatibility with older server versions. + req.links.extend(request_links) + return req async def cancel_activity(self, input: CancelActivityInput) -> None: @@ -733,6 +777,7 @@ async def describe_activity( is_local=False, ) ), + callbacks=resp.callbacks, ) def list_activities( diff --git a/temporalio/nexus/__init__.py b/temporalio/nexus/__init__.py index 003df44df..3abc9b0f2 100644 --- a/temporalio/nexus/__init__.py +++ b/temporalio/nexus/__init__.py @@ -25,6 +25,7 @@ wait_for_worker_shutdown_sync, ) from ._operation_handlers import ( + CancelActivityOptions, CancelUpdateWorkflowOptions, CancelWorkflowRunOptions, TemporalOperationHandler, @@ -34,6 +35,7 @@ __all__ = ( "workflow_run_operation", + "CancelActivityOptions", "CancelWorkflowRunOptions", "CancelUpdateWorkflowOptions", "Info", diff --git a/temporalio/nexus/_link_conversion.py b/temporalio/nexus/_link_conversion.py index 54ed3869a..e3ef3988b 100644 --- a/temporalio/nexus/_link_conversion.py +++ b/temporalio/nexus/_link_conversion.py @@ -23,6 +23,10 @@ r"^/namespaces/(?P[^/]+)/nexus-operations/(?P[^/]+)/(?P[^/]*)/details$" ) +_ACTIVITY_LINK_URL_PATH_REGEX = re.compile( + r"^/namespaces/(?P[^/]+)/activities/(?P[^/]+)/(?P[^/]+)/details$" +) + _WORKFLOW_LINK_URL_PATH_REGEX = re.compile( r"^/namespaces/(?P[^/]+)/workflows/(?P[^/]+)/(?P[^/]+)(?P/history)?$" ) @@ -32,6 +36,7 @@ class _LinkType(str, Enum): WORKFLOW_EVENT = temporalio.api.common.v1.Link.WorkflowEvent.DESCRIPTOR.full_name WORKFLOW = temporalio.api.common.v1.Link.Workflow.DESCRIPTOR.full_name NEXUS_OPERATION = temporalio.api.common.v1.Link.NexusOperation.DESCRIPTOR.full_name + ACTIVITY = temporalio.api.common.v1.Link.Activity.DESCRIPTOR.full_name LINK_EVENT_ID_PARAM_NAME = "eventID" @@ -88,6 +93,9 @@ def nexus_link_to_temporal_link( case _LinkType.NEXUS_OPERATION: return nexus_link_to_nexus_operation_link(nexus_link) + case _LinkType.ACTIVITY: + return nexus_link_to_activity_link(nexus_link) + def temporal_link_to_nexus_link( temporal_link: temporalio.api.common.v1.Link, @@ -106,10 +114,11 @@ def temporal_link_to_nexus_link( case "nexus_operation": return nexus_operation_to_nexus_link(temporal_link.nexus_operation) - case "activity" | "batch_job": - raise NotImplementedError( - "only workflow_event and nexus operation links are supported" - ) + case "activity": + return activity_link_to_nexus_link(temporal_link.activity) + + case "batch_job": + raise NotImplementedError("batch_job links are not supported") case None: logger.warning("Invalid Temporal link: missing variant") @@ -190,6 +199,17 @@ def nexus_operation_to_nexus_link( ) +def activity_link_to_nexus_link( + activity: temporalio.api.common.v1.Link.Activity, +) -> nexusrpc.Link: + """Convert an Activity link into a nexusrpc link.""" + namespace = urllib.parse.quote(activity.namespace, safe="") + activity_id = urllib.parse.quote(activity.activity_id, safe="") + run_id = urllib.parse.quote(activity.run_id, safe="") + path = f"/namespaces/{namespace}/activities/{activity_id}/{run_id}/details" + return nexusrpc.Link(url=_temporal_nexus_url(path), type=_LinkType.ACTIVITY.value) + + def _workflow_nexus_url( namespace: str, workflow_id: str, @@ -334,6 +354,28 @@ def nexus_link_to_nexus_operation_link( return temporalio.api.common.v1.Link(nexus_operation=nexus_op_link) +def nexus_link_to_activity_link( + nexus_link: nexusrpc.Link, +) -> temporalio.api.common.v1.Link | None: + """Convert a Nexus Activity link into a Temporal Activity link.""" + match = _ACTIVITY_LINK_URL_PATH_REGEX.match( + urllib.parse.urlparse(nexus_link.url).path + ) + if not match: + logger.warning( + f"Invalid Nexus link: {nexus_link}. Expected path to match {_ACTIVITY_LINK_URL_PATH_REGEX.pattern}" + ) + return None + groups = match.groupdict() + return temporalio.api.common.v1.Link( + activity=temporalio.api.common.v1.Link.Activity( + namespace=urllib.parse.unquote(groups["namespace"]), + activity_id=urllib.parse.unquote(groups["activity_id"]), + run_id=urllib.parse.unquote(groups["run_id"]), + ) + ) + + def _event_reference_to_query_params( event_ref: temporalio.api.common.v1.Link.WorkflowEvent.EventReference, ) -> str: diff --git a/temporalio/nexus/_operation_context.py b/temporalio/nexus/_operation_context.py index 7206e433e..7d391513e 100644 --- a/temporalio/nexus/_operation_context.py +++ b/temporalio/nexus/_operation_context.py @@ -46,7 +46,6 @@ from ._link_conversion import ( nexus_link_to_temporal_link, temporal_link_to_nexus_link, - workflow_event_to_nexus_link, workflow_execution_started_event_link_from_workflow_handle, ) from ._token import OperationToken, OperationTokenType, WorkflowHandle @@ -66,12 +65,12 @@ ContextVar("temporal-cancel-operation-context") ) -# A Nexus start handler might start zero or more workflows as usual using a Temporal client. In -# addition, it may start one "nexus-backing" workflow, using -# WorkflowRunOperationContext.start_workflow. This context is active while the latter is being done. +# A Nexus start handler might start zero or more async Temporal actions as usual using a Temporal client. In +# addition, it may start one "nexus-backing" async Temporal action, using +# WorkflowRunOperationContext.start_workflow or methods from TemporalNexusClient. This context is active while the latter is being done. # It is thus a narrower context than _temporal_start_operation_context. -_temporal_nexus_backing_workflow_start_context: ContextVar[bool] = ContextVar( - "temporal-nexus-backing-workflow-start-context" +_temporal_nexus_backing_start_context: ContextVar[bool] = ContextVar( + "temporal-nexus-backing-start-context" ) @@ -170,21 +169,21 @@ def _try_temporal_context() -> ( def _try_start_operation_context() -> _TemporalStartOperationContext | None: # pyright: ignore[reportUnusedFunction] - """The Nexus start-operation context if a handler is currently running, else None.""" + """Return the active Nexus start-operation context, if any.""" return _temporal_start_operation_context.get(None) @contextmanager -def _nexus_backing_workflow_start_context() -> Generator[None]: - token = _temporal_nexus_backing_workflow_start_context.set(True) +def _nexus_backing_start_context() -> Generator[None]: + token = _temporal_nexus_backing_start_context.set(True) try: yield finally: - _temporal_nexus_backing_workflow_start_context.reset(token) + _temporal_nexus_backing_start_context.reset(token) -def _in_nexus_backing_workflow_start_context() -> bool: # type:ignore[reportUnusedClass] - return _temporal_nexus_backing_workflow_start_context.get(False) +def _in_nexus_backing_start_context() -> bool: # type:ignore[reportUnusedClass] + return _temporal_nexus_backing_start_context.get(False) _OperationCtxT = TypeVar("_OperationCtxT", bound=OperationContext) @@ -254,11 +253,11 @@ def _get_request_links(self) -> list[temporalio.api.common.v1.Link]: ``links`` field so the callee's history event links back to whatever scheduled this Nexus operation. """ - event_links: list[temporalio.api.common.v1.Link] = [] + links: list[temporalio.api.common.v1.Link] = [] for inbound_link in self.nexus_context.inbound_links: if link := nexus_link_to_temporal_link(inbound_link): - event_links.append(link) - return event_links + links.append(link) + return links def _add_start_workflow_response_link( self, workflow_handle: temporalio.client.WorkflowHandle[Any, Any] @@ -302,17 +301,17 @@ def _add_response_link(self, link: temporalio.api.common.v1.Link | None) -> None """Append a response link returned by an RPC the operation handler issued. ``link`` is the ``common.v1.Link`` returned on a signal, signal-with-start, or start - response (or ``None`` against a server that did not return one). When present and of the - ``workflow_event`` variant, it is converted to a Nexus link and added to the operation's - outbound links so the caller workflow's Nexus history event links to the callee event. + response (or ``None`` against a server that did not return one). When present, it is + converted to a Nexus link and added to the operation's outbound links. This is only safe to call from the single thread/task that runs the operation handler. """ - if link is None or not link.HasField("workflow_event"): - return - self.nexus_context.outbound_links.append( - workflow_event_to_nexus_link(link.workflow_event) - ) + if link is not None: + try: + if response_link := temporal_link_to_nexus_link(link): + self.nexus_context.outbound_links.append(response_link) + except Exception as e: + logger.warning(f"Failed to create Nexus link from Temporal link: {e}") class WorkflowRunOperationContext(StartOperationContext): @@ -668,7 +667,7 @@ async def _start_nexus_backing_workflow( priority: temporalio.common.Priority = temporalio.common.Priority.default, versioning_override: temporalio.common.VersioningOverride | None = None, ) -> WorkflowHandle[ReturnType]: - # We must pass nexus_completion_callbacks, workflow_event_links, and request_id, + # We must pass nexus_completion_callbacks, links, and request_id, # but these are deliberately not exposed in overloads, hence the type-check # violation. @@ -677,7 +676,7 @@ async def _start_nexus_backing_workflow( # namespace to deliver the result to the caller namespace when the workflow reaches a # terminal state) and inbound links to the caller workflow (attached to history events of # the workflow started in the handler namespace, and displayed in the UI). - with _nexus_backing_workflow_start_context(): + with _nexus_backing_start_context(): token = OperationToken( type=OperationTokenType.WORKFLOW, namespace=temporal_context.client.namespace, diff --git a/temporalio/nexus/_operation_handlers.py b/temporalio/nexus/_operation_handlers.py index 0ad6c267c..36c44e96c 100644 --- a/temporalio/nexus/_operation_handlers.py +++ b/temporalio/nexus/_operation_handlers.py @@ -157,6 +157,16 @@ class CancelUpdateWorkflowOptions: """The workflow runID that accepted the update.""" +@dataclass(frozen=True) +class CancelActivityOptions: + """Options for cancelling the activity backing a Nexus operation.""" + + activity_id: str + """The activity ID to cancel.""" + run_id: str + """The run ID of the activity to cancel.""" + + class TemporalOperationHandler(OperationHandler[InputT, OutputT], ABC): """Operation handler for Nexus operations that interact with Temporal. Implementations override the start_operation method. @@ -200,16 +210,19 @@ async def cancel(self, ctx: CancelOperationContext, token: str) -> None: raise HandlerError( "Unable to decode operation token to cancel", type=HandlerErrorType.INTERNAL, + retryable_override=False, ) from err cancel_ctx = TemporalCancelOperationContext._from_cancel_operation_context(ctx) match operation_token.type: case OperationTokenType.WORKFLOW: + assert operation_token.workflow_id is not None options = CancelWorkflowRunOptions( workflow_id=operation_token.workflow_id ) await self.cancel_workflow_run(cancel_ctx, options) case OperationTokenType.UPDATE_WORKFLOW: + assert operation_token.workflow_id is not None assert operation_token.update_id is not None assert operation_token.run_id is not None cancel_options = CancelUpdateWorkflowOptions( @@ -218,6 +231,22 @@ async def cancel(self, ctx: CancelOperationContext, token: str) -> None: run_id=operation_token.run_id, ) await self.cancel_workflow_update(cancel_ctx, cancel_options) + case OperationTokenType.ACTIVITY: + assert operation_token.activity_id is not None + if not operation_token.run_id: + raise HandlerError( + "Expected operation token of type ACTIVITY to have a valid run id.", + type=HandlerErrorType.INTERNAL, + retryable_override=False, + ) + + await self.cancel_activity( + cancel_ctx, + CancelActivityOptions( + activity_id=operation_token.activity_id, + run_id=operation_token.run_id, + ), + ) async def cancel_workflow_run( self, @@ -253,3 +282,14 @@ async def cancel_workflow_update( """, type=HandlerErrorType.NOT_IMPLEMENTED, ) + + async def cancel_activity( + self, + ctx: TemporalCancelOperationContext, # pyright: ignore[reportUnusedParameter] + options: CancelActivityOptions, + ) -> None: + """Requests cancellation of the standalone activity backing the operation.""" + activity_handle = temporalio.nexus.client().get_activity_handle( + options.activity_id, run_id=options.run_id + ) + await activity_handle.cancel() diff --git a/temporalio/nexus/_temporal_client.py b/temporalio/nexus/_temporal_client.py index 479d1651b..393da2fae 100644 --- a/temporalio/nexus/_temporal_client.py +++ b/temporalio/nexus/_temporal_client.py @@ -21,11 +21,17 @@ import temporalio.common from temporalio.nexus._operation_context import ( + _nexus_backing_start_context, _start_nexus_backing_workflow, _start_nexus_operation_workflow_update, _TemporalStartOperationContext, ) +from temporalio.nexus._token import OperationToken, OperationTokenType from temporalio.types import ( + CallableAsyncNoParam, + CallableAsyncSingleParam, + CallableSyncNoParam, + CallableSyncSingleParam, MethodAsyncNoParam, MethodAsyncSingleParam, MultiParamSpec, @@ -34,8 +40,6 @@ SelfType, ) -from ._token import OperationToken, OperationTokenType - if TYPE_CHECKING: import temporalio.client import temporalio.workflow @@ -368,6 +372,207 @@ async def start_workflow_update( """ ... + # async no-param activity + @overload + async def start_activity( + self, + activity: CallableAsyncNoParam[ReturnType], + *, + id: str, + task_queue: str | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: ... + + # sync no-param activity + @overload + async def start_activity( + self, + activity: CallableSyncNoParam[ReturnType], + *, + id: str, + task_queue: str | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: ... + + # async single-param activity + @overload + async def start_activity( + self, + activity: CallableAsyncSingleParam[ParamType, ReturnType], + arg: ParamType, + *, + id: str, + task_queue: str | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: ... + + # sync single-param activity + @overload + async def start_activity( + self, + activity: CallableSyncSingleParam[ParamType, ReturnType], + arg: ParamType, + *, + id: str, + task_queue: str | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: ... + + # async multi-param activity + @overload + async def start_activity( + self, + activity: Callable[..., Awaitable[ReturnType]], + *, + args: Sequence[Any], + id: str, + task_queue: str | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: ... + + # sync multi-param activity + @overload + async def start_activity( + self, + activity: Callable[..., ReturnType], + *, + args: Sequence[Any], + id: str, + task_queue: str | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: ... + + # string-name activity + @overload + async def start_activity( + self, + activity: str, + arg: Any = temporalio.common._arg_unset, + *, + args: Sequence[Any] = [], + id: str, + task_queue: str | None = None, + result_type: type[ReturnType] | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: ... + + @abstractmethod + async def start_activity( + self, + activity: ( + str | Callable[..., Awaitable[ReturnType]] | Callable[..., ReturnType] + ), + arg: Any = temporalio.common._arg_unset, + *, + args: Sequence[Any] = [], + id: str, + task_queue: str | None = None, + result_type: type | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: + """Start a standalone activity that will deliver the Nexus operation result. + + If ``task_queue`` is not specified, the Nexus worker's task queue is used. + See :py:meth:`temporalio.client.Client.start_activity` for all other arguments. + """ + ... + class _TemporalNexusClient(TemporalNexusClient): # pyright: ignore[reportUnusedClass] """Nexus-aware wrapper around a Temporal Client. @@ -518,3 +723,72 @@ async def start_workflow_update( run_id=update_handle.workflow_run_id, ).encode() return TemporalOperationResult.async_token(token) + + async def start_activity( + self, + activity: ( + str | Callable[..., Awaitable[ReturnType]] | Callable[..., ReturnType] + ), + arg: Any = temporalio.common._arg_unset, + *, + args: Sequence[Any] = [], + id: str, + task_queue: str | None = None, + result_type: type | None = None, + schedule_to_start_timeout: timedelta | None = None, + schedule_to_close_timeout: timedelta | None = None, + start_to_close_timeout: timedelta | None = None, + heartbeat_timeout: timedelta | None = None, + start_delay: timedelta | None = None, + id_reuse_policy: temporalio.common.ActivityIDReusePolicy = temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy: temporalio.common.ActivityIDConflictPolicy = temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy: temporalio.common.RetryPolicy | None = None, + search_attributes: temporalio.common.TypedSearchAttributes | None = None, + summary: str | None = None, + priority: temporalio.common.Priority = temporalio.common.Priority.default, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> TemporalOperationResult[ReturnType]: + """Start a standalone activity that will deliver the Nexus operation result. + + If ``task_queue`` is not specified, the Nexus worker's task queue is used. + See :py:meth:`temporalio.client.Client.start_activity` for all other arguments. + """ + with self._reserve_async_start(): + # Here we are starting a "nexus-backing" standalone activity. The start request + # carries the Nexus completion callback so the activity result is delivered to + # the Nexus caller when the activity reaches a terminal state. + + with _nexus_backing_start_context(): + activity_handle: temporalio.client.ActivityHandle[ + ReturnType + ] = await self._temporal_context.client.start_activity( + activity=activity, # type: ignore + arg=arg, + args=args, + id=id, + task_queue=task_queue or self._temporal_context.info().task_queue, + result_type=result_type, + schedule_to_start_timeout=schedule_to_start_timeout, + schedule_to_close_timeout=schedule_to_close_timeout, + start_to_close_timeout=start_to_close_timeout, + heartbeat_timeout=heartbeat_timeout, + start_delay=start_delay, + id_reuse_policy=id_reuse_policy, + id_conflict_policy=id_conflict_policy, + retry_policy=retry_policy, + search_attributes=search_attributes, + summary=summary, + priority=priority, + rpc_metadata=rpc_metadata, + rpc_timeout=rpc_timeout, + ) + + activity_token = OperationToken( + type=OperationTokenType.ACTIVITY, + namespace=self._temporal_context.client.namespace, + activity_id=activity_handle.id, + run_id=activity_handle.run_id, + ) + + return TemporalOperationResult.async_token(activity_token.encode()) diff --git a/temporalio/nexus/_token.py b/temporalio/nexus/_token.py index fe0a466c9..a7c732f2a 100644 --- a/temporalio/nexus/_token.py +++ b/temporalio/nexus/_token.py @@ -14,6 +14,7 @@ class OperationTokenType(IntEnum): """Type discriminator for Nexus operation tokens.""" WORKFLOW = 1 + ACTIVITY = 2 UPDATE_WORKFLOW = 3 @@ -28,7 +29,8 @@ class OperationToken: version: int | None = None type: OperationTokenType namespace: str - workflow_id: str + workflow_id: str | None = None + activity_id: str | None = None run_id: str | None = None update_id: str | None = None @@ -37,8 +39,11 @@ def encode(self) -> str: token_details: dict[str, Any] = { "t": self.type, "ns": self.namespace, - "wid": self.workflow_id, } + if self.workflow_id is not None: + token_details["wid"] = self.workflow_id + if self.activity_id is not None: + token_details["aid"] = self.activity_id if self.version is not None: token_details["v"] = self.version if self.run_id is not None: @@ -90,7 +95,7 @@ def decode(cls, token: str) -> Self: ) workflow_id = token_details.get("wid") - if not isinstance(workflow_id, str): + if workflow_id is not None and not isinstance(workflow_id, str): raise TypeError( f"invalid token: expected workflow id to be a string, got {type(workflow_id)}" ) @@ -104,6 +109,17 @@ def decode(cls, token: str) -> Self: f"invalid token: expected non-empty workflow id for token type `{token_type.name}`" ) + activity_id = token_details.get("aid") + if activity_id is not None and not isinstance(activity_id, str): + raise TypeError( + f"invalid token: expected activity id to be a string, got {type(activity_id)}" + ) + + if token_type == OperationTokenType.ACTIVITY and not activity_id: + raise TypeError( + f"invalid token: expected non-empty activity id for token type `{token_type.name}`" + ) + update_id = token_details.get("uid") if not isinstance(update_id, str | None): raise TypeError( @@ -123,7 +139,6 @@ def decode(cls, token: str) -> Self: ) run_id = token_details.get("rid") - if not isinstance(run_id, str | None): raise TypeError( f"invalid token: expected run_id to be a string or None, got {type(run_id)}" @@ -133,6 +148,7 @@ def decode(cls, token: str) -> Self: type=OperationTokenType(token_type), namespace=namespace, workflow_id=workflow_id, + activity_id=activity_id, run_id=run_id, version=version, update_id=update_id, @@ -199,6 +215,9 @@ def from_token(cls, token: str) -> WorkflowHandle[OutputT]: f"invalid workflow token type: {op_token.type}, expected: {OperationTokenType.WORKFLOW}" ) + if not op_token.workflow_id: + raise TypeError("invalid workflow token: missing workflow id") + if op_token.version is not None and op_token.version != 0: raise TypeError( "invalid workflow token: 'v' field, if present, must be 0 or null/absent" diff --git a/tests/__init__.py b/tests/__init__.py index 4725d3a7e..af97849fe 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1 +1 @@ -DEV_SERVER_DOWNLOAD_VERSION = "v1.7.1-system-nexus-operations" +DEV_SERVER_DOWNLOAD_VERSION = "v1.7.4-standalone-nexus-operations" diff --git a/tests/conftest.py b/tests/conftest.py index bfe005ede..70dc6cff5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -143,6 +143,8 @@ async def env(env_type: str) -> AsyncGenerator[WorkflowEnvironment, None]: "history.enableSignalWithStartFromWorkflow=true", "--dynamic-config-value", "history.enableUpdateCallbacks=true", + "--dynamic-config-value", + "activity.enableCallbacks=true", ], dev_server_download_version=DEV_SERVER_DOWNLOAD_VERSION, ) diff --git a/tests/nexus/test_link_conversion.py b/tests/nexus/test_link_conversion.py index 4afe3367e..15fc8d77f 100644 --- a/tests/nexus/test_link_conversion.py +++ b/tests/nexus/test_link_conversion.py @@ -304,6 +304,19 @@ def test_link_conversion_workflow_to_link_and_back( url="temporal:///namespaces/ns/nexus-operations/op%2Fid//details", ), ), + ( + temporalio.api.common.v1.Link( + nexus_operation=temporalio.api.common.v1.Link.NexusOperation( + namespace="ns", + operation_id="op/id", + run_id="run/id", + ) + ), + nexusrpc.Link( + type=temporalio.api.common.v1.Link.NexusOperation.DESCRIPTOR.full_name, + url="temporal:///namespaces/ns/nexus-operations/op%2Fid/run%2Fid/details", + ), + ), ], ) def test_link_conversion_nexus_operation_to_link_and_back( @@ -342,6 +355,62 @@ def test_nexus_operation_link_with_unparseable_url_is_ignored(): assert temporalio.nexus._link_conversion.nexus_link_to_temporal_link(link) is None +@pytest.mark.parametrize( + ["link", "expected_link"], + [ + ( + nexusrpc.Link( + type=temporalio.api.common.v1.Link.Activity.DESCRIPTOR.full_name, + url="temporal:///namespaces/ns/activities/act-id/run-id/details", + ), + temporalio.api.common.v1.Link( + activity=temporalio.api.common.v1.Link.Activity( + namespace="ns", + activity_id="act-id", + run_id="run-id", + ), + ), + ), + ( + nexusrpc.Link( + type=temporalio.api.common.v1.Link.Activity.DESCRIPTOR.full_name, + url="temporal:///namespaces/ns%2F/activities/act-id%2F/run-id%3E/details", + ), + temporalio.api.common.v1.Link( + activity=temporalio.api.common.v1.Link.Activity( + namespace="ns/", + activity_id="act-id/", + run_id="run-id>", + ), + ), + ), + ], +) +def test_link_conversion_nexus_link_to_activity_link( + link: nexusrpc.Link, + expected_link: temporalio.api.common.v1.Link, +): + from_activity_link = temporalio.nexus._link_conversion.activity_link_to_nexus_link( + expected_link.activity + ) + assert link == from_activity_link + + from_temporal_link = temporalio.nexus._link_conversion.temporal_link_to_nexus_link( + expected_link + ) + assert link == from_temporal_link + + actual_activity = temporalio.nexus._link_conversion.nexus_link_to_activity_link( + link + ) + assert expected_link == actual_activity + + actual_temporal_link = ( + temporalio.nexus._link_conversion.nexus_link_to_temporal_link(link) + ) + assert expected_link == actual_temporal_link + + def test_link_conversion_utilities(): p2c = temporalio.nexus._link_conversion._event_type_pascal_case_to_constant_case c2p = temporalio.nexus._link_conversion._event_type_constant_case_to_pascal_case diff --git a/tests/nexus/test_signal_link_propagation.py b/tests/nexus/test_link_propagation.py similarity index 87% rename from tests/nexus/test_signal_link_propagation.py rename to tests/nexus/test_link_propagation.py index daff71c78..f34446563 100644 --- a/tests/nexus/test_signal_link_propagation.py +++ b/tests/nexus/test_link_propagation.py @@ -1,14 +1,14 @@ -"""Unit tests for Nexus signal-backlink propagation. +"""Unit tests for Nexus link propagation. -These exercise the in/out link propagation that happens when a Nexus operation handler issues a -signal, signal-with-start, or start-workflow RPC, against a mocked workflow service. -The corresponding end-to-end behavior requires a real server with EnableCHASMSignalBacklinks=true and is therefore -and is therefore not covered here. +These exercise link propagation when a Nexus operation handler signals or starts a +workflow or activity against a mocked workflow service. End-to-end signal backlinks +require a server with EnableCHASMSignalBacklinks enabled and are not covered here. """ from __future__ import annotations from collections.abc import Generator +from datetime import timedelta from typing import Any from unittest import mock @@ -32,9 +32,11 @@ import temporalio.converter import temporalio.nexus._link_conversion import temporalio.nexus._operation_context +import temporalio.nexus._token from temporalio.client._impl import _ClientImpl from temporalio.client._interceptor import ( SignalWorkflowInput, + StartActivityInput, StartWorkflowInput, ) from temporalio.nexus._operation_context import _TemporalStartOperationContext @@ -84,7 +86,7 @@ def nexus_ctx() -> Generator[_TemporalStartOperationContext]: operation="op", headers={}, request_id="req-id", - callback_url=None, + callback_url="https://callback.example", inbound_links=[inbound], callback_headers={}, task_cancellation=_NexusTaskCancellation(), @@ -163,6 +165,30 @@ def _start_input(start_signal: str | None = None) -> StartWorkflowInput: ) +def _start_activity_input() -> StartActivityInput: + return StartActivityInput( + activity_type="TestActivity", + args=[], + id="activity-target", + task_queue="tq", + result_type=None, + schedule_to_close_timeout=None, + start_to_close_timeout=timedelta(seconds=10), + schedule_to_start_timeout=None, + heartbeat_timeout=None, + id_reuse_policy=temporalio.common.ActivityIDReusePolicy.ALLOW_DUPLICATE, + id_conflict_policy=temporalio.common.ActivityIDConflictPolicy.FAIL, + retry_policy=None, + priority=temporalio.common.Priority.default, + search_attributes=None, + summary=None, + start_delay=None, + headers={}, + rpc_metadata={}, + rpc_timeout=None, + ) + + def _outbound_link_urls(ctx: Any) -> list[str]: return [link.url for link in ctx.nexus_context.outbound_links] @@ -433,7 +459,7 @@ async def test_backing_workflow_start_sets_on_conflict_options_without_duplicati # also re-add the context's request links. start_input = _start_input() start_input.links = [_inbound_nexus_link()] - with temporalio.nexus._operation_context._nexus_backing_workflow_start_context(): + with temporalio.nexus._operation_context._nexus_backing_start_context(): await impl.start_workflow(start_input) sent = workflow_service.start_workflow_execution.call_args.args[0] @@ -459,6 +485,43 @@ async def test_start_outside_nexus_context_leaves_on_conflict_options_unset() -> assert not sent.HasField("on_conflict_options") +# ── activity start ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.usefixtures("nexus_ctx") +async def test_activity_start_forwards_inbound_links() -> None: + impl = _make_client_impl(mock.MagicMock()) + + req = await impl._build_start_activity_execution_request(_start_activity_input()) + + assert len(req.links) == 1 + assert req.links[0] == _inbound_nexus_link() + assert not req.request_id + assert len(req.completion_callbacks) == 0 + + +@pytest.mark.usefixtures("nexus_ctx") +async def test_backing_activity_start_gets_nexus_request_fields() -> None: + impl = _make_client_impl(mock.MagicMock()) + + with temporalio.nexus._operation_context._nexus_backing_start_context(): + req = await impl._build_start_activity_execution_request( + _start_activity_input() + ) + + assert len(req.links) == 1 + assert req.links[0] == _inbound_nexus_link() + assert req.request_id == "req-id" + assert len(req.completion_callbacks) == 1 + operation_token = temporalio.nexus._token.OperationToken.decode( + req.completion_callbacks[0].nexus.header["nexus-operation-token"] + ) + assert operation_token.type is temporalio.nexus._token.OperationTokenType.ACTIVITY + assert operation_token.namespace == NAMESPACE + assert operation_token.activity_id == "activity-target" + assert list(req.completion_callbacks[0].links) == [_inbound_nexus_link()] + + # ── handler-level: backlinks land on the StartOperationResponse ────────────────────────────── # A response link that a handler stashes on ctx.outbound_links, mimicking what a signal RPC inside diff --git a/tests/nexus/test_nexus_type_errors.py b/tests/nexus/test_nexus_type_errors.py index 1486f9791..03af5d43f 100644 --- a/tests/nexus/test_nexus_type_errors.py +++ b/tests/nexus/test_nexus_type_errors.py @@ -11,7 +11,7 @@ import nexusrpc import temporalio.nexus -from temporalio import workflow +from temporalio import activity, workflow from temporalio.client import Client, NexusOperationHandle from temporalio.nexus import TemporalOperationStartHandlerFunc from temporalio.service import ServiceClient @@ -71,6 +71,40 @@ async def run( pass +@activity.defn +async def my_no_arg_activity() -> None: + pass + + +@activity.defn +async def my_one_arg_activity(_input: MyInput) -> None: + pass + + +@activity.defn +async def my_two_arg_activity(_input: MyInput, _arg2: int) -> None: + pass + + +@activity.defn +async def my_three_arg_activity(_input: MyInput, _arg2: int, _arg3: int) -> None: + pass + + +@activity.defn +async def my_four_arg_activity( + _input: MyInput, _arg2: int, _arg3: int, _arg4: int +) -> None: + pass + + +@activity.defn +async def my_five_arg_activity( + _input: MyInput, _arg2: int, _arg3: int, _arg4: int, _arg5: int +) -> None: + pass + + @nexusrpc.service class MyService: my_sync_operation: nexusrpc.Operation[MyInput, MyOutput] @@ -105,8 +139,9 @@ async def my_temporal_operation( input: int, ) -> temporalio.nexus.TemporalOperationResult[None]: """ - Typed proc workflow starts from a generic Temporal Nexus operation handler - infer TemporalOperationResult[None] for 0 to 5 workflow parameters. + Typed proc workflow and activity starts from a generic Temporal Nexus + operation handler infer TemporalOperationResult[None] for 0 to 5 + workflow or activity parameters. """ if input == 0: result_0: temporalio.nexus.TemporalOperationResult[ @@ -154,6 +189,78 @@ async def my_temporal_operation( id="proc-5", ) return result_5 + + # Typed activity starts infer TemporalOperationResult[None] for 0 to 5 + # activity parameters. Activities require a start_to_close_timeout. + if input == 6: + activity_result_0: temporalio.nexus.TemporalOperationResult[ + None + ] = await client.start_activity( + my_no_arg_activity, + id="activity-0", + start_to_close_timeout=timedelta(seconds=5), + ) + return activity_result_0 + if input == 7: + activity_result_1: temporalio.nexus.TemporalOperationResult[ + None + ] = await client.start_activity( + my_one_arg_activity, + MyInput(), + id="activity-1", + start_to_close_timeout=timedelta(seconds=5), + ) + return activity_result_1 + if input == 8: + activity_result_2: temporalio.nexus.TemporalOperationResult[ + None + ] = await client.start_activity( + my_two_arg_activity, + args=[MyInput(), 2], + id="activity-2", + start_to_close_timeout=timedelta(seconds=5), + ) + return activity_result_2 + if input == 9: + activity_result_3: temporalio.nexus.TemporalOperationResult[ + None + ] = await client.start_activity( + my_three_arg_activity, + args=[MyInput(), 2, 3], + id="activity-3", + start_to_close_timeout=timedelta(seconds=5), + ) + return activity_result_3 + if input == 10: + activity_result_4: temporalio.nexus.TemporalOperationResult[ + None + ] = await client.start_activity( + my_four_arg_activity, + args=[MyInput(), 2, 3, 4], + id="activity-4", + start_to_close_timeout=timedelta(seconds=5), + ) + return activity_result_4 + if input == 11: + activity_result_5: temporalio.nexus.TemporalOperationResult[ + None + ] = await client.start_activity( + my_five_arg_activity, + args=[MyInput(), 2, 3, 4, 5], + id="activity-5", + start_to_close_timeout=timedelta(seconds=5), + ) + return activity_result_5 + if input == 12: + # assert-type-error-pyright: 'No overloads for "start_activity" match' + return await client.start_activity( # type: ignore + my_one_arg_activity, + # assert-type-error-pyright: 'Argument of type .+ cannot be assigned to parameter' + "wrong-input-type", # type: ignore + id="activity-wrong-input", + start_to_close_timeout=timedelta(seconds=5), + ) + # assert-type-error-pyright: 'No overloads for "start_workflow" match' return await client.start_workflow( # type: ignore MyOneArgProcWorkflow.run, diff --git a/tests/nexus/test_operation_token.py b/tests/nexus/test_operation_token.py index 385f4f872..58d4a7859 100644 --- a/tests/nexus/test_operation_token.py +++ b/tests/nexus/test_operation_token.py @@ -8,6 +8,7 @@ OperationToken, OperationTokenType, WorkflowHandle, + _base64url_decode_no_padding, ) @@ -36,6 +37,39 @@ def test_operation_token_encode_decode_round_trip(): ) +def test_operation_token_activity_encode_decode_round_trip(): + token = OperationToken( + type=OperationTokenType.ACTIVITY, + namespace="default", + activity_id="activity-id", + run_id="run-id", + version=0, + ).encode() + + assert "=" not in token + assert OperationToken.decode(token) == OperationToken( + type=OperationTokenType.ACTIVITY, + namespace="default", + activity_id="activity-id", + run_id="run-id", + version=0, + ) + + +def test_operation_token_activity_encode_uses_activity_id_and_omits_workflow_id(): + token = OperationToken( + type=OperationTokenType.ACTIVITY, + namespace="default", + activity_id="activity-id", + ).encode() + + assert json.loads(_base64url_decode_no_padding(token)) == { + "t": 2, + "ns": "default", + "aid": "activity-id", + } + + def test_workflow_handle_to_from_token_round_trip(): handle = WorkflowHandle[str](namespace="default", workflow_id="workflow-id") @@ -80,6 +114,58 @@ def test_workflow_handle_to_from_token_round_trip(): version=0, ), ), + # Activity tokens + ( + _encode_json_token( + {"t": 2, "ns": "default", "aid": "activity-id", "rid": "run-id"} + ), + OperationToken( + type=OperationTokenType.ACTIVITY, + namespace="default", + activity_id="activity-id", + run_id="run-id", + ), + ), + ( + _encode_json_token( + {"t": 2, "ns": "", "aid": "activity-id", "rid": "run-id"} + ), + OperationToken( + type=OperationTokenType.ACTIVITY, + namespace="", + activity_id="activity-id", + run_id="run-id", + ), + ), + ( + _encode_json_token( + { + "t": 2, + "ns": "default", + "aid": "activity-id", + "rid": "run-id", + "v": None, + } + ), + OperationToken( + type=OperationTokenType.ACTIVITY, + namespace="default", + activity_id="activity-id", + run_id="run-id", + ), + ), + ( + _encode_json_token( + {"t": 2, "ns": "default", "aid": "activity-id", "rid": "run-id", "v": 0} + ), + OperationToken( + type=OperationTokenType.ACTIVITY, + namespace="default", + activity_id="activity-id", + run_id="run-id", + version=0, + ), + ), ], ) def test_operation_token_decode_accepts_valid_tokens( @@ -110,7 +196,7 @@ def test_operation_token_decode_accepts_valid_tokens( ), ( _encode_json_token({"t": 1, "ns": "default"}), - "expected workflow id to be a string", + "expected non-empty workflow id for token type `WORKFLOW`", ), ( _encode_json_token({"t": 1, "ns": "default", "wid": 123}), @@ -134,6 +220,41 @@ def test_operation_token_decode_accepts_valid_tokens( ), "expected version to be an int or null", ), + # Activity tokens + ( + _encode_json_token({"t": 2, "ns": "default"}), + "expected non-empty activity id for token type `ACTIVITY`", + ), + ( + _encode_json_token({"t": 2, "ns": "default", "aid": ""}), + "expected non-empty activity id for token type `ACTIVITY`", + ), + ( + _encode_json_token({"t": 2, "ns": "default", "aid": 123}), + "expected activity id to be a string", + ), + ( + _encode_json_token( + {"t": 2, "ns": "default", "aid": "activity-id", "rid": 123} + ), + "expected run_id to be a string", + ), + ( + _encode_json_token({"t": 2, "aid": "activity-id", "rid": "run-id"}), + "expected namespace to be a string", + ), + ( + _encode_json_token( + { + "t": 2, + "ns": "default", + "aid": "activity-id", + "rid": "run-id", + "v": "0", + } + ), + "expected version to be an int or null", + ), ], ) def test_operation_token_decode_rejects_invalid_tokens(token: str, message: str): diff --git a/tests/nexus/test_temporal_operation.py b/tests/nexus/test_temporal_operation.py index b6b1a92c7..3f16f5f34 100644 --- a/tests/nexus/test_temporal_operation.py +++ b/tests/nexus/test_temporal_operation.py @@ -6,14 +6,29 @@ import nexusrpc import pytest from nexusrpc import HandlerErrorType, Operation, service -from nexusrpc.handler import operation_handler, service_handler +from nexusrpc.handler import ( + CancelOperationContext, + OperationTaskCancellation, + operation_handler, + service_handler, +) from typing_extensions import override import temporalio.exceptions -from temporalio import nexus, workflow +from temporalio import activity, nexus, workflow from temporalio.api.common.v1 import Link -from temporalio.client import Client, WorkflowExecutionStatus, WorkflowFailureError -from temporalio.common import NexusOperationExecutionStatus, WorkflowIDConflictPolicy +from temporalio.client import ( + ActivityExecutionStatus, + Client, + NexusOperationFailureError, + WorkflowExecutionStatus, + WorkflowFailureError, +) +from temporalio.common import ( + NexusOperationExecutionStatus, + RetryPolicy, + WorkflowIDConflictPolicy, +) from temporalio.nexus._token import OperationToken, OperationTokenType from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker @@ -59,6 +74,29 @@ async def run(self, input: Input) -> str: return input.value +@activity.defn +async def echo_activity(input: Input) -> str: + return input.value + + +@activity.defn +async def raise_error_activity() -> None: + raise temporalio.exceptions.ApplicationError( + "test-activity-error-message", + type="test-activity-error-type", + non_retryable=True, + ) + + +@activity.defn +async def wait_for_cancel_activity() -> None: + # Heartbeat in a loop so the activity receives cancellation. Letting the + # resulting CancelledError bubble out transitions the activity to CANCELED. + while True: + await asyncio.sleep(0.3) + activity.heartbeat() + + @service class TestService: echo: Operation[Input, str] @@ -69,6 +107,12 @@ class TestService: sync_result: Operation[Input, str] custom_cancel: Operation[str, None] update_op: Operation[Input, str] + echo_activity: Operation[Input, str] + error_activity: Operation[Input, None] + blocking_activity: Operation[str, None] + custom_cancel_activity: Operation[str, None] + double_start_activity: Operation[Input, None] + mixed_start: Operation[Input, None] @service_handler(service=TestService) @@ -78,6 +122,8 @@ class TestServiceHandler: def __init__(self) -> None: self.started_custom_cancel_workflow = asyncio.Event() + self.started_custom_cancel_activity = asyncio.Event() + self.custom_cancel_activity_called = asyncio.Event() @nexus.temporal_operation async def echo( @@ -237,6 +283,130 @@ async def update_op( update_id=input.update_id, ) + @nexus.temporal_operation + async def echo_activity( + self, + _ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + input: Input, + ) -> nexus.TemporalOperationResult[str]: + return await client.start_activity( + echo_activity, + input, + id=f"echo_activity-{uuid.uuid4()}", + start_to_close_timeout=timedelta(seconds=5), + ) + + @nexus.temporal_operation + async def error_activity( + self, + _ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + _input: Input, + ) -> nexus.TemporalOperationResult[None]: + # The activity raises immediately. With a single permitted attempt it + # fails the backing activity, which in turn fails the Nexus operation. + return await client.start_activity( + raise_error_activity, + id=f"error_activity-{uuid.uuid4()}", + start_to_close_timeout=timedelta(seconds=5), + retry_policy=RetryPolicy(maximum_attempts=1), + ) + + @nexus.temporal_operation + async def blocking_activity( + self, + _ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + input: str, + ) -> nexus.TemporalOperationResult[None]: + return await client.start_activity( + wait_for_cancel_activity, + id=input, + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=1), + ) + + @nexus.temporal_operation + async def double_start_activity( + self, + _ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + input: Input, + ) -> nexus.TemporalOperationResult[None]: + await client.start_activity( + echo_activity, + input, + id=f"double-start-activity-{uuid.uuid4()}", + start_to_close_timeout=timedelta(seconds=5), + ) + await client.start_activity( + echo_activity, + input, + id=f"double-start-activity-{uuid.uuid4()}", + start_to_close_timeout=timedelta(seconds=5), + ) + return nexus.TemporalOperationResult.sync(None) + + @nexus.temporal_operation + async def mixed_start( + self, + _ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + input: Input, + ) -> nexus.TemporalOperationResult[None]: + # Starting a workflow reserves the single async start, so the subsequent + # start_activity must hit the same guard and raise a BAD_REQUEST error. + await client.start_workflow( + EchoWorkflow.run, input, id=f"mixed-start-{uuid.uuid4()}" + ) + await client.start_activity( + echo_activity, + input, + id=f"mixed-start-{uuid.uuid4()}", + start_to_close_timeout=timedelta(seconds=5), + ) + return nexus.TemporalOperationResult.sync(None) + + @operation_handler + def custom_cancel_activity(self) -> nexus.TemporalOperationHandler[str, None]: + started = self.started_custom_cancel_activity + cancel_called = self.custom_cancel_activity_called + + class CustomCancelActivityNexusOpHandler( + nexus.TemporalOperationHandler[str, None] + ): + @override + async def start_operation( + self, + ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + input: str, + ) -> nexus.TemporalOperationResult[None]: + result = await client.start_activity( + wait_for_cancel_activity, + id=input, + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=1), + ) + started.set() + return result + + @override + async def cancel_activity( + self, + ctx: nexus.TemporalCancelOperationContext, + options: nexus.CancelActivityOptions, + ): + # record that the custom override ran + cancel_called.set() + + # get a handle to the activity and cancel it + handle = nexus.client().get_activity_handle(options.activity_id) + await handle.cancel() + + return CustomCancelActivityNexusOpHandler() + @workflow.defn class EchoWorkflowCaller: @@ -556,6 +726,80 @@ async def test_temporal_operation_update_workflow_delayed( assert expected_backward_link in handler_links +async def test_temporal_operation_cancel_rejects_unknown_tokens(): + class FakeNexusTaskCancellation(OperationTaskCancellation): + def is_cancelled(self) -> bool: + return False + + def cancellation_reason(self) -> str | None: + return None + + def wait_until_cancelled_sync(self, timeout: float | None = None) -> bool: + return False + + async def wait_until_cancelled(self) -> None: + return None + + def cancel(self, _reason: str) -> bool: + return False + + cancel_ctx = CancelOperationContext( + service="TestService", + operation="echo", + headers={}, + task_cancellation=FakeNexusTaskCancellation(), + ) + + service_handler = TestServiceHandler() + + # Use a factory style operation form the handler to allow calling cancel directly + op_handler = service_handler.custom_cancel() + + # Invalid token type + token = OperationToken(type=30, namespace="default") # type: ignore + with pytest.raises(nexusrpc.HandlerError) as err: + await op_handler.cancel(cancel_ctx, token.encode()) + assert err.value.type == HandlerErrorType.INTERNAL + assert not err.value.retryable + underlying = err.value.__cause__ + assert isinstance(underlying, TypeError) + assert "unknown token type, got 30" in str(underlying) + + # Workflow ID missing from workflow type + token = OperationToken(type=OperationTokenType.WORKFLOW, namespace="default") + with pytest.raises(nexusrpc.HandlerError) as err: + await op_handler.cancel(cancel_ctx, token.encode()) + assert err.value.type == HandlerErrorType.INTERNAL + assert not err.value.retryable + underlying = err.value.__cause__ + assert isinstance(underlying, TypeError) + assert "expected non-empty workflow id for token type `WORKFLOW`" in str(underlying) + + # Activity ID missing from activity type + token = OperationToken(type=OperationTokenType.ACTIVITY, namespace="default") + with pytest.raises(nexusrpc.HandlerError) as err: + await op_handler.cancel(cancel_ctx, token.encode()) + assert err.value.type == HandlerErrorType.INTERNAL + assert not err.value.retryable + underlying = err.value.__cause__ + assert isinstance(underlying, TypeError) + assert "expected non-empty activity id for token type `ACTIVITY`" in str(underlying) + + activity_op_handler = service_handler.custom_cancel_activity() + for run_id in (None, ""): + token = OperationToken( + type=OperationTokenType.ACTIVITY, + namespace="default", + activity_id="activity-id", + run_id=run_id, + ) + with pytest.raises(nexusrpc.HandlerError) as err: + await activity_op_handler.cancel(cancel_ctx, token.encode()) + assert err.value.type == HandlerErrorType.INTERNAL + assert not err.value.retryable + assert not service_handler.custom_cancel_activity_called.is_set() + + @workflow.defn class BlockingWorkflow: def __init__(self) -> None: @@ -783,6 +1027,41 @@ async def test_temporal_operation_failed_start_allows_retry( await conflict_handle.cancel() +async def test_temporal_operation_mixed_start_raises_handler_err( + client: Client, env: WorkflowEnvironment +): + if env.supports_time_skipping: + pytest.skip( + "Standalone Nexus Operation tests don't work with time-skipping server" + ) + + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint_name, task_queue) + async with Worker( + env.client, + task_queue=task_queue, + nexus_service_handlers=[TestServiceHandler()], + workflows=[EchoWorkflow], + activities=[echo_activity], + ): + nexus_client = client.create_nexus_client(TestService, endpoint_name) + + with pytest.raises(NexusOperationFailureError) as err: + await nexus_client.execute_operation( + TestService.mixed_start, + Input(value="test", task_queue=task_queue), + id=str(uuid.uuid4()), + ) + + assert isinstance(err.value.cause, nexusrpc.HandlerError) + assert err.value.cause.type == HandlerErrorType.BAD_REQUEST + assert ( + "Only one async operation can be started per operation handler invocation" + in err.value.cause.message + ) + + @workflow.defn class SyncResultCaller: @workflow.run @@ -822,6 +1101,184 @@ async def test_temporal_operation_sync_result(client: Client, env: WorkflowEnvir ) +async def test_temporal_operation_start_activity( + client: Client, env: WorkflowEnvironment +): + if env.supports_time_skipping: + pytest.skip( + "Standalone Nexus Operation tests don't work with time-skipping server" + ) + + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint_name, task_queue) + async with Worker( + env.client, + task_queue=task_queue, + nexus_service_handlers=[TestServiceHandler()], + activities=[echo_activity], + ): + nexus_client = client.create_nexus_client(TestService, endpoint_name) + + result = await nexus_client.execute_operation( + TestService.echo_activity, + Input(value="test", task_queue=task_queue), + id=str(uuid.uuid4()), + ) + assert result == "test" + + +async def test_temporal_operation_start_activity_raises_error( + client: Client, env: WorkflowEnvironment +): + if env.supports_time_skipping: + pytest.skip( + "Standalone Nexus Operation tests don't work with time-skipping server" + ) + + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint_name, task_queue) + async with Worker( + env.client, + task_queue=task_queue, + nexus_service_handlers=[TestServiceHandler()], + activities=[raise_error_activity], + ): + nexus_client = client.create_nexus_client(TestService, endpoint_name) + + with pytest.raises(NexusOperationFailureError) as err: + await nexus_client.execute_operation( + TestService.error_activity, + Input(value="test", task_queue=task_queue), + id=str(uuid.uuid4()), + ) + + operation_err = err.value.__cause__ + assert isinstance(operation_err, temporalio.exceptions.ApplicationError) + assert operation_err.type == "OperationError" + assert "nexus operation completed unsuccessfully" in str(operation_err) + + application_err = operation_err.__cause__ + assert isinstance(application_err, temporalio.exceptions.ApplicationError) + + assert application_err.type == "test-activity-error-type" + assert "test-activity-error-message" in str(application_err) + assert application_err.__cause__ is None + + +async def test_temporal_operation_cancel_activity( + client: Client, env: WorkflowEnvironment +): + if env.supports_time_skipping: + pytest.skip( + "Standalone Nexus Operation tests don't work with time-skipping server" + ) + + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint_name, task_queue) + async with Worker( + env.client, + task_queue=task_queue, + nexus_service_handlers=[TestServiceHandler()], + activities=[wait_for_cancel_activity], + ): + nexus_client = client.create_nexus_client(TestService, endpoint_name) + + activity_id = f"blocking-activity-{uuid.uuid4()}" + op_handle = await nexus_client.start_operation( + TestService.blocking_activity, activity_id, id=str(uuid.uuid4()) + ) + + await op_handle.cancel() + + activity_handle = client.get_activity_handle(activity_id) + + async def check_cancelled(): + op_desc = await op_handle.describe() + assert op_desc.status is NexusOperationExecutionStatus.CANCELED + activity_desc = await activity_handle.describe() + assert activity_desc.status is ActivityExecutionStatus.CANCELED + + await assert_eventually(check_cancelled) + + +async def test_customized_temporal_operation_cancel_activity( + client: Client, env: WorkflowEnvironment +): + if env.supports_time_skipping: + pytest.skip( + "Standalone Nexus Operation tests don't work with time-skipping server" + ) + + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint_name, task_queue) + + service_handler = TestServiceHandler() + async with Worker( + env.client, + task_queue=task_queue, + nexus_service_handlers=[service_handler], + activities=[wait_for_cancel_activity], + ): + nexus_client = client.create_nexus_client(TestService, endpoint_name) + + activity_id = f"custom-cancel-activity-{uuid.uuid4()}" + op_handle = await nexus_client.start_operation( + TestService.custom_cancel_activity, activity_id, id=str(uuid.uuid4()) + ) + await service_handler.started_custom_cancel_activity.wait() + + await op_handle.cancel() + + activity_handle = client.get_activity_handle(activity_id) + + async def check_cancelled(): + assert service_handler.custom_cancel_activity_called.is_set() + op_desc = await op_handle.describe() + assert op_desc.status is NexusOperationExecutionStatus.CANCELED + activity_desc = await activity_handle.describe() + assert activity_desc.status is ActivityExecutionStatus.CANCELED + + await assert_eventually(check_cancelled) + + +async def test_temporal_operation_double_start_activity_raises_handler_err( + client: Client, env: WorkflowEnvironment +): + if env.supports_time_skipping: + pytest.skip( + "Standalone Nexus Operation tests don't work with time-skipping server" + ) + + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint_name, task_queue) + async with Worker( + env.client, + task_queue=task_queue, + nexus_service_handlers=[TestServiceHandler()], + activities=[echo_activity], + ): + nexus_client = client.create_nexus_client(TestService, endpoint_name) + + with pytest.raises(NexusOperationFailureError) as err: + await nexus_client.execute_operation( + TestService.double_start_activity, + Input(value="test", task_queue=task_queue), + id=str(uuid.uuid4()), + ) + + assert isinstance(err.value.cause, nexusrpc.HandlerError) + assert err.value.cause.type == HandlerErrorType.BAD_REQUEST + assert ( + "Only one async operation can be started per operation handler invocation" + in err.value.cause.message + ) + + @dataclass class TemporalOperationOverloadTestValue: value: int @@ -1067,3 +1524,64 @@ async def do_update(self, value: str) -> str: self.order_status = value update_result = f"Updated workflow status from {status} to {value}" return update_result + + +async def test_temporal_operation_includes_activity_token_in_callback( + client: Client, env: WorkflowEnvironment +): + if env.supports_time_skipping: + pytest.skip( + "Standalone Nexus Operation tests don't work with time-skipping server" + ) + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint_name, task_queue) + + @service_handler + class ActivityTokenHandler: + @nexus.temporal_operation + async def echo_activity( + self, + _ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + input: Input, + ) -> nexus.TemporalOperationResult[str]: + return await client.start_activity( + echo_activity, + input, + id=input.value, + start_to_close_timeout=timedelta(seconds=10), + start_delay=timedelta(milliseconds=100), + ) + + async with Worker( + client, + task_queue=task_queue, + nexus_service_handlers=[ActivityTokenHandler()], + activities=[echo_activity], + ): + input_value = f"test-{uuid.uuid4()}" + + nexus_client = client.create_nexus_client(ActivityTokenHandler, endpoint_name) + + result = await nexus_client.execute_operation( + ActivityTokenHandler.echo_activity, + Input(value=input_value, task_queue=task_queue), + id=str(uuid.uuid4()), + ) + assert result == input_value + + activity_handle = client.get_activity_handle(input_value) + + desc = await activity_handle.describe() + token = desc.raw_callbacks[0].info.callback.nexus.header[ + "nexus-operation-token" + ] + + expected_token = OperationToken( + type=OperationTokenType.ACTIVITY, + namespace=client.namespace, + activity_id=input_value, + ).encode() + + assert token == expected_token