Adding QueryWorkflowResponse Link for Nexus

CHANGELOG.md

@@ -20,6 +20,9 @@ to include examples, links to docs, or any other relevant information.
 
 ### Added
 
+- Added Nexus operation link propagation for Workflow Queries issued from operation handlers. The
+  queried Workflow link returned by the server is attached to the caller's Nexus operation event.
+
 ### Changed
 
 ### Deprecated

temporalio/client/_impl.py

@@ -472,6 +472,9 @@ async def query_workflow(self, input: QueryWorkflowInput) -> Any:
                 raise WorkflowQueryFailedError(err.message)
             else:
                 raise
+        nexus_ctx = temporalio.nexus._operation_context._try_start_operation_context()
+        if nexus_ctx is not None and resp.HasField("link"):
+            nexus_ctx._add_response_link(resp.link)
         if resp.HasField("query_rejected"):
             raise WorkflowQueryRejectedError(
                 WorkflowExecutionStatus(resp.query_rejected.status)

temporalio/nexus/_operation_context.py

@@ -300,9 +300,9 @@ def _add_start_workflow_response_link(
     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, it is
-        converted to a Nexus link and added to the operation's outbound links.
+        ``link`` is the ``common.v1.Link`` returned by a Temporal RPC (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.
         """

tests/nexus/test_link_propagation.py

@@ -1,8 +1,9 @@
 """Unit tests for Nexus link propagation.
 
-These exercise link propagation when a Nexus operation handler signals or starts a
-workflow or activity against a mocked workflow service. End-to-end signal backlinks
-require a server with EnableCHASMSignalBacklinks enabled and are not covered here.
+These exercise link propagation when a Nexus operation handler queries, signals, or
+starts a workflow or activity against a mocked workflow service. End-to-end signal
+backlinks require a server with EnableCHASMSignalBacklinks enabled and are not covered
+here.
 """
 
 from __future__ import annotations
@@ -35,6 +36,7 @@
 import temporalio.nexus._token
 from temporalio.client._impl import _ClientImpl
 from temporalio.client._interceptor import (
+    QueryWorkflowInput,
     SignalWorkflowInput,
     StartActivityInput,
     StartWorkflowInput,
@@ -63,6 +65,19 @@ def _workflow_event_link(
     )
 
 
+def _workflow_link(
+    workflow_id: str, run_id: str, *, reason: str
+) -> temporalio.api.common.v1.Link:
+    return temporalio.api.common.v1.Link(
+        workflow=temporalio.api.common.v1.Link.Workflow(
+            namespace=NAMESPACE,
+            workflow_id=workflow_id,
+            run_id=run_id,
+            reason=reason,
+        )
+    )
+
+
 def _inbound_nexus_link() -> temporalio.api.common.v1.Link:
     return _workflow_event_link(
         "caller-wf",
@@ -132,6 +147,20 @@ def _signal_input() -> SignalWorkflowInput:
     )
 
 
+def _query_input() -> QueryWorkflowInput:
+    return QueryWorkflowInput(
+        id=WORKFLOW_ID,
+        run_id=None,
+        query="query-done",
+        args=[],
+        reject_condition=None,
+        headers={},
+        ret_type=bool,
+        rpc_metadata={},
+        rpc_timeout=None,
+    )
+
+
 def _start_input(start_signal: str | None = None) -> StartWorkflowInput:
     return StartWorkflowInput(
         workflow="TestWorkflow",
@@ -193,9 +222,56 @@ def _outbound_link_urls(ctx: Any) -> list[str]:
     return [link.url for link in ctx.nexus_context.outbound_links]
 
 
+def test_response_link_captures_workflow_link(
+    nexus_ctx: _TemporalStartOperationContext,
+) -> None:
+    nexus_ctx._add_response_link(
+        _workflow_link(WORKFLOW_ID, "target-run", reason="Query processed")
+    )
+
+    assert nexus_ctx.nexus_context.outbound_links == [
+        nexusrpc.Link(
+            type=temporalio.api.common.v1.Link.Workflow.DESCRIPTOR.full_name,
+            url=(
+                "temporal:///namespaces/test-namespace/workflows/"
+                "wf-target/target-run?reason=Query+processed"
+            ),
+        )
+    ]
+
+
 # ── signal ────────────────────────────────────────────────────────────────────────────────
 
 
+# Query responses differ from Signal responses by linking to the Workflow rather than an event.
+async def test_query_captures_response_workflow_link(
+    nexus_ctx: _TemporalStartOperationContext,
+) -> None:
+    payloads = await temporalio.converter.DataConverter.default.encode([False])
+    workflow_service = mock.MagicMock()
+    workflow_service.query_workflow = mock.AsyncMock(
+        return_value=temporalio.api.workflowservice.v1.QueryWorkflowResponse(
+            query_result=temporalio.api.common.v1.Payloads(payloads=payloads),
+            link=_workflow_link(WORKFLOW_ID, "target-run", reason="Query processed"),
+        )
+    )
+    impl = _make_client_impl(workflow_service)
+
+    result = await impl.query_workflow(_query_input())
+
+    assert result is False
+    assert nexus_ctx.nexus_context.outbound_links == [
+        nexusrpc.Link(
+            type=temporalio.api.common.v1.Link.Workflow.DESCRIPTOR.full_name,
+            url=(
+                "temporal:///namespaces/test-namespace/workflows/"
+                "wf-target/target-run?reason=Query+processed"
+            ),
+        )
+    ]
+
+
+# Signal responses link to the event that accepted the Signal.
 async def test_signal_forwards_inbound_links_and_captures_response_backlink(
     nexus_ctx: _TemporalStartOperationContext,
 ) -> None:

tests/nexus/test_temporal_operation.py

@@ -116,6 +116,7 @@ class TestService:
     sync_result: Operation[Input, str]
     custom_cancel: Operation[str, None]
     update_op: Operation[Input, str]
+    query_op: Operation[str, bool]
     echo_activity: Operation[Input, str]
     error_activity: Operation[Input, None]
     blocking_activity: Operation[str, None]
@@ -292,6 +293,17 @@ async def update_op(
             update_id=input.update_id,
         )
 
+    @nexus.temporal_operation
+    async def query_op(
+        self,
+        _ctx: nexus.TemporalStartOperationContext,
+        client: nexus.TemporalNexusClient,
+        input: str,
+    ) -> nexus.TemporalOperationResult[bool]:
+        handle = client.client.get_workflow_handle(input)
+        result = await handle.query(BlockingWorkflow.query_done)
+        return nexus.TemporalOperationResult.sync(result)
+
     @nexus.temporal_operation
     async def echo_activity(
         self,
@@ -822,6 +834,74 @@ async def run(self) -> None:
     async def unblock(self):
         self.done = True
 
+    @workflow.query
+    def query_done(self) -> bool:
+        return self.done
+
+
+@workflow.defn
+class QueryWorkflowCaller:
+    @workflow.run
+    async def run(self, input: Input) -> bool:
+        client = workflow.create_nexus_client(
+            service=TestService, endpoint=make_nexus_endpoint_name(input.task_queue)
+        )
+        return await client.execute_operation(TestService.query_op, input.value)
+
+
+async def test_temporal_operation_query_workflow(
+    client: Client, env: WorkflowEnvironment
+) -> None:
+    task_queue = str(uuid.uuid4())
+    endpoint_name = make_nexus_endpoint_name(task_queue)
+    await env.create_nexus_endpoint(endpoint_name, task_queue)
+    target_workflow_id = f"query-target-{uuid.uuid4()}"
+
+    async with Worker(
+        env.client,
+        task_queue=task_queue,
+        nexus_service_handlers=[TestServiceHandler()],
+        workflows=[BlockingWorkflow, QueryWorkflowCaller],
+    ):
+        target_handle = await client.start_workflow(
+            BlockingWorkflow.run,
+            id=target_workflow_id,
+            task_queue=task_queue,
+        )
+        caller_handle = await client.start_workflow(
+            QueryWorkflowCaller.run,
+            Input(value=target_workflow_id, task_queue=task_queue),
+            id=f"query-caller-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+
+        try:
+            assert not await caller_handle.result()
+
+            caller_history = await caller_handle.fetch_history()
+            completed_event = next(
+                event
+                for event in caller_history.events
+                if event.event_type == EventType.EVENT_TYPE_NEXUS_OPERATION_COMPLETED
+            )
+
+            target_history = await target_handle.fetch_history()
+            assert not any(event.links for event in target_history.events)
+
+            if not completed_event.links:
+                pytest.skip("server did not return a Workflow Query response link")
+            assert target_handle.result_run_id is not None
+            assert Link(
+                workflow=Link.Workflow(
+                    namespace=client.namespace,
+                    workflow_id=target_workflow_id,
+                    run_id=target_handle.result_run_id,
+                    reason="Query processed",
+                )
+            ) in list(completed_event.links)
+        finally:
+            await target_handle.cancel()
+
 
 @workflow.defn
 class CancelBlockingWorkflowCaller: