[Bug] No clean way to cancel an activity and wait until it's cancelled

What are you really trying to do?

Start an activity from a workflow, cancel it, wait for clean cancellation acknowledgement, exit the workflow.

Describe the bug

There appear to be several issues.

1. If you call handle.cancel() on an activity handle from the workflow, the CancelledError will only be raised in the activity after an activity.heartbeat() call, followed by an await something() call. Otherwise the activity itself doesn't seem to receive CancelledError. This kinda sucks -- both that you need to call heartbeat() for cancellation requests to get through, and the fact that long-running coros do not get interrupted correctly. I don't understand the underlying implementation, but if the reason we need to heartbeat() is that workers don't want to poll for cancellation, still, when you call activity.heartbeat(), the client library could check for cancellation and immediately call task.cancel() on the task running the activity, no? Currently the work-around seems to be to litter activity code with asyncio.sleep(0.1).
2. From the workflow code, if you want to cleanly wait for the activity cancellation, you have to catch ActivityError, which is a broader exception type than I actually want to catch (presumably there are other ActivityErrors which I don't want to catch).
3. If you use WAIT_CANCELLATION_COMPLETED, if you try to cancel an activity that's currently awaiting a long-running coro (like asyncio.sleep(10)), the activity won't receive a CancelledError until it sends a heartbeat, which it can't do until the coro it's waiting on finishes, so your workflow won't finish until the activity's long-running coro finishes.

Minimal Reproduction

https://github.com/andmis/snippets/tree/temporal-python-sdk-cannot-cleanly-cancel-activities

Using python run_workflow.py with no args, the workflow exits promptly, but the activity never receives CancelledError, completes, and we get warning log spam:

2024-12-09 19:44:06.116611+00:00 (Workflow) Starting, use_sleep=False use_wait=False
2024-12-09 11:44:06.124587 (Activity) Running sandbox_activity
2024-12-09 19:44:07.130422+00:00 (Workflow) Cancelling activity
2024-12-09 19:44:07.130422+00:00 (Workflow) Activity cancelled
2024-12-09 19:44:07.130422+00:00 (Workflow) Exiting
2024-12-09 11:44:12.868185 (Activity) Completing sandbox_activity
2024-12-09T19:44:12.881460Z  WARN temporal_sdk_core::worker::activities: Activity not found on completion. This may happen if the activity has already been cancelled but completed anyway. task_token=TaskToken(CiRhOTc2ZWZkMy1iM2NiLTQwZmMtOWYzZi1jNjk0MzU2NjMzN2ESEHNhbmRib3gtd29ya2Zsb3caJDVhMjAyNWI3LWRhY2MtNDg4OC04NjA5LWI3NGQ3MzI2MzA3YyAFKAEyATFCEHNhbmRib3hfYWN0aXZpdHlKCAgBEIeKQBgB) details=Status { code: NotFound, message: "workflow execution already completed", details: b"\x08\x05\x12$workflow execution already completed\x1aB\n@type.googleapis.com/temporal.api.errordetails.v1.NotFoundFailure", metadata: MetadataMap { headers: {"content-type": "application/grpc"} }, source: None }

Using python run_workflow.py -w, the workflow waits until the activity's long-running sleep finishes, despite the activity being cancelled (note timestamps), and the activity completes rather than being cancelled:

2024-12-09 19:45:15.119427+00:00 (Workflow) Starting, use_sleep=False use_wait=True
2024-12-09 11:45:15.127826 (Activity) Running sandbox_activity
2024-12-09 19:45:16.130733+00:00 (Workflow) Cancelling activity
2024-12-09 11:45:24.371849 (Activity) Completing sandbox_activity
2024-12-09 19:45:24.384356+00:00 (Workflow) Exiting

Using python run_workflow.py -w -s, the workflow waits for the long-running coro in the activity to finish, which is bad, and the activity does cancel rather than completing, but -s sucks:

2024-12-09 19:47:00.283375+00:00 (Workflow) Starting, use_sleep=True use_wait=True
2024-12-09 11:47:00.291593 (Activity) Running sandbox_activity
2024-12-09 19:47:01.301447+00:00 (Workflow) Cancelling activity
2024-12-09 11:47:06.547612 (Activity) Cancelling sandbox_activity
2024-12-09 19:47:06.551643+00:00 (Workflow) Activity cancelled
2024-12-09 19:47:06.551643+00:00 (Workflow) Exiting

Using python run_workflow.py -s results in the workflow exiting cleanly and promptly (since we aren't using WAIT_CANCELLATION_COMPLETED), and the activity cancels rather than completing, but still waits on the long-running coro:

2024-12-09 19:48:40.102363+00:00 (Workflow) Starting, use_sleep=True use_wait=False
2024-12-09 11:48:40.110083 (Activity) Running sandbox_activity
2024-12-09 19:48:41.121725+00:00 (Workflow) Cancelling activity
2024-12-09 19:48:41.121725+00:00 (Workflow) Activity cancelled
2024-12-09 19:48:41.121725+00:00 (Workflow) Exiting
2024-12-09 11:48:47.629316 (Activity) Cancelling sandbox_activity

Environment/Versions

OS and processor: macOS, M1
Temporal version: 1.1.2
Python SDK version: 1.8.0
Are you using Docker or Kubernetes or building Temporal from source? No

Inizia dalla riproduzione minima collegata, in particolare da run_workflow.py e dai relativi percorsi di annullamento dell’attività. Esegui le varianti senza argomenti, -w, -s e -w -s per confrontare la tempistica dell’annullamento, il comportamento dell’heartbeat e WAIT_CANCELLATION_COMPLETED. Traccia la gestione dell’annullamento dell’attività dell’SDK a partire da questi punti di ingresso; il lavoro è completato quando l’annullamento interrompe la coroutine di lunga durata e il workflow può attendere un annullamento pulito senza una gestione ampia delle eccezioni.