From 1bc670cc73d47fae96cbd9f78d154052ecdc2d4f Mon Sep 17 00:00:00 2001
From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com>
Date: Sun, 16 Aug 2026 00:39:16 +0100
Subject: [PATCH] feat: add blueapi_task_id to metadata

---
 src/blueapi/worker/task.py                 | 2 +-
 src/blueapi/worker/task_worker.py          | 1 +
 tests/system_tests/test_blueapi_system.py  | 5 ++++-
 tests/unit_tests/service/test_interface.py | 5 ++++-
 4 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py
index 9ce373c769..e2efade0d7 100644
--- a/src/blueapi/worker/task.py
+++ b/src/blueapi/worker/task.py
@@ -19,7 +19,7 @@ class Task(BlueapiBaseModel):
     params: Mapping[str, Any] = Field(
         description="Values for parameters to plan, if any", default_factory=dict
     )
-    metadata: Mapping[str, Any] = Field(
+    metadata: dict[str, Any] = Field(
         description="Any metadata to apply to all runs within this task",
         default_factory=dict,
     )
diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py
index 4ada3c2a54..9da569f0c9 100644
--- a/src/blueapi/worker/task_worker.py
+++ b/src/blueapi/worker/task_worker.py
@@ -273,6 +273,7 @@ def submit_task(self, task: Task) -> str:
         """
         task.prepare_params(self._ctx)  # Will raise if parameters are invalid
         task_id: str = str(uuid.uuid4())
+        task.metadata["blueapi_task_id"] = task_id
         add_span_attributes({"TaskId": task_id})
         request_id = get_baggage("correlation_id")
         # If request id is not a string, we do not pass it into a TrackableTask
diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py
index a7e4c86413..5ed296a12e 100644
--- a/tests/system_tests/test_blueapi_system.py
+++ b/tests/system_tests/test_blueapi_system.py
@@ -342,7 +342,7 @@ def test_create_task_and_delete_task_by_id(
     rest_client.clear_task(create_task.task_id)
 
 
-def test_instrument_session_propagated(
+def test_task_metadata_propagated(
     rest_client: BlueapiRestClient, small_task: TaskRequest
 ):
     response = rest_client.create_task(small_task)
@@ -353,6 +353,7 @@ def test_instrument_session_propagated(
         "tiled_access_tags": [
             '{"proposal": 12345, "visit": 1, "beamline": "adsim"}',
         ],
+        "blueapi_task_id": response.task_id,
     }
 
 
@@ -624,6 +625,8 @@ def on_event(event: AnyEvent) -> None:
     assert start_metadata["instrument_session"] == VALID_INSTRUMENT_SESSION[user]
     assert "scan_id" in start_metadata
     assert start_metadata["scan_id"] == scan_id
+    assert "blueapi_task_id" in start_metadata
+    assert start_metadata["blueapi_task_id"] == final_event.task_id
     assert "detectors" in start_metadata
     assert "det" in start_metadata["detectors"]
 
diff --git a/tests/unit_tests/service/test_interface.py b/tests/unit_tests/service/test_interface.py
index ac756fe851..30a9afcaa8 100644
--- a/tests/unit_tests/service/test_interface.py
+++ b/tests/unit_tests/service/test_interface.py
@@ -368,6 +368,7 @@ def test_get_task_by_id(
 
     expected_metadata: dict[str, Any] = {
         "instrument_session": FAKE_INSTRUMENT_SESSION,
+        "blueapi_task_id": task_id,
     }
 
     if tiled_enabled:
@@ -408,13 +409,15 @@ def test_submit_task_inserts_metadata(context_mock: MagicMock):
         metadata,
     )
 
+    expected_metadata = {**metadata, "blueapi_task_id": task_id}
+
     assert interface.get_task_by_id(task_id) == TrackableTask.model_construct(
         task_id=task_id,
         request_id=ANY,
         task=Task(
             name="my_plan",
             params={},
-            metadata=metadata,
+            metadata=expected_metadata,
         ),
         is_complete=False,
         is_pending=True,