I've been thinking about this PR. Here's an example with the stimulus_id on the right after the # comment.

1. Client._send_to_scheduler({"op": "update-graph-hlg", "stimulus_id": "CHLG"}) # CHLG
2. Scheduler.update_graph_hlg # CHLG
3. Scheduler..update_graph # CHLG
4. Scheduler.worker_send({"op": "compute-task"}) # CHLG
5. Worker.handle_compute_task # CHLG
6. Worker.execute # CHLG
7. Worker._handle_instruction # "task-finished"
8. Worker.batched_stream({"op": "task-finished"}) # "task-finished"
9. Scheduler.handle_task_finished # "task-finished"

I've been thinking about this in a Cluster vs a LocalCluster approach.

1. Within a Distributed Cluster, the STIMULUS_ID ContextVar will be unset and must derived from a message from an external entity (Client, Scheduler, Worker).
2. Within a LocalCluster (Client, Scheduler, Worker in the same process), the STIMULUS_ID may already have been set.

I am considering how to write code to handle both cases and this is complicated by examples such as update_graph_hlg calling update_graph.

I also wonder if a Distributed ContextVar might be possible with the simple use of STIMULUS_ID.set(stimulus_id) at various handler boundaries and copy_context().run at thread boundaries.

I think there's a lot of complexity here (async, threads and inter-process communication) that I'm trying to get my head around. To improve my understanding I would like to model this with some minimal distributed.core.Server's in a test case.

A further thought on managing this complexity would be to add some kwargs to default_stimulus_id such as override (always override existing STIMULUS_ID) and require_empty (require STIMULUS_ID to not be present prior to setting it)