centralize and simplify start workflow response link handling. Move private start context accessor to nexus package. Update tests to reflect that response links are created when using plain start workflow

Author: VegetarianOrc

temporalio/client/_impl.py

@@ -206,19 +206,9 @@ async def start_workflow(
         # returned so the caller workflow's Nexus history event links to the signaled event. A
         # plain start does not capture a response link: it only forwards the inbound request links
         # onto the start request.
-        nexus_ctx = self._try_nexus_start_operation_context()
-        if (
-            nexus_ctx is not None
-            and not temporalio.nexus._operation_context._in_nexus_backing_workflow_start_context()
-            and isinstance(
-                resp,
-                temporalio.api.workflowservice.v1.SignalWithStartWorkflowExecutionResponse,
-            )
-        ):
-            # Server >= 1.31 with EnableCHASMSignalBacklinks returns signal_link pointing at
-            # the WorkflowExecutionSignaled event; older servers leave it unset.
-            if resp.HasField("signal_link"):
-                nexus_ctx._add_response_link(resp.signal_link)
+        nexus_ctx = temporalio.nexus._operation_context._try_start_operation_context()
+        if nexus_ctx is not None:
+            nexus_ctx._add_start_workflow_response_link(handle)
         return handle
 
     async def _build_start_workflow_execution_request(
@@ -256,7 +246,7 @@ async def _build_start_workflow_execution_request(
         # Links are duplicated on request for compatibility with older server versions.
         req.links.extend(links)
 
-        nexus_ctx = self._try_nexus_start_operation_context()
+        nexus_ctx = temporalio.nexus._operation_context._try_start_operation_context()
         if nexus_ctx is not None:
             # This start was issued from inside a Nexus operation handler. If the workflow ID
             # conflict policy is WORKFLOW_ID_CONFLICT_POLICY_USE_EXISTING and a conflict is
@@ -304,7 +294,9 @@ async def _build_signal_with_start_workflow_execution_request(
         # 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():
-            nexus_ctx = self._try_nexus_start_operation_context()
+            nexus_ctx = (
+                temporalio.nexus._operation_context._try_start_operation_context()
+            )
             if nexus_ctx is not None:
                 req.links.extend(nexus_ctx._get_request_links())
         return req
@@ -542,7 +534,7 @@ async def signal_workflow(self, input: SignalWorkflowInput) -> None:
             await self._apply_headers(input.headers, req.header.fields)
         # If this signal is issued from inside a Nexus operation handler, forward the inbound
         # Nexus task links so the WorkflowExecutionSignaled event links back to the caller.
-        nexus_ctx = self._try_nexus_start_operation_context()
+        nexus_ctx = temporalio.nexus._operation_context._try_start_operation_context()
         if nexus_ctx is not None:
             req.links.extend(nexus_ctx._get_request_links())
         resp = await self._client.workflow_service.signal_workflow_execution(
@@ -1685,17 +1677,6 @@ async def count_nexus_operations(
             )
         )
 
-    @staticmethod
-    def _try_nexus_start_operation_context() -> (
-        temporalio.nexus._operation_context._TemporalStartOperationContext | None
-    ):
-        """The Nexus start-operation context if a handler is currently running, else None."""
-        return (
-            temporalio.nexus._operation_context._temporal_start_operation_context.get(
-                None
-            )
-        )
-
     async def _apply_headers(
         self,
         source: Mapping[str, temporalio.api.common.v1.Payload] | None,

temporalio/nexus/_operation_context.py

@@ -23,6 +23,7 @@
     overload,
 )
 
+import nexusrpc
 from nexusrpc.handler import (
     CancelOperationContext,
     OperationContext,
@@ -42,8 +43,11 @@
     SelfType,
 )
 
+import temporalio.api.enums.v1
+
 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,
 )
@@ -167,6 +171,11 @@ def _try_temporal_context() -> (
     return start_ctx or cancel_ctx
 
 
+def _try_start_operation_context() -> _TemporalStartOperationContext | None:  # pyright: ignore[reportUnusedFunction]
+    """The Nexus start-operation context if a handler is currently running, else None."""
+    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)
@@ -253,35 +262,43 @@ def _get_request_links(self) -> list[temporalio.api.common.v1.Link]:
                 event_links.append(link)
         return event_links
 
-    def _add_backing_workflow_response_link(
+    def _add_start_workflow_response_link(
         self, workflow_handle: temporalio.client.WorkflowHandle[Any, Any]
     ):
-        # If links were not sent in StartWorkflowExecutionResponse then construct them.
-        wf_event_links: list[temporalio.api.common.v1.Link.WorkflowEvent] = []
-        try:
-            if isinstance(
-                workflow_handle._start_workflow_response,
-                temporalio.api.workflowservice.v1.StartWorkflowExecutionResponse,
-            ):
-                if workflow_handle._start_workflow_response.HasField("link"):
-                    if link := workflow_handle._start_workflow_response.link:
-                        if link.HasField("workflow_event"):
-                            wf_event_links.append(link.workflow_event)
-            if not wf_event_links:
-                wf_event_links = [
-                    workflow_execution_started_event_link_from_workflow_handle(
+        response = workflow_handle._start_workflow_response
+
+        nexus_link: nexusrpc.Link | None = None
+        if isinstance(
+            response, temporalio.api.workflowservice.v1.StartWorkflowExecutionResponse
+        ):
+            if response.HasField("link"):
+                nexus_link = temporal_link_to_nexus_link(response.link)
+            else:
+                # If a link was not sent in response then construct it.
+                link = temporalio.api.common.v1.Link(
+                    workflow_event=workflow_execution_started_event_link_from_workflow_handle(
                         workflow_handle,
                         self.nexus_context.request_id,
                     )
-                ]
-            self.nexus_context.outbound_links.extend(
-                workflow_event_to_nexus_link(link) for link in wf_event_links
-            )
+                )
+                nexus_link = temporal_link_to_nexus_link(link)
+
+        elif isinstance(
+            response,
+            temporalio.api.workflowservice.v1.SignalWithStartWorkflowExecutionResponse,
+        ):
+            # Server >= 1.31 with EnableCHASMSignalBacklinks returns signal_link pointing at
+            # the WorkflowExecutionSignaled event; older servers leave it unset.
+            if response.HasField("signal_link"):
+                nexus_link = temporal_link_to_nexus_link(response.signal_link)
+
+        try:
+            if nexus_link is not None:
+                self.nexus_context.outbound_links.append(nexus_link)
         except Exception as e:
             logger.warning(
-                f"Failed to create WorkflowExecutionStarted event links for workflow {workflow_handle}: {e}"
+                f"Failed to create event links for workflow {workflow_handle}: {e}"
             )
-        return workflow_handle
 
     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.
@@ -699,6 +716,4 @@ async def _start_nexus_backing_workflow(
             request_id=temporal_context.nexus_context.request_id,
         )
 
-    temporal_context._add_backing_workflow_response_link(wf_handle)
-
     return WorkflowHandle[ReturnType]._unsafe_from_client_workflow_handle(wf_handle)

tests/nexus/test_signal_link_propagation.py

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Generator
 from typing import Any
 from unittest import mock
 
@@ -31,6 +32,7 @@
 import temporalio.converter
 import temporalio.nexus._link_conversion
 import temporalio.nexus._operation_context
+from temporalio.nexus._operation_context import _TemporalStartOperationContext
 from temporalio.client._impl import _ClientImpl
 from temporalio.client._interceptor import (
     SignalWorkflowInput,
@@ -68,7 +70,7 @@ def _inbound_nexus_link() -> temporalio.api.common.v1.Link:
 
 
 @pytest.fixture
-def nexus_ctx() -> Any:
+def nexus_ctx() -> Generator[_TemporalStartOperationContext]:
     """Install a Nexus start-operation context with a single inbound link.
 
     The inbound link is provided in nexusrpc.Link form, exactly as the worker populates it from
@@ -169,7 +171,7 @@ def _outbound_link_urls(ctx: Any) -> list[str]:
 
 
 async def test_signal_forwards_inbound_links_and_captures_response_backlink(
-    nexus_ctx: Any,
+    nexus_ctx: _TemporalStartOperationContext,
 ) -> None:
     response_link = _workflow_event_link(
         WORKFLOW_ID,
@@ -197,7 +199,7 @@ async def test_signal_forwards_inbound_links_and_captures_response_backlink(
 
 
 async def test_signal_against_older_server_captures_no_backlink(
-    nexus_ctx: Any,
+    nexus_ctx: _TemporalStartOperationContext,
 ) -> None:
     workflow_service = mock.MagicMock()
     workflow_service.signal_workflow_execution = mock.AsyncMock(
@@ -215,7 +217,9 @@ async def test_signal_against_older_server_captures_no_backlink(
     assert nexus_ctx.nexus_context.outbound_links == []
 
 
-async def test_multiple_signals_accumulate_all_backlinks(nexus_ctx: Any) -> None:
+async def test_multiple_signals_accumulate_all_backlinks(
+    nexus_ctx: _TemporalStartOperationContext,
+) -> None:
     first = _workflow_event_link(
         "callee-a",
         "run-a",
@@ -265,7 +269,7 @@ async def test_signal_outside_nexus_context_does_not_touch_links() -> None:
 
 
 async def test_signal_with_start_forwards_inbound_links_and_captures_backlink(
-    nexus_ctx: Any,
+    nexus_ctx: _TemporalStartOperationContext,
 ) -> None:
     response_link = _workflow_event_link(
         WORKFLOW_ID,
@@ -294,7 +298,7 @@ async def test_signal_with_start_forwards_inbound_links_and_captures_backlink(
 
 
 async def test_signal_with_start_against_older_server_captures_no_backlink(
-    nexus_ctx: Any,
+    nexus_ctx: _TemporalStartOperationContext,
 ) -> None:
     workflow_service = mock.MagicMock()
     workflow_service.signal_with_start_workflow_execution = mock.AsyncMock(
@@ -314,8 +318,8 @@ async def test_signal_with_start_against_older_server_captures_no_backlink(
 # ── start ─────────────────────────────────────────────────────────────────────────────────
 
 
-async def test_start_forwards_inbound_links_and_captures_no_backlink(
-    nexus_ctx: Any,
+async def test_start_forwards_inbound_links_and_captures_backlink(
+    nexus_ctx: _TemporalStartOperationContext,
 ) -> None:
     server_link = _workflow_event_link(
         WORKFLOW_ID,
@@ -338,12 +342,13 @@ async def test_start_forwards_inbound_links_and_captures_no_backlink(
     assert len(sent.links) == 1
     assert sent.links[0] == _inbound_nexus_link()
 
-    # Backward: a plain start does not capture a backlink, even when the server returns one.
-    assert nexus_ctx.nexus_context.outbound_links == []
+    # Backward: a plain start captures a backlink
+    assert len(nexus_ctx.nexus_context.outbound_links) == 1
+    assert "wf-target" in _outbound_link_urls(nexus_ctx)[0]
 
 
 async def test_start_against_older_server_captures_no_backlink(
-    nexus_ctx: Any,
+    nexus_ctx: _TemporalStartOperationContext,
 ) -> None:
     workflow_service = mock.MagicMock()
     workflow_service.start_workflow_execution = mock.AsyncMock(
@@ -360,8 +365,9 @@ async def test_start_against_older_server_captures_no_backlink(
     assert len(sent.links) == 1
     assert sent.links[0] == _inbound_nexus_link()
 
-    # Backward: a plain start never fabricates a backlink.
-    assert nexus_ctx.nexus_context.outbound_links == []
+    # Backward: a plain start fabricates a backlink when the server doesn't return one.
+    assert len(nexus_ctx.nexus_context.outbound_links) == 1
+    assert "wf-target" in _outbound_link_urls(nexus_ctx)[0]
 
 
 async def test_start_outside_nexus_context_does_not_touch_links() -> None: