On start callback, test interceptor

Author: dandavison

temporalio/client.py

@@ -992,16 +992,22 @@ async def _start_update_with_start(
             update_workflow_input=update_input,
         )
 
+        def on_start(
+            start_response: temporalio.api.workflowservice.v1.StartWorkflowExecutionResponse,
+        ):
+            start_workflow_operation._workflow_handle.set_result(
+                WorkflowHandle(
+                    self,
+                    start_workflow_operation._start_workflow_input.id,
+                    first_execution_run_id=start_response.run_id,
+                )
+            )
+
+        setattr(input, "_on_start_callback", on_start)
+
         update_handle = await self._impl.start_workflow_update_with_start(input)
         # TODO https://github.com/temporalio/sdk-python/issues/682
         assert update_handle.workflow_run_id, "In Client.start_workflow_update why don't we use the run ID from the update response?"
-        start_workflow_operation._workflow_handle.set_result(
-            WorkflowHandle(
-                self,
-                update_handle.workflow_id,
-                first_execution_run_id=update_handle.workflow_run_id,
-            )
-        )
         return update_handle
 
     def list_workflows(
@@ -5878,6 +5884,14 @@ async def start_workflow_update_with_start(
                 ),
             ],
         )
+
+        on_start = getattr(input, "_on_start_callback", None)
+        if not on_start:
+            raise RuntimeError(
+                "Missing on_start callback. Please report this as a bug at https://github.com/temporalio/sdk-python/issues."
+            )
+        seen_start = False
+
         # Repeatedly try to invoke ExecuteMultiOperation until the update
         # reaches user-provided wait stage or is at least ACCEPTED (as of the
         # time of this writing, the user cannot specify sooner than ACCEPTED)
@@ -5943,6 +5957,9 @@ async def start_workflow_update_with_start(
                 raise RuntimeError("Invalid ExecuteMultiOperationResponse")
             [start_response] = start_responses
             [update_response] = update_responses
+            if not seen_start:
+                on_start(start_response)
+                seen_start = True
             known_outcome = (
                 update_response.outcome if update_response.HasField("outcome") else None
             )

tests/worker/test_update_with_start.py

@@ -4,19 +4,24 @@
 from contextlib import contextmanager
 from datetime import timedelta
 from enum import Enum
-from typing import Iterator, NoReturn
+from typing import Any, Iterator, NoReturn
 from unittest.mock import patch
 
 import pytest
 
+import temporalio.api.common.v1
 import temporalio.api.errordetails.v1
 import temporalio.worker
 from temporalio import activity, workflow
 from temporalio.client import (
     Client,
+    Interceptor,
+    OutboundInterceptor,
     RPCError,
+    StartWorkflowUpdateWithStartInput,
     WithStartWorkflowOperation,
     WorkflowUpdateFailedError,
+    WorkflowUpdateHandle,
     WorkflowUpdateStage,
 )
 from temporalio.common import (
@@ -430,8 +435,8 @@ async def test_workflow_update_poll_loop(client: Client):
     )
     with patch.object(
         client.workflow_service,
-        "update_workflow_execution",
-        wraps=client.workflow_service.update_workflow_execution,
+        "execute_multi_operation",
+        wraps=client.workflow_service.execute_multi_operation,
     ) as workflow_service_method:
         try:
             await client.execute_update_with_start(
@@ -440,6 +445,68 @@ async def test_workflow_update_poll_loop(client: Client):
             )
         except:
             print(
-                f"update_workflow_execution was called {workflow_service_method.call_count} times"
+                f"execute_multi_operation was called {workflow_service_method.call_count} times"
             )
             raise
+
+
+class SimpleClientInterceptor(Interceptor):
+    def intercept_client(self, next: OutboundInterceptor) -> OutboundInterceptor:
+        return SimpleClientOutboundInterceptor(super().intercept_client(next))
+
+
+class SimpleClientOutboundInterceptor(OutboundInterceptor):
+    def __init__(self, next: OutboundInterceptor) -> None:
+        super().__init__(next)
+
+    async def start_workflow_update_with_start(
+        self, input: StartWorkflowUpdateWithStartInput
+    ) -> WorkflowUpdateHandle[Any]:
+        input.start_workflow_input.args = ["intercepted-workflow-arg"]
+        input.update_workflow_input.args = ["intercepted-update-arg"]
+        return await super().start_workflow_update_with_start(input)
+
+
+@workflow.defn
+class UpdateWithStartInterceptorWorkflow:
+    def __init__(self) -> None:
+        self.received_update = False
+
+    @workflow.run
+    async def run(self, arg: str) -> str:
+        await workflow.wait_condition(lambda: self.received_update)
+        return arg
+
+    @workflow.update
+    async def my_update(self, arg: str) -> str:
+        self.received_update = True
+        await workflow.wait_condition(lambda: self.received_update)
+        return arg
+
+
+async def test_update_with_start_client_outbound_interceptor(
+    client: Client,
+):
+    interceptor = SimpleClientInterceptor()
+    client = Client(**{**client.config(), "interceptors": [interceptor]})
+
+    async with new_worker(
+        client,
+        UpdateWithStartInterceptorWorkflow,
+    ) as worker:
+        start_op = WithStartWorkflowOperation(
+            UpdateWithStartInterceptorWorkflow.run,
+            "original-workflow-arg",
+            id=f"wf-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+            id_conflict_policy=WorkflowIDConflictPolicy.FAIL,
+        )
+        update_result = await client.execute_update_with_start(
+            UpdateWithStartInterceptorWorkflow.my_update,
+            "original-update-arg",
+            start_workflow_operation=start_op,
+        )
+        assert update_result == "intercepted-update-arg"
+
+        wf_handle = await start_op.workflow_handle()
+        assert await wf_handle.result() == "intercepted-workflow-arg"