From 62ea4941e2084b8c51952e83b1ff2c6fb1f637b4 Mon Sep 17 00:00:00 2001
From: Dan Davison <dan.davison@temporal.io>
Date: Sat, 12 Jul 2025 20:58:47 -0400
Subject: [PATCH] Use links from StartWorkflowExecutionResponse if present

---
 temporalio/client.py                   |  8 ++++++
 temporalio/nexus/_link_conversion.py   |  2 +-
 temporalio/nexus/_operation_context.py | 36 +++++++++++++++-----------
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/temporalio/client.py b/temporalio/client.py
index 3740ab335..3eb3f35e1 100644
--- a/temporalio/client.py
+++ b/temporalio/client.py
@@ -1542,6 +1542,12 @@ def __init__(
         result_run_id: Optional[str] = None,
         first_execution_run_id: Optional[str] = None,
         result_type: Optional[Type] = None,
+        start_workflow_response: Optional[
+            Union[
+                temporalio.api.workflowservice.v1.StartWorkflowExecutionResponse,
+                temporalio.api.workflowservice.v1.SignalWithStartWorkflowExecutionResponse,
+            ]
+        ] = None,
     ) -> None:
         """Create workflow handle."""
         self._client = client
@@ -1550,6 +1556,7 @@ def __init__(
         self._result_run_id = result_run_id
         self._first_execution_run_id = first_execution_run_id
         self._result_type = result_type
+        self._start_workflow_response = start_workflow_response
         self.__temporal_eagerly_started = False
 
     @property
@@ -5832,6 +5839,7 @@ async def start_workflow(
             result_run_id=resp.run_id,
             first_execution_run_id=first_execution_run_id,
             result_type=input.ret_type,
+            start_workflow_response=resp,
         )
         setattr(handle, "__temporal_eagerly_started", eagerly_started)
         return handle
diff --git a/temporalio/nexus/_link_conversion.py b/temporalio/nexus/_link_conversion.py
index 87027333b..9df56b9bf 100644
--- a/temporalio/nexus/_link_conversion.py
+++ b/temporalio/nexus/_link_conversion.py
@@ -23,7 +23,7 @@
 LINK_EVENT_TYPE_PARAM_NAME = "eventType"
 
 
-def workflow_handle_to_workflow_execution_started_event_link(
+def workflow_execution_started_event_link_from_workflow_handle(
     handle: temporalio.client.WorkflowHandle[Any, Any],
 ) -> temporalio.api.common.v1.Link.WorkflowEvent:
     """Create a WorkflowEvent link corresponding to a started workflow"""
diff --git a/temporalio/nexus/_operation_context.py b/temporalio/nexus/_operation_context.py
index 9ff676dc9..e0f28b28f 100644
--- a/temporalio/nexus/_operation_context.py
+++ b/temporalio/nexus/_operation_context.py
@@ -18,6 +18,7 @@
 from typing_extensions import Concatenate
 
 import temporalio.api.common.v1
+import temporalio.api.workflowservice.v1
 import temporalio.client
 import temporalio.common
 from temporalio.nexus import _link_conversion
@@ -142,25 +143,30 @@ def _get_workflow_event_links(
     def _add_outbound_links(
         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:
-            link = _link_conversion.workflow_event_to_nexus_link(
-                _link_conversion.workflow_handle_to_workflow_execution_started_event_link(
-                    workflow_handle
-                )
+            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 = [
+                    _link_conversion.workflow_execution_started_event_link_from_workflow_handle(
+                        workflow_handle
+                    )
+                ]
+            self.nexus_context.outbound_links.extend(
+                _link_conversion.workflow_event_to_nexus_link(link)
+                for link in wf_event_links
             )
         except Exception as e:
             logger.warning(
-                f"Failed to create WorkflowExecutionStarted event link for workflow {id}: {e}"
-            )
-        else:
-            self.nexus_context.outbound_links.append(
-                # TODO(nexus-prerelease): Before, WorkflowRunOperation was generating an EventReference
-                # link to send back to the caller. Now, it checks if the server returned
-                # the link in the StartWorkflowExecutionResponse, and if so, send the link
-                # from the response to the caller. Fallback to generating the link for
-                # backwards compatibility. PR reference in Go SDK:
-                # https://github.com/temporalio/sdk-go/pull/1934
-                link
+                f"Failed to create WorkflowExecutionStarted event links for workflow {workflow_handle}: {e}"
             )
         return workflow_handle