diff --git a/docs/discovery.md b/docs/discovery.md index dc7b8ea..ec49fd8 100644 --- a/docs/discovery.md +++ b/docs/discovery.md @@ -107,7 +107,7 @@ function(estop) fleet emergency-st ### Discovery -#### `discover(selector, offset=0, limit=200)` +#### `discover(selector, offset=0, limit=200, *, where=None)` Resolves a selector to matched entities. Returns devices, function tuples, or event tuples depending on the selector scope. The response includes a @@ -119,6 +119,65 @@ and switches to a name-and-labels summary above `DEVICE_CONNECT_FUNCTION_THRESHOLD` (default 20). The threshold is configurable via environment variable. +##### State predicates + +Pass a CEL `where` expression to query stored device state without invoking +a function. This also works for reporting-only sensors and instruments: + +```python +discover( + "device(model_id:ophyd_async:SimStage)", + where="status.x_readback < 51.99 && status.x_setpoint == 2.0", +) +# {"scope": "device_only", "matched": 1, "returned": 1, "offset": 0, +# "next_offset": None, +# "results": [{"device_id": "stage-0051", +# "device_type": "ophyd_async:SimStage", "location": "lab-A"}], +# "label_histogram": {...}} +``` + +The registry uses the same CEL evaluator as `broadcast`. Discovery binds: + +| Variable | Source | +| --- | --- | +| `status` | Latest stored heartbeat status, including driver-specific fields | +| `identity` | Registered identity plus `device_id` | +| `labels` | Device capability labels, with legacy `status.location` and `identity.device_type` defaults for `location` and `type`; declared labels take precedence | + +Discovery does not provide broadcast's `bindings` payload. Missing fields +and evaluation type errors make that device a non-match, as in broadcast. +Malformed, empty or non-string expressions produce JSON-RPC `-32602` +(invalid params), exposed by the tool as `error.code = "invalid_predicate"`. + +`discovery/listDevices` accepts optional `where` alongside `device_type`, +`location`, `offset` and `limit`. It evaluates the predicate before paging; +`total_matched` and `next_offset` refer to state matches. A successful +predicate query includes `where_applied: true`. For predicate queries, +visibility ACLs filter the matches before pagination and counting, so +neither the total nor the cursor exposes hidden devices' state. +The agent tool resolves its selector over +visible state matches, then calculates `matched`, the label histogram and +the requested page. For function/event selectors, `where` filters the +owning devices before resolving their functions/events. + +With `where`, device-only tool results contain only `device_id`, +`device_type` and `location`; full records remain inside the registry and +tool process. A fleet of 5,000 devices with one state match therefore +requires one registry page and returns one compact tool row. The registry +still scans its tenant's stored snapshot; CPU cost scales with fleet size. + +State queries bypass the SDK's local fleet cache. The registry's existing +snapshot cache still applies (`DC_FLEET_CACHE_TTL`, default 2 seconds), and +state freshness also depends on device heartbeat frequency. Pages are not +a transactionally frozen view if devices change state between requests. + +The server installs CEL support by default; agent tools need no predicate +extra for registry discovery. Upgrade the server and agent's edge SDK +together: clients reject replies without `where_applied: true` rather than +returning an unfiltered fleet. There is no client fallback for older +servers, and `where` requires registry mode (D2D is unsupported). +Calls without `where` retain their existing behavior and response shape. + #### `discover_labels(key=None, offset=0, limit=50)` Returns the fleet label vocabulary. Use this first when you do not know diff --git a/packages/device-connect-agent-tools/device_connect_agent_tools/adapters/claude.py b/packages/device-connect-agent-tools/device_connect_agent_tools/adapters/claude.py index f4a2883..bb0a029 100644 --- a/packages/device-connect-agent-tools/device_connect_agent_tools/adapters/claude.py +++ b/packages/device-connect-agent-tools/device_connect_agent_tools/adapters/claude.py @@ -91,8 +91,19 @@ async def discover_labels(args: dict[str, Any]) -> dict[str, Any]: "'device(category:camera, location:zone-A/*)', " "'device(*).function(direction:write)', 'event(modality:motion)'. " "Response includes a label_histogram (per-key vocabulary across the " - "matched set) so the agent can narrow next.", - {"selector": str, "offset": int, "limit": int}, + "matched set) so the agent can narrow next. Optional CEL where filters " + "registry-stored status, identity and labels; device-only results are " + "compact ids, types and locations, including devices with no functions.", + { + "type": "object", + "properties": { + "selector": {"type": "string"}, + "offset": {"type": "integer"}, + "limit": {"type": "integer"}, + "where": {"type": "string", "description": "CEL predicate over status, identity and labels"}, + }, + "required": ["selector", "offset", "limit"], + }, ) async def discover(args: dict[str, Any]) -> dict[str, Any]: return _text( @@ -100,6 +111,7 @@ async def discover(args: dict[str, Any]) -> dict[str, Any]: selector=args["selector"], offset=int(args.get("offset", 0)), limit=int(args.get("limit", 200)), + where=args.get("where"), ) ) diff --git a/packages/device-connect-agent-tools/device_connect_agent_tools/connection.py b/packages/device-connect-agent-tools/device_connect_agent_tools/connection.py index ad9148c..7b49e9a 100644 --- a/packages/device-connect-agent-tools/device_connect_agent_tools/connection.py +++ b/packages/device-connect-agent-tools/device_connect_agent_tools/connection.py @@ -432,20 +432,31 @@ def list_devices( self, device_type: Optional[str] = None, location: Optional[str] = None, + *, + where: Optional[str] = None, ) -> List[Dict[str, Any]]: """List devices via the discovery provider (D2D or registry).""" - return self._run(self._async_list_devices(device_type, location)) + return self._run(self._async_list_devices(device_type, location, where=where)) async def _async_list_devices( self, device_type: Optional[str] = None, location: Optional[str] = None, + *, + where: Optional[str] = None, ) -> List[Dict[str, Any]]: if self._provider is None: raise RuntimeError("Not connected — call connect() first") - devices = await self._provider.list_devices( - device_type=device_type, location=location, - ) + if where is not None: + if not isinstance(self._provider, _SDKRegistryClient): + raise RuntimeError("Discovery where predicates require a registry connection; D2D is unsupported") + devices = await self._provider.list_devices( + device_type=device_type, location=location, where=where, + ) + else: + devices = await self._provider.list_devices( + device_type=device_type, location=location, + ) return [flatten_device(d) for d in devices] def invalidate_cache(self) -> None: diff --git a/packages/device-connect-agent-tools/device_connect_agent_tools/tools.py b/packages/device-connect-agent-tools/device_connect_agent_tools/tools.py index 4b9e87e..b619469 100644 --- a/packages/device-connect-agent-tools/device_connect_agent_tools/tools.py +++ b/packages/device-connect-agent-tools/device_connect_agent_tools/tools.py @@ -218,6 +218,8 @@ def discover( selector: str, offset: int = 0, limit: int = DEFAULT_DISCOVER_LIMIT, + *, + where: str | None = None, ) -> dict[str, Any]: """Resolve a selector to matched devices, functions, or events. @@ -241,6 +243,11 @@ def discover( selector: A selector expression string. offset: Pagination offset (rows skipped). limit: Max rows per page (capped at DISCOVER_HARD_LIMIT). + where: Optional CEL predicate over registry-stored ``status``, + ``identity`` and ``labels``. Filters devices before resolving + functions/events. Device-only results contain device_id, + device_type and location. No callable function is required. + Requires a registry with state discovery support. Returns: A response envelope: @@ -279,7 +286,15 @@ def discover( try: conn = get_connection() - devices = conn.list_devices() + if where is None: + devices = conn.list_devices() + else: + from device_connect_edge.predicate import PredicateCompileError + + try: + devices = conn.list_devices(where=where) + except PredicateCompileError as e: + return _empty_envelope(scope=sel.scope.value, error=_error("invalid_predicate", str(e))) except Exception as e: logger.error("discover(%r) failed loading fleet: %s", selector, e) return _empty_envelope( @@ -298,7 +313,13 @@ def discover( total = len(matched_devices) page_devices, next_offset = _paginate(matched_devices, safe_offset, safe_limit) expand = SMALL_FLEET_THRESHOLD > 0 and total <= SMALL_FLEET_THRESHOLD - results = [_device_summary_for_discover(d, expand) for d in page_devices] + if where is not None: + results = [ + {key: d.get(key) for key in ("device_id", "device_type", "location")} + for d in page_devices + ] + else: + results = [_device_summary_for_discover(d, expand) for d in page_devices] histogram, multivalued, unique = label_histogram(matched_devices, count_unique=True) formatted_histogram = _format_label_histogram(histogram, multivalued, unique) return { diff --git a/packages/device-connect-agent-tools/tests/test_claude_adapter.py b/packages/device-connect-agent-tools/tests/test_claude_adapter.py index 4960a49..ce9fb29 100644 --- a/packages/device-connect-agent-tools/tests/test_claude_adapter.py +++ b/packages/device-connect-agent-tools/tests/test_claude_adapter.py @@ -78,6 +78,18 @@ def _mock_sdk_and_connection(): class TestClaudeAdapterExports: + @pytest.mark.asyncio + @pytest.mark.parametrize("where_args", [{}, {"where": "status.online"}]) + async def test_discover_exposes_and_forwards_optional_where(self, where_args): + from device_connect_agent_tools.adapters import claude as adapter + + schema = adapter.discover._tool_schema + assert schema["properties"]["where"]["type"] == "string" + assert "where" not in schema["required"] + with patch.object(adapter, "_discover", return_value={"matched": 0}) as discover: + await adapter.discover.__wrapped__({"selector": "device(*)", **where_args}) + discover.assert_called_once_with(selector="device(*)", offset=0, limit=200, where=where_args.get("where")) + def test_module_exports_all_tools(self): from device_connect_agent_tools.adapters import claude as adapter diff --git a/packages/device-connect-agent-tools/tests/test_discover.py b/packages/device-connect-agent-tools/tests/test_discover.py index 3a5702c..e675942 100644 --- a/packages/device-connect-agent-tools/tests/test_discover.py +++ b/packages/device-connect-agent-tools/tests/test_discover.py @@ -13,6 +13,54 @@ import pytest from device_connect_agent_tools import tools as tools_mod +from device_connect_edge.predicate import PredicateCompileError + + +def test_discover_where_returns_compact_rows_and_paginates_matching_devices(): + conn = MagicMock() + conn.list_devices.return_value = [ + {"device_id": f"stage-{i}", "device_type": "stage", "location": "lab-A", + "labels": {"model_id": "ophyd_async:SimStage"}, + "functions": [{"name": "huge_schema", "parameters": {"description": "x" * 10000}}]} + for i in range(3) + ] + [{"device_id": "sensor", "labels": {"model_id": "sensor"}}] + with patch.object(tools_mod, "get_connection", return_value=conn): + result = tools_mod.discover( + "device(model_id:ophyd_async:SimStage)", offset=1, limit=1, where="status.x_readback < 51.99", + ) + conn.list_devices.assert_called_once_with(where="status.x_readback < 51.99") + assert (result["matched"], result["returned"], result["next_offset"]) == (3, 1, 2) + assert result["results"] == [{"device_id": "stage-1", "device_type": "stage", "location": "lab-A"}] + assert result["label_histogram"]["model_id"]["values"] == {"ophyd_async:SimStage": 3} + + +def test_discover_invalid_where_is_a_structured_error(): + conn = MagicMock() + conn.list_devices.side_effect = PredicateCompileError("failed to compile where") + with patch.object(tools_mod, "get_connection", return_value=conn): + result = tools_mod.discover("device(*)", where="status.x > > 1") + assert result["error"]["code"] == "invalid_predicate" + assert result["matched"] == 0 + assert result["results"] == [] + + +def test_discover_where_does_not_mislabel_connection_value_errors(): + conn = MagicMock() + conn.list_devices.side_effect = ValueError("malformed registry response") + with patch.object(tools_mod, "get_connection", return_value=conn): + result = tools_mod.discover("device(*)", where="status.online") + assert result["error"]["code"] == "connection_error" + + +def test_discover_where_also_filters_function_and_event_scopes(): + conn = MagicMock() + conn.list_devices.return_value = [SAMPLE_DEVICES[0]] + with patch.object(tools_mod, "get_connection", return_value=conn): + functions = tools_mod.discover("function(*)", where="status.online") + events = tools_mod.discover("event(*)", where="status.online") + assert functions["matched"] == 1 + assert events["matched"] == 2 + assert all(row["device_id"] == "cam-001" for row in functions["results"] + events["results"]) # -- Fixture: labeled fleet --------------------------------------- diff --git a/packages/device-connect-edge/device_connect_edge/device.py b/packages/device-connect-edge/device_connect_edge/device.py index c3e2f51..8e2f033 100644 --- a/packages/device-connect-edge/device_connect_edge/device.py +++ b/packages/device-connect-edge/device_connect_edge/device.py @@ -1463,7 +1463,14 @@ def _run() -> bool: deadline = loop.time() + timeout_s try: - await asyncio.wait_for(self._where_eval_semaphore.acquire(), timeout_s) + if self._where_eval_semaphore.locked(): + await asyncio.wait_for(self._where_eval_semaphore.acquire(), timeout_s) + else: + # acquire() completes immediately when capacity is available. + # wait_for() would schedule another task and yield even then, + # letting unrelated device work consume the whole deadline + # before this predicate is submitted to its idle executor. + await self._where_eval_semaphore.acquire() except asyncio.TimeoutError: self._logger.warning( "Broadcast %s: where predicate timed out after %.3fs (skipping)", diff --git a/packages/device-connect-edge/device_connect_edge/registry_client.py b/packages/device-connect-edge/device_connect_edge/registry_client.py index aa4ec0f..e1fac3c 100644 --- a/packages/device-connect-edge/device_connect_edge/registry_client.py +++ b/packages/device-connect-edge/device_connect_edge/registry_client.py @@ -37,6 +37,7 @@ from device_connect_edge.messaging.base import MessagingClient from device_connect_edge.messaging.exceptions import RequestTimeoutError +from device_connect_edge.predicate import PredicateCompileError logger = logging.getLogger(__name__) @@ -49,6 +50,14 @@ _DEFAULT_LIST_PAGE_SIZE = int(os.getenv("DEVICE_CONNECT_LIST_PAGE_SIZE", "100")) +class RegistryError(RuntimeError): + """JSON-RPC error from the registry, preserving its machine-readable code.""" + + def __init__(self, code: int, message: str): + self.code = code + super().__init__(f"Registry error ({code}): {message}") + + class RegistryClient: """JSON-RPC client for the device registry service. @@ -116,10 +125,7 @@ async def _request( if "error" in response: error = response["error"] - raise RuntimeError( - f"Registry error ({error.get('code', -1)}): " - f"{error.get('message', 'Unknown error')}" - ) + raise RegistryError(error.get("code", -1), error.get("message", "Unknown error")) return response.get("result") except RequestTimeoutError as e: last_err = e @@ -141,6 +147,7 @@ async def list_devices( location: Optional[str] = None, capabilities: Optional[List[str]] = None, timeout: Optional[float] = None, + where: Optional[str] = None, ) -> List[Dict[str, Any]]: """List devices from the registry service. @@ -153,12 +160,14 @@ async def list_devices( location: Filter by location. capabilities: Filter by required capabilities. timeout: Override default timeout. + where: CEL predicate evaluated by the registry over stored state. + Bypasses the local fleet cache. Requires a supporting server. Returns: List of device dictionaries with full registration data. """ # Check cache - if self._cache_ttl > 0 and self._cache is not None: + if where is None and self._cache_ttl > 0 and self._cache is not None: age = time.time() - self._cache_time if age < self._cache_ttl: logger.debug("Using cached device list (age: %.1fs)", age) @@ -182,6 +191,7 @@ async def list_devices( offset=offset, limit=_DEFAULT_LIST_PAGE_SIZE, timeout=timeout, + where=where, ) devices.extend(page) if next_offset is None: @@ -208,6 +218,7 @@ async def list_devices( and device_type is None and location is None and not capabilities + and where is None ): self._cache = devices self._cache_time = time.time() @@ -223,6 +234,7 @@ async def list_devices_page( location: Optional[str] = None, capabilities: Optional[List[str]] = None, timeout: Optional[float] = None, + where: Optional[str] = None, ) -> Tuple[List[Dict[str, Any]], Optional[int], int]: """Fetch a single page of devices with pagination metadata. @@ -230,12 +242,15 @@ async def list_devices_page( most callers should stick with :meth:`list_devices`, which loops internally and returns the full fleet. + ``where`` filters stored device state before pagination. Older + servers that do not acknowledge the predicate raise an error. + Returns: ``(devices, next_offset, total_matched)`` where ``next_offset`` is ``None`` on the final page. ACL caveat: - When the registry has ACLs enabled, server-side filtering + Without ``where``, when the registry has ACLs enabled, filtering runs *after* slicing. As a result ``len(devices)`` for a given page may be smaller than ``limit`` even when more pages follow, and ``total_matched`` is the unfiltered total @@ -243,6 +258,8 @@ async def list_devices_page( ``total_matched`` as an upper bound on what the caller will ever see, and must not assume ``len(devices) == limit`` implies a full page. + With ``where``, visibility filtering precedes pagination and + ``total_matched`` counts only visible state matches. """ return await self._list_devices_page( device_type=device_type, @@ -251,6 +268,7 @@ async def list_devices_page( offset=offset, limit=limit, timeout=timeout, + where=where, ) async def _list_devices_page( @@ -262,6 +280,7 @@ async def _list_devices_page( offset: int, limit: int, timeout: Optional[float], + where: Optional[str] = None, ) -> Tuple[List[Dict[str, Any]], Optional[int], int]: subject = f"device-connect.{self._tenant}.discovery" params: Dict[str, Any] = {"offset": int(offset), "limit": int(limit)} @@ -271,10 +290,19 @@ async def _list_devices_page( params["location"] = location if capabilities: params["capabilities"] = capabilities + if where is not None: + params["where"] = where - result = await self._request( - subject, "discovery/listDevices", params, timeout, - ) + try: + result = await self._request( + subject, "discovery/listDevices", params, timeout, + ) + except RegistryError as e: + if where is not None and e.code == -32602: + raise PredicateCompileError(str(e)) from e + raise + if where is not None and result.get("where_applied") is not True: + raise RuntimeError("Registry does not support discovery where predicates; upgrade the server") devices = result.get("devices", []) next_offset = result.get("next_offset") total = result.get("total_matched", len(devices)) diff --git a/packages/device-connect-edge/tests/test_device_where.py b/packages/device-connect-edge/tests/test_device_where.py index e58c242..b248580 100644 --- a/packages/device-connect-edge/tests/test_device_where.py +++ b/packages/device-connect-edge/tests/test_device_where.py @@ -25,6 +25,34 @@ def evaluate(self, context): raise RuntimeError("bad predicate") +@pytest.mark.asyncio +async def test_available_where_worker_is_not_delayed_by_other_runtime_work(): + """An idle executor must start before other devices occupy the event loop.""" + runtime = DeviceRuntime(device_id="where-free-slot-test") + runtime._logger = Mock() + finished = threading.Event() + + class _FastPredicate: + def evaluate(self, context): + finished.set() + return True + + try: + evaluation = asyncio.create_task(runtime._evaluate_where_with_timeout( + _FastPredicate(), {}, "corr-free-slot", timeout_s=1.0, + )) + await asyncio.sleep(0) + # The worker must start without another loop turn, even if another + # runtime occupies this loop. Stay within the evaluation deadline; + # asyncio's handling of already-expired waits differs by Python version. + submitted_without_another_loop_turn = finished.wait(timeout=0.1) + assert await evaluation is True + assert submitted_without_another_loop_turn + runtime._logger.warning.assert_not_called() + finally: + runtime._shutdown_where_eval_executor() + + @pytest.mark.asyncio async def test_where_eval_timeout_fails_closed_and_warns(): runtime = DeviceRuntime(device_id="where-timeout-test") diff --git a/packages/device-connect-edge/tests/test_registry_client.py b/packages/device-connect-edge/tests/test_registry_client.py index 5817e64..75047c4 100644 --- a/packages/device-connect-edge/tests/test_registry_client.py +++ b/packages/device-connect-edge/tests/test_registry_client.py @@ -14,6 +14,7 @@ from device_connect_edge.messaging.exceptions import RequestTimeoutError from device_connect_edge.registry_client import RegistryClient +from device_connect_edge.predicate import PredicateCompileError def _make_client(mock_messaging=None, **kwargs): @@ -103,6 +104,42 @@ async def test_request_raises_after_all_retries_exhausted(self, mock_sleep): class TestListDevicesPagination: """Verify list_devices transparently pages through the registry.""" + @pytest.mark.asyncio + async def test_where_walk_bypasses_cache_and_preserves_unfiltered_cache(self): + client, messaging = _make_client(cache_ttl=30) + messaging.request.side_effect = [ + _success_response({"devices": [{"device_id": "unfiltered"}]}), + _success_response({"devices": [{"device_id": "match-a"}], "next_offset": 1, + "total_matched": 2, "where_applied": True}), + _success_response({"devices": [{"device_id": "match-b"}], "next_offset": None, + "total_matched": 2, "where_applied": True}), + ] + assert await client.list_devices() == [{"device_id": "unfiltered"}] + assert await client.list_devices(where="status.online") == [ + {"device_id": "match-a"}, {"device_id": "match-b"}, + ] + assert await client.list_devices() == [{"device_id": "unfiltered"}] + requests = [json.loads(c.args[1])["params"] for c in messaging.request.call_args_list] + assert "where" not in requests[0] + assert [(p["where"], p["offset"]) for p in requests[1:]] == [ + ("status.online", 0), ("status.online", 1), + ] + + @pytest.mark.asyncio + async def test_where_page_rejects_older_server_that_ignores_predicate(self): + client, messaging = _make_client() + messaging.request.return_value = _success_response({"devices": [{"device_id": "unfiltered"}]}) + with pytest.raises(RuntimeError, match="does not support.*where"): + await client.list_devices_page(where="status.online") + + @pytest.mark.asyncio + async def test_where_surfaces_registry_validation_as_predicate_error(self): + client, messaging = _make_client() + messaging.request.return_value = _error_response(-32602, "failed to compile where") + with pytest.raises(PredicateCompileError, match="failed to compile where"): + await client.list_devices(where="status.x > > 1") + assert messaging.request.call_count == 1 + @staticmethod def _paged_responses(total: int, page_size: int): """Build the sequence of NATS reply bytes the server would emit.""" diff --git a/packages/device-connect-server/device_connect_server/registry/client.py b/packages/device-connect-server/device_connect_server/registry/client.py index 23f84b6..53c259b 100644 --- a/packages/device-connect-server/device_connect_server/registry/client.py +++ b/packages/device-connect-server/device_connect_server/registry/client.py @@ -182,6 +182,7 @@ async def list_devices( location: Optional[str] = None, capabilities: Optional[List[str]] = None, timeout: Optional[float] = None, + where: Optional[str] = None, ) -> List[Dict[str, Any]]: """List all registered devices. @@ -190,6 +191,7 @@ async def list_devices( location: Filter by location capabilities: Filter by required capabilities timeout: Request timeout + where: CEL predicate over stored state; requires a supporting server. Returns: List of device dictionaries with full registration data @@ -215,6 +217,7 @@ async def list_devices( offset=offset, limit=_DEFAULT_LIST_PAGE_SIZE, timeout=timeout, + where=where, ) devices.extend(page) if next_offset is None: @@ -242,19 +245,25 @@ async def list_devices_page( location: Optional[str] = None, capabilities: Optional[List[str]] = None, timeout: Optional[float] = None, + where: Optional[str] = None, ) -> Tuple[List[Dict[str, Any]], Optional[int], int]: """Fetch one page of devices with pagination metadata. Returns ``(devices, next_offset, total_matched)``; ``next_offset`` is ``None`` on the final page. + ``where`` filters stored state before pagination. A server that + ignores the predicate raises an error instead of returning its fleet. + ACL caveat: - When the registry has ACLs enabled, filtering runs *after* + Without ``where``, when the registry has ACLs enabled, filtering runs *after* slicing. ``len(devices)`` for a page may be smaller than ``limit`` even when more pages follow, and ``total_matched`` is the unfiltered total (before the caller's ACL applies). Callers should treat ``total_matched`` as an upper bound and must not infer "full page" from ``len(devices) == limit``. + With ``where``, visibility filtering precedes pagination and + ``total_matched`` counts only visible state matches. """ return await self._list_devices_page( device_type=device_type, @@ -263,6 +272,7 @@ async def list_devices_page( offset=offset, limit=limit, timeout=timeout, + where=where, ) async def _list_devices_page( @@ -274,6 +284,7 @@ async def _list_devices_page( offset: int, limit: int, timeout: Optional[float], + where: Optional[str] = None, ) -> Tuple[List[Dict[str, Any]], Optional[int], int]: subject = f"device-connect.{self._tenant}.discovery" params: Dict[str, Any] = {"offset": int(offset), "limit": int(limit)} @@ -283,10 +294,14 @@ async def _list_devices_page( params["location"] = location if capabilities: params["capabilities"] = capabilities + if where is not None: + params["where"] = where result = await self._request( subject, "discovery/listDevices", params, timeout, ) + if where is not None and result.get("where_applied") is not True: + raise RuntimeError("Registry does not support discovery where predicates; upgrade the server") devices = result.get("devices", []) next_offset = result.get("next_offset") total = result.get("total_matched", len(devices)) diff --git a/packages/device-connect-server/device_connect_server/registry/service/Dockerfile b/packages/device-connect-server/device_connect_server/registry/service/Dockerfile index 8f0eb20..ecaebdf 100644 --- a/packages/device-connect-server/device_connect_server/registry/service/Dockerfile +++ b/packages/device-connect-server/device_connect_server/registry/service/Dockerfile @@ -9,7 +9,7 @@ RUN pip install --no-cache-dir --only-binary eclipse-zenoh -r requirements.txt # Install device-connect-edge from sibling package COPY packages/device-connect-edge/ /tmp/device-connect-edge/ -RUN pip install --no-cache-dir /tmp/device-connect-edge/ && rm -rf /tmp/device-connect-edge/ +RUN pip install --no-cache-dir '/tmp/device-connect-edge[predicate]' && rm -rf /tmp/device-connect-edge/ # Copy device-connect-server library COPY packages/device-connect-server/device_connect_server/ ./device_connect_server/ diff --git a/packages/device-connect-server/device_connect_server/registry/service/main.py b/packages/device-connect-server/device_connect_server/registry/service/main.py index a2c60c6..ce6f86b 100644 --- a/packages/device-connect-server/device_connect_server/registry/service/main.py +++ b/packages/device-connect-server/device_connect_server/registry/service/main.py @@ -400,6 +400,8 @@ def _make_list_handler( ``requester_id`` field in the RPC params. """ + from device_connect_edge.predicate import PredicateCompileError + async def rpc_discovery(data: bytes, reply: Optional[str]): if not reply: logger.debug("[device-registry] discovery request with no reply address; ignoring") @@ -418,6 +420,8 @@ async def rpc_discovery(data: bytes, reply: Optional[str]): if method == "discovery/listDevices": device_type = params.get("device_type") location = params.get("location") + where = params.get("where") + predicate_params = {"where": where} if where is not None else {} # Pagination: ``offset`` and ``limit`` are optional. # # If ``limit`` is absent the caller is on the legacy @@ -527,23 +531,42 @@ async def rpc_discovery(data: bytes, reply: Optional[str]): "unintended.", requested_limit_int, _LIST_DEVICES_MAX_LIMIT, ) - page, next_offset, total = await asyncio.to_thread( - registry.list_devices_page, tenant, - device_type=device_type, - location=location, - offset=offset_val, - limit=effective_limit, - ) + if where is not None and acl_manager: + # State predicates can probe arbitrary status fields. + # Count only visible matches so metadata cannot reveal + # the state of a device hidden from this requester. + matches = await asyncio.to_thread( + registry.list_devices, tenant, + device_type=device_type, location=location, + **predicate_params, + ) + matches = acl_manager.filter_visible_devices( + params.get("requester_id", ""), matches, tenant=tenant, + ) + total = len(matches) + end = offset_val + effective_limit + page = matches[offset_val:end] + next_offset = end if end < total else None + else: + page, next_offset, total = await asyncio.to_thread( + registry.list_devices_page, tenant, + device_type=device_type, + location=location, + offset=offset_val, + limit=effective_limit, + **predicate_params, + ) else: page = await asyncio.to_thread( registry.list_devices, tenant, device_type=device_type, location=location, + **predicate_params, ) # next_offset / total are unused on the legacy reply # path (see the ``if paged`` branch below); the # legacy shape is just ``{"devices": page}``. - if acl_manager: + if acl_manager and not (paged and where is not None): requester_id = params.get("requester_id", "") # ACL filtering runs after pagination — devices the # caller is not allowed to see are dropped from the @@ -570,6 +593,9 @@ async def rpc_discovery(data: bytes, reply: Optional[str]): # ``next_offset`` so old clients that ignore unknown # keys aren't surprised by new metadata. response_result = {"devices": page} + if where is not None: + # Clients must detect older registries that ignore where. + response_result["where_applied"] = True await messaging.publish( reply, build_rpc_response(payload.get("id"), response_result), @@ -609,6 +635,10 @@ async def rpc_discovery(data: bytes, reply: Optional[str]): ) else: return # Not a discovery method — ignore + except PredicateCompileError as e: + await messaging.publish( + reply, build_rpc_error(payload.get("id"), -32602, str(e)), + ) except Exception as e: await messaging.publish( reply, diff --git a/packages/device-connect-server/device_connect_server/registry/service/registry.py b/packages/device-connect-server/device_connect_server/registry/service/registry.py index 4655562..d59309c 100644 --- a/packages/device-connect-server/device_connect_server/registry/service/registry.py +++ b/packages/device-connect-server/device_connect_server/registry/service/registry.py @@ -26,6 +26,8 @@ import etcd3gw from requests.adapters import HTTPAdapter +from device_connect_edge.predicate import PredicateEvalError, compile_where + _logger = logging.getLogger(__name__) ETCD_HOST = os.getenv("ETCD_HOST", "localhost") @@ -248,6 +250,8 @@ def list_devices( tenant: str, device_type: str | None = None, location: str | None = None, + *, + where: str | None = None, ) -> List[dict]: """Return registered device payloads for ``tenant``, optionally filtered. @@ -255,7 +259,11 @@ def list_devices( tenant: Tenant namespace. device_type: Filter by device type (case-insensitive substring match). location: Filter by device location (case-insensitive substring match). + where: CEL predicate over stored status, identity and device labels. """ + # Compile once, including for empty fleets, so invalid expressions + # always produce an error. Evaluation runs in the handler's worker. + predicate = compile_where(where) if where is not None else None # Shallow-copy the cached snapshot so the filter rebinds below # never mutate the shared cache entry (the dicts are read-only). devices: List[dict] = list(self._decoded_fleet(tenant)) @@ -272,6 +280,25 @@ def list_devices( d for d in devices if loc in (d.get("status", {}).get("location") or "").lower() ] + if predicate is not None: + matched = [] + for device in devices: + identity = {**(device.get("identity") or {}), "device_id": device.get("device_id")} + status = device.get("status") or {} + labels = dict((device.get("capabilities") or {}).get("labels") or {}) + # Match the legacy-label defaults used by agent discovery. + if status.get("location"): + labels.setdefault("location", status["location"]) + if identity.get("device_type"): + labels.setdefault("type", identity["device_type"]) + try: + if predicate.evaluate({"identity": identity, "labels": labels, "status": status}): + matched.append(device) + except PredicateEvalError: + # Heterogeneous fleets may omit fields or use different + # types. Like broadcast, those devices do not match. + continue + devices = matched return devices def list_devices_page( @@ -280,6 +307,7 @@ def list_devices_page( *, device_type: str | None = None, location: str | None = None, + where: str | None = None, offset: int = 0, limit: int | None = None, ) -> Tuple[List[dict], int | None, int]: @@ -304,7 +332,7 @@ def list_devices_page( (devices_page, next_offset, total_matched). ``next_offset`` is None when the page reaches the end of the filtered list. ``total_matched`` is the size after the - ``device_type``/``location`` filters and before pagination. + ``device_type``/``location``/``where`` filters and before pagination. ACL filtering, when enabled at the handler layer, runs after this method returns and can further shrink the page. @@ -322,7 +350,7 @@ def list_devices_page( fleets materially larger than the current ~1400 devices. """ all_devices = self.list_devices( - tenant, device_type=device_type, location=location, + tenant, device_type=device_type, location=location, where=where, ) total = len(all_devices) safe_offset = max(0, int(offset or 0)) @@ -414,9 +442,11 @@ def list_devices( tenant: str, device_type: str | None = None, location: str | None = None, + *, + where: str | None = None, ) -> List[dict]: """Return a list of registered devices for ``tenant``, optionally filtered.""" - return _REGISTRY.list_devices(tenant, device_type=device_type, location=location) + return _REGISTRY.list_devices(tenant, device_type=device_type, location=location, where=where) def list_devices_page( @@ -424,6 +454,7 @@ def list_devices_page( *, device_type: str | None = None, location: str | None = None, + where: str | None = None, offset: int = 0, limit: int | None = None, ) -> Tuple[List[dict], int | None, int]: @@ -432,6 +463,7 @@ def list_devices_page( tenant, device_type=device_type, location=location, + where=where, offset=offset, limit=limit, ) diff --git a/packages/device-connect-server/pyproject.toml b/packages/device-connect-server/pyproject.toml index 479df0c..3edc014 100644 --- a/packages/device-connect-server/pyproject.toml +++ b/packages/device-connect-server/pyproject.toml @@ -29,7 +29,7 @@ classifiers = [ "Topic :: Software Development :: Libraries :: Python Modules", ] dependencies = [ - "device-connect-edge>=0.2.0", + "device-connect-edge[predicate]>=0.2.5", # nats-py, pydantic, nkeys, pyyaml are transitive via device-connect-edge ] diff --git a/packages/device-connect-server/tests/device_connect_server/test_discovery_where.py b/packages/device-connect-server/tests/device_connect_server/test_discovery_where.py new file mode 100644 index 0000000..233b135 --- /dev/null +++ b/packages/device-connect-server/tests/device_connect_server/test_discovery_where.py @@ -0,0 +1,117 @@ +# Copyright (c) 2024-2026, Arm Limited and Contributors. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""State discovery uses stored records and paginates only the matches.""" + +import json +from unittest.mock import AsyncMock, patch + +import pytest + +from device_connect_server.registry.service import registry +from device_connect_server.registry.service.main import _make_list_handler +from device_connect_server.security.acl import ACLManager, DeviceACL + + +def stage(index, **status): + return { + "device_id": f"stage-{index:04d}", + "identity": {"device_type": "ophyd_async:SimStage"}, + "capabilities": {"labels": {"model_id": "ophyd_async:SimStage"}, "functions": []}, + "status": {"location": "lab-A", "x_readback": float(index), "x_setpoint": 2.0, **status}, + } + + +@pytest.fixture +def fleet(monkeypatch): + records = [stage(i) for i in range(6)] + monkeypatch.setattr(registry._REGISTRY, "_decoded_fleet", lambda tenant: records) + return records + + +def test_where_filters_before_pagination(fleet): + page, cursor, total = registry.list_devices_page( + "default", where="status.x_readback >= 2.0", offset=1, limit=2, + ) + assert [d["device_id"] for d in page] == ["stage-0003", "stage-0004"] + assert (cursor, total) == (3, 4) + page, cursor, total = registry.list_devices_page( + "default", where="status.x_readback >= 2.0", offset=3, limit=2, + ) + assert [d["device_id"] for d in page] == ["stage-0005"] + assert (cursor, total) == (None, 4) + + +def test_where_binds_identity_labels_and_status_without_mutating_records(fleet): + original = json.dumps(fleet) + page = registry.list_devices( + "default", device_type="SimStage", location="lab", + where=("identity.device_id == 'stage-0002' && " + "identity.device_type == 'ophyd_async:SimStage' && " + "labels.model_id == 'ophyd_async:SimStage' && " + "labels.location == 'lab-A' && labels.type == identity.device_type && " + "status.x_setpoint == 2.0"), + ) + assert [d["device_id"] for d in page] == ["stage-0002"] + assert json.dumps(fleet) == original + + +def test_where_missing_or_incompatible_status_does_not_match(fleet): + fleet[0]["status"] = {} + fleet[1]["status"]["x_readback"] = "unknown" + page = registry.list_devices("default", where="status.x_readback < 3.0") + assert [d["device_id"] for d in page] == ["stage-0002"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("where", ["status.x_readback > > 1", "", " ", 42, False, {}]) +async def test_malformed_where_returns_invalid_params_even_for_empty_fleet(monkeypatch, where): + monkeypatch.setattr(registry._REGISTRY, "_decoded_fleet", lambda tenant: []) + messaging = AsyncMock() + request = {"id": "bad-where", "method": "discovery/listDevices", "params": {"where": where, "limit": 10}} + await _make_list_handler("default", messaging)(json.dumps(request).encode(), "reply") + response = json.loads(messaging.publish.call_args.args[1]) + assert response["id"] == "bad-where" + assert response["error"]["code"] == -32602 + assert "where" in response["error"]["message"] + + +@pytest.mark.asyncio +async def test_five_thousand_functionless_devices_return_one_small_reply(monkeypatch): + records = [stage(i, x_setpoint=2.0 if i == 51 else 0.0) for i in range(5000)] + monkeypatch.setattr(registry._REGISTRY, "_decoded_fleet", lambda tenant: records) + messaging = AsyncMock() + request = { + "id": "state-query", "method": "discovery/listDevices", + "params": {"where": "status.x_readback < 51.99 && status.x_setpoint == 2.0", "limit": 200}, + } + with patch.object(registry, "compile_where", wraps=registry.compile_where) as compile_predicate: + await _make_list_handler("default", messaging)(json.dumps(request).encode(), "reply") + compile_predicate.assert_called_once() + response_bytes = messaging.publish.call_args.args[1] + result = json.loads(response_bytes)["result"] + assert result["total_matched"] == 1 + assert result["next_offset"] is None + assert result["where_applied"] is True + assert [d["device_id"] for d in result["devices"]] == ["stage-0051"] + assert len(response_bytes) < 1000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("where,offset,expected_ids,cursor,total", [ + ("status.x_readback >= 0.0", 0, ["stage-0001", "stage-0003"], 2, 3), + ("status.x_readback >= 0.0", 2, ["stage-0005"], None, 3), + ("identity.device_id == 'stage-0000' && status.x_setpoint == 2.0", 0, [], None, 0), +]) +async def test_where_counts_and_pages_only_acl_visible_matches(fleet, where, offset, expected_ids, cursor, total): + acl = ACLManager() + for index in (0, 2, 4): + acl.set_acl(DeviceACL(device_id=f"stage-{index:04d}", tenant="default", hidden_from=["observer"])) + messaging = AsyncMock() + request = {"id": "private-state", "method": "discovery/listDevices", + "params": {"where": where, "offset": offset, "limit": 2, "requester_id": "observer"}} + await _make_list_handler("default", messaging, acl)(json.dumps(request).encode(), "reply") + result = json.loads(messaging.publish.call_args.args[1])["result"] + assert [d["device_id"] for d in result["devices"]] == expected_ids + assert (result["next_offset"], result["total_matched"]) == (cursor, total) diff --git a/packages/device-connect-server/tests/device_connect_server/test_registry_client.py b/packages/device-connect-server/tests/device_connect_server/test_registry_client.py index 151bd3f..2ee9287 100644 --- a/packages/device-connect-server/tests/device_connect_server/test_registry_client.py +++ b/packages/device-connect-server/tests/device_connect_server/test_registry_client.py @@ -88,6 +88,25 @@ async def test_context_manager(self): class TestListDevices: + @pytest.mark.asyncio + @pytest.mark.parametrize("paged", [False, True]) + async def test_where_is_forwarded_and_acknowledged(self, paged): + mc = _mock_messaging({"result": {"devices": SAMPLE_DEVICES[:1], "where_applied": True, + "next_offset": None, "total_matched": 1}}) + client = RegistryClient(mc) + await client.connect() + result = await (client.list_devices_page(where="status.online") if paged + else client.list_devices(where="status.online")) + assert (result[0] if paged else result) == SAMPLE_DEVICES[:1] + assert json.loads(mc.request.call_args.args[1])["params"]["where"] == "status.online" + + @pytest.mark.asyncio + async def test_where_rejects_an_unfiltered_legacy_response(self): + client = RegistryClient(_mock_messaging({"result": {"devices": SAMPLE_DEVICES}})) + await client.connect() + with pytest.raises(RuntimeError, match="does not support.*where"): + await client.list_devices(where="status.online") + @pytest.mark.asyncio async def test_list_all(self): mc = _mock_messaging({"result": {"devices": SAMPLE_DEVICES}}) diff --git a/packages/device-connect-server/tests/device_connect_server/test_registry_service.py b/packages/device-connect-server/tests/device_connect_server/test_registry_service.py index 2bc492e..a07baaa 100644 --- a/packages/device-connect-server/tests/device_connect_server/test_registry_service.py +++ b/packages/device-connect-server/tests/device_connect_server/test_registry_service.py @@ -789,7 +789,7 @@ def test_module_refresh(self, mock_reg): def test_module_list_devices(self, mock_reg): mock_reg.list_devices.return_value = [{"id": "cam-001"}] result = list_devices("default") - mock_reg.list_devices.assert_called_once_with("default", device_type=None, location=None) + mock_reg.list_devices.assert_called_once_with("default", device_type=None, location=None, where=None) assert result == [{"id": "cam-001"}] @patch("device_connect_server.registry.service.registry._REGISTRY") diff --git a/tests/tests/test_tools_large_fleet_broadcast.py b/tests/tests/test_tools_large_fleet_broadcast.py index 99124fb..644eaa4 100644 --- a/tests/tests/test_tools_large_fleet_broadcast.py +++ b/tests/tests/test_tools_large_fleet_broadcast.py @@ -5,6 +5,7 @@ """Slow NATS-backed large-fleet tests for broadcast reply fan-out.""" import asyncio +from functools import partialmethod import time import uuid @@ -99,7 +100,7 @@ async def test_broadcast_large_fan_out_returns_correlation_and_target_count( async def test_broadcast_where_self_election_narrows_large_candidates( - messaging_backend, messaging_url, clear_registry, device_spawner + messaging_backend, messaging_url, clear_registry, device_spawner, monkeypatch ): """A broad broadcast can be narrowed by edge-side where self-election.""" if messaging_backend != "nats": @@ -107,6 +108,23 @@ async def test_broadcast_where_self_election_narrows_large_candidates( pytest.importorskip("celpy") fleet_size = scale_fleet_size() + from device_connect_edge import DeviceRuntime + + # These devices share one interpreter/event loop, unlike a physical fleet. + # Budget for the whole batch so this self-election test does not measure + # host scheduling delays against each edge's 50 ms production deadline. + # test_device_where.py separately verifies timeouts and worker bounds. + monkeypatch.setattr( + DeviceRuntime, + "_evaluate_where_with_timeout", + partialmethod( + DeviceRuntime._evaluate_where_with_timeout, + timeout_s=min( + _reply_timeout(fleet_size), + max(1.0, fleet_size * DeviceRuntime._WHERE_EVAL_TIMEOUT_S), + ), + ), + ) prefix = f"itest-bcwhere-{uuid.uuid4().hex[:8]}" selected_location = f"{prefix}-selected" other_location = f"{prefix}-other"