Add SAA TemporalNexusOperation handling

temporalio/client/_activity.py

@@ -309,13 +309,17 @@ 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,
         info: temporalio.api.activity.v1.ActivityExecutionInfo,
         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,
         )
 
 

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(

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",

temporalio/nexus/_link_conversion.py

@@ -23,6 +23,10 @@
     r"^/namespaces/(?P<namespace>[^/]+)/nexus-operations/(?P<operation_id>[^/]+)/(?P<run_id>[^/]*)/details$"
 )
 
+_ACTIVITY_LINK_URL_PATH_REGEX = re.compile(
+    r"^/namespaces/(?P<namespace>[^/]+)/activities/(?P<activity_id>[^/]+)/(?P<run_id>[^/]+)/details$"
+)
+
 _WORKFLOW_LINK_URL_PATH_REGEX = re.compile(
     r"^/namespaces/(?P<namespace>[^/]+)/workflows/(?P<workflow_id>[^/]+)/(?P<run_id>[^/]+)(?P<history>/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:

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,