From f83052de6374e92ba09d8e55676f8e3cd38d88ff Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 2 Nov 2026 16:01:25 +0000
Subject: [PATCH 1/6] Remove client protocols

---
 src/blueapi/client/client.py               | 21 ++++++++++----
 src/blueapi/client/device_cache.py         |  4 +--
 src/blueapi/client/plan_cache.py           | 23 ++++++----------
 src/blueapi/client/protocols.py            | 32 ----------------------
 src/blueapi/client/rest.py                 |  6 +++-
 tests/unit_tests/client/test_client.py     | 12 ++++----
 tests/unit_tests/client/test_plan_cache.py | 22 ++++++---------
 7 files changed, 46 insertions(+), 74 deletions(-)
 delete mode 100644 src/blueapi/client/protocols.py

diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index b0258141f..4dc9f4fbc 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -1,12 +1,12 @@
 import itertools
 import logging
 import time
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from concurrent.futures import Future
 from contextlib import suppress
 from functools import cached_property
 from pathlib import Path
-from typing import Self
+from typing import Any, Self
 
 from bluesky_stomp.messaging import MessageContext, StompClient
 from bluesky_stomp.models import Broker
@@ -43,7 +43,6 @@
 from .device_cache import DeviceCache
 from .event_bus import AnyEvent, EventBusClient, OnAnyEvent
 from .plan_cache import PlanCache
-from .protocols import ClientProtocol
 from .rest import (
     BlueapiRestClient,
     BlueskyRemoteControlError,
@@ -61,7 +60,7 @@ class MissingInstrumentSessionError(Exception):
     pass
 
 
-class BlueapiClient(ClientProtocol):
+class BlueapiClient:
     """Unified client for controlling blueapi"""
 
     _rest: BlueapiRestClient
@@ -120,7 +119,7 @@ def for_host(cls, host: str | None = None) -> Self:
     @cached_property
     @start_as_current_span(TRACER)
     def plans(self) -> PlanCache:
-        return PlanCache(self, self._rest.get_plans().plans)
+        return PlanCache(self.run_plan, self._rest.get_plans().plans)
 
     @cached_property
     @start_as_current_span(TRACER)
@@ -321,6 +320,18 @@ def get_active_task(self) -> WorkerTask:
 
         return self.active_task
 
+    @start_as_current_span(TRACER, "name", "args", "params")
+    def run_plan(
+        self, name: str, args: tuple[Any, ...], kwargs: Mapping[str, Any]
+    ) -> TaskStatus:
+        req = TaskRequest(
+            name=name,
+            args=list(args),
+            params=kwargs,
+            instrument_session=self.instrument_session,
+        )
+        return self.run_task(req)
+
     @start_as_current_span(TRACER, "task")
     def run_task(
         self,
diff --git a/src/blueapi/client/device_cache.py b/src/blueapi/client/device_cache.py
index db305e397..447983d6b 100644
--- a/src/blueapi/client/device_cache.py
+++ b/src/blueapi/client/device_cache.py
@@ -1,13 +1,13 @@
+from blueapi.client.rest import BlueapiRestClient, ClientObjectRef
 from blueapi.service.model import DeviceModel
 
-from .protocols import ClientObjectRef, RestClientProtocol
 from .rest import NotFoundError
 
 
 class DeviceCache:
     """Cache and lazily resolve devices from a Blueapi server."""
 
-    def __init__(self, rest: RestClientProtocol):
+    def __init__(self, rest: BlueapiRestClient):
         self._rest = rest
         self._cache = {
             model.name: DeviceRef(name=model.name, cache=self, model=model)
diff --git a/src/blueapi/client/plan_cache.py b/src/blueapi/client/plan_cache.py
index acf5b27ee..3e3993f63 100644
--- a/src/blueapi/client/plan_cache.py
+++ b/src/blueapi/client/plan_cache.py
@@ -1,12 +1,11 @@
 import logging
+from collections.abc import Callable
 from typing import Any
 
 from blueapi.client.rest import InvalidParametersError
-from blueapi.service.model import PlanModel, TaskRequest
+from blueapi.service.model import PlanModel
 from blueapi.worker.event import TaskError, TaskResult, TaskStatus
 
-from .protocols import ClientProtocol
-
 _REPR_MAX_LENGTH = 100
 _REPR_MAX_ARGS_INLINE = 3
 _JSON_TYPE_MAP = {
@@ -19,6 +18,8 @@
 
 log = logging.getLogger(__name__)
 
+PlanRunner = Callable[[str, tuple[Any, ...], dict[str, Any]], TaskStatus]
+
 
 class PlanFailedError(Exception):
     def __init__(self, typ: str, message: str):
@@ -29,9 +30,9 @@ def __init__(self, typ: str, message: str):
 class PlanCache:
     """Collection of plans that can be accessed by name or attribute."""
 
-    def __init__(self, client: ClientProtocol, plans: list[PlanModel]):
+    def __init__(self, runner: PlanRunner, plans: list[PlanModel]):
         self._cache = {
-            model.name: Plan(name=model.name, model=model, client=client)
+            model.name: Plan(name=model.name, model=model, runner=runner)
             for model in plans
         }
         for name, plan in self._cache.items():
@@ -55,10 +56,10 @@ def __repr__(self) -> str:
 class Plan:
     """Callable client-side reference to a registered BlueAPI plan."""
 
-    def __init__(self, name: str, model: PlanModel, client: ClientProtocol):
+    def __init__(self, name: str, model: PlanModel, runner: PlanRunner):
         self.name = name
         self.model = model
-        self._client = client
+        self._runner = runner
         self.__doc__ = model.description
 
     def __call__(self, *args, **kwargs) -> Any:
@@ -75,14 +76,8 @@ def __call__(self, *args, **kwargs) -> Any:
             TypeError: If the supplied arguments do not match the plan
                 parameter schema.
         """
-        req = TaskRequest(
-            name=self.name,
-            args=list(args),
-            params=kwargs,
-            instrument_session=self._client.instrument_session,
-        )
         try:
-            match self._client.run_task(req):
+            match self._runner(self.name, args, kwargs):
                 case TaskStatus(result=TaskResult(result=res)):
                     return res
                 case TaskStatus(result=TaskError(type=typ, message=msg)):
diff --git a/src/blueapi/client/protocols.py b/src/blueapi/client/protocols.py
deleted file mode 100644
index dbb923af6..000000000
--- a/src/blueapi/client/protocols.py
+++ /dev/null
@@ -1,32 +0,0 @@
-from typing import Protocol
-
-from blueapi.service.model import DeviceModel, DeviceResponse, TaskRequest
-from blueapi.worker.event import TaskStatus
-
-from .event_bus import OnAnyEvent
-
-# This module is needed to break circular dependencies. The Device and Plan cache
-# reference the Clients. The Clients reference the Device and Plan cache.
-
-
-class ClientProtocol(Protocol):
-    def run_task(
-        self,
-        task: TaskRequest,
-        on_event: OnAnyEvent | None = None,
-        timeout: float | None = None,
-    ) -> TaskStatus: ...
-
-    @property
-    def instrument_session(self) -> str: ...
-
-
-class RestClientProtocol(Protocol):
-    def get_devices(self) -> DeviceResponse: ...
-
-    def get_device(self, name: str) -> DeviceModel: ...
-
-
-class ClientObjectRef:
-    name: str
-    model: DeviceModel
diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index f07f6b659..a070b0579 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -16,7 +16,6 @@
 from websockets.sync.client import connect
 
 from blueapi import __version__
-from blueapi.client.protocols import ClientObjectRef
 from blueapi.config import RestConfig
 from blueapi.core.bluesky_types import DataEvent
 from blueapi.service.authentication import JWTAuth, SessionManager
@@ -201,6 +200,11 @@ def _response_json(response: requests.Response) -> Any:
         ) from exc
 
 
+class ClientObjectRef:
+    name: str
+    model: DeviceModel
+
+
 class BlueapiRestClient:
     _config: RestConfig
     session_manager: SessionManager | None
diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py
index e7e31f930..a662f2aa0 100644
--- a/tests/unit_tests/client/test_client.py
+++ b/tests/unit_tests/client/test_client.py
@@ -497,8 +497,8 @@ def callback(on_event: Callable[[AnyEvent, MessageContext], None]):
 
 
 def test_scripting_interface_returns_result():
-    client = Mock(spec=BlueapiClient, instrument_session="cm12345-1")
-    client.run_task.return_value = TaskStatus(
+    runner = Mock()
+    runner.return_value = TaskStatus(
         task_id="foobar",
         task_complete=True,
         task_failed=False,
@@ -506,15 +506,15 @@ def test_scripting_interface_returns_result():
     )
     demo_plan = Plan(
         name="demo",
-        client=client,
+        runner=runner,
         model=PlanModel(name="demo", description="Demo plan", schema={}),
     )
     assert demo_plan() == 42
 
 
 def test_scripting_interface_raises_exceptions():
-    client = Mock(spec=BlueapiClient, instrument_session="cm12345-1")
-    client.run_task.return_value = TaskStatus(
+    runner = Mock()
+    runner.return_value = TaskStatus(
         task_id="foobar",
         task_complete=True,
         task_failed=True,
@@ -522,7 +522,7 @@ def test_scripting_interface_raises_exceptions():
     )
     demo_plan = Plan(
         name="demo",
-        client=client,
+        runner=runner,
         model=PlanModel(name="demo", description="Demo plan", schema={}),
     )
     with pytest.raises(PlanFailedError, match="Plan failed"):
diff --git a/tests/unit_tests/client/test_plan_cache.py b/tests/unit_tests/client/test_plan_cache.py
index 6b3624693..99eb07052 100644
--- a/tests/unit_tests/client/test_plan_cache.py
+++ b/tests/unit_tests/client/test_plan_cache.py
@@ -6,7 +6,7 @@
 from blueapi.client.client import PlanCache
 from blueapi.client.plan_cache import Plan
 from blueapi.client.rest import InvalidParametersError, ParameterError
-from blueapi.service.model import PlanModel, TaskRequest
+from blueapi.service.model import PlanModel
 
 FULL_PLAN = PlanModel(
     name="foobar",
@@ -143,22 +143,16 @@ def test_plan_empty_fallback_help_text(client):
     ],
 )
 def test_plan_param_forwarding(args, kwargs):
-    client = Mock()
-    client.instrument_session = "cm12345-1"
-    plan = Plan("foobar", FULL_PLAN, client)
+    runner = Mock()
+    plan = Plan("foobar", FULL_PLAN, runner)
 
     plan(*args, **kwargs)
-    client.run_task.assert_called_once_with(
-        TaskRequest(
-            name="foobar", instrument_session="cm12345-1", args=args, params=kwargs
-        )
-    )
+    runner.assert_called_once_with("foobar", args, kwargs)
 
 
 def test_plan_invalid_param_mapping():
-    client = Mock()
-    client.instrument_session = "cm12345-1"
-    client.run_task.side_effect = InvalidParametersError(
+    runner = Mock()
+    runner.side_effect = InvalidParametersError(
         [
             ParameterError(
                 msg="foo is not valid",
@@ -168,8 +162,8 @@ def test_plan_invalid_param_mapping():
             )
         ]
     )
-    plan = Plan("foobar", FULL_PLAN, client)
+    plan = Plan("foobar", FULL_PLAN, runner)
 
     with pytest.raises(TypeError, match="foo is not valid"):
         plan(1, 2, foo="bar")
-    client.run_task.assert_called_once()
+    runner.assert_called_once()

From 4df1f58988df3acc3efc16ceab34fd3eadc0a7a9 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 2 Nov 2026 18:44:09 +0000
Subject: [PATCH 2/6] As restored from previous PR

---
 pyproject.toml                                |   1 +
 src/blueapi/cli/cli.py                        |  13 ++
 src/blueapi/cli/stubgen.py                    | 117 ++++++++++
 src/blueapi/client/device_cache.py            |   2 +
 src/blueapi/client/plan_cache.py              |   2 +
 .../stubs/templates/cache_template.pyi.jinja  |  72 ++++++
 tests/unit_tests/cli/test_stubgen.py          | 214 ++++++++++++++++++
 uv.lock                                       |  68 +++---
 8 files changed, 456 insertions(+), 33 deletions(-)
 create mode 100644 src/blueapi/cli/stubgen.py
 create mode 100644 src/blueapi/stubs/templates/cache_template.pyi.jinja
 create mode 100644 tests/unit_tests/cli/test_stubgen.py

diff --git a/pyproject.toml b/pyproject.toml
index 5c5008e6b..384aa2665 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,7 @@ dependencies = [
     "graypy>=2.1.0",
     "httpx>=0.28.1",
     "aiohttp>=3.13.5",
+    "jinja2>=3.1.6",
 ]
 dynamic = ["version"]
 license.file = "LICENSE"
diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py
index 8424c3c4e..a4da2766b 100644
--- a/src/blueapi/cli/cli.py
+++ b/src/blueapi/cli/cli.py
@@ -239,6 +239,19 @@ def start_application(obj: dict):
     start(config)
 
 
+@main.command()
+@click.pass_obj
+@click.argument("target", type=click.Path(file_okay=False))
+def generate_stubs(obj: dict, target: Path):
+    click.echo(f"Writing stubs to {target}")
+
+    config: ApplicationConfig = obj["config"]
+    bc = BlueapiClient.from_config(config)
+    from . import stubgen
+
+    stubgen.generate_stubs(Path(target), list(bc.plans), list(bc.devices))
+
+
 @main.group()
 @click.option(
     "-o",
diff --git a/src/blueapi/cli/stubgen.py b/src/blueapi/cli/stubgen.py
new file mode 100644
index 000000000..6f6fbf4bd
--- /dev/null
+++ b/src/blueapi/cli/stubgen.py
@@ -0,0 +1,117 @@
+import logging
+from dataclasses import dataclass
+from inspect import cleandoc
+from pathlib import Path
+from textwrap import dedent
+from typing import Self, TextIO
+
+from jinja2 import Environment, PackageLoader
+
+from blueapi.client.cache import DeviceRef, Plan
+from blueapi.core import context
+from blueapi.core.bluesky_types import BLUESKY_PROTOCOLS
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class ArgSpec:
+    name: str
+    type: str
+    optional: bool
+
+
+@dataclass
+class PlanSpec:
+    name: str
+    docs: str
+    args: list[ArgSpec]
+
+    @classmethod
+    def from_plan(cls, plan: Plan) -> Self:
+        req = set(plan.required)
+        args = [
+            ArgSpec(arg, _type_string(spec), arg not in req)
+            for arg, spec in plan.model.parameter_schema.get("properties", {}).items()
+        ]
+        return cls(plan.name, plan.help_text, args)
+
+
+BLUESKY_PROTOCOL_NAMES = {context.qualified_name(proto) for proto in BLUESKY_PROTOCOLS}
+
+
+def _type_string(spec) -> str:
+    """Best effort attempt at making useful type hints for plans"""
+    match spec.get("type"):
+        case "array":
+            return f"list[{_type_string(spec.get('items'))}]"
+        case "integer":
+            return "int"
+        case "number":
+            return "float"
+        case proto if proto in BLUESKY_PROTOCOL_NAMES:
+            return "DeviceRef"
+        case "object":
+            return "dict[str, Any]"
+        case "string":
+            return "str"
+        case "boolean":
+            return "bool"
+        case None if opts := spec.get("anyOf"):
+            return " | ".join(_type_string(opt) for opt in opts)
+        case _:
+            return "Any"
+
+
+def generate_stubs(target: Path, plans: list[Plan], devices: list[DeviceRef]):
+    log.info("Generating stubs for %d plans and %d devices", len(plans), len(devices))
+    target.mkdir(parents=True, exist_ok=True)
+    client_dir = target / "src" / "blueapi-stubs" / "client"
+
+    log.debug("Making project structure: %s", client_dir)
+    client_dir.mkdir(parents=True, exist_ok=True)
+
+    stub_file = client_dir / "cache.pyi"
+    project_file = target / "pyproject.toml"
+    py_typed = target / "src" / "blueapi-stubs" / "py.typed"
+
+    log.debug("Writing pyproject.toml to %s", project_file)
+    with open(project_file, "w") as out:
+        out.write(
+            dedent("""
+                [project]
+                name = "blueapi-stubs"
+                version = "0.1.0"
+                description = "Generated client stubs for a running server"
+                readme = "README.md"
+                requires-python = ">=3.11"
+
+                dependencies = [
+                    "blueapi"
+                ]
+                """)
+        )
+
+    log.debug("Writing py.typed file to %s", py_typed)
+    with open(py_typed, "w") as out:
+        out.write("partial\n")
+
+    log.debug("Writing stub file to %s", stub_file)
+    with open(stub_file, "w") as out:
+        render_stub_file(out, plans, devices)
+
+
+def _docstring(text: str) -> str:
+    # """Convert a docstring to a format that can be inserted into the template"""
+    return cleandoc(text).replace('"""', '\\"""')
+
+
+def render_stub_file(
+    stub_file: TextIO, plan_models: list[Plan], devices: list[DeviceRef]
+):
+    plans = [PlanSpec.from_plan(p) for p in plan_models]
+
+    env = Environment(loader=PackageLoader("blueapi", package_path="stubs/templates"))
+    env.filters["docstring"] = _docstring
+    tmpl = env.get_template("cache_template.pyi.jinja")
+    stub_file.write(tmpl.render(plans=plans, devices=devices))
diff --git a/src/blueapi/client/device_cache.py b/src/blueapi/client/device_cache.py
index 447983d6b..bf5fec11b 100644
--- a/src/blueapi/client/device_cache.py
+++ b/src/blueapi/client/device_cache.py
@@ -3,6 +3,8 @@
 
 from .rest import NotFoundError
 
+# This file should be kept in sync with the type stub template in stubs/templates
+
 
 class DeviceCache:
     """Cache and lazily resolve devices from a Blueapi server."""
diff --git a/src/blueapi/client/plan_cache.py b/src/blueapi/client/plan_cache.py
index 3e3993f63..6c01a0f9e 100644
--- a/src/blueapi/client/plan_cache.py
+++ b/src/blueapi/client/plan_cache.py
@@ -18,6 +18,8 @@
 
 log = logging.getLogger(__name__)
 
+# This file should be kept in sync with the type stub template in stubs/templates
+
 PlanRunner = Callable[[str, tuple[Any, ...], dict[str, Any]], TaskStatus]
 
 
diff --git a/src/blueapi/stubs/templates/cache_template.pyi.jinja b/src/blueapi/stubs/templates/cache_template.pyi.jinja
new file mode 100644
index 000000000..b06ef3638
--- /dev/null
+++ b/src/blueapi/stubs/templates/cache_template.pyi.jinja
@@ -0,0 +1,72 @@
+from collections.abc import Callable
+from typing import Any
+from blueapi.client.rest import BlueapiRestClient
+from blueapi.service.model import DeviceModel, PlanModel
+from blueapi.worker.event import WorkerEvent
+
+{#-
+ This file is based on the cache.py file in blueapi/client/cache.py and should
+ be kept in sync with changes there.
+#}
+
+# This file is auto-generated for a live server and should not be modified directly
+
+PlanRunner = Callable[[str, dict[str, Any]], WorkerEvent]
+
+class PlanCache:
+    def __init__(self, runner: PlanRunner, plans: list[PlanModel]) -> None: ...
+    def __getitem__(self, name: str) -> Plan: ...
+    def __iter__(self):  # -> Iterator[Plan]:
+        ...
+    def __repr__(self) -> str: ...
+
+### Generated plans
+{%- for item in plans %}
+    def {{ item.name }}(self,{% for arg in item.args %}
+            {{ arg.name }}: {{ arg.type }}{% if arg.optional %} | None = None{% endif %},
+            {%- endfor %}
+    ) -> WorkerEvent:
+        """
+        {{ item.docs | docstring | indent(8) }}
+        """
+        ...
+{%- endfor %}
+### End
+
+
+class Plan:
+    model: PlanModel
+    def __init__(self, model: PlanModel, runner: PlanRunner) -> None: ...
+    def __call__(self, *args, **kwargs):  # -> None:
+        ...
+
+    @property
+    def name(self) -> str: ...
+    @property
+    def help_text(self) -> str: ...
+    @property
+    def properties(self) -> set[str]: ...
+    @property
+    def required(self) -> list[str]: ...
+    def __repr__(self) -> str: ...
+
+
+class DeviceRef(str):
+    model: DeviceModel
+    _cache: DeviceCache
+    def __new__(cls, name: str, cache: DeviceCache, model: DeviceModel): ...
+    def __getattr__(self, name) -> DeviceRef: ...
+    def __repr__(self) -> str: ...
+
+class DeviceCache:
+    def __init__(self, rest: BlueapiRestClient) -> None: ...
+    def __getitem__(self, name: str) -> DeviceRef: ...
+    def __iter__(self):  # -> Iterator[DeviceRef]:
+        ...
+    def __repr__(self) -> str: ...
+
+### Generated devices
+    {%- for item in devices %}
+    {{ item }}: DeviceRef
+    {%- endfor %}
+### End
diff --git a/tests/unit_tests/cli/test_stubgen.py b/tests/unit_tests/cli/test_stubgen.py
new file mode 100644
index 000000000..27c6ad239
--- /dev/null
+++ b/tests/unit_tests/cli/test_stubgen.py
@@ -0,0 +1,214 @@
+from io import StringIO
+from textwrap import dedent
+from types import FunctionType
+from unittest.mock import Mock
+
+import pytest
+from blueapi.client.cache import DeviceRef, Plan
+
+from blueapi.cli.stubgen import (
+    _docstring,
+    _type_string,
+    generate_stubs,
+    render_stub_file,
+)
+from blueapi.service.model import DeviceModel, PlanModel
+
+
+def single_line():
+    """Single line docstring"""
+
+
+def single_line_new_line():
+    """
+    Single line docstring
+    """
+
+
+def multi_line_inline():
+    """First line
+    Second line"""
+
+
+def multi_line_new_line():
+    """
+    First line
+    Second line
+    """
+
+
+def indented_multi_line():
+    """
+    First line
+        indented
+    """
+
+
+@pytest.mark.parametrize(
+    "input,expected",
+    [
+        (single_line, "Single line docstring"),
+        (single_line_new_line, "Single line docstring"),
+        (multi_line_inline, "First line\nSecond line"),
+        (multi_line_new_line, "First line\nSecond line"),
+        (indented_multi_line, "First line\n    indented"),
+    ],
+)
+def test_docstring_filter(input: FunctionType, expected: str):
+    assert input.__doc__
+    assert _docstring(input.__doc__) == expected
+
+
+@pytest.mark.parametrize(
+    "typ,expected",
+    [
+        ({"type": "string"}, "str"),
+        ({"type": "number"}, "float"),
+        ({"type": "integer"}, "int"),
+        ({"type": "object"}, "dict[str, Any]"),
+        ({"type": "boolean"}, "bool"),
+        ({"type": "array", "items": {"type": "integer"}}, "list[int]"),
+        ({"type": "array", "items": {"type": "object"}}, "list[dict[str, Any]]"),
+        (
+            {
+                "type": "array",
+                "items": {"anyOf": [{"type": "integer"}, {"type": "boolean"}]},
+            },
+            "list[int | bool]",
+        ),
+        ({"anyOf": [{"type": "object"}, {"type": "string"}]}, "dict[str, Any] | str"),
+        ({"type": "unknown.other.Type"}, "Any"),
+        # Special case the bluesky protocols to require device references
+        ({"type": "bluesky.protocols.Readable"}, "DeviceRef"),
+        ({}, "Any"),
+    ],
+    ids=lambda param: param.get("type") if isinstance(param, dict) else param,
+)
+def test_type_string(typ: dict, expected: str):
+    assert _type_string(typ) == expected
+
+
+def test_render_empty():
+    output = StringIO()
+
+    render_stub_file(output, [], [])
+    plan_text, device_text = _extract_rendered(output)
+
+    assert plan_text == ""
+    assert device_text == ""
+
+
+FOO = PlanModel(name="empty", description="Doc string for empty", schema={})
+
+BAR = PlanModel(
+    name="two_args",
+    description="Doc string for two_args",
+    schema={
+        "properties": {
+            "one": {"type": "integer"},
+            "two": {"type": "string"},
+        },
+        "required": ["one"],
+    },
+)
+
+
+def test_render_empty_plan_function():
+    output = StringIO()
+    plans = [Plan(model=FOO, runner=Mock())]
+    render_stub_file(output, plans, [])
+    plan_text, device_text = _extract_rendered(output)
+
+    assert device_text == ""
+
+    assert (
+        plan_text
+        == """\
+    def empty(self,
+    ) -> WorkerEvent:
+        \"""
+        Doc string for empty
+        \"""
+        ...\n"""
+    )
+
+
+def test_render_multiple_plan_functions():
+    output = StringIO()
+    runner = Mock()
+    plans = [Plan(FOO, runner), Plan(BAR, runner)]
+    render_stub_file(output, plans, [])
+    plan_text, device_text = _extract_rendered(output)
+    assert device_text == ""
+
+    assert (
+        plan_text
+        == """\
+    def empty(self,
+    ) -> WorkerEvent:
+        \"""
+        Doc string for empty
+        \"""
+        ...
+    def two_args(self,
+            one: int,
+            two: str | None = None,
+    ) -> WorkerEvent:
+        \"""
+        Doc string for two_args
+        \"""
+        ...\n"""
+    )
+
+
+def test_device_fields():
+    output = StringIO()
+    cache = Mock()
+    devices = [
+        DeviceRef("one", cache, DeviceModel(name="one", protocols=[])),
+        DeviceRef("two", cache, DeviceModel(name="two", protocols=[])),
+    ]
+    render_stub_file(output, [], devices)
+
+    plan_text, device_text = _extract_rendered(output)
+    assert plan_text == ""
+    assert device_text == "    one: DeviceRef\n    two: DeviceRef\n"
+
+
+def test_package_creation(tmp_path):
+    generate_stubs(tmp_path / "blueapi-stubs", [], [])
+    with open(tmp_path / "blueapi-stubs" / "pyproject.toml") as pyproj:
+        assert pyproj.read().startswith(
+            dedent("""
+            [project]
+            name = "blueapi-stubs"
+            version = "0.1.0"
+            """)
+        )
+    with open(
+        tmp_path / "blueapi-stubs" / "src" / "blueapi-stubs" / "py.typed"
+    ) as typed:
+        assert typed.read() == "partial\n"
+
+    assert (
+        tmp_path / "blueapi-stubs" / "src" / "blueapi-stubs" / "client" / "cache.pyi"
+    ).exists()
+
+
+def _extract_rendered(src: StringIO) -> tuple[str, str]:
+    src.seek(0)
+    _read_until_line(src, "### Generated plans")
+    plan_text = _read_until_line(src, "### End")
+    _read_until_line(src, "### Generated devices")
+    device_text = _read_until_line(src, "### End")
+    return plan_text, device_text
+
+
+def _read_until_line(src: StringIO, match: str) -> str:
+    text = ""
+    for line in src:
+        if line.startswith(match):
+            break
+        text += line
+
+    return text
diff --git a/uv.lock b/uv.lock
index 62a751355..96000e90e 100644
--- a/uv.lock
+++ b/uv.lock
@@ -3,16 +3,16 @@ revision = 5
 requires-python = ">=3.11"
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
-    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
+    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.12.*' and sys_platform != 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
-    "python_full_version < '3.11.5' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
-    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 
@@ -491,6 +491,7 @@ dependencies = [
     { name = "gitpython" },
     { name = "graypy" },
     { name = "httpx" },
+    { name = "jinja2" },
     { name = "observability-utils" },
     { name = "opentelemetry-distro" },
     { name = "opentelemetry-instrumentation-fastapi" },
@@ -560,6 +561,7 @@ requires-dist = [
     { name = "gitpython", specifier = ">=3.1.58" },
     { name = "graypy", specifier = ">=2.1.0" },
     { name = "httpx", specifier = ">=0.28.1" },
+    { name = "jinja2", specifier = ">=3.1.6" },
     { name = "observability-utils", specifier = ">=0.1.4" },
     { name = "opentelemetry-distro", specifier = ">=0.48b0" },
     { name = "opentelemetry-instrumentation-fastapi", specifier = ">=0.48b0" },
@@ -1082,10 +1084,10 @@ version = "1.3.3"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
-    "python_full_version < '3.11.5' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
-    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 dependencies = [
@@ -1172,10 +1174,10 @@ version = "1.4.0"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
-    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
+    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.12.*' and sys_platform != 'emscripten'",
 ]
 dependencies = [
@@ -3713,10 +3715,10 @@ version = "3.6.1"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
-    "python_full_version < '3.11.5' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
-    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" }
@@ -3730,10 +3732,10 @@ version = "3.7"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
-    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
+    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.12.*' and sys_platform != 'emscripten'",
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/dc/76/3af777226b63a5e64a6b36b1ec5855c14e2b94a37096d4760e595fc43511/networkx-3.7.tar.gz", hash = "sha256:fd77a511bd90f39f3d016351345b52cf5319b813bdca01de3f755d3cca62e96a", size = 1866482, upload-time = "2026-09-21T16:45:16.974Z" }
@@ -3862,10 +3864,10 @@ version = "2.4.6"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
-    "python_full_version < '3.11.5' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
-    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" }
@@ -3949,10 +3951,10 @@ version = "2.5.4"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
-    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
+    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.12.*' and sys_platform != 'emscripten'",
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/95/b0/c7453d0b6e2073c3264468b106ee1563750cecc910965e67357e3698c83e/numpy-2.5.4.tar.gz", hash = "sha256:9a94cf751c9ad8ebaa835bcd3d40dacf8534ad086b88c38029b65123c7999d2a", size = 20866315, upload-time = "2026-10-10T20:05:31.422Z" }
@@ -4642,10 +4644,10 @@ version = "0.25.3"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
-    "python_full_version < '3.11.5' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
-    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 dependencies = [
@@ -4665,10 +4667,10 @@ version = "0.26.2"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
-    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
+    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.12.*' and sys_platform != 'emscripten'",
 ]
 dependencies = [
@@ -5259,10 +5261,10 @@ version = "13.0.1"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
-    "python_full_version < '3.11.5' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
-    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 dependencies = [
@@ -5279,10 +5281,10 @@ version = "14.0.0"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
-    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
+    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.12.*' and sys_platform != 'emscripten'",
 ]
 dependencies = [
@@ -6190,10 +6192,10 @@ version = "3.0.4"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
-    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
+    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.12.*' and sys_platform != 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
@@ -6224,10 +6226,10 @@ version = "9.0.4"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
-    "python_full_version < '3.11.5' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
-    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 dependencies = [
@@ -6292,8 +6294,8 @@ version = "9.2.0"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
 ]
 dependencies = [
@@ -7160,10 +7162,10 @@ version = "3.3"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'",
-    "python_full_version < '3.11.5' and sys_platform == 'win32'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'",
-    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'win32'",
+    "python_full_version < '3.11.5' and sys_platform == 'emscripten'",
     "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 dependencies = [
@@ -7186,10 +7188,10 @@ version = "3.7"
 source = { registry = "https://pypi.org/simple" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'emscripten'",
-    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
-    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten'",
+    "python_full_version == '3.13.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.13.*' and sys_platform != 'emscripten'",
+    "python_full_version == '3.12.*' and sys_platform == 'emscripten'",
     "python_full_version == '3.12.*' and sys_platform != 'emscripten'",
 ]
 dependencies = [

From 2776a3a66497c68542c25447df5b7a6f828715d4 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 2 Nov 2026 19:18:54 +0000
Subject: [PATCH 3/6] Update stubgen for changes to main

---
 src/blueapi/cli/stubgen.py                    | 27 +++---
 .../templates/device_cache_template.pyi.jinja | 32 ++++++++
 ...yi.jinja => plan_cache_template.pyi.jinja} | 36 ++------
 tests/unit_tests/cli/test_stubgen.py          | 82 ++++++++++---------
 4 files changed, 96 insertions(+), 81 deletions(-)
 create mode 100644 src/blueapi/stubs/templates/device_cache_template.pyi.jinja
 rename src/blueapi/stubs/templates/{cache_template.pyi.jinja => plan_cache_template.pyi.jinja} (54%)

diff --git a/src/blueapi/cli/stubgen.py b/src/blueapi/cli/stubgen.py
index 6f6fbf4bd..2f3a7926e 100644
--- a/src/blueapi/cli/stubgen.py
+++ b/src/blueapi/cli/stubgen.py
@@ -3,11 +3,12 @@
 from inspect import cleandoc
 from pathlib import Path
 from textwrap import dedent
-from typing import Self, TextIO
+from typing import Any, Self, TextIO
 
 from jinja2 import Environment, PackageLoader
 
-from blueapi.client.cache import DeviceRef, Plan
+from blueapi.client.device_cache import DeviceRef
+from blueapi.client.plan_cache import Plan
 from blueapi.core import context
 from blueapi.core.bluesky_types import BLUESKY_PROTOCOLS
 
@@ -71,7 +72,6 @@ def generate_stubs(target: Path, plans: list[Plan], devices: list[DeviceRef]):
     log.debug("Making project structure: %s", client_dir)
     client_dir.mkdir(parents=True, exist_ok=True)
 
-    stub_file = client_dir / "cache.pyi"
     project_file = target / "pyproject.toml"
     py_typed = target / "src" / "blueapi-stubs" / "py.typed"
 
@@ -96,9 +96,14 @@ def generate_stubs(target: Path, plans: list[Plan], devices: list[DeviceRef]):
     with open(py_typed, "w") as out:
         out.write("partial\n")
 
-    log.debug("Writing stub file to %s", stub_file)
-    with open(stub_file, "w") as out:
-        render_stub_file(out, plans, devices)
+    log.debug("Writing stub files to %s", client_dir)
+
+    plan_model = [PlanSpec.from_plan(p) for p in plans]
+    with open(client_dir / "plan_cache.pyi", "w") as out:
+        render_stub_file(out, "plan_cache_template.pyi.jinja", plans=plan_model)
+
+    with open(client_dir / "device_cache.pyi", "w") as out:
+        render_stub_file(out, "device_cache_template.pyi.jinja", devices=devices)
 
 
 def _docstring(text: str) -> str:
@@ -106,12 +111,8 @@ def _docstring(text: str) -> str:
     return cleandoc(text).replace('"""', '\\"""')
 
 
-def render_stub_file(
-    stub_file: TextIO, plan_models: list[Plan], devices: list[DeviceRef]
-):
-    plans = [PlanSpec.from_plan(p) for p in plan_models]
-
+def render_stub_file(stub_file: TextIO, template: str, **model: Any):
     env = Environment(loader=PackageLoader("blueapi", package_path="stubs/templates"))
     env.filters["docstring"] = _docstring
-    tmpl = env.get_template("cache_template.pyi.jinja")
-    stub_file.write(tmpl.render(plans=plans, devices=devices))
+    tmpl = env.get_template(template)
+    stub_file.write(tmpl.render(**model))
diff --git a/src/blueapi/stubs/templates/device_cache_template.pyi.jinja b/src/blueapi/stubs/templates/device_cache_template.pyi.jinja
new file mode 100644
index 000000000..e2341277f
--- /dev/null
+++ b/src/blueapi/stubs/templates/device_cache_template.pyi.jinja
@@ -0,0 +1,32 @@
+from collections.abc import Callable
+from typing import Any
+from blueapi.client.rest import BlueapiRestClient
+from blueapi.service.model import DeviceModel, PlanModel
+from blueapi.worker.event import WorkerEvent
+
+{#-
+ This file is based on the devices_cache.py file in
+ blueapi/client/device_cache.py and should be kept in sync with changes there.
+#}
+
+# This file is auto-generated for a live server and should not be modified directly
+
+class DeviceRef(str):
+    model: DeviceModel
+    _cache: DeviceCache
+    def __new__(cls, name: str, cache: DeviceCache, model: DeviceModel): ...
+    def __getattr__(self, name) -> DeviceRef: ...
+    def __repr__(self) -> str: ...
+
+class DeviceCache:
+    def __init__(self, rest: BlueapiRestClient) -> None: ...
+    def __getitem__(self, name: str) -> DeviceRef: ...
+    def __iter__(self):  # -> Iterator[DeviceRef]:
+        ...
+    def __repr__(self) -> str: ...
+
+### Generated
+    {%- for item in devices %}
+    {{ item.name }}: DeviceRef
+    {%- endfor %}
+### End
diff --git a/src/blueapi/stubs/templates/cache_template.pyi.jinja b/src/blueapi/stubs/templates/plan_cache_template.pyi.jinja
similarity index 54%
rename from src/blueapi/stubs/templates/cache_template.pyi.jinja
rename to src/blueapi/stubs/templates/plan_cache_template.pyi.jinja
index b06ef3638..1f31ea255 100644
--- a/src/blueapi/stubs/templates/cache_template.pyi.jinja
+++ b/src/blueapi/stubs/templates/plan_cache_template.pyi.jinja
@@ -1,17 +1,16 @@
 from collections.abc import Callable
 from typing import Any
-from blueapi.client.rest import BlueapiRestClient
-from blueapi.service.model import DeviceModel, PlanModel
-from blueapi.worker.event import WorkerEvent
+from blueapi.service.model import PlanModel
+from blueapi.worker.event import TaskStatus
 
 {#-
- This file is based on the cache.py file in blueapi/client/cache.py and should
- be kept in sync with changes there.
+ This file is based on the plan_cache.py file in blueapi/client/plan_cache.py
+ and should be kept in sync with changes there.
 #}
 
 # This file is auto-generated for a live server and should not be modified directly
 
-PlanRunner = Callable[[str, dict[str, Any]], WorkerEvent]
+PlanRunner = Callable[[str, dict[str, Any]], TaskStatus]
 
 class PlanCache:
     def __init__(self, runner: PlanRunner, plans: list[PlanModel]) -> None: ...
@@ -20,12 +19,12 @@ class PlanCache:
         ...
     def __repr__(self) -> str: ...
 
-### Generated plans
+### Generated
 {%- for item in plans %}
     def {{ item.name }}(self,{% for arg in item.args %}
             {{ arg.name }}: {{ arg.type }}{% if arg.optional %} | None = None{% endif %},
             {%- endfor %}
-    ) -> WorkerEvent:
+    ) -> TaskStatus:
         """
         {{ item.docs | docstring | indent(8) }}
         """
@@ -49,24 +48,3 @@ class Plan:
     @property
     def required(self) -> list[str]: ...
     def __repr__(self) -> str: ...
-
-
-class DeviceRef(str):
-    model: DeviceModel
-    _cache: DeviceCache
-    def __new__(cls, name: str, cache: DeviceCache, model: DeviceModel): ...
-    def __getattr__(self, name) -> DeviceRef: ...
-    def __repr__(self) -> str: ...
-
-class DeviceCache:
-    def __init__(self, rest: BlueapiRestClient) -> None: ...
-    def __getitem__(self, name: str) -> DeviceRef: ...
-    def __iter__(self):  # -> Iterator[DeviceRef]:
-        ...
-    def __repr__(self) -> str: ...
-
-### Generated devices
-    {%- for item in devices %}
-    {{ item }}: DeviceRef
-    {%- endfor %}
-### End
diff --git a/tests/unit_tests/cli/test_stubgen.py b/tests/unit_tests/cli/test_stubgen.py
index 27c6ad239..d525839e0 100644
--- a/tests/unit_tests/cli/test_stubgen.py
+++ b/tests/unit_tests/cli/test_stubgen.py
@@ -1,17 +1,20 @@
 from io import StringIO
 from textwrap import dedent
 from types import FunctionType
+from typing import Any
 from unittest.mock import Mock
 
 import pytest
-from blueapi.client.cache import DeviceRef, Plan
 
 from blueapi.cli.stubgen import (
+    PlanSpec,
     _docstring,
     _type_string,
     generate_stubs,
     render_stub_file,
 )
+from blueapi.client.device_cache import DeviceRef
+from blueapi.client.plan_cache import Plan
 from blueapi.service.model import DeviceModel, PlanModel
 
 
@@ -88,14 +91,18 @@ def test_type_string(typ: dict, expected: str):
     assert _type_string(typ) == expected
 
 
-def test_render_empty():
+@pytest.mark.parametrize(
+    "template,model",
+    [
+        ("device_cache_template.pyi.jinja", {"devices": []}),
+        ("plan_cache_template.pyi.jinja", {"plans": []}),
+    ],
+)
+def test_render_empty(template: str, model: dict[str, Any]):
     output = StringIO()
 
-    render_stub_file(output, [], [])
-    plan_text, device_text = _extract_rendered(output)
-
-    assert plan_text == ""
-    assert device_text == ""
+    render_stub_file(output, template, **model)
+    assert _extract_rendered(output) == ""
 
 
 FOO = PlanModel(name="empty", description="Doc string for empty", schema={})
@@ -115,17 +122,16 @@ def test_render_empty():
 
 def test_render_empty_plan_function():
     output = StringIO()
-    plans = [Plan(model=FOO, runner=Mock())]
-    render_stub_file(output, plans, [])
-    plan_text, device_text = _extract_rendered(output)
-
-    assert device_text == ""
+    plans = [Plan(name="empty", model=FOO, runner=Mock())]
+    plan_model = [PlanSpec.from_plan(p) for p in plans]
+    render_stub_file(output, "plan_cache_template.pyi.jinja", plans=plan_model)
+    rendered = _extract_rendered(output)
 
     assert (
-        plan_text
+        rendered
         == """\
     def empty(self,
-    ) -> WorkerEvent:
+    ) -> TaskStatus:
         \"""
         Doc string for empty
         \"""
@@ -136,16 +142,16 @@ def empty(self,
 def test_render_multiple_plan_functions():
     output = StringIO()
     runner = Mock()
-    plans = [Plan(FOO, runner), Plan(BAR, runner)]
-    render_stub_file(output, plans, [])
-    plan_text, device_text = _extract_rendered(output)
-    assert device_text == ""
+    plans = [Plan("empty", FOO, runner), Plan("two_args", BAR, runner)]
+    plan_model = [PlanSpec.from_plan(p) for p in plans]
+    render_stub_file(output, "plan_cache_template.pyi.jinja", plans=plan_model)
+    rendered = _extract_rendered(output)
 
     assert (
-        plan_text
+        rendered
         == """\
     def empty(self,
-    ) -> WorkerEvent:
+    ) -> TaskStatus:
         \"""
         Doc string for empty
         \"""
@@ -153,7 +159,7 @@ def empty(self,
     def two_args(self,
             one: int,
             two: str | None = None,
-    ) -> WorkerEvent:
+    ) -> TaskStatus:
         \"""
         Doc string for two_args
         \"""
@@ -168,16 +174,17 @@ def test_device_fields():
         DeviceRef("one", cache, DeviceModel(name="one", protocols=[])),
         DeviceRef("two", cache, DeviceModel(name="two", protocols=[])),
     ]
-    render_stub_file(output, [], devices)
+    render_stub_file(output, "device_cache_template.pyi.jinja", devices=devices)
 
-    plan_text, device_text = _extract_rendered(output)
-    assert plan_text == ""
-    assert device_text == "    one: DeviceRef\n    two: DeviceRef\n"
+    assert _extract_rendered(output) == "    one: DeviceRef\n    two: DeviceRef\n"
 
 
 def test_package_creation(tmp_path):
-    generate_stubs(tmp_path / "blueapi-stubs", [], [])
-    with open(tmp_path / "blueapi-stubs" / "pyproject.toml") as pyproj:
+
+    stub_root = tmp_path / "blueapi-stubs"
+
+    generate_stubs(stub_root, [], [])
+    with open(stub_root / "pyproject.toml") as pyproj:
         assert pyproj.read().startswith(
             dedent("""
             [project]
@@ -185,23 +192,20 @@ def test_package_creation(tmp_path):
             version = "0.1.0"
             """)
         )
-    with open(
-        tmp_path / "blueapi-stubs" / "src" / "blueapi-stubs" / "py.typed"
-    ) as typed:
+
+    stub_src = stub_root / "src" / "blueapi-stubs"
+
+    with open(stub_src / "py.typed") as typed:
         assert typed.read() == "partial\n"
 
-    assert (
-        tmp_path / "blueapi-stubs" / "src" / "blueapi-stubs" / "client" / "cache.pyi"
-    ).exists()
+    assert (stub_src / "client" / "device_cache.pyi").exists()
+    assert (stub_src / "client" / "plan_cache.pyi").exists()
 
 
-def _extract_rendered(src: StringIO) -> tuple[str, str]:
+def _extract_rendered(src: StringIO) -> str:
     src.seek(0)
-    _read_until_line(src, "### Generated plans")
-    plan_text = _read_until_line(src, "### End")
-    _read_until_line(src, "### Generated devices")
-    device_text = _read_until_line(src, "### End")
-    return plan_text, device_text
+    _read_until_line(src, "### Generated")
+    return _read_until_line(src, "### End")
 
 
 def _read_until_line(src: StringIO, match: str) -> str:

From 366789f6c5f5a1771ee7480517bd8aefd6fcce1f Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 2 Nov 2026 19:20:53 +0000
Subject: [PATCH 4/6] Extract template names to constants

---
 src/blueapi/cli/stubgen.py           |  7 +++++--
 tests/unit_tests/cli/test_stubgen.py | 12 +++++++-----
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/blueapi/cli/stubgen.py b/src/blueapi/cli/stubgen.py
index 2f3a7926e..24ec261b8 100644
--- a/src/blueapi/cli/stubgen.py
+++ b/src/blueapi/cli/stubgen.py
@@ -14,6 +14,9 @@
 
 log = logging.getLogger(__name__)
 
+PLAN_TMPL = "plan_cache_template.pyi.jinja"
+DEV_TMPL = "device_cache_template.pyi.jinja"
+
 
 @dataclass
 class ArgSpec:
@@ -100,10 +103,10 @@ def generate_stubs(target: Path, plans: list[Plan], devices: list[DeviceRef]):
 
     plan_model = [PlanSpec.from_plan(p) for p in plans]
     with open(client_dir / "plan_cache.pyi", "w") as out:
-        render_stub_file(out, "plan_cache_template.pyi.jinja", plans=plan_model)
+        render_stub_file(out, PLAN_TMPL, plans=plan_model)
 
     with open(client_dir / "device_cache.pyi", "w") as out:
-        render_stub_file(out, "device_cache_template.pyi.jinja", devices=devices)
+        render_stub_file(out, DEV_TMPL, devices=devices)
 
 
 def _docstring(text: str) -> str:
diff --git a/tests/unit_tests/cli/test_stubgen.py b/tests/unit_tests/cli/test_stubgen.py
index d525839e0..9deb6109a 100644
--- a/tests/unit_tests/cli/test_stubgen.py
+++ b/tests/unit_tests/cli/test_stubgen.py
@@ -7,6 +7,8 @@
 import pytest
 
 from blueapi.cli.stubgen import (
+    DEV_TMPL,
+    PLAN_TMPL,
     PlanSpec,
     _docstring,
     _type_string,
@@ -94,8 +96,8 @@ def test_type_string(typ: dict, expected: str):
 @pytest.mark.parametrize(
     "template,model",
     [
-        ("device_cache_template.pyi.jinja", {"devices": []}),
-        ("plan_cache_template.pyi.jinja", {"plans": []}),
+        (DEV_TMPL, {"devices": []}),
+        (PLAN_TMPL, {"plans": []}),
     ],
 )
 def test_render_empty(template: str, model: dict[str, Any]):
@@ -124,7 +126,7 @@ def test_render_empty_plan_function():
     output = StringIO()
     plans = [Plan(name="empty", model=FOO, runner=Mock())]
     plan_model = [PlanSpec.from_plan(p) for p in plans]
-    render_stub_file(output, "plan_cache_template.pyi.jinja", plans=plan_model)
+    render_stub_file(output, PLAN_TMPL, plans=plan_model)
     rendered = _extract_rendered(output)
 
     assert (
@@ -144,7 +146,7 @@ def test_render_multiple_plan_functions():
     runner = Mock()
     plans = [Plan("empty", FOO, runner), Plan("two_args", BAR, runner)]
     plan_model = [PlanSpec.from_plan(p) for p in plans]
-    render_stub_file(output, "plan_cache_template.pyi.jinja", plans=plan_model)
+    render_stub_file(output, PLAN_TMPL, plans=plan_model)
     rendered = _extract_rendered(output)
 
     assert (
@@ -174,7 +176,7 @@ def test_device_fields():
         DeviceRef("one", cache, DeviceModel(name="one", protocols=[])),
         DeviceRef("two", cache, DeviceModel(name="two", protocols=[])),
     ]
-    render_stub_file(output, "device_cache_template.pyi.jinja", devices=devices)
+    render_stub_file(output, DEV_TMPL, devices=devices)
 
     assert _extract_rendered(output) == "    one: DeviceRef\n    two: DeviceRef\n"
 

From 3ecb238a25934fe8d469744f9ecffb3af1303ca1 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 2 Nov 2026 20:14:07 +0000
Subject: [PATCH 5/6] Support tuple parameters

---
 src/blueapi/cli/stubgen.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/cli/stubgen.py b/src/blueapi/cli/stubgen.py
index 24ec261b8..9552dfabd 100644
--- a/src/blueapi/cli/stubgen.py
+++ b/src/blueapi/cli/stubgen.py
@@ -47,8 +47,11 @@ def from_plan(cls, plan: Plan) -> Self:
 def _type_string(spec) -> str:
     """Best effort attempt at making useful type hints for plans"""
     match spec.get("type"):
-        case "array":
+        case "array" if "items" in spec:
             return f"list[{_type_string(spec.get('items'))}]"
+        case "array" if "prefixItems" in spec:
+            items = spec.get("prefixItems")
+            return f"tuple[{', '.join(_type_string(item) for item in items)}]"
         case "integer":
             return "int"
         case "number":

From b3c3d8626c007753fe07b62fd42226c3d89d0971 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 2 Nov 2026 20:14:31 +0000
Subject: [PATCH 6/6] Update templates

---
 src/blueapi/client/client.py                               | 2 +-
 .../stubs/templates/device_cache_template.pyi.jinja        | 7 ++-----
 src/blueapi/stubs/templates/plan_cache_template.pyi.jinja  | 3 ++-
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index 4dc9f4fbc..7685932ee 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -320,7 +320,7 @@ def get_active_task(self) -> WorkerTask:
 
         return self.active_task
 
-    @start_as_current_span(TRACER, "name", "args", "params")
+    @start_as_current_span(TRACER, "name", "args", "kwargs")
     def run_plan(
         self, name: str, args: tuple[Any, ...], kwargs: Mapping[str, Any]
     ) -> TaskStatus:
diff --git a/src/blueapi/stubs/templates/device_cache_template.pyi.jinja b/src/blueapi/stubs/templates/device_cache_template.pyi.jinja
index e2341277f..3cfd508ed 100644
--- a/src/blueapi/stubs/templates/device_cache_template.pyi.jinja
+++ b/src/blueapi/stubs/templates/device_cache_template.pyi.jinja
@@ -1,8 +1,5 @@
-from collections.abc import Callable
-from typing import Any
 from blueapi.client.rest import BlueapiRestClient
-from blueapi.service.model import DeviceModel, PlanModel
-from blueapi.worker.event import WorkerEvent
+from blueapi.service.model import DeviceModel
 
 {#-
  This file is based on the devices_cache.py file in
@@ -11,7 +8,7 @@ from blueapi.worker.event import WorkerEvent
 
 # This file is auto-generated for a live server and should not be modified directly
 
-class DeviceRef(str):
+class DeviceRef:
     model: DeviceModel
     _cache: DeviceCache
     def __new__(cls, name: str, cache: DeviceCache, model: DeviceModel): ...
diff --git a/src/blueapi/stubs/templates/plan_cache_template.pyi.jinja b/src/blueapi/stubs/templates/plan_cache_template.pyi.jinja
index 1f31ea255..e63108f15 100644
--- a/src/blueapi/stubs/templates/plan_cache_template.pyi.jinja
+++ b/src/blueapi/stubs/templates/plan_cache_template.pyi.jinja
@@ -2,6 +2,7 @@ from collections.abc import Callable
 from typing import Any
 from blueapi.service.model import PlanModel
 from blueapi.worker.event import TaskStatus
+from blueapi.client.device_cache import DeviceRef
 
 {#-
  This file is based on the plan_cache.py file in blueapi/client/plan_cache.py
@@ -10,7 +11,7 @@ from blueapi.worker.event import TaskStatus
 
 # This file is auto-generated for a live server and should not be modified directly
 
-PlanRunner = Callable[[str, dict[str, Any]], TaskStatus]
+PlanRunner = Callable[[str, tuple[Any, ...], dict[str, Any]], TaskStatus]
 
 class PlanCache:
     def __init__(self, runner: PlanRunner, plans: list[PlanModel]) -> None: ...