[Bug] TaskGroup cancellation of mixed Temporal operations can break replay

"""An asyncio.TaskGroup that is aborted while it holds several kinds of Temporal operation: does the recorded history replay?

Needs Python 3.11+ (asyncio.TaskGroup). It starts a local dev server with WorkflowEnvironment.start_local().

    python tgreplay_repro.py [HISTORIES] [REPLAYS_PER_HISTORY]
"""

import asyncio
import sys
import uuid
from datetime import timedelta

import temporalio
from temporalio import activity, workflow
from temporalio.api.enums.v1 import EventType
from temporalio.client import Client
from temporalio.common import RetryPolicy
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Replayer, Worker


@activity.defn
async def wait_forever_activity() -> None:
    while True:
        activity.heartbeat()
        await asyncio.sleep(0.1)


@workflow.defn
class WaitingChild:
    @workflow.run
    async def run(self) -> None:
        await asyncio.Future()


@workflow.defn
class TaskGroupAbort:
    @workflow.run
    async def run(self, kinds: list[str]) -> None:
        async def fail_soon() -> None:
            await workflow.sleep(0.3)
            raise ValueError("planned abort")

        try:
            async with asyncio.TaskGroup() as group:
                for index in range(2):
                    if "timer" in kinds:
                        group.create_task(workflow.sleep(1000))
                    if "activity" in kinds:
                        group.create_task(
                            workflow.execute_activity(
                                wait_forever_activity,
                                start_to_close_timeout=timedelta(seconds=60),
                                heartbeat_timeout=timedelta(seconds=10),
                                retry_policy=RetryPolicy(maximum_attempts=1),
                            )
                        )
                    if "child" in kinds:
                        group.create_task(
                            workflow.execute_child_workflow(
                                WaitingChild.run,
                                id=f"{workflow.info().workflow_id}-child-{index}",
                            )
                        )
                group.create_task(fail_soon())
        except Exception:
            pass  # the workflow handles the group's error and carries on
        await workflow.sleep(1)


# Which command the server recorded for each cancelled operation, in history order.
CANCEL_EVENTS = {
    EventType.EVENT_TYPE_TIMER_CANCELED: "T",
    EventType.EVENT_TYPE_ACTIVITY_TASK_CANCEL_REQUESTED: "A",
    EventType.EVENT_TYPE_REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION_INITIATED: "C",
}


async def check(client: Client, kinds: list[str], histories: int, replays: int) -> None:
    task_queue = str(uuid.uuid4())
    async with Worker(
        client,
        task_queue=task_queue,
        workflows=[TaskGroupAbort, WaitingChild],
        activities=[wait_forever_activity],
    ):
        handles = [
            await client.start_workflow(
                TaskGroupAbort.run,
                kinds,
                id=str(uuid.uuid4()),
                task_queue=task_queue,
            )
            for _ in range(histories)
        ]
        await asyncio.wait_for(
            asyncio.gather(*(handle.result() for handle in handles)), 120
        )
        recorded = [await handle.fetch_history() for handle in handles]

    replayer = Replayer(workflows=[TaskGroupAbort, WaitingChild])
    failed_histories = failed_replays = 0
    first_failure = None
    for history in recorded:
        failures = []
        for _ in range(replays):
            result = await replayer.replay_workflow(
                history, raise_on_replay_failure=False
            )
            if result.replay_failure:
                failures.append(result.replay_failure)
        failed_histories += bool(failures)
        failed_replays += len(failures)
        first_failure = first_failure or (failures[0] if failures else None)

    orders = {
        "".join(CANCEL_EVENTS[e.event_type] for e in h.events if e.event_type in CANCEL_EVENTS)
        for h in recorded
    }
    print(
        f"kinds={'+'.join(kinds):<22} {histories} workflows completed; "
        f"histories failing replay: {failed_histories} of {histories}; "
        f"replays failing: {failed_replays} of {histories * replays}; "
        f"distinct recorded cancel orders: {len(orders)}"
    )
    if first_failure:
        print(f"  first failure: {type(first_failure).__name__}: {str(first_failure)[:220]}")


async def main() -> None:
    histories = int(sys.argv[1]) if len(sys.argv) > 1 else 20
    replays = int(sys.argv[2]) if len(sys.argv) > 2 else 3
    print(f"temporalio {temporalio.__version__}, python {sys.version.split()[0]}")
    async with await WorkflowEnvironment.start_local() as env:
        for kinds in (
            ["timer", "activity", "child"],  # mixed: the report
            ["timer", "activity"],
            ["timer", "child"],
            ["activity", "child"],
            ["timer"],  # single kind: controls
            ["activity"],
            ["child"],
        ):
            await check(env.client, kinds, histories, replays)


if __name__ == "__main__":
    asyncio.run(main())

Named trace (children named s0 to s5: s0 and s3 are timers, s1 and s4 activities, s2 and s5 child workflows; 3 histories, 4 replays each, with a wrapper around TaskGroup._abort that prints [t.get_name() for t in self._tasks]). In all 3 histories the recorded cancel events followed the live _abort() iteration order; all 12 replays walked a different order and all 12 failed. First history:

history 0: cancel events recorded in order CATTAC (T timer, A activity, C child)
  live _abort() iteration order:   ['s5', 's1', 's3', 's0', 's4', 's2']
  replay 0 _abort() iteration order: ['s0', 's4', 's1', 's3', 's5', 's2'] -> FAILS (nondeterminism)
  replay 1 _abort() iteration order: ['s2', 's5', 's1', 's4', 's0', 's3'] -> FAILS (nondeterminism)

Plain asyncio, no Temporal: 200 aborts of the same six child tasks, three runs per interpreter, gave 18, 14, 14 distinct cancellation orders on CPython 3.11.17; 19, 19, 18 on 3.12.3; 15, 15, 15 on 3.12.15; 7, 6, 6 on 3.13.17; and 48, 18, 18 on 3.14.9.

New processes: replaying 20 other saved histories of the three-kind group, 19 of 20 failed in each of three new processes.

In a variant that waits 8 s after the abort and does not set the activity retry policy (10 workflows per case): After restart, 9/10 result waiters remained pending after 30 seconds, with 18 recorded workflow-task failures. With caching disabled, 10/10 remained pending, with 29 failures. All single-kind control waiters finished within the observation window. The 47 failing-workflow-task log lines of that run (18 + 29) all carry force_cause: NonDeterministicError.