fix: Pass headers into TiledWriter

Author: DiamondJoseph

src/blueapi/core/context.py

@@ -7,6 +7,7 @@
 from types import ModuleType, NoneType, UnionType
 from typing import Any, Generic, TypeVar, Union, get_args, get_origin, get_type_hints
 
+from bluesky.callbacks.tiled_writer import TiledWriter
 from bluesky.protocols import HasName
 from bluesky.run_engine import RunEngine
 from dodal.common.beamlines.beamline_utils import get_path_provider, set_path_provider
@@ -16,6 +17,8 @@
 from pydantic.fields import FieldInfo
 from pydantic.json_schema import JsonSchemaValue, SkipJsonSchema
 from pydantic_core import CoreSchema, core_schema
+from tiled.client import from_uri
+from tiled.client.base import BaseClient
 
 from blueapi import utils
 from blueapi.client.numtracker import NumtrackerClient
@@ -108,6 +111,7 @@ class BlueskyContext:
     run_engine: RunEngine = field(
         default_factory=lambda: RunEngine(context_managers=[])
     )
+    tiled_client: BaseClient | None = field(default=None, init=False, repr=False)
     numtracker: NumtrackerClient | None = field(default=None, init=False, repr=False)
     plans: dict[str, Plan] = field(default_factory=dict)
     devices: dict[str, Device] = field(default_factory=dict)
@@ -157,6 +161,12 @@ def _update_scan_num(md: dict[str, Any]) -> int:
                 "the devices. Remove this path provider to use numtracker."
             )
 
+        if (tiled_conf := configuration.tiled) is not None and tiled_conf.enabled:
+            self.tiled_client = client = from_uri(
+                str(tiled_conf.url), api_key=tiled_conf.api_key
+            )
+            self.run_engine.subscribe(TiledWriter(client))
+
     def find_device(self, addr: str | list[str]) -> Device | None:
         """
         Find a device in this context, allows for recursive search.

src/blueapi/service/interface.py

@@ -1,14 +1,13 @@
 from collections.abc import Mapping
 from functools import cache
+from queue import Full
 from typing import Any
 
-from bluesky.callbacks.tiled_writer import TiledWriter
 from bluesky_stomp.messaging import StompClient
 from bluesky_stomp.models import Broker, DestinationBase, MessageTopic
-from tiled.client import from_uri
 
 from blueapi.cli.scratch import get_python_environment
-from blueapi.config import ApplicationConfig, OIDCConfig, StompConfig, TiledConfig
+from blueapi.config import ApplicationConfig, OIDCConfig, StompConfig
 from blueapi.core.context import BlueskyContext
 from blueapi.core.event import EventStream
 from blueapi.log import set_up_logging
@@ -20,7 +19,7 @@
     TaskRequest,
     WorkerTask,
 )
-from blueapi.worker.event import TaskStatusEnum, WorkerState
+from blueapi.worker.event import TaskStatusEnum, WorkerEvent, WorkerState
 from blueapi.worker.task import Task
 from blueapi.worker.task_worker import TaskWorker, TrackableTask
 
@@ -87,16 +86,6 @@ def stomp_client() -> StompClient | None:
         return None
 
 
-@cache
-def tiled_writer() -> TiledWriter | None:
-    tiled_config: TiledConfig = config().tiled
-    if tiled_config.enabled:
-        client = from_uri(str(tiled_config.url), api_key=tiled_config.api_key)
-        return TiledWriter(client, batch_size=1)
-    else:
-        return None
-
-
 def setup(config: ApplicationConfig) -> None:
     """Creates and starts a worker with supplied config"""
     set_config(config)
@@ -105,8 +94,6 @@ def setup(config: ApplicationConfig) -> None:
     # Eagerly initialize worker and messaging connection
     worker()
     stomp_client()
-    if writer := tiled_writer():
-        context().run_engine.subscribe(writer)
 
 
 def teardown() -> None:
@@ -116,7 +103,6 @@ def teardown() -> None:
     context.cache_clear()
     worker.cache_clear()
     stomp_client.cache_clear()
-    tiled_writer.cache_clear()
 
 
 def _publish_event_streams(
@@ -161,7 +147,10 @@ def submit_task(task_request: TaskRequest) -> str:
     task = Task(
         name=task_request.name,
         params=task_request.params,
-        metadata={"instrument_session": task_request.instrument_session},
+        metadata={
+            "instrument_session": task_request.instrument_session,
+            "tiled_access_tags": [task_request.instrument_session],
+        },
     )
     return worker().submit_task(task)
 
@@ -175,8 +164,38 @@ def begin_task(
     task: WorkerTask, pass_through_headers: Mapping[str, str] | None = None
 ) -> WorkerTask:
     """Trigger a task. Will fail if the worker is busy"""
+    if worker().get_active_task() is not None:
+        raise Full()
     if nt := context().numtracker:
         nt.set_headers(pass_through_headers or {})
+
+        def unset_headers_when_task_finished(
+            event: WorkerEvent, correlation_id: str | None
+        ) -> None:
+            if (
+                event.task_status
+                and event.task_status.task_id == task.task_id
+                and event.task_status.task_complete
+            ):
+                nt.set_headers({})
+
+        worker().worker_events.subscribe(unset_headers_when_task_finished)
+    if tiled_client := context().tiled_client:
+        tiled_client.context.http_client.headers.update(pass_through_headers or {})
+
+        def unset_headers_when_task_finished(
+            event: WorkerEvent, correlation_id: str | None
+        ) -> None:
+            if (
+                event.task_status
+                and event.task_status.task_id == task.task_id
+                and event.task_status.task_complete
+            ):
+                for header in pass_through_headers or {}:
+                    del tiled_client.context.http_client.headers[header]
+
+        worker().worker_events.subscribe(unset_headers_when_task_finished)
+
     if task.task_id is not None:
         worker().begin_task(task.task_id)
     return task

tests/unit_tests/core/test_context.py

@@ -30,7 +30,14 @@
 from pydantic.json_schema import SkipJsonSchema
 from pytest import LogCaptureFixture
 
-from blueapi.config import EnvironmentConfig, MetadataConfig, Source, SourceKind
+from blueapi.config import (
+    ApplicationConfig,
+    EnvironmentConfig,
+    MetadataConfig,
+    Source,
+    SourceKind,
+    TiledConfig,
+)
 from blueapi.core import BlueskyContext, is_bluesky_compatible_device
 from blueapi.core.context import DefaultFactory, generic_bounds, qualified_name
 from blueapi.utils.connect_devices import _establish_device_connections
@@ -693,3 +700,30 @@ def demo_plan(foo: int | None) -> MsgGenerator:
         "foo": {"title": "Foo", "anyOf": [{"type": "integer"}, {"type": "null"}]}
     }
     assert "foo" in schema.get("required", [])
+
+
+def test_setup_without_tiled_not_makes_tiled_inserter():
+    with patch("blueapi.core.context.from_uri") as from_uri:
+        BlueskyContext(ApplicationConfig())
+
+        assert from_uri.call_count == 0
+
+
+def test_setup_with_tiled_makes_tiled_inserter():
+    with patch("blueapi.core.context.from_uri") as from_uri:
+        BlueskyContext(ApplicationConfig(tiled=TiledConfig(enabled=True)))
+
+        assert from_uri.call_count == 1
+        assert from_uri.call_args.args == ("http://localhost:8407/",)
+        assert from_uri.call_args.kwargs == {"api_key": None}
+
+
+def test_setup_with_tiled_api_key_makes_tiled_inserter():
+    with patch("blueapi.core.context.from_uri") as from_uri:
+        BlueskyContext(
+            ApplicationConfig(tiled=TiledConfig(enabled=True, api_key="foobarbaz"))
+        )
+
+        assert from_uri.call_count == 1
+        assert from_uri.call_args.args == ("http://localhost:8407/",)
+        assert from_uri.call_args.kwargs == {"api_key": "foobarbaz"}

tests/unit_tests/service/test_interface.py

@@ -26,7 +26,6 @@
     Source,
     SourceKind,
     StompConfig,
-    TiledConfig,
 )
 from blueapi.core.context import BlueskyContext
 from blueapi.service import interface
@@ -510,34 +509,6 @@ def test_setup_with_numtracker_makes_start_document_provider():
     clear_path_provider()
 
 
-def test_setup_without_tiled_not_makes_tiled_inserter():
-    with patch("blueapi.service.interface.from_uri") as from_uri:
-        conf = ApplicationConfig()
-        interface.setup(conf)
-
-        assert from_uri.call_count == 0
-
-
-def test_setup_with_tiled_makes_tiled_inserter():
-    with patch("blueapi.service.interface.from_uri") as from_uri:
-        conf = ApplicationConfig(tiled=TiledConfig(enabled=True))
-        interface.setup(conf)
-
-        assert from_uri.call_count == 1
-        assert from_uri.call_args.args == ("http://localhost:8407/",)
-        assert from_uri.call_args.kwargs == {"api_key": None}
-
-
-def test_setup_with_tiled_api_key_makes_tiled_inserter():
-    with patch("blueapi.service.interface.from_uri") as from_uri:
-        conf = ApplicationConfig(tiled=TiledConfig(enabled=True, api_key="foobarbaz"))
-        interface.setup(conf)
-
-        assert from_uri.call_count == 1
-        assert from_uri.call_args.args == ("http://localhost:8407/",)
-        assert from_uri.call_args.kwargs == {"api_key": "foobarbaz"}
-
-
 def test_setup_with_numtracker_raises_if_provider_is_defined_in_device_module():
     conf = ApplicationConfig(
         env=EnvironmentConfig(