Upgrade on CAN (#1311)

* upgrade on CAN (needs test)

* fix CAN overloads (add initial_version_behavior), add test

* Add PR suggestions, add test to check CAN suggested reasons persist across WFTs

* add docs to UNSPECIFIED enums for suggested CAN reasons and initial versioning behavior

* remove unused _to_proto (casting instead)

* formattign

* add missing proto.nexus import

* fix lint - remove misc print

* add param to v2, add unused param suppression

* get again formatting

* Update to use new version upgraded flag

* Formatting

* Remove continue-as-new reasons exposure in favor of new bool API

The continue-as-new reasons exposure functionality is no longer needed
with the new bool API. This commit removes:

- get_suggested_continue_as_new_reasons() method from workflow.Info
- SuggestContinueAsNewReason enum class
- Related implementation in _workflow_instance.py
- Abstract method declaration from _Runtime

Keeps is_target_worker_deployment_version_changed() functionality for
upgrade-on-continue-as-new feature.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Revert protobuf files to main branch versions

Reset all pb2.py and pb2.pyi files to their main branch versions to remove
continue-as-new reason exposures that are no longer needed with the new
bool API.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Remove invalid test

---------

Co-authored-by: Tim Conley <timothy.conley@temporal.io>
Co-authored-by: Claude <noreply@anthropic.com>

Author: 3 people

temporalio/worker/_interceptor.py

@@ -21,7 +21,7 @@
 import temporalio.nexus
 import temporalio.nexus._util
 import temporalio.workflow
-from temporalio.workflow import VersioningIntent
+from temporalio.workflow import ContinueAsNewVersioningBehavior, VersioningIntent
 
 
 class Interceptor:
@@ -172,6 +172,7 @@ class ContinueAsNewInput:
     )
     headers: Mapping[str, temporalio.api.common.v1.Payload]
     versioning_intent: VersioningIntent | None
+    initial_versioning_behavior: ContinueAsNewVersioningBehavior | None
     # The types may be absent
     arg_types: list[type] | None
 

temporalio/worker/_workflow_instance.py

@@ -251,6 +251,7 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
         self._current_history_length = 0
         self._current_history_size = 0
         self._continue_as_new_suggested = False
+        self._target_worker_deployment_version_changed = False
         # Lazily loaded
         self._untyped_converted_memo: MutableMapping[str, Any] | None = None
         # Handles which are ready to run on the next event loop iteration
@@ -405,6 +406,9 @@ def activate(
         self._current_history_length = act.history_length
         self._current_history_size = act.history_size_bytes
         self._continue_as_new_suggested = act.continue_as_new_suggested
+        self._target_worker_deployment_version_changed = (
+            act.target_worker_deployment_version_changed
+        )
         self._time_ns = act.timestamp.ToNanoseconds()
         self._is_replaying = act.is_replaying
         self._current_thread_id = threading.get_ident()
@@ -1131,6 +1135,8 @@ def workflow_continue_as_new(
             temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
         ),
         versioning_intent: temporalio.workflow.VersioningIntent | None,
+        initial_versioning_behavior: temporalio.workflow.ContinueAsNewVersioningBehavior
+        | None,
     ) -> NoReturn:
         self._assert_not_read_only("continue as new")
         # Use definition if callable
@@ -1158,6 +1164,7 @@ def workflow_continue_as_new(
                 headers={},
                 arg_types=arg_types,
                 versioning_intent=versioning_intent,
+                initial_versioning_behavior=initial_versioning_behavior,
             )
         )
 
@@ -1227,6 +1234,9 @@ def workflow_instance(self) -> Any:
     def workflow_is_continue_as_new_suggested(self) -> bool:
         return self._continue_as_new_suggested
 
+    def workflow_is_target_worker_deployment_version_changed(self) -> bool:
+        return self._target_worker_deployment_version_changed
+
     def workflow_is_replaying(self) -> bool:
         return self._is_replaying
 
@@ -3428,6 +3438,11 @@ def _apply_command(self) -> None:
             )
         if self._input.versioning_intent:
             v.versioning_intent = self._input.versioning_intent._to_proto()
+        if self._input.initial_versioning_behavior:
+            v.initial_versioning_behavior = cast(
+                "temporalio.api.enums.v1.ContinueAsNewVersioningBehavior.ValueType",
+                int(self._input.initial_versioning_behavior),
+            )
 
 
 def _encode_search_attributes(

temporalio/workflow.py

@@ -50,6 +50,8 @@
 )
 
 import temporalio.api.common.v1
+import temporalio.api.enums
+import temporalio.api.enums.v1
 import temporalio.bridge.proto.child_workflow
 import temporalio.bridge.proto.common
 import temporalio.bridge.proto.nexus
@@ -601,6 +603,16 @@ def is_continue_as_new_suggested(self) -> bool:
         """
         return _Runtime.current().workflow_is_continue_as_new_suggested()
 
+    def is_target_worker_deployment_version_changed(self) -> bool:
+        """Check whether the target worker deployment version has changed.
+
+        Note: Upgrade-on-Continue-as-New is currently experimental.
+
+        Returns:
+            True if the target worker deployment version has changed.
+        """
+        return _Runtime.current().workflow_is_target_worker_deployment_version_changed()
+
 
 @dataclass(frozen=True)
 class ParentInfo:
@@ -690,6 +702,7 @@ def workflow_continue_as_new(
             temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
         ),
         versioning_intent: VersioningIntent | None,
+        initial_versioning_behavior: ContinueAsNewVersioningBehavior | None,
     ) -> NoReturn: ...
 
     @abstractmethod
@@ -735,6 +748,9 @@ def workflow_instance(self) -> Any: ...
     @abstractmethod
     def workflow_is_continue_as_new_suggested(self) -> bool: ...
 
+    @abstractmethod
+    def workflow_is_target_worker_deployment_version_changed(self) -> bool: ...
+
     @abstractmethod
     def workflow_is_replaying(self) -> bool: ...
 
@@ -4800,6 +4816,7 @@ def continue_as_new(
         temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
     ) = None,
     versioning_intent: VersioningIntent | None = None,
+    initial_versioning_behavior: ContinueAsNewVersioningBehavior | None = None,
 ) -> NoReturn: ...
 
 
@@ -4818,6 +4835,7 @@ def continue_as_new(
         temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
     ) = None,
     versioning_intent: VersioningIntent | None = None,
+    initial_versioning_behavior: ContinueAsNewVersioningBehavior | None = None,
 ) -> NoReturn: ...
 
 
@@ -4837,6 +4855,7 @@ def continue_as_new(
         temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
     ) = None,
     versioning_intent: VersioningIntent | None = None,
+    initial_versioning_behavior: ContinueAsNewVersioningBehavior | None = None,
 ) -> NoReturn: ...
 
 
@@ -4856,6 +4875,7 @@ def continue_as_new(
         temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
     ) = None,
     versioning_intent: VersioningIntent | None = None,
+    initial_versioning_behavior: ContinueAsNewVersioningBehavior | None = None,
 ) -> NoReturn: ...
 
 
@@ -4875,6 +4895,7 @@ def continue_as_new(
         temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
     ) = None,
     versioning_intent: VersioningIntent | None = None,
+    initial_versioning_behavior: ContinueAsNewVersioningBehavior | None = None,
 ) -> NoReturn: ...
 
 
@@ -4893,6 +4914,7 @@ def continue_as_new(
         temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
     ) = None,
     versioning_intent: VersioningIntent | None = None,
+    initial_versioning_behavior: ContinueAsNewVersioningBehavior | None = None,
 ) -> NoReturn:
     """Stop the workflow immediately and continue as new.
 
@@ -4934,6 +4956,7 @@ def continue_as_new(
         memo=memo,
         search_attributes=search_attributes,
         versioning_intent=versioning_intent,
+        initial_versioning_behavior=initial_versioning_behavior,
     )
 
 
@@ -5350,6 +5373,34 @@ def _to_proto(self) -> temporalio.bridge.proto.common.VersioningIntent.ValueType
         return temporalio.bridge.proto.common.VersioningIntent.UNSPECIFIED
 
 
+class ContinueAsNewVersioningBehavior(IntEnum):
+    """Experimental. Optionally decide the versioning behavior that the first task of the new run should use.
+    For example, choose to AutoUpgrade on continue-as-new instead of inheriting the pinned version
+    of the previous run.
+    """
+
+    UNSPECIFIED = int(
+        temporalio.api.enums.v1.ContinueAsNewVersioningBehavior.CONTINUE_AS_NEW_VERSIONING_BEHAVIOR_UNSPECIFIED
+    )
+    """An initial versioning behavior is not set, follow the existing continue-as-new inheritance semantics.
+    See https://docs.temporal.io/worker-versioning#inheritance-semantics for more detail.
+    """
+
+    AUTO_UPGRADE = int(
+        temporalio.api.enums.v1.ContinueAsNewVersioningBehavior.CONTINUE_AS_NEW_VERSIONING_BEHAVIOR_AUTO_UPGRADE
+    )
+    """Start the new run with AutoUpgrade behavior. Use the Target Version of the workflow's task queue at
+    start-time, as AutoUpgrade workflows do. After the first workflow task completes, use whatever
+    Versioning Behavior the workflow is annotated with in the workflow code.
+    
+    Note that if the previous workflow had a Pinned override, that override will be inherited by the
+    new workflow run regardless of the ContinueAsNewVersioningBehavior specified in the continue-as-new
+    command. If a Pinned override is inherited by the new run, and the new run starts with AutoUpgrade
+    behavior, the base version of the new run will be the Target Version as described above, but the
+    effective version will be whatever is specified by the Versioning Override until the override is removed.
+    """
+
+
 ServiceT = TypeVar("ServiceT")
 
 

tests/worker/test_worker.py

@@ -27,6 +27,7 @@
 )
 from temporalio.client import (
     Client,
+    WorkflowHandle,
 )
 from temporalio.common import PinnedVersioningOverride, RawValue, VersioningBehavior
 from temporalio.runtime import (
@@ -1162,6 +1163,54 @@ async def set_ramping_version(
     return response
 
 
+async def wait_for_worker_deployment_routing_config_propagation(
+    client: Client,
+    deployment_name: str,
+    expected_current_build_id: str,
+    expected_ramping_build_id: str = "",
+) -> None:
+    """Wait for routing config to be propagated to all task queues."""
+    import temporalio.api.enums.v1
+
+    async def check() -> bool:
+        resp = await client.workflow_service.describe_worker_deployment(
+            DescribeWorkerDeploymentRequest(
+                namespace=client.namespace,
+                deployment_name=deployment_name,
+            )
+        )
+        routing_config = resp.worker_deployment_info.routing_config
+        if (
+            routing_config.current_deployment_version.build_id
+            != expected_current_build_id
+        ):
+            return False
+        if (
+            routing_config.ramping_deployment_version.build_id
+            != expected_ramping_build_id
+        ):
+            return False
+        state = resp.worker_deployment_info.routing_config_update_state
+        if (
+            state
+            == temporalio.api.enums.v1.RoutingConfigUpdateState.ROUTING_CONFIG_UPDATE_STATE_COMPLETED
+        ):
+            return True
+        if (
+            state
+            == temporalio.api.enums.v1.RoutingConfigUpdateState.ROUTING_CONFIG_UPDATE_STATE_UNSPECIFIED
+        ):
+            return True  # unimplemented
+        if (
+            state
+            == temporalio.api.enums.v1.RoutingConfigUpdateState.ROUTING_CONFIG_UPDATE_STATE_IN_PROGRESS
+        ):
+            return False
+        return False
+
+    await assert_eventually(check)
+
+
 def create_worker(
     client: Client,
     on_fatal_error: Callable[[BaseException], Awaitable[None]] | None = None,
@@ -1316,3 +1365,124 @@ async def capture_client_activity() -> None:
     assert len(captured_clients) == 2
     assert captured_clients[0] is client
     assert captured_clients[1] is client2  # This will fail before the fix
+
+
+@workflow.defn(
+    name="ContinueAsNewWithVersionUpgrade",
+    versioning_behavior=VersioningBehavior.PINNED,
+)
+class ContinueAsNewWithVersionUpgradeV1:
+    @workflow.run
+    async def run(self, attempt: int) -> str:
+        if attempt > 0:
+            return "v1.0"
+
+        # Loop waiting for CAN suggestion with version changed
+        while True:
+            # Trigger a WFT when timer expires, thereby refreshing the continue-as-new-suggested flag
+            await asyncio.sleep(0.01)
+            info = workflow.info()
+            if info.is_target_worker_deployment_version_changed():
+                workflow.continue_as_new(
+                    arg=attempt + 1,
+                    initial_versioning_behavior=workflow.ContinueAsNewVersioningBehavior.AUTO_UPGRADE,
+                )
+
+
+@workflow.defn(
+    name="ContinueAsNewWithVersionUpgrade",
+    versioning_behavior=VersioningBehavior.PINNED,
+)
+class ContinueAsNewWithVersionUpgradeV2:
+    @workflow.run
+    async def run(self, attempt: int) -> str:  # type:ignore[reportUnusedParameter]
+        return "v2.0"
+
+
+async def wait_for_workflow_running_on_version(
+    handle: WorkflowHandle[Any, Any], expected_build_id: str
+) -> None:
+    """Wait until workflow is RUNNING with expected build ID."""
+
+    async def check() -> bool:
+        desc = await handle.describe()
+        if (
+            desc.status
+            != temporalio.api.enums.v1.WorkflowExecutionStatus.WORKFLOW_EXECUTION_STATUS_RUNNING
+        ):
+            return False
+        versioning_info = desc.raw_description.workflow_execution_info.versioning_info
+        if not versioning_info.HasField("deployment_version"):
+            return False
+        return versioning_info.deployment_version.build_id == expected_build_id
+
+    await assert_eventually(check)
+
+
+async def test_continue_as_new_with_version_upgrade(
+    client: Client, env: WorkflowEnvironment
+):
+    if env.supports_time_skipping:
+        pytest.skip("Test Server doesn't support worker deployments")
+
+    deployment_name = f"deployment-can-upgrade-{uuid.uuid4()}"
+    v1 = WorkerDeploymentVersion(deployment_name=deployment_name, build_id="1.0")
+    v2 = WorkerDeploymentVersion(deployment_name=deployment_name, build_id="2.0")
+
+    async with (
+        new_worker(
+            client,
+            ContinueAsNewWithVersionUpgradeV1,
+            deployment_config=WorkerDeploymentConfig(
+                version=v1,
+                use_worker_versioning=True,
+            ),
+        ) as w1,
+        new_worker(
+            client,
+            ContinueAsNewWithVersionUpgradeV2,
+            deployment_config=WorkerDeploymentConfig(
+                version=v2,
+                use_worker_versioning=True,
+            ),
+            task_queue=w1.task_queue,
+        ),
+    ):
+        # Wait for the deployment to be ready
+        describe_resp = await wait_until_worker_deployment_visible(client, v1)
+
+        # Set version 1.0 as current
+        resp2 = await set_current_deployment_version(
+            client, describe_resp.conflict_token, v1
+        )
+
+        # Wait for v1.0-as-Current routing config to be propagated
+        await wait_for_worker_deployment_routing_config_propagation(
+            client, deployment_name, v1.build_id
+        )
+
+        # Start workflow with v1 as current
+        handle = await client.start_workflow(
+            "ContinueAsNewWithVersionUpgrade",
+            0,
+            id=f"test-can-version-upgrade-{uuid.uuid4()}",
+            task_queue=w1.task_queue,
+        )
+
+        # Wait for workflow to complete one WFT on v1.0
+        await wait_for_workflow_running_on_version(handle, v1.build_id)
+
+        # Wait for version 2.0 to be ready
+        await wait_until_worker_deployment_visible(client, v2)
+
+        # Set version 2.0 as current
+        await set_current_deployment_version(client, resp2.conflict_token, v2)
+
+        # Wait for v2.0-as-Current routing config to be propagated
+        await wait_for_worker_deployment_routing_config_propagation(
+            client, deployment_name, v2.build_id
+        )
+
+        # Expect workflow to return "v2.0", indicating that it continued-as-new and completed on v2
+        result = await handle.result()
+        assert result == "v2.0"