From 1f79370a39be761c5af57a7a4c3d55cfe26345bc Mon Sep 17 00:00:00 2001
From: LindseyZ1205 <yingzizhang1205@gmail.com>
Date: Fri, 23 Oct 2026 16:11:29 -0400
Subject: [PATCH] Stamp the time-skipping config on update-with-start workflows

TimeSkipper's client interceptor only set time_skipping_config in
start_workflow(), so workflows started with update-with-start ran
without time skipping.
---
 temporalio/testing/_timeskipping.py   |  9 +++++
 tests/testing/test_timeskipping_v2.py | 52 ++++++++++++++++++++++++++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/temporalio/testing/_timeskipping.py b/temporalio/testing/_timeskipping.py
index ab5a091d3a..a5add56da7 100644
--- a/temporalio/testing/_timeskipping.py
+++ b/temporalio/testing/_timeskipping.py
@@ -361,3 +361,12 @@ async def start_workflow(
         if self._skipper._ts_enabled:
             input.time_skipping_config = self._skipper.config._to_proto()
         return await super().start_workflow(input)
+
+    async def start_update_with_start_workflow(
+        self, input: temporalio.client.StartWorkflowUpdateWithStartInput
+    ) -> temporalio.client.WorkflowUpdateHandle[Any]:
+        if self._skipper._ts_enabled:
+            input.start_workflow_input.time_skipping_config = (
+                self._skipper.config._to_proto()
+            )
+        return await super().start_update_with_start_workflow(input)
diff --git a/tests/testing/test_timeskipping_v2.py b/tests/testing/test_timeskipping_v2.py
index 9411ecbbec..047a501474 100644
--- a/tests/testing/test_timeskipping_v2.py
+++ b/tests/testing/test_timeskipping_v2.py
@@ -10,7 +10,8 @@
 
 from temporalio import workflow
 from temporalio.api.enums.v1 import event_type_pb2 as _event_type
-from temporalio.common import RetryPolicy
+from temporalio.client import WithStartWorkflowOperation
+from temporalio.common import RetryPolicy, WorkflowIDConflictPolicy
 from temporalio.exceptions import ApplicationError
 from temporalio.testing import TimeSkipper, TimeSkippingConfig, WorkflowEnvironment
 from tests import DEV_SERVER_DOWNLOAD_VERSION
@@ -522,6 +523,55 @@ async def test_signal_with_start_stamps_time_skipping_config(
     await assert_time_was_skipped(handle)
 
 
+@workflow.defn
+class UpdateWithStartTargetWorkflow:
+    """Accepts a ``ping`` update, then does a long sleep."""
+
+    @workflow.run
+    async def run(self, sleep_seconds: float) -> dict[str, float]:
+        t_start = workflow.now().timestamp()
+        await workflow.sleep(timedelta(seconds=sleep_seconds))
+        t_end = workflow.now().timestamp()
+        return {"start": t_start, "end": t_end}
+
+    @workflow.update
+    def ping(self) -> str:
+        return "pong"
+
+
+async def test_update_with_start_stamps_time_skipping_config(
+    ts_env: WorkflowEnvironment,
+) -> None:
+    """Timeskip an update-with-start workflow."""
+    async with new_worker(ts_env.client, UpdateWithStartTargetWorkflow) as worker:
+        wall_start = monotonic()
+        start_op = WithStartWorkflowOperation(
+            UpdateWithStartTargetWorkflow.run,
+            3600.0,
+            id=f"wf-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+            id_conflict_policy=WorkflowIDConflictPolicy.FAIL,
+        )
+        update_result = await ts_env.client.execute_update_with_start_workflow(
+            UpdateWithStartTargetWorkflow.ping, start_workflow_operation=start_op
+        )
+        assert update_result == "pong"
+        handle = await start_op.workflow_handle()
+        # Fail here rather than wait for the timer in real time
+        tsi = await ts_env.get_time_skipping_info(handle)
+        assert tsi is not None
+        assert tsi.effective_config.enabled, (
+            "expected time skipping to be enabled (env-stamped)"
+        )
+        result = await handle.result()
+        wall_elapsed = monotonic() - wall_start
+
+    assert wall_elapsed < 10
+    virtual_elapsed = result["end"] - result["start"]
+    assert_duration_same(3600, virtual_elapsed, tolerance=50)
+    await assert_time_was_skipped(handle)
+
+
 async def test_get_time_skipping_info_during_workflow(
     ts_env: WorkflowEnvironment,
 ) -> None: