[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

Start with the linked minimal reproduction, especially run_workflow.py and its activity cancellation paths. Run the no-argument, -w, -s, and -w -s variants to compare cancellation timing, heartbeat behavior, and WAIT_CANCELLATION_COMPLETED. Trace the SDK's activity cancellation handling from these entry points; done means cancellation interrupts the long-running coroutine and the workflow can wait for clean cancellation without broad exception handling.