Add SAA Temporal Nexus operation handling

Expose customizable Temporal Nexus operation handlers and client support.

Add activity link conversion, operation token handling, and related tests.

Author: VegetarianOrc

temporalio/client/_activity.py

@@ -691,6 +691,8 @@ def __init__(
         *,
         run_id: str | None = None,
         result_type: type | None = None,
+        start_activity_response: None
+        | temporalio.api.workflowservice.v1.StartActivityExecutionResponse = None,
     ) -> None:
         """Create activity handle."""
         self._client = client
@@ -700,6 +702,7 @@ def __init__(
         self._known_outcome: (
             temporalio.api.activity.v1.ActivityExecutionOutcome | None
         ) = None
+        self._start_activity_response = start_activity_response
 
     @functools.cached_property
     def _data_converter(self) -> temporalio.converter.DataConverter:

temporalio/client/_client.py

@@ -1484,6 +1484,12 @@ async def start_activity(
         start_delay: timedelta | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
         rpc_timeout: timedelta | None = None,
+        # The following options should not be considered part of the public API. They
+        # are deliberately not exposed in overloads, and are not subject to any
+        # backwards compatibility guarantees.
+        callbacks: Sequence[Callback] = [],
+        links: Sequence[temporalio.api.common.v1.Link] = [],
+        request_id: str | None = None,
     ) -> ActivityHandle[ReturnType]:
         """Start an activity and return its handle.
 
@@ -1542,6 +1548,9 @@ async def start_activity(
                 rpc_metadata=rpc_metadata,
                 rpc_timeout=rpc_timeout,
                 priority=priority,
+                callbacks=callbacks,
+                links=links,
+                request_id=request_id,
             )
         )
 

temporalio/client/_impl.py

@@ -237,7 +237,7 @@ async def _build_start_workflow_execution_request(
         # Links are duplicated on request for compatibility with older server versions.
         req.links.extend(links)
 
-        if temporalio.nexus._operation_context._in_nexus_backing_workflow_start_context():
+        if temporalio.nexus._operation_context._in_nexus_backing_start_context():
             req.on_conflict_options.attach_request_id = True
             req.on_conflict_options.attach_completion_callbacks = True
             req.on_conflict_options.attach_links = True
@@ -566,6 +566,7 @@ async def start_activity(self, input: StartActivityInput) -> ActivityHandle[Any]
             input.id,
             run_id=resp.run_id,
             result_type=input.result_type,
+            start_activity_response=resp,
         )
 
     async def _build_start_activity_execution_request(
@@ -609,6 +610,8 @@ async def _build_start_activity_execution_request(
             ),
         )
 
+        if input.request_id:
+            req.request_id = input.request_id
         if input.schedule_to_close_timeout is not None:
             req.schedule_to_close_timeout.FromTimedelta(input.schedule_to_close_timeout)
         if input.start_to_close_timeout is not None:
@@ -644,6 +647,23 @@ async def _build_start_activity_execution_request(
         # Set priority
         req.priority.CopyFrom(input.priority._to_proto())
 
+        req.completion_callbacks.extend(
+            temporalio.api.common.v1.Callback(
+                nexus=temporalio.api.common.v1.Callback.Nexus(
+                    url=callback.url,
+                    header=callback.headers,
+                ),
+                links=input.links,
+            )
+            for callback in input.callbacks
+        )
+        req.links.extend(input.links)
+
+        if temporalio.nexus._operation_context._in_nexus_backing_start_context():
+            req.on_conflict_options.attach_request_id = True
+            req.on_conflict_options.attach_completion_callbacks = True
+            req.on_conflict_options.attach_links = True
+
         return req
 
     async def cancel_activity(self, input: CancelActivityInput) -> None:

temporalio/client/_interceptor.py

@@ -230,6 +230,10 @@ class StartActivityInput:
     headers: Mapping[str, temporalio.api.common.v1.Payload]
     rpc_metadata: Mapping[str, str | bytes]
     rpc_timeout: timedelta | None
+    # The following options are experimental and unstable.
+    callbacks: Sequence[Callback]
+    links: Sequence[temporalio.api.common.v1.Link]
+    request_id: str | None
 
 
 @dataclass

temporalio/nexus/__init__.py

@@ -25,6 +25,7 @@
     wait_for_worker_shutdown_sync,
 )
 from ._operation_handlers import (
+    CancelActivityOptions,
     CancelWorkflowRunOptions,
     TemporalNexusOperationHandler,
 )
@@ -33,6 +34,7 @@
 
 __all__ = (
     "workflow_run_operation",
+    "CancelActivityOptions",
     "CancelWorkflowRunOptions",
     "Info",
     "LoggerAdapter",

temporalio/nexus/_link_conversion.py

@@ -23,6 +23,10 @@
     r"^/namespaces/(?P<namespace>[^/]+)/nexus-operations/(?P<operation_id>[^/]+)$"
 )
 
+_ACTIVITY_LINK_URL_PATH_REGEX = re.compile(
+    r"^/namespaces/(?P<namespace>[^/]+)/activities/(?P<activity_id>[^/]+)$"
+)
+
 _WORFKLOW_LINK_URL_PATH_REGEX = re.compile(
     r"^/namespaces/(?P<namespace>[^/]+)/workflows/(?P<workflow_id>[^/]+)/(?P<run_id>[^/]+)/history$"
 )
@@ -31,6 +35,7 @@
 class _LinkType(str, Enum):
     WORKFLOW = temporalio.api.common.v1.Link.WorkflowEvent.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"
@@ -84,6 +89,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,
@@ -99,7 +107,10 @@ def temporal_link_to_nexus_link(
         case "nexus_operation":
             return nexus_operation_to_nexus_link(temporal_link.nexus_operation)
 
-        case "activity" | "batch_job":
+        case "activity":
+            return activity_link_to_nexus_link(temporal_link.activity)
+
+        case "batch_job":
             raise NotImplementedError("only workflow links are supported")
 
         case None:
@@ -165,6 +176,25 @@ def nexus_operation_to_nexus_link(
     return nexusrpc.Link(url=url, type=_LinkType.NEXUS_OPERATION.value)
 
 
+def activity_link_to_nexus_link(
+    activity: temporalio.api.common.v1.Link.Activity,
+) -> nexusrpc.Link:
+    """Convert an Activity link into a nexusrpc link."""
+    scheme = "temporal"
+    namespace = urllib.parse.quote(activity.namespace, safe="")
+    activity_id = urllib.parse.quote(activity.activity_id, safe="")
+    path = f"/namespaces/{namespace}/activities/{activity_id}"
+
+    if activity.run_id:
+        query_params = urllib.parse.urlencode({LINK_RUN_ID_PARAM_NAME: activity.run_id})
+    else:
+        query_params = ""
+
+    url = f"{scheme}://{urllib.parse.urlunparse(('', '', path, '', query_params, ''))}"
+
+    return nexusrpc.Link(url=url, type=_LinkType.ACTIVITY.value)
+
+
 def nexus_link_to_workflow_event_link(
     link: nexusrpc.Link,
 ) -> temporalio.api.common.v1.Link | None:
@@ -250,6 +280,38 @@ 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 link into a Temporal Activity link."""
+    url = urllib.parse.urlparse(nexus_link.url)
+    match = _ACTIVITY_LINK_URL_PATH_REGEX.match(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
+
+    query_params = urllib.parse.parse_qs(url.query, keep_blank_values=True)
+
+    match query_params.get(LINK_RUN_ID_PARAM_NAME):
+        case [run_id_param]:
+            run_id = run_id_param
+        case _:
+            logger.warning(
+                f"Invalid Nexus link: {nexus_link}. Expected {LINK_RUN_ID_PARAM_NAME} to have exactly 1 value"
+            )
+            return None
+
+    groups = match.groupdict()
+    activity_link = temporalio.api.common.v1.Link.Activity(
+        namespace=urllib.parse.unquote(groups["namespace"]),
+        activity_id=urllib.parse.unquote(groups["activity_id"]),
+        run_id=run_id,
+    )
+    return temporalio.api.common.v1.Link(activity=activity_link)
+
+
 def _event_reference_to_query_params(
     event_ref: temporalio.api.common.v1.Link.WorkflowEvent.EventReference,
 ) -> str:

temporalio/nexus/_operation_context.py

@@ -43,6 +43,7 @@
 )
 
 from ._link_conversion import (
+    activity_link_to_nexus_link,
     nexus_link_to_temporal_link,
     workflow_event_to_nexus_link,
     workflow_execution_started_event_link_from_workflow_handle,
@@ -64,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 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"
 )
 
 
@@ -168,16 +169,16 @@ def _try_temporal_context() -> (
 
 
 @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)
@@ -243,13 +244,13 @@ def _get_callbacks(
     def _get_links(
         self,
     ) -> list[temporalio.api.common.v1.Link]:
-        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_outbound_links(
+    def _add_outbound_workflow_links(
         self, workflow_handle: temporalio.client.WorkflowHandle[Any, Any]
     ):
         # If links were not sent in StartWorkflowExecutionResponse then construct them.
@@ -279,6 +280,45 @@ def _add_outbound_links(
             )
         return workflow_handle
 
+    def _add_outbound_activity_links(
+        self, activity_handle: temporalio.client.ActivityHandle[Any]
+    ):
+        activity_links: list[temporalio.api.common.v1.Link.Activity] = []
+        try:
+            if isinstance(
+                activity_handle._start_activity_response,
+                temporalio.api.workflowservice.v1.StartActivityExecutionResponse,
+            ):
+                if activity_handle._start_activity_response.HasField("link"):
+                    if activity_handle._start_activity_response.link.HasField(
+                        "activity"
+                    ):
+                        activity_links.append(
+                            activity_handle._start_activity_response.link.activity
+                        )
+            if not activity_links:
+                activity_run_id = activity_handle.run_id
+                if activity_run_id is None:
+                    raise ValueError(
+                        f"Activity handle {activity_handle} has no run ID. "
+                        "Cannot create Activity link."
+                    )
+                activity_links.append(
+                    temporalio.api.common.v1.Link.Activity(
+                        namespace=activity_handle._client.namespace,
+                        activity_id=activity_handle.id,
+                        run_id=activity_run_id,
+                    )
+                )
+            self.nexus_context.outbound_links.extend(
+                activity_link_to_nexus_link(link) for link in activity_links
+            )
+        except Exception as e:
+            logger.warning(
+                f"Failed to create Activity link for activity {activity_handle}: {e}"
+            )
+        return activity_handle
+
 
 class WorkflowRunOperationContext(StartOperationContext):
     """Context received by a workflow run operation."""
@@ -642,7 +682,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():
         wf_handle = await temporal_context.client.start_workflow(  # type: ignore
             workflow=workflow,
             arg=arg,
@@ -674,6 +714,6 @@ async def _start_nexus_backing_workflow(
             request_id=temporal_context.nexus_context.request_id,
         )
 
-    temporal_context._add_outbound_links(wf_handle)
+    temporal_context._add_outbound_workflow_links(wf_handle)
 
     return WorkflowHandle[ReturnType]._unsafe_from_client_workflow_handle(wf_handle)