refactor(entities): migrate entities consumers to typed EntitiesClient

plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/api/v2/events.py

@@ -24,6 +24,8 @@
 from nemo_iron_swarm_plugin.entities import IRON_SWARM_RUN_TYPE
 from nemo_iron_swarm_plugin.filesets import download_fileset
 from nemo_platform_plugin.authz import CallerKind, path_rule
+from nemo_platform_plugin.client.adapter import client_from_platform
+from nemo_platform_plugin.entities.client import EntitiesClient
 from pydantic import BaseModel
 from starlette.concurrency import run_in_threadpool
 
@@ -153,10 +155,14 @@ def _fileset_fallback(workspace: str, name: str, stream: Any, after: int) -> lis
         sdk = _get_sdk()
         # get_entity_by_name returns a generic Entity — its domain fields live under `.data`
         # (same access pattern as sdk.py::_run_to_dict), not as top-level attributes.
-        run = sdk.entities.get_entity_by_name(
-            name=name,
-            entity_type=IRON_SWARM_RUN_TYPE,
-            workspace=workspace,
+        run = (
+            client_from_platform(sdk, EntitiesClient)
+            .get_entity_by_name(
+                name=name,
+                entity_type=IRON_SWARM_RUN_TYPE,
+                workspace=workspace,
+            )
+            .data()
         )
         fileset_ref = (getattr(run, "data", None) or {}).get("events_fileset")
         if fileset_ref:

plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/manifest.py

@@ -19,6 +19,9 @@
 from nemo_iron_swarm_plugin.entities import IRON_SWARM_MANIFEST_TYPE
 from nemo_iron_swarm_plugin.filesets import download_and_extract_project, upload_project_dir
 from nemo_iron_swarm_plugin.jobs.errors import CATEGORY_FILESET, CATEGORY_MANIFEST, IronSwarmRunError
+from nemo_platform_plugin.client.adapter import client_from_platform
+from nemo_platform_plugin.entities.client import EntitiesClient
+from nemo_platform_plugin.entities.types import EntityUpdate
 from nemo_platform_plugin.job_context import JobContext
 
 logger = logging.getLogger(__name__)
@@ -147,8 +150,10 @@ def _materialize_manifest(
         raise IronSwarmRunError(
             CATEGORY_MANIFEST, "running a saved manifest requires the platform SDK (submit the job, don't run locally)."
         )
-    record = sdk.entities.get_entity_by_name(
-        name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace
+    record = (
+        client_from_platform(sdk, EntitiesClient)
+        .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace)
+        .data()
     )
     data = {**(getattr(record, "data", {}) or {}), **(config_overrides or {})}
     manifest_dir = ctx.storage.persistent
@@ -255,8 +260,11 @@ def _persist_upgraded_bundle(sdk: Any, manifest_id: str, ctx: JobContext, record
         updated = {**(getattr(record, "data", {}) or {})}
         updated["agent_fileset"] = fileset
         updated["manifest_yaml"] = yaml.safe_dump(resolved.manifest, sort_keys=False)
-        sdk.entities.update_entity_by_name(
-            name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace, data=updated
+        client_from_platform(sdk, EntitiesClient).update_entity_by_name(
+            name=manifest_id,
+            entity_type=IRON_SWARM_MANIFEST_TYPE,
+            workspace=ctx.workspace,
+            body=EntityUpdate(data=updated),
         )
     except Exception:  # the war-game matters more than the upgrade; it retries next run
         logger.warning("could not freeze manifest %s on this run", manifest_id, exc_info=True)

plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/records.py

@@ -21,6 +21,9 @@
 )
 from nemo_iron_swarm_plugin.jobs import benign_suite
 from nemo_iron_swarm_plugin.jobs.errors import RunFailure
+from nemo_platform_plugin.client.adapter import client_from_platform
+from nemo_platform_plugin.entities.client import EntitiesClient
+from nemo_platform_plugin.entities.types import EntityCreateInput, EntityUpdate
 from nemo_platform_plugin.entity_client import NemoEntitiesClient
 from nemo_platform_plugin.job_context import JobContext
 
@@ -74,7 +77,11 @@ def _create_run(sdk: Any, *, workspace: str, data: dict[str, Any]) -> str | None
     if sdk is None or not hasattr(sdk, "entities"):
         return None
     try:
-        entity = sdk.entities.create(IRON_SWARM_RUN_TYPE, workspace=workspace, data=data)
+        entity = (
+            client_from_platform(sdk, EntitiesClient)
+            .create_entity(entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace, body=EntityCreateInput(data=data))
+            .data()
+        )
         return getattr(entity, "name", None)
     except Exception:  # recording is best-effort, not part of the war-game
         logger.warning("failed to persist IronSwarmRun record", exc_info=True)
@@ -120,7 +127,11 @@ def _run_facts(sdk: Any, *, workspace: str, name: str) -> tuple[str, int]:
     if sdk is None or not hasattr(sdk, "entities"):
         return "", 0
     try:
-        record = sdk.entities.get_entity_by_name(name=name, entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace)
+        record = (
+            client_from_platform(sdk, EntitiesClient)
+            .get_entity_by_name(name=name, entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace)
+            .data()
+        )
         data = getattr(record, "data", {}) or {}
         port = data.get("port")
         return str(data.get("agent") or ""), int(port) if isinstance(port, int) else 0
@@ -134,7 +145,9 @@ def _update_run(sdk: Any, *, workspace: str, name: str, data: dict[str, Any]) ->
     if sdk is None or not hasattr(sdk, "entities"):
         return
     try:
-        sdk.entities.update_entity_by_name(name=name, entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace, data=data)
+        client_from_platform(sdk, EntitiesClient).update_entity_by_name(
+            name=name, entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace, body=EntityUpdate(data=data)
+        ).data()
     except Exception:  # recording is best-effort, not part of the war-game
         logger.warning("failed to update IronSwarmRun record", exc_info=True)
 
@@ -144,8 +157,10 @@ def _manifest_rounds(sdk: Any, manifest_id: str, ctx: JobContext) -> int:
     if sdk is None or not hasattr(sdk, "entities"):
         return 1
     try:
-        record = sdk.entities.get_entity_by_name(
-            name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace
+        record = (
+            client_from_platform(sdk, EntitiesClient)
+            .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace)
+            .data()
         )
         rounds = (getattr(record, "data", {}) or {}).get("rounds")
         return rounds if isinstance(rounds, int) and rounds >= 1 else 1
@@ -159,8 +174,10 @@ def _manifest_models(sdk: Any, manifest_id: str, ctx: JobContext) -> dict[str, A
     if sdk is None or not hasattr(sdk, "entities"):
         return {}
     try:
-        record = sdk.entities.get_entity_by_name(
-            name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace
+        record = (
+            client_from_platform(sdk, EntitiesClient)
+            .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace)
+            .data()
         )
         models = (getattr(record, "data", {}) or {}).get("models")
         return models if isinstance(models, dict) else {}
@@ -174,8 +191,10 @@ def _cached_benign_suite(sdk: Any, manifest_id: str, ctx: JobContext) -> list[di
     if sdk is None or not hasattr(sdk, "entities"):
         return []
     try:
-        record = sdk.entities.get_entity_by_name(
-            name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace
+        record = (
+            client_from_platform(sdk, EntitiesClient)
+            .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace)
+            .data()
         )
         suite = (getattr(record, "data", {}) or {}).get("benign_suite") or []
         return [row for row in suite if isinstance(row, dict)]
@@ -216,15 +235,17 @@ def _persist_benign_suite(
     if sdk is None or not hasattr(sdk, "entities") or not suite:
         return
     try:
-        record = sdk.entities.get_entity_by_name(
-            name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=workspace
+        record = (
+            client_from_platform(sdk, EntitiesClient)
+            .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=workspace)
+            .data()
         )
         data = dict(getattr(record, "data", {}) or {})
         data["benign_suite"] = suite
         if interview:
             data["benign_interview"] = interview
-        sdk.entities.update_entity_by_name(
-            name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=workspace, data=data
+        client_from_platform(sdk, EntitiesClient).update_entity_by_name(
+            name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=workspace, body=EntityUpdate(data=data)
         )
     except Exception:  # caching is best-effort, not part of the war-game
         logger.warning("failed to cache benign suite on manifest %s", manifest_id, exc_info=True)

plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/sdk.py

@@ -25,6 +25,9 @@
 from nemo_iron_swarm_plugin.jobs.run import IronSwarmRunJob
 from nemo_iron_swarm_plugin.jobs.synth_benign import IronSwarmSynthBenignJob
 from nemo_platform import AsyncNeMoPlatform, NeMoPlatform
+from nemo_platform_plugin.client.adapter import client_from_platform
+from nemo_platform_plugin.entities.client import EntitiesClient
+from nemo_platform_plugin.entities.types import ListEntitiesQueryParams
 from nemo_platform_plugin.scheduler import NemoJobScheduler
 from nemo_platform_plugin.sdk import NemoPluginSDKResources
 
@@ -108,8 +111,16 @@ def _list_newest(platform: NeMoPlatform, entity_type: str, *, workspace: str, li
     ``page_size`` bounds the *page*, not the total — iterating it walks the entire history. We ask for
     one page of *limit* and take only that page's items, which is a single request.
     """
-    page = platform.entities.list(entity_type, workspace=workspace, sort="-created_at", page_size=limit)
-    return [_run_to_dict(item) for item in itertools.islice(page, limit)]
+    page = (
+        client_from_platform(platform, EntitiesClient)
+        .list_entities(
+            entity_type=entity_type,
+            workspace=workspace,
+            query_params=ListEntitiesQueryParams(sort="-created_at", page_size=limit),
+        )
+        .page()
+    )
+    return [_run_to_dict(item) for item in itertools.islice(page.items, limit)]
 
 
 class _RunsResource:

plugins/nemo-iron-swarm/tests/unit/conftest.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared test doubles for the iron-swarm plugin unit tests.
+
+The records/manifest/sdk/events modules talk to the entity store through the typed
+:class:`EntitiesClient` via :func:`client_from_platform`. Unit tests fake the entity store
+as a ``SimpleNamespace(entities=...)`` shape, so ``client_from_platform`` must be patched
+at each consuming module's boundary to route the typed-client calls back onto that fake.
+This keeps the per-test ``entities`` fakes (and the capturing run-service doubles) unchanged.
+
+The typed client method calls are translated to the fake's shape here, once:
+create_entity/get_entity_by_name/update_entity_by_name/list_entities -> entities.create/
+get_entity_by_name/update_entity_by_name/list, with the request ``body``/``query_params``
+flattened back to the ``data``/``sort``/``page_size`` kwargs the fakes expect.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+from nemo_iron_swarm_plugin import sdk as sdk_module
+from nemo_iron_swarm_plugin.api.v2 import events as events_module
+from nemo_iron_swarm_plugin.jobs import manifest as manifest_module
+from nemo_iron_swarm_plugin.jobs import records as records_module
+
+
+def _data(body: Any) -> Any:
+    """Flatten the request body (EntityCreateInput/EntityUpdate) to its data dict."""
+    return body.data
+
+
+def _ok(value: Any) -> Any:
+    """Typed-client responses expose ``.data()``; wrap the fake so call sites can call it."""
+    return SimpleNamespace(data=lambda: value)
+
+
+def _page(value: Any) -> Any:
+    """Typed paginated responses expose ``.page().items``; surface the fake's iteration through it."""
+    return SimpleNamespace(page=lambda: SimpleNamespace(items=iter(value)))
+
+
+def _fake_entities_client(platform: Any, _client_cls: Any) -> Any:
+    """Build a typed-client-shaped stub that delegates to the fake ``platform.entities`` namespace."""
+    entities = platform.entities
+    return SimpleNamespace(
+        create_entity=lambda *, entity_type, workspace, body: _ok(
+            entities.create(entity_type, workspace=workspace, data=_data(body))
+        ),
+        get_entity_by_name=lambda *, name, entity_type, workspace: _ok(
+            entities.get_entity_by_name(name=name, entity_type=entity_type, workspace=workspace)
+        ),
+        update_entity_by_name=lambda *, name, entity_type, workspace, body: _ok(
+            entities.update_entity_by_name(name=name, entity_type=entity_type, workspace=workspace, data=_data(body))
+        ),
+        list_entities=lambda *, entity_type, workspace, query_params=None: _page(
+            entities.list(
+                entity_type,
+                sort=query_params.get("sort") if query_params else None,
+                workspace=workspace,
+                page_size=query_params.get("page_size") if query_params else None,
+            )
+        ),
+    )
+
+
+@pytest.fixture(autouse=True)
+def _fake_entities_client_boundary(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Route typed-client entity calls onto the fake ``entities`` namespace in every consuming module."""
+    for mod in (records_module, manifest_module, events_module):
+        monkeypatch.setattr(mod, "client_from_platform", _fake_entities_client)
+    monkeypatch.setattr(sdk_module, "client_from_platform", _fake_entities_client)

services/core/entities/tests/integration/test_workspaces_crud_with_auth.py

@@ -19,6 +19,9 @@
 
 import pytest
 from nemo_platform import ConflictError, NeMoPlatform, PermissionDeniedError
+from nemo_platform_plugin.client.adapter import client_from_platform
+from nemo_platform_plugin.entities.client import EntitiesClient
+from nemo_platform_plugin.entities.types import EntityCreateInput
 from nmp.core.entities.service import EntitiesService
 from nmp.testing import TEST_USER_EMAIL, create_test_client, short_unique_name
 
@@ -560,12 +563,14 @@ def test_delete_workspace_with_entities_marks_for_deletion(self, sdk: NeMoPlatfo
 
         # Create entity as service principal (generic entities API requires service credentials)
         with as_service(sdk, "entities"):
-            sdk.entities.create(
+            client_from_platform(sdk, EntitiesClient).create_entity(
                 workspace=workspace_name,
                 entity_type="test-entity-type",
-                name="test-entity",
-                data={"key": "value"},
-            )
+                body=EntityCreateInput(
+                    name="test-entity",
+                    data={"key": "value"},
+                ),
+            ).data()
 
         with as_user(sdk, admin_email):
             sdk.workspaces.delete(workspace_name)

tests/agentic-use/entities-basic-cli-easy/tests/test_outputs.py

@@ -14,29 +14,29 @@
 import os
 
 import pytest
-from nemo_platform import NeMoPlatform
+from nemo_platform_plugin.entities.client import EntitiesClient
 
 WORKSPACE = "default"
 
 
 @pytest.fixture
-def client() -> NeMoPlatform:
+def client() -> EntitiesClient:
     nmp_base_url = os.environ.get("NMP_BASE_URL", "http://localhost:8080")
-    return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE)
+    return EntitiesClient(base_url=nmp_base_url, workspace=WORKSPACE)
 
 
-def test_harbor_test_model_deleted(client: NeMoPlatform) -> None:
+def test_harbor_test_model_deleted(client: EntitiesClient) -> None:
     """Test that harbor-test-model was deleted after CRUD operations."""
-    response = client.entities.list(entity_type="model")
-    entity_names = [e.name for e in response.data]
+    response = client.list_entities(entity_type="model")
+    entity_names = [e.name for e in response.page().items]
     assert "harbor-test-model" not in entity_names, (
         f"Entity 'harbor-test-model' should have been deleted but still exists! Found: {entity_names}"
     )
 
 
-def test_harbor_final_dataset_exists(client: NeMoPlatform) -> None:
+def test_harbor_final_dataset_exists(client: EntitiesClient) -> None:
     """Test that harbor-final-dataset was created and has correct data."""
-    response = client.entities.get_entity_by_name(name="harbor-final-dataset", entity_type="dataset")
+    response = client.get_entity_by_name(name="harbor-final-dataset", entity_type="dataset").data()
     assert response.name == "harbor-final-dataset", (
         f"Expected entity name 'harbor-final-dataset', got '{response.name}'"
     )

tests/agentic-use/entities-basic-cli/tests/test_outputs.py

@@ -14,29 +14,29 @@
 import os
 
 import pytest
-from nemo_platform import NeMoPlatform
+from nemo_platform_plugin.entities.client import EntitiesClient
 
 WORKSPACE = "default"
 
 
 @pytest.fixture
-def client() -> NeMoPlatform:
+def client() -> EntitiesClient:
     nmp_base_url = os.environ.get("NMP_BASE_URL", "http://localhost:8080")
-    return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE)
+    return EntitiesClient(base_url=nmp_base_url, workspace=WORKSPACE)
 
 
-def test_harbor_test_model_deleted(client: NeMoPlatform) -> None:
+def test_harbor_test_model_deleted(client: EntitiesClient) -> None:
     """Test that harbor-test-model was deleted after CRUD operations."""
-    response = client.entities.list(entity_type="model")
-    entity_names = [e.name for e in response.data]
+    response = client.list_entities(entity_type="model")
+    entity_names = [e.name for e in response.page().items]
     assert "harbor-test-model" not in entity_names, (
         f"Entity 'harbor-test-model' should have been deleted but still exists! Found: {entity_names}"
     )
 
 
-def test_harbor_final_dataset_exists(client: NeMoPlatform) -> None:
+def test_harbor_final_dataset_exists(client: EntitiesClient) -> None:
     """Test that harbor-final-dataset was created and has correct data."""
-    response = client.entities.get_entity_by_name(name="harbor-final-dataset", entity_type="dataset")
+    response = client.get_entity_by_name(name="harbor-final-dataset", entity_type="dataset").data()
     assert response.name == "harbor-final-dataset", (
         f"Expected entity name 'harbor-final-dataset', got '{response.name}'"
     )