Support RequestIdReference in Nexus Links (#1228)

* WIP add Nexus request id reference to links

* Small cleanups. Move to request ID reference for workflow_execution_started_event_link_from_workflow_handle

* Update tests/nexus/test_link_conversion.py

Co-authored-by: James Watkins-Harvey <[email protected]>

---------

Co-authored-by: James Watkins-Harvey <[email protected]>

Author: VegetarianOrc

temporalio/nexus/_link_conversion.py

@@ -24,26 +24,31 @@
 )
 LINK_EVENT_ID_PARAM_NAME = "eventID"
 LINK_EVENT_TYPE_PARAM_NAME = "eventType"
+LINK_REQUEST_ID_PARAM_NAME = "requestID"
+LINK_REFERENCE_TYPE_PARAM_NAME = "referenceType"
+
+EVENT_REFERENCE_TYPE = "EventReference"
+REQUEST_ID_REFERENCE_TYPE = "RequestIdReference"
 
 
 def workflow_execution_started_event_link_from_workflow_handle(
-    handle: temporalio.client.WorkflowHandle[Any, Any],
+    handle: temporalio.client.WorkflowHandle[Any, Any], request_id: str
 ) -> temporalio.api.common.v1.Link.WorkflowEvent:
     """Create a WorkflowEvent link corresponding to a started workflow"""
     if handle.first_execution_run_id is None:
         raise ValueError(
             f"Workflow handle {handle} has no first execution run ID. "
             f"Cannot create WorkflowExecutionStarted event link."
         )
+
     return temporalio.api.common.v1.Link.WorkflowEvent(
         namespace=handle._client.namespace,
         workflow_id=handle.id,
         run_id=handle.first_execution_run_id,
-        event_ref=temporalio.api.common.v1.Link.WorkflowEvent.EventReference(
-            event_id=1,
+        request_id_ref=temporalio.api.common.v1.Link.WorkflowEvent.RequestIdReference(
+            request_id=request_id,
             event_type=temporalio.api.enums.v1.EventType.EVENT_TYPE_WORKFLOW_EXECUTION_STARTED,
         ),
-        # TODO(nexus-preview): RequestIdReference
     )
 
 
@@ -60,9 +65,21 @@ def workflow_event_to_nexus_link(
     workflow_id = urllib.parse.quote(workflow_event.workflow_id)
     run_id = urllib.parse.quote(workflow_event.run_id)
     path = f"/namespaces/{namespace}/workflows/{workflow_id}/{run_id}/history"
-    query_params = _event_reference_to_query_params(workflow_event.event_ref)
+
+    query_params = None
+    match workflow_event.WhichOneof("reference"):
+        case "event_ref":
+            query_params = _event_reference_to_query_params(workflow_event.event_ref)
+        case "request_id_ref":
+            query_params = _request_id_reference_to_query_params(
+                workflow_event.request_id_ref
+            )
+
+    # urllib will omit '//' from the url if netloc is empty so we add the scheme manually
+    url = f"{scheme}://{urllib.parse.urlunparse(('', '', path, '', query_params, ''))}"
+
     return nexusrpc.Link(
-        url=urllib.parse.urlunparse((scheme, "", path, "", query_params, "")),
+        url=url,
         type=workflow_event.DESCRIPTOR.full_name,
     )
 
@@ -83,7 +100,20 @@ def nexus_link_to_workflow_event(
         )
         return None
     try:
-        event_ref = _query_params_to_event_reference(url.query)
+        query_params = urllib.parse.parse_qs(url.query)
+
+        request_id_ref = None
+        event_ref = None
+        match query_params.get(LINK_REFERENCE_TYPE_PARAM_NAME):
+            case ["EventReference"]:
+                event_ref = _query_params_to_event_reference(query_params)
+            case ["RequestIdReference"]:
+                request_id_ref = _query_params_to_request_id_reference(query_params)
+            case _:
+                raise ValueError(
+                    f"Invalid Nexus link: {link}. Expected {LINK_REFERENCE_TYPE_PARAM_NAME} to be '{EVENT_REFERENCE_TYPE}' or '{REQUEST_ID_REFERENCE_TYPE}'"
+                )
+
     except ValueError as err:
         logger.warning(
             f"Failed to parse event reference from Nexus link URL query parameters: {link} ({err})"
@@ -96,6 +126,7 @@ def nexus_link_to_workflow_event(
         workflow_id=urllib.parse.unquote(groups["workflow_id"]),
         run_id=urllib.parse.unquote(groups["run_id"]),
         event_ref=event_ref,
+        request_id_ref=request_id_ref,
     )
 
 
@@ -109,36 +140,58 @@ def _event_reference_to_query_params(
         )
     return urllib.parse.urlencode(
         {
-            "eventID": event_ref.event_id,
-            "eventType": event_type_name,
-            "referenceType": "EventReference",
+            LINK_EVENT_ID_PARAM_NAME: event_ref.event_id,
+            LINK_EVENT_TYPE_PARAM_NAME: event_type_name,
+            LINK_REFERENCE_TYPE_PARAM_NAME: EVENT_REFERENCE_TYPE,
         }
     )
 
 
+def _request_id_reference_to_query_params(
+    request_id_ref: temporalio.api.common.v1.Link.WorkflowEvent.RequestIdReference,
+) -> str:
+    params = {
+        LINK_REFERENCE_TYPE_PARAM_NAME: REQUEST_ID_REFERENCE_TYPE,
+    }
+
+    if request_id_ref.request_id:
+        params[LINK_REQUEST_ID_PARAM_NAME] = request_id_ref.request_id
+
+    event_type_name = temporalio.api.enums.v1.EventType.Name(request_id_ref.event_type)
+    if event_type_name.startswith("EVENT_TYPE_"):
+        event_type_name = _event_type_constant_case_to_pascal_case(
+            event_type_name.removeprefix("EVENT_TYPE_")
+        )
+    params[LINK_EVENT_TYPE_PARAM_NAME] = event_type_name
+
+    return urllib.parse.urlencode(params)
+
+
 def _query_params_to_event_reference(
-    raw_query_params: str,
+    query_params: dict[str, list[str]],
 ) -> temporalio.api.common.v1.Link.WorkflowEvent.EventReference:
     """Return an EventReference from the query params or raise ValueError."""
-    query_params = urllib.parse.parse_qs(raw_query_params)
-
-    [reference_type] = query_params.get("referenceType") or [""]
-    if reference_type != "EventReference":
+    [reference_type] = query_params.get(LINK_REFERENCE_TYPE_PARAM_NAME) or [""]
+    if reference_type != EVENT_REFERENCE_TYPE:
         raise ValueError(
             f"Expected Nexus link URL query parameter referenceType to be EventReference but got: {reference_type}"
         )
+
     # event type
-    [raw_event_type_name] = query_params.get(LINK_EVENT_TYPE_PARAM_NAME) or [""]
-    if not raw_event_type_name:
-        raise ValueError(f"query params do not contain event type: {query_params}")
-    if raw_event_type_name.startswith("EVENT_TYPE_"):
-        event_type_name = raw_event_type_name
-    elif re.match("[A-Z][a-z]", raw_event_type_name):
-        event_type_name = "EVENT_TYPE_" + _event_type_pascal_case_to_constant_case(
-            raw_event_type_name
-        )
-    else:
-        raise ValueError(f"Invalid event type name: {raw_event_type_name}")
+    match query_params.get(LINK_EVENT_TYPE_PARAM_NAME):
+        case None:
+            raise ValueError(f"query params do not contain event type: {query_params}")
+
+        case [raw_event_type_name] if raw_event_type_name.startswith("EVENT_TYPE_"):
+            event_type_name = raw_event_type_name
+
+        case [raw_event_type_name] if re.match("[A-Z][a-z]", raw_event_type_name):
+            event_type_name = "EVENT_TYPE_" + _event_type_pascal_case_to_constant_case(
+                raw_event_type_name
+            )
+
+        case raw_event_type_name:
+            raise ValueError(f"Invalid event type name: {raw_event_type_name}")
 
     # event id
     event_id = 0
@@ -155,6 +208,34 @@ def _query_params_to_event_reference(
     )
 
 
+def _query_params_to_request_id_reference(
+    query_params: dict[str, list[str]],
+) -> temporalio.api.common.v1.Link.WorkflowEvent.RequestIdReference:
+    """Return an EventReference from the query params or raise ValueError."""
+    # event type
+    match query_params.get(LINK_EVENT_TYPE_PARAM_NAME):
+        case None:
+            raise ValueError(f"query params do not contain event type: {query_params}")
+
+        case [raw_event_type_name] if raw_event_type_name.startswith("EVENT_TYPE_"):
+            event_type_name = raw_event_type_name
+
+        case [raw_event_type_name] if re.match("[A-Z][a-z]", raw_event_type_name):
+            event_type_name = "EVENT_TYPE_" + _event_type_pascal_case_to_constant_case(
+                raw_event_type_name
+            )
+
+        case raw_event_type_name:
+            raise ValueError(f"Invalid event type name: {raw_event_type_name}")
+
+    [request_id] = query_params.get(LINK_REQUEST_ID_PARAM_NAME, [""])
+
+    return temporalio.api.common.v1.Link.WorkflowEvent.RequestIdReference(
+        request_id=request_id,
+        event_type=temporalio.api.enums.v1.EventType.Value(event_type_name),
+    )
+
+
 def _event_type_constant_case_to_pascal_case(s: str) -> str:
     """Convert a CONSTANT_CASE string to PascalCase.
 

temporalio/nexus/_operation_context.py

@@ -183,7 +183,8 @@ def _add_outbound_links(
             if not wf_event_links:
                 wf_event_links = [
                     _link_conversion.workflow_execution_started_event_link_from_workflow_handle(
-                        workflow_handle
+                        workflow_handle,
+                        self.nexus_context.request_id,
                     )
                 ]
             self.nexus_context.outbound_links.extend(

tests/nexus/test_link_conversion.py

@@ -1,6 +1,7 @@
 import urllib.parse
 from typing import Any
 
+import nexusrpc
 import pytest
 
 import temporalio.api.common.v1
@@ -39,8 +40,9 @@
 def test_query_params_to_event_reference(
     query_param_str: str, expected_event_ref: dict[str, Any]
 ):
+    query_params = urllib.parse.parse_qs(query_param_str)
     event_ref = temporalio.nexus._link_conversion._query_params_to_event_reference(
-        query_param_str
+        query_params
     )
     for k, v in expected_event_ref.items():
         assert getattr(event_ref, k) == v
@@ -72,6 +74,118 @@ def test_event_reference_to_query_params(
     assert query_params == expected_query_params
 
 
+@pytest.mark.parametrize(
+    ["query_param_str", "expected_event_ref"],
+    [
+        (
+            "eventType=NexusOperationScheduled&referenceType=RequestIdReference&requestID=req-123",
+            {
+                "event_type": temporalio.api.enums.v1.EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED,
+                "request_id": "req-123",
+            },
+        ),
+        # event ID is optional in query params; we leave it unset in the ref if missing
+        (
+            "eventType=NexusOperationScheduled&referenceType=RequestIdReference",
+            {
+                "event_type": temporalio.api.enums.v1.EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED,
+                "request_id": "",
+            },
+        ),
+        # Older server sends EVENT_TYPE_CONSTANT_CASE event type name
+        (
+            "eventType=EVENT_TYPE_NEXUS_OPERATION_SCHEDULED&referenceType=RequestIdReference&requestID=req-123",
+            {
+                "event_type": temporalio.api.enums.v1.EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED,
+                "request_id": "req-123",
+            },
+        ),
+    ],
+)
+def test_query_params_to_request_id_reference(
+    query_param_str: str, expected_event_ref: dict[str, Any]
+):
+    query_params = urllib.parse.parse_qs(query_param_str)
+    event_ref = temporalio.nexus._link_conversion._query_params_to_request_id_reference(
+        query_params
+    )
+    for k, v in expected_event_ref.items():
+        assert getattr(event_ref, k) == v
+
+
+@pytest.mark.parametrize(
+    ["event_ref", "expected_query_param_str"],
+    [
+        # We always send PascalCase event type names (no EventType prefix)
+        (
+            {
+                "event_type": temporalio.api.enums.v1.EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED,
+                "request_id": "req-123",
+            },
+            "eventType=NexusOperationScheduled&referenceType=RequestIdReference&requestID=req-123",
+        ),
+    ],
+)
+def test_request_id_reference_to_query_params(
+    event_ref: dict[str, Any], expected_query_param_str: str
+):
+    query_params_str = (
+        temporalio.nexus._link_conversion._request_id_reference_to_query_params(
+            temporalio.api.common.v1.Link.WorkflowEvent.RequestIdReference(**event_ref)
+        )
+    )
+    query_params = urllib.parse.parse_qs(query_params_str)
+    expected_query_params = urllib.parse.parse_qs(expected_query_param_str)
+    assert query_params == expected_query_params
+
+
+@pytest.mark.parametrize(
+    ["event", "expected_link"],
+    [
+        (
+            temporalio.api.common.v1.Link.WorkflowEvent(
+                namespace="ns",
+                workflow_id="wid",
+                run_id="rid",
+                request_id_ref=temporalio.api.common.v1.Link.WorkflowEvent.RequestIdReference(
+                    event_type=temporalio.api.enums.v1.event_type_pb2.EVENT_TYPE_WORKFLOW_TASK_COMPLETED,
+                    request_id="req-123",
+                ),
+            ),
+            nexusrpc.Link(
+                type=temporalio.api.common.v1.Link.WorkflowEvent.DESCRIPTOR.full_name,
+                url="temporal:///namespaces/ns/workflows/wid/rid/history?referenceType=RequestIdReference&requestID=req-123&eventType=WorkflowTaskCompleted",
+            ),
+        ),
+        (
+            temporalio.api.common.v1.Link.WorkflowEvent(
+                namespace="ns2",
+                workflow_id="wid2",
+                run_id="rid2",
+                event_ref=temporalio.api.common.v1.Link.WorkflowEvent.EventReference(
+                    event_id=42,
+                    event_type=temporalio.api.enums.v1.event_type_pb2.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED,
+                ),
+            ),
+            nexusrpc.Link(
+                type=temporalio.api.common.v1.Link.WorkflowEvent.DESCRIPTOR.full_name,
+                url="temporal:///namespaces/ns2/workflows/wid2/rid2/history?eventID=42&eventType=WorkflowExecutionCompleted&referenceType=EventReference",
+            ),
+        ),
+    ],
+)
+def test_link_conversion_workflow_event_to_link_and_back(
+    event: temporalio.api.common.v1.Link.WorkflowEvent, expected_link: nexusrpc.Link
+):
+    actual_link = temporalio.nexus._link_conversion.workflow_event_to_nexus_link(event)
+    assert expected_link == actual_link
+
+    actual_event = temporalio.nexus._link_conversion.nexus_link_to_workflow_event(
+        actual_link
+    )
+    assert event == actual_event
+
+
 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