feat(rollout): remote inference draft

tests/policy_server/conftest.py

@@ -0,0 +1,275 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared fixtures for the remote-inference test suite.
+
+The mock policy is deterministic: chunk[t, j] = state[j] + 0.01 * t (so
+tests can predict exact values), accepts the RTC kwargs, and records
+every call for assertions.  Pipelines mimic the
+``PolicyProcessorPipeline`` surface the server uses (``__call__``,
+``reset``, ``steps``); the mock postprocessor doubles actions so tests
+can tell model-space from robot-space chunks.
+"""
+
+from __future__ import annotations
+
+import socket
+from dataclasses import dataclass, field
+from threading import Event
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.policy_server.manifest import ModelSpec, PolicyServerManifest, ZenohSpec
+from lerobot.policy_server.validation import PolicyClassification, ServingClass
+
+ACTION_DIM = 6
+CHUNK_SIZE = 20
+STATE_DIM = 6
+IMG_H, IMG_W = 48, 64
+ACTION_NAMES = [f"joint_{i}.pos" for i in range(ACTION_DIM)]
+TASK = "test task"
+MODEL_ID = "mock/model"
+
+
+# ---------------------------------------------------------------------------
+# Mock policy & config
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class MockPolicyConfig:
+    type: str = "mockchunk"
+    pretrained_path: str = MODEL_ID
+    chunk_size: int = CHUNK_SIZE
+    action_feature_names: list[str] = field(default_factory=lambda: list(ACTION_NAMES))
+    input_features: dict = field(
+        default_factory=lambda: {
+            "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)),
+            "observation.images.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMG_H, IMG_W)),
+        }
+    )
+    rtc_config: object | None = None
+
+
+class MockChunkPolicy:
+    """Deterministic chunk policy with the RTC kwargs surface."""
+
+    name = "mockchunk"
+
+    def __init__(self, config: MockPolicyConfig | None = None):
+        self.config = config or MockPolicyConfig()
+        self.calls: list[dict] = []
+        self.reset_count = 0
+        self.rtc_initialized = False
+
+    # nn.Module surface the server touches
+    def to(self, *args, **kwargs):
+        return self
+
+    def eval(self):
+        return self
+
+    def reset(self):
+        self.reset_count += 1
+
+    def init_rtc_processor(self):
+        self.rtc_initialized = True
+
+    def predict_action_chunk(self, batch, inference_delay=None, prev_chunk_left_over=None):
+        state = batch["observation.state"]
+        if state.ndim == 1:
+            state = state.unsqueeze(0)
+        self.calls.append(
+            {
+                "state": state.detach().clone(),
+                "inference_delay": inference_delay,
+                "prev_chunk_left_over": None
+                if prev_chunk_left_over is None
+                else prev_chunk_left_over.detach().clone(),
+                "task": batch.get("task"),
+            }
+        )
+        steps = torch.arange(CHUNK_SIZE, dtype=torch.float32).unsqueeze(1) * 0.01
+        return (state[:, :ACTION_DIM].unsqueeze(1) + steps.unsqueeze(0)).clone()
+
+
+# ---------------------------------------------------------------------------
+# Mock processor pipelines
+# ---------------------------------------------------------------------------
+
+
+class MockPipeline:
+    """Mimics the PolicyProcessorPipeline surface used by the server."""
+
+    def __init__(self, transform=None, steps=()):
+        self._transform = transform
+        self.steps = list(steps)
+        self.reset_count = 0
+        self.call_count = 0
+
+    def __call__(self, x):
+        self.call_count += 1
+        return self._transform(x) if self._transform is not None else x
+
+    def reset(self):
+        self.reset_count += 1
+
+
+def make_mock_processors():
+    """Identity preprocessor + doubling postprocessor (model vs robot space)."""
+    return MockPipeline(), MockPipeline(transform=lambda actions: actions * 2.0)
+
+
+# ---------------------------------------------------------------------------
+# Server fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def mock_policy():
+    return MockChunkPolicy()
+
+
+@pytest.fixture
+def shared_rtc_classification():
+    return PolicyClassification(
+        ServingClass.SHARED, supports_rtc=True, needs_queue_population=False, reason="mock"
+    )
+
+
+def make_manifest(**overrides) -> PolicyServerManifest:
+    kwargs = {
+        "model": ModelSpec(repo_or_path=MODEL_ID, device="cpu"),
+        "zenoh": ZenohSpec(mode="peer"),
+        "default_task": TASK,
+        "max_sessions": 4,
+        "warmup_inferences": 0,
+        "trained_fps": 30.0,
+        "health_port": 0,
+    }
+    kwargs.update(overrides)
+    return PolicyServerManifest(**kwargs)
+
+
+@pytest.fixture
+def manifest():
+    return make_manifest()
+
+
+def make_logic_server(
+    manifest: PolicyServerManifest | None = None,
+    policy: MockChunkPolicy | None = None,
+    classification: PolicyClassification | None = None,
+    processor_factory=None,
+):
+    """A PolicyServer with everything injected and no zenoh transport."""
+    from lerobot.policy_server.server import PolicyServer
+
+    policy = policy or MockChunkPolicy()
+    factory_calls = []
+
+    def default_factory():
+        pair = make_mock_processors()
+        factory_calls.append(pair)
+        return pair
+
+    server = PolicyServer(
+        manifest or make_manifest(),
+        policy=policy,
+        policy_cfg=policy.config,
+        processor_factory=processor_factory or default_factory,
+        classification=classification
+        or PolicyClassification(
+            ServingClass.SHARED, supports_rtc=True, needs_queue_population=False, reason="mock"
+        ),
+    )
+    server.load_policy()
+    server.factory_calls = factory_calls
+    return server
+
+
+# ---------------------------------------------------------------------------
+# Client-side fixtures (hw features, observations)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def hw_features():
+    return {
+        "observation.state": {
+            "dtype": "float32",
+            "shape": (STATE_DIM,),
+            "names": list(ACTION_NAMES),
+        },
+        "observation.images.front": {
+            "dtype": "video",
+            "shape": (IMG_H, IMG_W, 3),
+            "names": ["height", "width", "channels"],
+        },
+    }
+
+
+def make_robot_obs(seed: float = 1.0) -> dict:
+    obs = {name: seed + 0.1 * i for i, name in enumerate(ACTION_NAMES)}
+    rng = np.random.default_rng(int(seed * 10))
+    obs["front"] = rng.integers(0, 255, size=(IMG_H, IMG_W, 3), dtype=np.uint8)
+    return obs
+
+
+@pytest.fixture
+def shutdown_event():
+    return Event()
+
+
+# ---------------------------------------------------------------------------
+# Loopback helpers
+# ---------------------------------------------------------------------------
+
+
+def free_tcp_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def make_loopback_manifest(port: int, **overrides) -> PolicyServerManifest:
+    return make_manifest(
+        zenoh=ZenohSpec(mode="peer", listen_endpoints=[f"tcp/127.0.0.1:{port}"]),
+        **overrides,
+    )
+
+
+def make_remote_config(port: int, **overrides):
+    """RemoteInferenceConfig dialing a loopback server (fast watchdogs)."""
+    from lerobot.rollout.inference.factory import RemoteInferenceConfig
+
+    kwargs = {
+        "connect_endpoint": f"tcp/127.0.0.1:{port}",
+        "zenoh_mode": "peer",
+        "service_model_id": MODEL_ID,
+        "service_task": TASK,
+        "jpeg_quality": 0,  # raw images: byte-exact loopback
+        "buffer_time_s": 0.2,
+        "handshake_timeout_s": 2.0,
+        "request_timeout_s": 1.0,
+        "degraded_after_s": 0.3,
+        "reconnect_initial_backoff_s": 0.1,
+        "reconnect_max_backoff_s": 0.5,
+        "max_offline_s": 8.0,
+    }
+    kwargs.update(overrides)
+    return RemoteInferenceConfig(**kwargs)

tests/policy_server/test_codec.py

@@ -0,0 +1,412 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the MessagePack wire codecs (tensors, images, messages)."""
+
+import numpy as np
+import pytest
+
+msgpack = pytest.importorskip("msgpack")
+
+from lerobot.policy_server.codec import (  # noqa: E402
+    decode_action_chunk,
+    decode_image,
+    decode_observation,
+    decode_raw,
+    decode_reset,
+    decode_reset_ack,
+    decode_session_ack,
+    decode_session_close,
+    decode_session_open,
+    decode_status,
+    decode_tensor,
+    encode_action_chunk,
+    encode_image,
+    encode_observation,
+    encode_reset,
+    encode_reset_ack,
+    encode_session_ack,
+    encode_session_close,
+    encode_session_open,
+    encode_status,
+    encode_tensor,
+)
+from lerobot.policy_server.schema import (  # noqa: E402
+    IMAGE_CODEC_JPEG,
+    IMAGE_CODEC_RAW,
+    ActionChunkMsg,
+    ObservationMsg,
+    ResetAckMsg,
+    ResetMsg,
+    SessionAckMsg,
+    SessionCloseMsg,
+    SessionOpenMsg,
+    StatusMsg,
+)
+
+# ---------------------------------------------------------------------------
+# Tensor codec
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "arr",
+    [
+        np.array([1.5, -2.25, 3.0], dtype=np.float32),
+        np.arange(12, dtype=np.float64).reshape(3, 4),
+        np.array([[1, -2], [3, 4]], dtype=np.int64),
+        np.array([[True, False], [False, True]], dtype=np.bool_),
+        np.zeros((0,), dtype=np.float32),  # empty 1-d
+        np.zeros((0, 6), dtype=np.float64),  # empty 2-d
+        np.arange(24, dtype=np.int64).reshape(2, 3, 4),
+    ],
+    ids=["f32_1d", "f64_2d", "i64_2d", "bool_2d", "f32_empty", "f64_empty_2d", "i64_3d"],
+)
+def test_tensor_roundtrip(arr):
+    out = decode_tensor(encode_tensor(arr))
+    assert out.dtype == arr.dtype
+    assert out.shape == arr.shape
+    np.testing.assert_array_equal(out, arr)
+
+
+def test_tensor_roundtrip_0d_preserves_value():
+    # KNOWN QUIRK: np.ascontiguousarray inside encode_tensor promotes
+    # 0-d arrays to shape (1,), so the round-trip is value-preserving
+    # but not shape-preserving for scalars.
+    arr = np.array(3.5, dtype=np.float32)
+    out = decode_tensor(encode_tensor(arr))
+    assert out.dtype == arr.dtype
+    assert out.shape in ((), (1,))
+    assert float(np.squeeze(out)) == 3.5
+
+
+def test_tensor_none_passthrough():
+    assert encode_tensor(None) is None
+    assert decode_tensor(None) is None
+
+
+def test_tensor_big_endian_input_values_identical():
+    be = np.array([1.0, 2.5, -3.75], dtype=">f4")
+    enc = encode_tensor(be)
+    assert np.dtype(enc["dtype"]).byteorder != ">"
+    out = decode_tensor(enc)
+    np.testing.assert_array_equal(out, be.astype("<f4"))
+    np.testing.assert_array_equal(out, np.array([1.0, 2.5, -3.75], dtype=np.float32))
+
+
+def test_tensor_decoded_writable_and_contiguous():
+    arr = np.arange(6, dtype=np.float32).reshape(2, 3)
+    out = decode_tensor(encode_tensor(arr))
+    assert out.flags.writeable
+    assert out.flags.c_contiguous
+    out[0, 0] = 99.0  # must not raise
+    assert out[0, 0] == 99.0
+
+
+def test_tensor_decode_refuses_object_dtype():
+    with pytest.raises(ValueError, match="object dtype"):
+        decode_tensor({"dtype": "|O", "shape": [1], "data": b"\x00" * 8})
+
+
+def test_tensor_roundtrip_through_msgpack():
+    arr = np.arange(10, dtype=np.float32)
+    packed = msgpack.packb(encode_tensor(arr), use_bin_type=True)
+    out = decode_tensor(msgpack.unpackb(packed, raw=False))
+    np.testing.assert_array_equal(out, arr)
+
+
+# ---------------------------------------------------------------------------
+# Image codec
+# ---------------------------------------------------------------------------
+
+
+def _gradient_image(h: int = 32, w: int = 48) -> np.ndarray:
+    """Smooth RGB gradient: JPEG-friendly, deterministic."""
+    img = np.zeros((h, w, 3), dtype=np.uint8)
+    img[..., 0] = np.linspace(0, 255, w, dtype=np.uint8)[None, :]
+    img[..., 1] = np.linspace(0, 255, h, dtype=np.uint8)[:, None]
+    img[..., 2] = 128
+    return img
+
+
+def test_image_raw_roundtrip_byte_exact():
+    img = _gradient_image()
+    enc = encode_image(img, jpeg_quality=0)
+    assert enc["codec"] == IMAGE_CODEC_RAW
+    out = decode_image(enc)
+    assert out.dtype == np.uint8
+    assert out.shape == img.shape
+    np.testing.assert_array_equal(out, img)
+
+
+def test_image_jpeg_roundtrip_approximately_equal():
+    img = _gradient_image()
+    enc = encode_image(img, jpeg_quality=95)
+    assert enc["codec"] == IMAGE_CODEC_JPEG
+    out = decode_image(enc)
+    assert out.dtype == np.uint8
+    assert out.shape == img.shape
+    err = np.abs(out.astype(np.int32) - img.astype(np.int32)).mean()
+    assert err < 5.0, f"JPEG round-trip too lossy: mean abs error {err}"
+
+
+def test_image_jpeg_rgb_order_regression_pure_red_stays_red():
+    # A silent BGR swap would poison every VLA in a fleet: pure red must
+    # come back red-dominant, not blue-dominant.
+    img = np.zeros((32, 32, 3), dtype=np.uint8)
+    img[..., 0] = 255  # RGB: red channel
+    out = decode_image(encode_image(img, jpeg_quality=90))
+    red_mean = out[..., 0].astype(np.float64).mean()
+    blue_mean = out[..., 2].astype(np.float64).mean()
+    assert red_mean > 200, f"red channel lost: mean {red_mean}"
+    assert blue_mean < 50, f"blue channel gained: mean {blue_mean}"
+    assert red_mean > blue_mean
+
+
+def test_encode_image_rejects_float_arrays():
+    with pytest.raises(ValueError, match="uint8 HWC RGB"):
+        encode_image(np.zeros((8, 8, 3), dtype=np.float32))
+
+
+@pytest.mark.parametrize(
+    "shape", [(3, 16, 24), (16, 24), (16, 24, 4), (16, 24, 1)], ids=["chw", "hw", "hwc4", "hwc1"]
+)
+def test_encode_image_rejects_non_hwc(shape):
+    with pytest.raises(ValueError, match="uint8 HWC RGB"):
+        encode_image(np.zeros(shape, dtype=np.uint8))
+
+
+def test_decode_image_rejects_unknown_codec():
+    with pytest.raises(ValueError, match="Unknown image codec"):
+        decode_image({"codec": "webp", "data": b""})
+
+
+# ---------------------------------------------------------------------------
+# Data-plane messages
+# ---------------------------------------------------------------------------
+
+
+def test_observation_roundtrip_full():
+    rng = np.random.default_rng(0)
+    state = np.array([0.1, -0.2, 0.3, 0.4], dtype=np.float32)
+    front = rng.integers(0, 255, size=(16, 24, 3), dtype=np.uint8)
+    wrist = rng.integers(0, 255, size=(8, 12, 3), dtype=np.uint8)
+    prefix_model = rng.standard_normal((5, 4)).astype(np.float32)
+    prefix_robot = (prefix_model * 2.0).astype(np.float32)
+    msg = ObservationMsg(
+        state=state,
+        images={"front": front, "wrist": wrist},
+        task="fold the towel",
+        inference_delay_steps=3,
+        prefix_model=prefix_model,
+        prefix_robot=prefix_robot,
+        episode_start=True,
+        jpeg_quality=0,  # raw: byte-exact images
+    )
+    out = decode_observation(encode_observation(msg))
+    np.testing.assert_array_equal(out.state, state)
+    assert set(out.images) == {"front", "wrist"}
+    np.testing.assert_array_equal(out.images["front"], front)
+    np.testing.assert_array_equal(out.images["wrist"], wrist)
+    assert out.task == "fold the towel"
+    assert out.inference_delay_steps == 3
+    np.testing.assert_array_equal(out.prefix_model, prefix_model)
+    np.testing.assert_array_equal(out.prefix_robot, prefix_robot)
+    assert out.episode_start is True
+
+
+def test_observation_roundtrip_minimal_defaults():
+    out = decode_observation(encode_observation(ObservationMsg()))
+    assert out.state is None
+    assert out.images == {}
+    assert out.task == ""
+    assert out.inference_delay_steps == 0
+    assert out.prefix_model is None
+    assert out.prefix_robot is None
+    assert out.episode_start is False
+
+
+def test_action_chunk_roundtrip():
+    chunk_model = np.arange(12, dtype=np.float32).reshape(3, 4)
+    chunk_robot = chunk_model * 2.0
+    msg = ActionChunkMsg(
+        seq_id_echo=17,
+        client_mono_ns_echo=123456789,
+        episode_id_echo=2,
+        chunk_model=chunk_model,
+        chunk_robot=chunk_robot,
+        queue_wait_ms=1.5,
+        inference_ms=12.25,
+        superseded_seqs=4,
+        server_load=0.75,
+    )
+    out = decode_action_chunk(encode_action_chunk(msg))
+    assert out.seq_id_echo == 17
+    assert out.client_mono_ns_echo == 123456789
+    assert out.episode_id_echo == 2
+    np.testing.assert_array_equal(out.chunk_model, chunk_model)
+    np.testing.assert_array_equal(out.chunk_robot, chunk_robot)
+    assert out.queue_wait_ms == 1.5
+    assert out.inference_ms == 12.25
+    assert out.superseded_seqs == 4
+    assert out.server_load == 0.75
+
+
+# ---------------------------------------------------------------------------
+# Control-plane messages
+# ---------------------------------------------------------------------------
+
+
+def test_session_open_roundtrip():
+    msg = SessionOpenMsg(
+        client_uuid="uuid-1",
+        robot_type="so101_follower",
+        policy_type="pi0",
+        fps=30.0,
+        action_names=["j0.pos", "j1.pos"],
+        camera_names=["front", "wrist"],
+        state_dim=6,
+        rtc_enabled=True,
+        task="fold",
+        tags={"site": "lab-3"},
+    )
+    out = decode_session_open(encode_session_open(msg))
+    assert out == msg
+
+
+def test_session_ack_roundtrip():
+    msg = SessionAckMsg(
+        accepted=True,
+        warnings=["fps mismatch"],
+        session_id="sess-1",
+        model_repo="org/model",
+        model_revision="main",
+        policy_type="pi0",
+        action_names=["j0.pos"],
+        expected_cameras=["front"],
+        state_dim=6,
+        chunk_size=50,
+        trained_fps=30.0,
+        supports_rtc=True,
+        rtc_execution_horizon=25,
+        serving_mode="shared",
+        warmed_up=True,
+        server_load=0.5,
+    )
+    out = decode_session_ack(encode_session_ack(msg))
+    assert out == msg
+
+
+def test_status_roundtrip():
+    msg = StatusMsg(
+        model_repo="org/model",
+        model_revision="abc123",
+        policy_type="act",
+        action_names=["j0.pos", "j1.pos"],
+        expected_cameras=["front"],
+        state_dim=6,
+        chunk_size=100,
+        trained_fps=30.0,
+        supports_rtc=False,
+        rtc_execution_horizon=0,
+        serving_mode="exclusive",
+        warmed_up=False,
+        active_sessions=2,
+        max_sessions=4,
+    )
+    out = decode_status(encode_status(msg))
+    assert out == msg
+
+
+def test_reset_and_reset_ack_roundtrip():
+    out = decode_reset(encode_reset(ResetMsg(client_uuid="uuid-1", episode_id=5)))
+    assert out == ResetMsg(client_uuid="uuid-1", episode_id=5)
+    out_ack = decode_reset_ack(encode_reset_ack(ResetAckMsg(ok=False, error="busy")))
+    assert out_ack == ResetAckMsg(ok=False, error="busy")
+
+
+def test_session_close_roundtrip():
+    msg = SessionCloseMsg(client_uuid="uuid-1", session_id="sess-1")
+    out = decode_session_close(encode_session_close(msg))
+    assert out == msg
+
+
+# ---------------------------------------------------------------------------
+# Schema evolution (additive-only contract)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("encoded", "decoder", "expected"),
+    [
+        (
+            encode_session_ack(SessionAckMsg(accepted=True, session_id="s")),
+            decode_session_ack,
+            SessionAckMsg(accepted=True, session_id="s"),
+        ),
+        (
+            encode_reset(ResetMsg(client_uuid="u", episode_id=1)),
+            decode_reset,
+            ResetMsg(client_uuid="u", episode_id=1),
+        ),
+        (
+            encode_session_open(SessionOpenMsg(client_uuid="u")),
+            decode_session_open,
+            SessionOpenMsg(client_uuid="u"),
+        ),
+    ],
+    ids=["session_ack", "reset", "session_open"],
+)
+def test_unknown_keys_ignored(encoded, decoder, expected):
+    obj = msgpack.unpackb(encoded, raw=False)
+    obj["a_future_field"] = {"nested": [1, 2, 3]}
+    out = decoder(msgpack.packb(obj, use_bin_type=True))
+    assert out == expected
+
+
+def test_missing_optional_keys_take_defaults():
+    minimal = msgpack.packb({"accepted": True}, use_bin_type=True)
+    out = decode_session_ack(minimal)
+    assert out.accepted is True
+    assert out.error == ""
+    assert out.warnings == []
+    assert out.chunk_size == 0
+    assert out.server_load == 0.0
+
+    out_chunk = decode_action_chunk(msgpack.packb({"seq_id_echo": 9}, use_bin_type=True))
+    assert out_chunk.seq_id_echo == 9
+    assert out_chunk.chunk_model is None
+    assert out_chunk.chunk_robot is None
+    assert out_chunk.queue_wait_ms == 0.0
+
+    out_obs = decode_observation(msgpack.packb({"task": "t"}, use_bin_type=True))
+    assert out_obs.task == "t"
+    assert out_obs.state is None
+    assert out_obs.images == {}
+    assert out_obs.episode_start is False
+
+
+# ---------------------------------------------------------------------------
+# decode_raw
+# ---------------------------------------------------------------------------
+
+
+def test_decode_raw_returns_plain_dict_with_op():
+    open_obj = decode_raw(encode_session_open(SessionOpenMsg(client_uuid="u")))
+    assert isinstance(open_obj, dict)
+    assert open_obj["op"] == "open"
+
+    close_obj = decode_raw(encode_session_close(SessionCloseMsg(client_uuid="u")))
+    assert isinstance(close_obj, dict)
+    assert close_obj["op"] == "close"

tests/policy_server/test_golden_parity.py

@@ -0,0 +1,235 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Golden parity contract test: remote request path == local RTC compute path.
+
+The local side replicates exactly what ``RTCInferenceEngine._rtc_loop``
+(rtc.py) does per iteration; the remote side runs the same observation
+through the wire codec (encode -> decode), ``PolicyServer.run_inference_request``,
+and the action-chunk codec — no network, no threads.  With the same
+deterministic policy and identical inputs, both ActionQueues must stay
+byte-identical merge after merge.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+pytest.importorskip("msgpack")
+
+from lerobot.policies.rtc import ActionQueue  # noqa: E402
+from lerobot.policies.rtc.configuration_rtc import RTCConfig  # noqa: E402
+from lerobot.policies.utils import prepare_observation_for_inference  # noqa: E402
+from lerobot.policy_server import codec  # noqa: E402
+from lerobot.policy_server.schema import MsgHeader, ObservationMsg  # noqa: E402
+from lerobot.policy_server.server import _normalize_prev_actions_length  # noqa: E402
+from lerobot.policy_server.session import Session  # noqa: E402
+from lerobot.utils.constants import OBS_STATE, OBS_STR  # noqa: E402
+from lerobot.utils.feature_utils import build_dataset_frame  # noqa: E402
+from tests.policy_server.conftest import (  # noqa: E402
+    ACTION_NAMES,
+    CHUNK_SIZE,
+    STATE_DIM,
+    TASK,
+    MockChunkPolicy,
+    make_logic_server,
+    make_mock_processors,
+    make_robot_obs,
+)
+
+# Must match make_manifest()'s default RTCConfig (enabled=True, horizon=10).
+EXECUTION_HORIZON = 10
+ROBOT_TYPE = "mock_robot"
+# Fixed per-cycle inference-delay hints; cycle 2 exercises a non-zero delay.
+DELAYS = [0, 2, 1]
+# Actions consumed from both queues between cycles (makes prefixes non-trivial).
+CONSUME_K = 4
+
+STATE_ONLY_FEATURES = {
+    OBS_STATE: {
+        "dtype": "float32",
+        "shape": (STATE_DIM,),
+        "names": list(ACTION_NAMES),
+    },
+}
+
+
+def _make_queue() -> ActionQueue:
+    return ActionQueue(RTCConfig(enabled=True, execution_horizon=EXECUTION_HORIZON))
+
+
+def _local_cycle(policy, pre, post, queue, features, obs, delay) -> None:
+    """Replicates the loop body of RTCInferenceEngine._rtc_loop (rtc.py)."""
+    idx_before = queue.get_action_index()
+    prev_actions = queue.get_left_over()
+
+    obs_batch = build_dataset_frame(features, obs, prefix=OBS_STR)
+    obs_batch = prepare_observation_for_inference(obs_batch, torch.device("cpu"), TASK, ROBOT_TYPE)
+    obs_batch["task"] = [TASK]
+    preprocessed = pre(obs_batch)
+
+    if prev_actions is not None:
+        prev_actions = _normalize_prev_actions_length(prev_actions, target_steps=EXECUTION_HORIZON)
+
+    actions = policy.predict_action_chunk(
+        preprocessed, inference_delay=delay, prev_chunk_left_over=prev_actions
+    )
+    original = actions.squeeze(0).clone()
+    processed = post(actions).squeeze(0)
+    queue.merge(original, processed, delay, idx_before)
+
+
+def _remote_cycle(server, session, queue, features, obs, delay, seq_id) -> None:
+    """Replicates RemoteInferenceEngine._request_cycle with the wire codec in
+    the loop (encode -> decode on both legs) but no network or threads."""
+    idx_before = queue.get_action_index()
+
+    obs_frame = build_dataset_frame(features, obs, prefix=OBS_STR)
+    state = obs_frame.pop(OBS_STATE, None)
+    images = {k: v for k, v in obs_frame.items() if isinstance(v, np.ndarray) and v.ndim == 3}
+
+    prefix_model: np.ndarray | None = None
+    prefix_robot: np.ndarray | None = None
+    left_over = queue.get_left_over()
+    if left_over is not None and left_over.numel():
+        prefix_model = left_over[:EXECUTION_HORIZON].to(torch.float32).numpy()
+    processed_left_over = queue.get_processed_left_over()
+    if processed_left_over is not None and processed_left_over.numel():
+        prefix_robot = processed_left_over[:EXECUTION_HORIZON].to(torch.float32).numpy()
+
+    msg = ObservationMsg(
+        state=state,
+        images=images,
+        task=TASK,
+        inference_delay_steps=delay,
+        prefix_model=prefix_model,
+        prefix_robot=prefix_robot,
+        jpeg_quality=0,  # raw image codec: byte-exact transport
+    )
+    decoded = codec.decode_observation(codec.encode_observation(msg))
+
+    # The float32 wire dtype must introduce zero drift: byte-exact roundtrip.
+    assert decoded.state.dtype == np.float32
+    assert decoded.state.tobytes() == np.ascontiguousarray(state).tobytes()
+    if prefix_model is not None:
+        assert decoded.prefix_model.tobytes() == np.ascontiguousarray(prefix_model).tobytes()
+    if prefix_robot is not None:
+        assert decoded.prefix_robot.tobytes() == np.ascontiguousarray(prefix_robot).tobytes()
+    for name, img in images.items():
+        assert np.array_equal(decoded.images[name], img)
+
+    reply = server.run_inference_request(session, MsgHeader(seq_id=seq_id), decoded)
+    chunk = codec.decode_action_chunk(codec.encode_action_chunk(reply))
+
+    # Reply leg is byte-exact too (float32 in, float32 on the wire).
+    assert chunk.chunk_model.tobytes() == np.ascontiguousarray(reply.chunk_model).tobytes()
+    assert chunk.chunk_robot.tobytes() == np.ascontiguousarray(reply.chunk_robot).tobytes()
+
+    queue.merge(
+        torch.from_numpy(np.ascontiguousarray(chunk.chunk_model)),
+        torch.from_numpy(np.ascontiguousarray(chunk.chunk_robot)),
+        delay,
+        idx_before,
+    )
+
+
+def _drive_parity(features) -> tuple[MockChunkPolicy, MockChunkPolicy]:
+    """Run DELAYS cycles through both paths, asserting queue parity after
+    each merge and consuming CONSUME_K actions from both queues between
+    cycles. Returns (local_policy, remote_policy) for call-level checks."""
+    policy_local = MockChunkPolicy()
+    pre_local, post_local = make_mock_processors()
+    queue_local = _make_queue()
+
+    policy_remote = MockChunkPolicy()
+    server = make_logic_server(policy=policy_remote)
+    pre_remote, post_remote = make_mock_processors()
+    session = Session(
+        session_id="parity",
+        client_uuid="parity-client",
+        task=TASK,
+        robot_type=ROBOT_TYPE,
+        rtc_enabled=True,
+        preprocessor=pre_remote,
+        postprocessor=post_remote,
+    )
+    queue_remote = _make_queue()
+
+    for cycle, delay in enumerate(DELAYS):
+        obs = make_robot_obs(seed=float(cycle + 1))
+        _local_cycle(policy_local, pre_local, post_local, queue_local, features, obs, delay)
+        _remote_cycle(server, session, queue_remote, features, obs, delay, seq_id=cycle + 1)
+
+        assert queue_local.queue is not None and queue_remote.queue is not None
+        assert torch.equal(queue_local.queue, queue_remote.queue), (
+            f"robot-space queues diverged (cycle {cycle})"
+        )
+        assert torch.equal(queue_local.original_queue, queue_remote.original_queue), (
+            f"model-space queues diverged (cycle {cycle})"
+        )
+        assert queue_local.queue.shape == (CHUNK_SIZE - min(delay, CHUNK_SIZE), len(ACTION_NAMES))
+
+        # Consume the same k actions on both sides so the next cycle's RTC
+        # prefixes are non-trivial (and identical).
+        for _ in range(CONSUME_K):
+            action_local = queue_local.get()
+            action_remote = queue_remote.get()
+            assert action_local is not None and action_remote is not None
+            assert torch.equal(action_local, action_remote)
+
+    return policy_local, policy_remote
+
+
+def test_remote_path_matches_local_rtc_path_state_only():
+    """3 cycles, state-only features: queues stay byte-identical."""
+    _drive_parity(STATE_ONLY_FEATURES)
+
+
+def test_remote_path_matches_local_rtc_path_with_images(hw_features):
+    """Images in the loop (raw codec) must not perturb state-driven outputs."""
+    _drive_parity(hw_features)
+
+
+def test_policy_inputs_identical_across_paths(hw_features):
+    """The strongest contract: both policies saw byte-identical inputs."""
+    policy_local, policy_remote = _drive_parity(hw_features)
+
+    assert len(policy_local.calls) == len(policy_remote.calls) == len(DELAYS)
+    for i, (local_call, remote_call) in enumerate(zip(policy_local.calls, policy_remote.calls, strict=True)):
+        assert torch.equal(local_call["state"], remote_call["state"]), f"state diverged (call {i})"
+        assert local_call["state"].dtype == remote_call["state"].dtype == torch.float32
+        assert local_call["inference_delay"] == remote_call["inference_delay"] == DELAYS[i]
+        if local_call["prev_chunk_left_over"] is None:
+            assert remote_call["prev_chunk_left_over"] is None
+        else:
+            assert torch.equal(local_call["prev_chunk_left_over"], remote_call["prev_chunk_left_over"])
+            assert local_call["prev_chunk_left_over"].shape == (EXECUTION_HORIZON, len(ACTION_NAMES))
+
+    # First cycle has no leftover; later cycles must carry a real prefix.
+    assert policy_local.calls[0]["prev_chunk_left_over"] is None
+    assert all(call["prev_chunk_left_over"] is not None for call in policy_local.calls[1:])
+
+
+def test_float32_wire_dtype_is_byte_exact():
+    """Round-tripping non-dyadic float32 values through the tensor codec
+    must reproduce the exact bytes (no dtype casts, no re-quantization)."""
+    rng = np.random.default_rng(7)
+    arr = (rng.standard_normal((CHUNK_SIZE, STATE_DIM)) * 0.1).astype(np.float32)
+    decoded = codec.decode_tensor(codec.encode_tensor(arr))
+    assert decoded.dtype == np.float32
+    assert decoded.shape == arr.shape
+    assert decoded.tobytes() == arr.tobytes()
+    assert torch.equal(torch.from_numpy(decoded), torch.from_numpy(arr))

tests/policy_server/test_loopback.py

@@ -0,0 +1,366 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Real zenoh peer-to-peer loopback tests: PolicyServer ↔ RemoteInferenceEngine.
+
+The server listens on a fresh loopback TCP port per test; the engine
+dials it directly (``mode=peer``, no router).  Mock policy values are
+deterministic — chunk_robot[t, j] = 2 * (state[j] + 0.01 t) — so first
+actions identify which client's observation produced them (the
+per-session isolation regression).  Chaos tests kill/restart the server
+mid-episode and assert the engine degrades and recovers without ever
+raising on the control thread.
+"""
+
+import time
+from threading import Event
+
+import pytest
+import torch
+
+pytest.importorskip("msgpack")
+zenoh = pytest.importorskip("zenoh")
+
+from lerobot.policy_server.schema import MsgHeader, obs_key  # noqa: E402
+from lerobot.policy_server.zenoh_utils import build_zenoh_config  # noqa: E402
+from lerobot.rollout.inference.factory import FallbackMode  # noqa: E402
+from lerobot.rollout.inference.remote import ClientState, RemoteInferenceEngine  # noqa: E402
+from tests.policy_server.conftest import (  # noqa: E402
+    ACTION_DIM,
+    ACTION_NAMES,
+    TASK,
+    free_tcp_port,
+    make_logic_server,
+    make_loopback_manifest,
+    make_remote_config,
+    make_robot_obs,
+)
+
+_FPS = 30.0
+_TICK_S = 1.0 / _FPS
+# Settle time after server.start() for zenoh declarations to propagate.
+_DECLARATION_SETTLE_S = 0.5
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _wait_until(predicate, timeout_s: float, interval_s: float = 0.05) -> bool:
+    """Poll ``predicate`` until true or the deadline passes (never a fixed sleep)."""
+    deadline = time.monotonic() + timeout_s
+    while time.monotonic() < deadline:
+        if predicate():
+            return True
+        time.sleep(interval_s)
+    return bool(predicate())
+
+
+def _start_loopback_server(port: int, attempts: int = 3):
+    """Start a fully-injected PolicyServer listening on tcp/127.0.0.1:<port>."""
+    last_error: Exception | None = None
+    for _ in range(attempts):
+        server = make_logic_server(make_loopback_manifest(port))
+        try:
+            server.start()
+        except Exception as e:  # noqa: BLE001 — e.g. lingering socket on restart
+            last_error = e
+            server.stop()
+            time.sleep(0.5)
+            continue
+        time.sleep(_DECLARATION_SETTLE_S)
+        return server
+    raise last_error
+
+
+def _make_engine(port: int, server, hw_features: dict, **config_overrides) -> RemoteInferenceEngine:
+    return RemoteInferenceEngine(
+        config=make_remote_config(port, **config_overrides),
+        policy_config=server._policy_cfg,
+        hw_features=hw_features,
+        ordered_action_keys=list(ACTION_NAMES),
+        task=TASK,
+        fps=_FPS,
+        robot_type="mock",
+        shutdown_event=Event(),
+    )
+
+
+def _start_engine(engine: RemoteInferenceEngine, attempts: int = 4) -> None:
+    """Engine start with handshake retries (declarations may still be settling)."""
+    last_error: Exception | None = None
+    for _ in range(attempts):
+        try:
+            engine.start()
+            return
+        except ConnectionError as e:
+            last_error = e
+            engine.stop()
+            time.sleep(0.3)
+    raise last_error
+
+
+def _collect_actions(engine: RemoteInferenceEngine, n: int, timeout_s: float) -> list[torch.Tensor]:
+    """Poll ``get_action`` at ~30 Hz until ``n`` actions arrive or the deadline passes."""
+    actions: list[torch.Tensor] = []
+    deadline = time.monotonic() + timeout_s
+    while len(actions) < n and time.monotonic() < deadline:
+        action = engine.get_action(None)
+        if action is not None:
+            actions.append(action)
+        time.sleep(_TICK_S)
+    return actions
+
+
+# ---------------------------------------------------------------------------
+# Happy path
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.timeout(60)
+def test_end_to_end_chunks(hw_features):
+    port = free_tcp_port()
+    server = _start_loopback_server(port)
+    engine = _make_engine(port, server, hw_features)
+    try:
+        _start_engine(engine)
+        engine.notify_observation(make_robot_obs(2.0))
+
+        actions = _collect_actions(engine, n=20, timeout_s=15.0)
+        assert len(actions) >= 20
+
+        # chunk_robot[t, j] = 2 * (2.0 + 0.1 j + 0.01 t); the queue head is
+        # trimmed by the (small, loopback) delay → within 0.1 of the t=0 value.
+        first = actions[0]
+        assert first.shape == (ACTION_DIM,)
+        for j in range(ACTION_DIM):
+            expected = 2.0 * (2.0 + 0.1 * j)
+            assert abs(first[j].item() - expected) < 0.1, f"joint {j}: {first[j].item()} vs {expected}"
+
+        assert engine.state == ClientState.STREAMING
+        assert engine.ready
+        assert engine.failed is False
+        assert engine.stats["merges"] >= 1
+    finally:
+        engine.stop()
+        server.stop()
+
+
+@pytest.mark.timeout(60)
+def test_multi_client_no_cross_contamination(hw_features):
+    port = free_tcp_port()
+    server = _start_loopback_server(port)
+    engine_a = _make_engine(port, server, hw_features, client_uuid="client-a")
+    engine_b = _make_engine(port, server, hw_features, client_uuid="client-b")
+    try:
+        _start_engine(engine_a)
+        _start_engine(engine_b)
+        engine_a.notify_observation(make_robot_obs(2.0))
+        engine_b.notify_observation(make_robot_obs(7.0))
+
+        actions_a = _collect_actions(engine_a, n=1, timeout_s=10.0)
+        actions_b = _collect_actions(engine_b, n=1, timeout_s=10.0)
+        assert actions_a, "engine A produced no actions"
+        assert actions_b, "engine B produced no actions"
+
+        # Each engine's first action must reflect ITS OWN observation seed:
+        # 2*(2.0) = 4.0 for A, 2*(7.0) = 14.0 for B (gap 10.0 ≫ tolerance).
+        first_a = actions_a[0][0].item()
+        first_b = actions_b[0][0].item()
+        assert abs(first_a - 4.0) < 0.3, f"engine A got {first_a} (cross-contamination?)"
+        assert abs(first_b - 14.0) < 0.3, f"engine B got {first_b} (cross-contamination?)"
+    finally:
+        engine_a.stop()
+        engine_b.stop()
+        server.stop()
+
+
+@pytest.mark.timeout(60)
+def test_reset_roundtrip(hw_features):
+    port = free_tcp_port()
+    server = _start_loopback_server(port)
+    engine = _make_engine(port, server, hw_features)
+    try:
+        _start_engine(engine)
+        engine.notify_observation(make_robot_obs(2.0))
+        assert _collect_actions(engine, n=3, timeout_s=10.0), "no actions before reset"
+
+        merges_before = engine.stats["merges"]
+        engine.reset()
+        engine.notify_observation(make_robot_obs(5.0))
+
+        # New merges land after the reset (worker keeps cycling).
+        assert _wait_until(lambda: engine.stats["merges"] > merges_before, timeout_s=10.0)
+        # The queue refills with post-reset chunks.
+        assert _collect_actions(engine, n=1, timeout_s=5.0), "queue did not refill after reset"
+
+        # Server-side session advanced to the new episode (via the acked
+        # reset query, or via the episode bump in the next obs header).
+        def _episode_advanced() -> bool:
+            sessions = server.registry.snapshot()
+            return bool(sessions) and sessions[0].episode_id >= 1
+
+        assert _wait_until(_episode_advanced, timeout_s=8.0), "server session episode_id never advanced"
+        assert engine.failed is False
+    finally:
+        engine.stop()
+        server.stop()
+
+
+# ---------------------------------------------------------------------------
+# Chaos: server death / restart
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.timeout(60)
+def test_server_death_is_safe(hw_features):
+    port = free_tcp_port()
+    server = _start_loopback_server(port)
+    engine = _make_engine(port, server, hw_features)
+    try:
+        _start_engine(engine)
+        engine.notify_observation(make_robot_obs(2.0))
+        assert _collect_actions(engine, n=5, timeout_s=10.0), "no actions before server death"
+
+        server.stop()
+
+        # Keep ticking at 30 Hz for ~2 s: get_action must never raise.
+        results = []
+        deadline = time.monotonic() + 2.0
+        while time.monotonic() < deadline:
+            results.append(engine.get_action(None))
+            time.sleep(_TICK_S)
+
+        # The local queue drains and HOLD fallback yields None.
+        assert all(result is None for result in results[-5:]), "queue never drained to HOLD fallback"
+        # max_offline_s=8 not reached → not failed, in a degraded-but-alive state.
+        assert engine.failed is False
+        assert engine.state in {ClientState.DEGRADED, ClientState.STALLED, ClientState.RECONNECTING}
+    finally:
+        engine.stop()
+        server.stop()
+
+
+@pytest.mark.timeout(60)
+def test_server_restart_recovery(hw_features):
+    port = free_tcp_port()
+    server = _start_loopback_server(port)
+    engine = _make_engine(port, server, hw_features, max_offline_s=45.0)
+    new_server = None
+    try:
+        _start_engine(engine)
+        engine.notify_observation(make_robot_obs(2.0))
+        assert _collect_actions(engine, n=3, timeout_s=10.0), "no actions before server death"
+
+        server.stop()
+        # Let the engine notice the death (liveliness drop / request timeout).
+        _wait_until(lambda: engine.state != ClientState.STREAMING, timeout_s=5.0)
+
+        new_server = _start_loopback_server(port)
+
+        # Re-handshake: bounded by the engine backoff and zenoh's TCP
+        # reconnect period — poll generously rather than sleeping.
+        reconnected = _wait_until(lambda: engine.stats["reconnects"] >= 1, timeout_s=25.0, interval_s=0.1)
+        assert reconnected, f"engine never re-handshook (state={engine.state})"
+
+        engine.notify_observation(make_robot_obs(2.0))
+        actions = _collect_actions(engine, n=3, timeout_s=10.0)
+        assert len(actions) >= 3, "no actions after server restart"
+        assert abs(actions[-1][0].item() - 4.0) < 0.3
+        assert engine.failed is False
+    finally:
+        engine.stop()
+        server.stop()
+        if new_server is not None:
+            new_server.stop()
+
+
+# ---------------------------------------------------------------------------
+# Robustness: unknown clients, fallback modes
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.timeout(60)
+def test_unknown_client_dropped(hw_features):
+    port = free_tcp_port()
+    server = _start_loopback_server(port)
+    intruder = None
+    engine = None
+    try:
+        intruder = zenoh.open(build_zenoh_config(mode="peer", connect_endpoints=[f"tcp/127.0.0.1:{port}"]))
+        key = obs_key(server.prefix, "intruder-uuid")
+        header_bytes = MsgHeader().pack()  # valid header, garbage body, no session
+
+        deadline = time.monotonic() + 8.0
+        while server.metrics["dropped_unknown_client_total"] < 1 and time.monotonic() < deadline:
+            intruder.put(key, b"\xde\xad\xbe\xef", attachment=header_bytes)
+            time.sleep(0.1)
+        assert server.metrics["dropped_unknown_client_total"] >= 1
+        assert len(server.registry) == 0
+
+        # The server stays healthy: a legitimate engine still works.
+        engine = _make_engine(port, server, hw_features)
+        _start_engine(engine)
+        engine.notify_observation(make_robot_obs(2.0))
+        actions = _collect_actions(engine, n=1, timeout_s=10.0)
+        assert actions, "legitimate engine got no actions after garbage traffic"
+        assert abs(actions[0][0].item() - 4.0) < 0.3
+    finally:
+        if engine is not None:
+            engine.stop()
+        if intruder is not None:
+            intruder.close()
+        server.stop()
+
+
+@pytest.mark.timeout(60)
+def test_fallback_zero(hw_features):
+    port = free_tcp_port()
+    server = _start_loopback_server(port)
+    engine = _make_engine(port, server, hw_features, fallback=FallbackMode.ZERO)
+    try:
+        _start_engine(engine)
+        engine.notify_observation(make_robot_obs(2.0))
+        # With ZERO fallback an empty queue already yields zeros, so wait
+        # for a *streamed* (nonzero ~4.0) action to prove chunks flowed.
+        streamed = False
+        deadline = time.monotonic() + 10.0
+        while time.monotonic() < deadline:
+            action = engine.get_action(None)
+            if action is not None and torch.count_nonzero(action) > 0:
+                streamed = True
+                break
+            time.sleep(_TICK_S)
+        assert streamed, "no streamed (nonzero) actions before server death"
+
+        server.stop()
+
+        # Drain the local queue; once dry, ZERO fallback must return an
+        # explicit zero command (never None) of the action dimension.
+        saw_zero = False
+        deadline = time.monotonic() + 6.0
+        while time.monotonic() < deadline:
+            action = engine.get_action(None)
+            assert action is not None, "FallbackMode.ZERO returned None"
+            if torch.count_nonzero(action) == 0:
+                assert action.shape == (len(ACTION_NAMES),)
+                saw_zero = True
+                break
+            time.sleep(_TICK_S)
+        assert saw_zero, "queue never drained to the zero fallback"
+        assert engine.failed is False
+    finally:
+        engine.stop()
+        server.stop()

tests/policy_server/test_manifest.py

@@ -0,0 +1,190 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the policy-server manifest (defaults + __post_init__ validation)."""
+
+from __future__ import annotations
+
+import textwrap
+
+import draccus
+import pytest
+
+from lerobot.policy_server.manifest import (
+    SERVING_MODE_AUTO,
+    ModelSpec,
+    PolicyServerManifest,
+    ZenohSpec,
+)
+
+
+def _manifest(**overrides) -> PolicyServerManifest:
+    kwargs: dict = {"model": ModelSpec(repo_or_path="mock/model")}
+    kwargs.update(overrides)
+    return PolicyServerManifest(**kwargs)
+
+
+# ---------------------------------------------------------------------------
+# Defaults
+# ---------------------------------------------------------------------------
+
+
+def test_defaults_parse():
+    # The default zenoh mode is "client", which requires a router address;
+    # peer mode is the minimal valid transport for a defaults-only manifest.
+    manifest = _manifest(zenoh=ZenohSpec(mode="peer"))
+    assert manifest.model.repo_or_path == "mock/model"
+    assert manifest.model.revision == "main"
+    assert manifest.model.device == "cuda"
+    assert manifest.serving_mode == SERVING_MODE_AUTO
+    assert manifest.max_sessions == 5
+    assert manifest.warmup_inferences == 2
+    assert manifest.trained_fps == 30.0
+    assert manifest.strict_fps is False
+    assert manifest.pin_task is False
+    assert manifest.session_idle_timeout_s == 300.0
+    assert manifest.health_port == 9100
+    assert manifest.debug.capture_dir is None
+    assert manifest.rtc.enabled is True
+    # Bare ZenohSpec defaults (validated only when embedded in a manifest).
+    assert ZenohSpec().mode == "client"
+    assert ZenohSpec().connect_endpoints == []
+
+
+# ---------------------------------------------------------------------------
+# __post_init__ validation
+# ---------------------------------------------------------------------------
+
+
+def test_missing_repo_or_path_raises():
+    with pytest.raises(ValueError, match="repo_or_path"):
+        PolicyServerManifest()
+
+
+def test_bad_serving_mode_raises():
+    with pytest.raises(ValueError, match="serving_mode"):
+        _manifest(serving_mode="multiplexed")
+
+
+@pytest.mark.parametrize("max_sessions", [0, -3])
+def test_max_sessions_below_one_raises(max_sessions):
+    with pytest.raises(ValueError, match="max_sessions"):
+        _manifest(max_sessions=max_sessions)
+
+
+def test_zenoh_client_mode_requires_connect_endpoints():
+    with pytest.raises(ValueError, match="connect_endpoints"):
+        _manifest(zenoh=ZenohSpec(mode="client", connect_endpoints=[]))
+
+
+def test_zenoh_client_mode_with_endpoints_ok():
+    manifest = _manifest(zenoh=ZenohSpec(mode="client", connect_endpoints=["tcp/router:7447"]))
+    assert manifest.zenoh.connect_endpoints == ["tcp/router:7447"]
+
+
+def test_zenoh_peer_mode_without_endpoints_ok():
+    manifest = _manifest(zenoh=ZenohSpec(mode="peer"))
+    assert manifest.zenoh.mode == "peer"
+    assert manifest.zenoh.connect_endpoints == []
+
+
+def test_partial_tls_triple_raises():
+    with pytest.raises(ValueError, match="TLS"):
+        _manifest(zenoh=ZenohSpec(mode="peer", tls_root_ca_certificate="/certs/ca.pem"))
+
+
+def test_full_tls_triple_ok():
+    manifest = _manifest(
+        zenoh=ZenohSpec(
+            mode="peer",
+            tls_root_ca_certificate="/certs/ca.pem",
+            tls_connect_certificate="/certs/cert.pem",
+            tls_connect_private_key="/certs/key.pem",
+        )
+    )
+    assert manifest.zenoh.tls_connect_private_key == "/certs/key.pem"
+
+
+# ---------------------------------------------------------------------------
+# Draccus round-trip (YAML manifest → dataclass)
+# ---------------------------------------------------------------------------
+
+
+def test_draccus_yaml_round_trip(tmp_path):
+    yaml_path = tmp_path / "server.yaml"
+    yaml_path.write_text(
+        textwrap.dedent(
+            """\
+            model:
+              repo_or_path: mock/model
+              revision: v2.0
+              device: cpu
+            zenoh:
+              mode: peer
+              listen_endpoints:
+                - tcp/127.0.0.1:7447
+            default_task: pick the cube
+            pin_task: true
+            serving_mode: exclusive
+            max_sessions: 1
+            trained_fps: 25.0
+            strict_fps: true
+            health_port: 0
+            """
+        )
+    )
+
+    manifest = draccus.parse(PolicyServerManifest, config_path=str(yaml_path), args=[])
+
+    assert isinstance(manifest, PolicyServerManifest)
+    assert manifest.model.repo_or_path == "mock/model"
+    assert manifest.model.revision == "v2.0"
+    assert manifest.model.device == "cpu"
+    assert manifest.zenoh.mode == "peer"
+    assert manifest.zenoh.listen_endpoints == ["tcp/127.0.0.1:7447"]
+    assert manifest.default_task == "pick the cube"
+    assert manifest.pin_task is True
+    assert manifest.serving_mode == "exclusive"
+    assert manifest.max_sessions == 1
+    assert manifest.trained_fps == 25.0
+    assert manifest.strict_fps is True
+    assert manifest.health_port == 0
+    # Untouched fields keep their defaults.
+    assert manifest.warmup_inferences == 2
+    assert manifest.session_idle_timeout_s == 300.0
+
+
+def test_draccus_cli_override_on_top_of_yaml(tmp_path):
+    yaml_path = tmp_path / "server.yaml"
+    yaml_path.write_text(
+        textwrap.dedent(
+            """\
+            model:
+              repo_or_path: mock/model
+              device: cpu
+            zenoh:
+              mode: peer
+            """
+        )
+    )
+
+    manifest = draccus.parse(
+        PolicyServerManifest,
+        config_path=str(yaml_path),
+        args=["--max_sessions", "3", "--model.revision", "exp-1"],
+    )
+
+    assert manifest.max_sessions == 3
+    assert manifest.model.revision == "exp-1"
+    assert manifest.model.repo_or_path == "mock/model"

tests/policy_server/test_scheduler.py

@@ -0,0 +1,117 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the RoundRobinScheduler fairness guarantees.
+
+The scheduler only reads ``client_uuid`` from sessions, so minimal fakes
+stand in for real Session objects.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+
+from lerobot.policy_server.scheduler import RoundRobinScheduler
+
+
+class FakeSession:
+    """The scheduler touches nothing but client_uuid."""
+
+    def __init__(self, client_uuid: str):
+        self.client_uuid = client_uuid
+
+
+def picks(scheduler: RoundRobinScheduler, ready: list[FakeSession], n: int) -> list[str]:
+    served = []
+    for _ in range(n):
+        chosen = scheduler.select(list(ready))
+        assert len(chosen) == 1
+        served.append(chosen[0].client_uuid)
+    return served
+
+
+def test_empty_ready_returns_empty():
+    scheduler = RoundRobinScheduler()
+    assert scheduler.select([]) == []
+
+
+def test_empty_ready_after_serving_returns_empty():
+    scheduler = RoundRobinScheduler()
+    scheduler.select([FakeSession("a")])
+    assert scheduler.select([]) == []
+
+
+def test_single_session_picked_repeatedly():
+    scheduler = RoundRobinScheduler()
+    only = FakeSession("solo")
+    for _ in range(5):
+        assert scheduler.select([only]) == [only]
+
+
+def test_three_sessions_served_fairly_in_sorted_uuid_order():
+    scheduler = RoundRobinScheduler()
+    a, b, c = FakeSession("a"), FakeSession("b"), FakeSession("c")
+    # Pass ready in non-sorted order: the ring is sorted by uuid internally.
+    served = picks(scheduler, [c, a, b], 9)
+
+    assert served == ["a", "b", "c", "a", "b", "c", "a", "b", "c"]
+    assert Counter(served) == {"a": 3, "b": 3, "c": 3}
+
+
+def test_session_leaving_between_calls_keeps_fairness():
+    scheduler = RoundRobinScheduler()
+    a, b, c = FakeSession("a"), FakeSession("b"), FakeSession("c")
+    assert scheduler.select([a, b, c]) == [a]
+
+    # 'a' leaves; remaining sessions alternate with no crash or starvation.
+    served = picks(scheduler, [b, c], 4)
+    assert served == ["b", "c", "b", "c"]
+
+
+def test_departed_last_served_uuid_resumes_after_it():
+    scheduler = RoundRobinScheduler()
+    a, b, c = FakeSession("a"), FakeSession("b"), FakeSession("c")
+    picks(scheduler, [a, b, c], 2)  # last served is 'b'
+
+    # 'b' leaves; the next pick is the first uuid greater than 'b'.
+    assert scheduler.select([a, c]) == [c]
+    assert scheduler.select([a, c]) == [a]
+
+
+def test_wraparound_from_last_uuid_back_to_first():
+    scheduler = RoundRobinScheduler()
+    a, b, c = FakeSession("a"), FakeSession("b"), FakeSession("c")
+    assert scheduler.select([c]) == [c]  # last served is the highest uuid
+
+    # Everyone is <= last served: wrap back to the first sorted uuid.
+    assert scheduler.select([a, b, c]) == [a]
+
+
+def test_newly_ready_session_joins_ring_fairly():
+    scheduler = RoundRobinScheduler()
+    a, c = FakeSession("a"), FakeSession("c")
+    served = picks(scheduler, [a, c], 2)
+    assert served == ["a", "c"]
+
+    # 'b' becomes ready; wrap-around lands on 'a', then 'b' gets its turn.
+    b = FakeSession("b")
+    served = picks(scheduler, [a, b, c], 3)
+    assert served == ["a", "b", "c"]
+
+
+def test_no_starvation_over_many_cycles():
+    scheduler = RoundRobinScheduler()
+    ready = [FakeSession(f"u{i:02d}") for i in range(5)]
+    served = picks(scheduler, ready, 50)
+    assert Counter(served) == {f"u{i:02d}": 10 for i in range(5)}

tests/policy_server/test_schema.py

@@ -0,0 +1,197 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the wire schema: packed header and key-expression layout."""
+
+import pytest
+
+from lerobot.policy_server.schema import (
+    HEADER_SIZE,
+    MSG_TYPE_CHUNK,
+    MSG_TYPE_EVENT,
+    MSG_TYPE_OBS,
+    RESERVED_SEGMENTS,
+    SCHEMA_VERSION,
+    MsgHeader,
+    action_key,
+    client_alive_key,
+    client_alive_wildcard,
+    client_uuid_from_key,
+    obs_key,
+    obs_wildcard,
+    reset_key,
+    reset_wildcard,
+    sanitize_key_segment,
+    server_alive_key,
+    service_prefix,
+    session_key,
+    status_key,
+)
+
+# ---------------------------------------------------------------------------
+# MsgHeader
+# ---------------------------------------------------------------------------
+
+
+def test_header_roundtrip_all_fields():
+    hdr = MsgHeader(
+        schema_version=3,
+        msg_type=MSG_TYPE_CHUNK,
+        seq_id=123456789,
+        episode_id=42,
+        client_mono_ns=987654321012345,
+        session_epoch=7,
+    )
+    out = MsgHeader.unpack(hdr.pack())
+    assert out == hdr
+
+
+def test_header_defaults_roundtrip():
+    out = MsgHeader.unpack(MsgHeader().pack())
+    assert out.schema_version == SCHEMA_VERSION
+    assert out.msg_type == MSG_TYPE_OBS
+    assert out.seq_id == 0
+    assert out.episode_id == 0
+    assert out.client_mono_ns == 0
+    assert out.session_epoch == 0
+
+
+def test_header_negative_client_mono_ns():
+    hdr = MsgHeader(msg_type=MSG_TYPE_EVENT, client_mono_ns=-123456789)
+    out = MsgHeader.unpack(hdr.pack())
+    assert out.client_mono_ns == -123456789
+
+
+def test_header_max_u64_seq_id():
+    max_u64 = 2**64 - 1
+    hdr = MsgHeader(seq_id=max_u64)
+    out = MsgHeader.unpack(hdr.pack())
+    assert out.seq_id == max_u64
+
+
+def test_header_size_constant_matches_pack():
+    assert len(MsgHeader().pack()) == HEADER_SIZE
+
+
+def test_header_unpack_rejects_wrong_length():
+    packed = MsgHeader().pack()
+    with pytest.raises(ValueError, match="Bad header length"):
+        MsgHeader.unpack(packed[:-1])
+    with pytest.raises(ValueError, match="Bad header length"):
+        MsgHeader.unpack(packed + b"\x00")
+    with pytest.raises(ValueError, match="Bad header length"):
+        MsgHeader.unpack(b"")
+
+
+# ---------------------------------------------------------------------------
+# sanitize_key_segment
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("bad_char", ["/", "*", "$", "?", "#", " "])
+def test_sanitize_folds_unsafe_chars_to_dash(bad_char):
+    assert sanitize_key_segment(f"a{bad_char}b") == "a-b"
+
+
+def test_sanitize_folds_runs_to_single_dash():
+    assert sanitize_key_segment("a/*$?# b") == "a-b"
+
+
+def test_sanitize_strips_whitespace_and_edge_dashes():
+    assert sanitize_key_segment("  hello  ") == "hello"
+    assert sanitize_key_segment("/leading/trailing/") == "leading-trailing"
+
+
+def test_sanitize_preserves_allowed_chars():
+    assert sanitize_key_segment("Az09_.-x") == "Az09_.-x"
+
+
+@pytest.mark.parametrize("empty_input", ["", "   ", "***", "/?#$*"])
+def test_sanitize_raises_on_empty_after_sanitize(empty_input):
+    with pytest.raises(ValueError, match="empty"):
+        sanitize_key_segment(empty_input)
+
+
+@pytest.mark.parametrize(
+    "reserved", sorted(["status", "session", "server", "obs", "action", "reset", "alive"])
+)
+def test_sanitize_raises_on_reserved_segments(reserved):
+    assert reserved in RESERVED_SEGMENTS
+    with pytest.raises(ValueError, match="reserved"):
+        sanitize_key_segment(reserved)
+
+
+# ---------------------------------------------------------------------------
+# service_prefix
+# ---------------------------------------------------------------------------
+
+
+def test_service_prefix_example():
+    prefix = service_prefix("lerobot/pi0_towels", "main", "fold the towel")
+    assert prefix == "@lerobot/lerobot-pi0_towels/main/fold-the-towel"
+
+
+def test_service_prefix_defaults_for_empty_revision_and_task():
+    prefix = service_prefix("org/model", "", "")
+    assert prefix == "@lerobot/org-model/main/default"
+
+
+# ---------------------------------------------------------------------------
+# Key builders and wildcards
+# ---------------------------------------------------------------------------
+
+PREFIX = "@lerobot/org-model/main/default"
+UUID = "client-uuid-1234"
+
+
+def test_per_client_keys():
+    assert obs_key(PREFIX, UUID) == f"{PREFIX}/{UUID}/obs"
+    assert action_key(PREFIX, UUID) == f"{PREFIX}/{UUID}/action"
+    assert reset_key(PREFIX, UUID) == f"{PREFIX}/{UUID}/reset"
+    assert client_alive_key(PREFIX, UUID) == f"{PREFIX}/{UUID}/alive"
+
+
+def test_service_level_keys():
+    assert status_key(PREFIX) == f"{PREFIX}/status"
+    assert session_key(PREFIX) == f"{PREFIX}/session"
+    assert server_alive_key(PREFIX) == f"{PREFIX}/server/alive"
+
+
+def test_wildcards_are_single_depth():
+    assert obs_wildcard(PREFIX) == f"{PREFIX}/*/obs"
+    assert reset_wildcard(PREFIX) == f"{PREFIX}/*/reset"
+    assert client_alive_wildcard(PREFIX) == f"{PREFIX}/*/alive"
+    assert "**" not in obs_wildcard(PREFIX)
+
+
+def test_key_builders_sanitize_client_uuid():
+    assert obs_key(PREFIX, "bad uuid") == f"{PREFIX}/bad-uuid/obs"
+    with pytest.raises(ValueError):
+        obs_key(PREFIX, "status")
+
+
+# ---------------------------------------------------------------------------
+# client_uuid_from_key
+# ---------------------------------------------------------------------------
+
+
+def test_client_uuid_from_obs_reset_alive_keys():
+    assert client_uuid_from_key(obs_key(PREFIX, UUID)) == UUID
+    assert client_uuid_from_key(reset_key(PREFIX, UUID)) == UUID
+    assert client_uuid_from_key(client_alive_key(PREFIX, UUID)) == UUID
+
+
+def test_client_uuid_from_key_rejects_keys_without_client_chunk():
+    with pytest.raises(ValueError, match="no client chunk"):
+        client_uuid_from_key("obs")

tests/policy_server/test_server_logic.py

@@ -0,0 +1,405 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Pure-logic PolicyServer tests (no zenoh transport).
+
+Covers status snapshots, session open/reject/re-handshake, the
+per-request inference path (determinism, RTC forwarding, echo fields,
+supersession), episode boundaries in ``_serve_one``, warmup, and the
+error/metrics accounting.
+"""
+
+import numpy as np
+import pytest
+
+pytest.importorskip("msgpack")
+
+from lerobot.policy_server import codec  # noqa: E402
+from lerobot.policy_server.schema import MsgHeader, ObservationMsg, SessionOpenMsg  # noqa: E402
+from lerobot.policy_server.validation import PolicyClassification, ServingClass  # noqa: E402
+from tests.policy_server.conftest import (  # noqa: E402
+    ACTION_DIM,
+    ACTION_NAMES,
+    CHUNK_SIZE,
+    IMG_H,
+    IMG_W,
+    MODEL_ID,
+    STATE_DIM,
+    TASK,
+    MockChunkPolicy,
+    make_logic_server,
+    make_manifest,
+)
+
+CAMERA_KEY = "observation.images.front"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def make_open_msg(client_uuid: str = "client-a", **overrides) -> SessionOpenMsg:
+    kwargs = {
+        "client_uuid": client_uuid,
+        "robot_type": "so101",
+        "policy_type": "mockchunk",
+        "fps": 30.0,
+        "action_names": list(ACTION_NAMES),
+        "camera_names": [CAMERA_KEY],
+        "state_dim": STATE_DIM,
+        "rtc_enabled": True,
+        "task": TASK,
+    }
+    kwargs.update(overrides)
+    return SessionOpenMsg(**kwargs)
+
+
+def make_obs(**overrides) -> ObservationMsg:
+    kwargs = {
+        "state": np.arange(STATE_DIM, dtype=np.float32),
+        "images": {CAMERA_KEY: np.zeros((IMG_H, IMG_W, 3), dtype=np.uint8)},
+        "task": TASK,
+        "jpeg_quality": 0,
+    }
+    kwargs.update(overrides)
+    return ObservationMsg(**kwargs)
+
+
+def open_session(server, client_uuid: str = "client-a", **overrides):
+    ack = server._handle_session_open(make_open_msg(client_uuid=client_uuid, **overrides))
+    assert ack.accepted, ack.error
+    return server.registry.get(client_uuid), ack
+
+
+def deposit(session, obs: ObservationMsg, header: MsgHeader | None = None) -> None:
+    session.deposit(header or MsgHeader(episode_id=session.episode_id), codec.encode_observation(obs))
+
+
+def make_exclusive_server(policy=None):
+    return make_logic_server(
+        manifest=make_manifest(serving_mode="exclusive"),
+        policy=policy,
+        classification=PolicyClassification(
+            ServingClass.EXCLUSIVE, supports_rtc=False, needs_queue_population=False, reason="x"
+        ),
+    )
+
+
+def expected_chunk(state: np.ndarray) -> np.ndarray:
+    steps = np.arange(CHUNK_SIZE, dtype=np.float32)[:, None] * np.float32(0.01)
+    return state[None, :ACTION_DIM] + steps
+
+
+# ---------------------------------------------------------------------------
+# status_snapshot
+# ---------------------------------------------------------------------------
+
+
+def test_status_snapshot_capabilities():
+    server = make_logic_server()
+    snap = server.status_snapshot()
+    assert snap.model_repo == MODEL_ID
+    assert snap.policy_type == "mockchunk"
+    assert snap.action_names == ACTION_NAMES
+    assert snap.state_dim == STATE_DIM
+    assert snap.chunk_size == CHUNK_SIZE
+    assert snap.expected_cameras == [CAMERA_KEY]
+    assert snap.supports_rtc is True
+    assert snap.warmed_up is True
+    assert snap.serving_mode == "shared"
+    assert snap.active_sessions == 0
+    assert snap.max_sessions == 4  # make_manifest default
+
+
+# ---------------------------------------------------------------------------
+# _handle_session_open
+# ---------------------------------------------------------------------------
+
+
+def test_session_open_happy_path():
+    server = make_logic_server()
+    session, ack = open_session(server, "client-a")
+    assert ack.accepted is True
+    assert ack.error == ""
+    assert ack.session_id == session.session_id
+    assert ack.session_id != ""
+    assert ack.model_repo == MODEL_ID
+    assert ack.policy_type == "mockchunk"
+    assert ack.action_names == ACTION_NAMES
+    assert ack.expected_cameras == [CAMERA_KEY]
+    assert ack.state_dim == STATE_DIM
+    assert ack.chunk_size == CHUNK_SIZE
+    assert ack.supports_rtc is True
+    assert ack.serving_mode == "shared"
+    assert ack.warmed_up is True
+    assert len(server.registry) == 1
+    assert session.rtc_enabled is True
+    assert session.task == TASK
+    assert server.metrics["sessions_opened_total"] == 1
+
+
+def test_session_open_fresh_processor_pair_per_session():
+    server = make_logic_server()
+    assert len(server.factory_calls) == 0  # warmup_inferences=0: no warmup pair
+    session_a, _ = open_session(server, "client-a")
+    assert len(server.factory_calls) == 1
+    session_b, _ = open_session(server, "client-b")
+    assert len(server.factory_calls) == 2
+    assert session_a.preprocessor is not session_b.preprocessor
+    assert session_a.postprocessor is not session_b.postprocessor
+
+
+def test_session_open_rejects_action_order_mismatch():
+    server = make_logic_server()
+    ack = server._handle_session_open(make_open_msg(action_names=list(reversed(ACTION_NAMES))))
+    assert ack.accepted is False
+    assert "action" in ack.error
+    assert len(server.registry) == 0
+
+
+def test_session_open_rejects_at_capacity():
+    server = make_logic_server()
+    for i in range(4):  # make_manifest max_sessions=4
+        open_session(server, f"client-{i}")
+    ack = server._handle_session_open(make_open_msg(client_uuid="client-overflow"))
+    assert ack.accepted is False
+    assert "sessions" in ack.error
+    assert len(server.registry) == 4
+
+
+def test_rehandshake_replaces_session_without_counting_against_capacity():
+    server = make_logic_server()
+    first_ack = None
+    for i in range(4):
+        _, ack = open_session(server, f"client-{i}")
+        if i == 0:
+            first_ack = ack
+    # Server is full; the same client re-handshakes and must be accepted.
+    session, ack = open_session(server, "client-0")
+    assert ack.accepted is True
+    assert len(server.registry) == 4
+    assert ack.session_id != first_ack.session_id
+    assert server.registry.get("client-0").session_id == session.session_id
+
+
+def test_session_open_rtc_downgrade():
+    server = make_logic_server(
+        classification=PolicyClassification(
+            ServingClass.SHARED, supports_rtc=False, needs_queue_population=False, reason="x"
+        )
+    )
+    session, ack = open_session(server, "client-a", rtc_enabled=True)
+    assert ack.accepted is True
+    assert ack.supports_rtc is False
+    assert session.rtc_enabled is False
+    assert any("RTC" in w for w in ack.warnings)
+
+
+# ---------------------------------------------------------------------------
+# run_inference_request
+# ---------------------------------------------------------------------------
+
+
+def test_inference_deterministic_chunks():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    state = np.arange(STATE_DIM, dtype=np.float32)
+    reply = server.run_inference_request(session, MsgHeader(), make_obs(state=state))
+    assert reply.chunk_model.shape == (CHUNK_SIZE, ACTION_DIM)
+    assert reply.chunk_robot.shape == (CHUNK_SIZE, ACTION_DIM)
+    np.testing.assert_allclose(reply.chunk_model[0], state, rtol=0, atol=0)
+    np.testing.assert_allclose(reply.chunk_model, expected_chunk(state), rtol=1e-6)
+    np.testing.assert_allclose(reply.chunk_robot, 2.0 * reply.chunk_model, rtol=0, atol=0)
+
+
+def test_inference_delay_forwarded_to_policy():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    policy = server._policy
+    server.run_inference_request(session, MsgHeader(), make_obs(inference_delay_steps=3))
+    assert policy.calls[-1]["inference_delay"] == 3
+    assert policy.calls[-1]["prev_chunk_left_over"] is None
+
+
+def test_prefix_model_forwarded_padded_to_execution_horizon():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    policy = server._policy
+    prefix = (np.arange(3 * ACTION_DIM, dtype=np.float32).reshape(3, ACTION_DIM)) + 1.0
+    server.run_inference_request(session, MsgHeader(), make_obs(prefix_model=prefix))
+    received = policy.calls[-1]["prev_chunk_left_over"]
+    assert received is not None
+    horizon = server._manifest.rtc.execution_horizon  # 10 by default
+    assert received.shape == (horizon, ACTION_DIM)
+    np.testing.assert_allclose(received[:3].numpy(), prefix, rtol=0, atol=0)
+    np.testing.assert_allclose(received[3:].numpy(), np.zeros((horizon - 3, ACTION_DIM)), rtol=0, atol=0)
+
+
+def test_prefix_model_truncated_to_execution_horizon():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    policy = server._policy
+    horizon = server._manifest.rtc.execution_horizon
+    prefix = np.ones((horizon + 5, ACTION_DIM), dtype=np.float32)
+    server.run_inference_request(session, MsgHeader(), make_obs(prefix_model=prefix))
+    assert policy.calls[-1]["prev_chunk_left_over"].shape == (horizon, ACTION_DIM)
+
+
+def test_reply_echo_fields():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    header = MsgHeader(seq_id=7, episode_id=2, client_mono_ns=123_456_789)
+    reply = server.run_inference_request(session, header, make_obs())
+    assert reply.seq_id_echo == 7
+    assert reply.episode_id_echo == 2
+    assert reply.client_mono_ns_echo == 123_456_789
+
+
+def test_superseded_seqs_reported_then_reset():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    deposit(session, make_obs(), MsgHeader(seq_id=1))
+    deposit(session, make_obs(), MsgHeader(seq_id=2))  # supersedes seq 1
+    item = session.take()
+    assert item.header.seq_id == 2  # latest-only mailbox
+    reply = server.run_inference_request(session, item.header, codec.decode_observation(item.payload))
+    assert reply.superseded_seqs == 1
+    deposit(session, make_obs(), MsgHeader(seq_id=3))
+    item = session.take()
+    reply = server.run_inference_request(session, item.header, codec.decode_observation(item.payload))
+    assert reply.superseded_seqs == 0
+
+
+# ---------------------------------------------------------------------------
+# _serve_one
+# ---------------------------------------------------------------------------
+
+
+def test_serve_one_episode_boundary_resets_session_pipelines():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    # Fresh sessions start at the -1 sentinel so their first request
+    # always lands on the episode-boundary branch (mid-episode reconnects
+    # can never inherit stale state).
+    assert session.episode_id == -1
+    deposit(session, make_obs(episode_start=True), MsgHeader(episode_id=1))
+    server._serve_one(session)
+    assert session.preprocessor.reset_count == 1
+    assert session.postprocessor.reset_count == 1
+    assert session.episode_id == 1
+    # Shared mode never resets the policy itself.
+    assert server._policy.reset_count == 0
+
+
+def test_serve_one_no_boundary_no_reset():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    # First request always resets (the -1 sentinel) and syncs the episode.
+    deposit(session, make_obs(), MsgHeader(episode_id=0))
+    server._serve_one(session)
+    assert session.preprocessor.reset_count == 1
+    assert session.episode_id == 0
+    # Same-episode follow-up: no further reset.
+    deposit(session, make_obs(), MsgHeader(episode_id=0, seq_id=2))
+    server._serve_one(session)
+    assert session.preprocessor.reset_count == 1
+    assert session.postprocessor.reset_count == 1
+
+
+def test_serve_one_exclusive_mode_resets_policy_on_boundary():
+    policy = MockChunkPolicy()
+    server = make_exclusive_server(policy=policy)
+    assert server.status_snapshot().serving_mode == "exclusive"
+    assert server.status_snapshot().max_sessions == 1  # exclusive forces 1
+    session, _ = open_session(server, rtc_enabled=False)
+    # Exclusive session open already resets the policy to fresh state.
+    base_resets = policy.reset_count
+    assert base_resets >= 1
+    deposit(session, make_obs(episode_start=True), MsgHeader(episode_id=1))
+    server._serve_one(session)
+    assert policy.reset_count == base_resets + 1
+    assert session.episode_id == 1
+
+
+def test_serve_one_inference_error_counted_not_propagated(monkeypatch):
+    server = make_logic_server()
+    session, _ = open_session(server)
+
+    def boom(*args, **kwargs):
+        raise RuntimeError("boom")
+
+    monkeypatch.setattr(server._policy, "predict_action_chunk", boom)
+    deposit(session, make_obs())
+    server._serve_one(session)  # must not raise
+    assert server.metrics["errors_total"] == 1
+    assert server.metrics["requests_total"] == 0
+    assert session.stats.errors == 1
+
+
+def test_serve_one_increments_requests_total():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    for seq in (1, 2):
+        deposit(session, make_obs(), MsgHeader(seq_id=seq))
+        server._serve_one(session)
+    assert server.metrics["requests_total"] == 2
+    assert server.metrics["errors_total"] == 0
+    assert session.stats.requests == 2
+
+
+# ---------------------------------------------------------------------------
+# Episode reset semantics (session-level, as used by _on_reset_query)
+# ---------------------------------------------------------------------------
+
+
+def test_session_reset_episode_clears_state():
+    server = make_logic_server()
+    session, _ = open_session(server)
+    deposit(session, make_obs())
+    assert session.has_pending()
+    session.reset_episode(5)
+    assert not session.has_pending()  # mailbox cleared
+    assert session.episode_id == 5
+    assert session.preprocessor.reset_count == 1
+    assert session.postprocessor.reset_count == 1
+    session.reset_episode()  # no explicit id: increments
+    assert session.episode_id == 6
+
+
+# ---------------------------------------------------------------------------
+# Warmup
+# ---------------------------------------------------------------------------
+
+
+def test_warmup_runs_before_any_session():
+    policy = MockChunkPolicy()
+    server = make_logic_server(manifest=make_manifest(warmup_inferences=2), policy=policy)
+    assert len(policy.calls) == 2
+    assert len(server.registry) == 0  # warmup session is not registered
+    for call in policy.calls:
+        assert tuple(call["state"].shape) == (1, STATE_DIM)
+        assert float(call["state"].abs().sum()) == 0.0  # synthetic zeros
+    assert server.status_snapshot().warmed_up is True
+
+
+def test_synthetic_observation_matches_input_features():
+    server = make_logic_server()
+    obs = server._synthetic_observation()
+    assert obs.state.shape == (STATE_DIM,)
+    assert obs.state.dtype == np.float32
+    assert set(obs.images) == {CAMERA_KEY}
+    assert obs.images[CAMERA_KEY].shape == (IMG_H, IMG_W, 3)
+    assert obs.images[CAMERA_KEY].dtype == np.uint8
+    assert obs.jpeg_quality == 0

tests/policy_server/test_session.py

@@ -0,0 +1,298 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for Session (latest-only mailbox, episode reset, close,
+processor-step introspection) and SessionRegistry (thread-safe map)."""
+
+from __future__ import annotations
+
+import threading
+
+from lerobot.policy_server.schema import MsgHeader
+from lerobot.policy_server.session import Session, SessionRegistry
+from lerobot.processor import NormalizerProcessorStep, RelativeActionsProcessorStep
+from tests.policy_server.conftest import TASK, MockPipeline, make_mock_processors
+
+
+def make_session(
+    client_uuid: str = "client-a",
+    preprocessor: MockPipeline | None = None,
+    postprocessor: MockPipeline | None = None,
+    publisher=None,
+) -> Session:
+    default_pre, default_post = make_mock_processors()
+    return Session(
+        session_id=f"sess-{client_uuid}",
+        client_uuid=client_uuid,
+        task=TASK,
+        robot_type="mock_robot",
+        rtc_enabled=False,
+        preprocessor=preprocessor if preprocessor is not None else default_pre,
+        postprocessor=postprocessor if postprocessor is not None else default_post,
+        action_publisher=publisher,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Mailbox: latest-only deposit / take
+# ---------------------------------------------------------------------------
+
+
+def test_deposit_then_take_returns_item():
+    session = make_session()
+    header = MsgHeader(seq_id=7)
+    session.deposit(header, b"payload-7")
+
+    item = session.take()
+    assert item is not None
+    assert item.header.seq_id == 7
+    assert item.payload == b"payload-7"
+    assert item.recv_mono > 0
+
+
+def test_second_deposit_supersedes_and_take_returns_newer():
+    session = make_session()
+    session.deposit(MsgHeader(seq_id=1), b"old")
+    session.deposit(MsgHeader(seq_id=2), b"new")
+
+    assert session.stats.superseded == 1
+    assert session.stats.superseded_since_reply == 1
+
+    item = session.take()
+    assert item is not None
+    assert item.header.seq_id == 2
+    assert item.payload == b"new"
+
+
+def test_deposit_after_take_is_not_superseded():
+    session = make_session()
+    session.deposit(MsgHeader(seq_id=1), b"one")
+    session.take()
+    session.deposit(MsgHeader(seq_id=2), b"two")
+
+    assert session.stats.superseded == 0
+    assert session.stats.superseded_since_reply == 0
+
+
+def test_take_clears_mailbox_second_take_is_none():
+    session = make_session()
+    session.deposit(MsgHeader(seq_id=1), b"one")
+
+    assert session.take() is not None
+    assert session.take() is None
+
+
+def test_has_pending_transitions():
+    session = make_session()
+    assert not session.has_pending()
+
+    session.deposit(MsgHeader(seq_id=1), b"one")
+    assert session.has_pending()
+
+    session.take()
+    assert not session.has_pending()
+
+
+def test_deposit_marks_alive_and_clears_token_drop():
+    session = make_session()
+    session.alive = False
+    session.token_dropped_mono = 123.4
+    before = session.last_seen_mono
+
+    session.deposit(MsgHeader(seq_id=1), b"one")
+
+    assert session.alive is True
+    assert session.token_dropped_mono is None
+    assert session.last_seen_mono >= before
+
+
+# ---------------------------------------------------------------------------
+# Episode boundary
+# ---------------------------------------------------------------------------
+
+
+def test_reset_episode_resets_pipelines_clears_mailbox_and_increments():
+    pre, post = make_mock_processors()
+    session = make_session(preprocessor=pre, postprocessor=post)
+    session.deposit(MsgHeader(seq_id=1), b"stale")
+    assert session.episode_id == 0
+
+    session.reset_episode()
+
+    assert not session.has_pending()
+    assert pre.reset_count == 1
+    assert post.reset_count == 1
+    assert session.episode_id == 1
+
+    session.reset_episode()
+    assert session.episode_id == 2
+    assert pre.reset_count == 2
+    assert post.reset_count == 2
+
+
+def test_reset_episode_with_explicit_id():
+    session = make_session()
+    session.reset_episode(episode_id=7)
+    assert session.episode_id == 7
+
+
+# ---------------------------------------------------------------------------
+# close()
+# ---------------------------------------------------------------------------
+
+
+class FakePublisher:
+    def __init__(self, raise_on_undeclare: bool = False):
+        self.undeclare_calls = 0
+        self.raise_on_undeclare = raise_on_undeclare
+
+    def undeclare(self):
+        self.undeclare_calls += 1
+        if self.raise_on_undeclare:
+            raise RuntimeError("transport already closed")
+
+
+def test_close_clears_mailbox_and_undeclares_publisher_exactly_once():
+    publisher = FakePublisher()
+    session = make_session(publisher=publisher)
+    session.deposit(MsgHeader(seq_id=1), b"stale")
+
+    session.close()
+
+    assert not session.has_pending()
+    assert publisher.undeclare_calls == 1
+    assert session.action_publisher is None
+
+    # Idempotent: a second close must not undeclare again.
+    session.close()
+    assert publisher.undeclare_calls == 1
+
+
+def test_close_tolerates_undeclare_raising():
+    publisher = FakePublisher(raise_on_undeclare=True)
+    session = make_session(publisher=publisher)
+    session.deposit(MsgHeader(seq_id=1), b"stale")
+
+    session.close()  # must not raise
+
+    assert publisher.undeclare_calls == 1
+    assert not session.has_pending()
+    assert session.action_publisher is None
+
+
+def test_close_without_publisher_is_noop():
+    session = make_session(publisher=None)
+    session.close()  # must not raise
+    assert session.action_publisher is None
+
+
+# ---------------------------------------------------------------------------
+# Processor-step introspection
+# ---------------------------------------------------------------------------
+
+
+def test_relative_and_normalizer_steps_detected():
+    relative = RelativeActionsProcessorStep(enabled=True)
+    normalizer = NormalizerProcessorStep(features={}, norm_map={})
+    pre = MockPipeline(steps=[relative, normalizer])
+    session = make_session(preprocessor=pre)
+
+    assert session.relative_step is relative
+    assert session.normalizer_step is normalizer
+
+
+def test_disabled_relative_step_is_not_detected():
+    relative = RelativeActionsProcessorStep(enabled=False)
+    pre = MockPipeline(steps=[relative])
+    session = make_session(preprocessor=pre)
+
+    assert session.relative_step is None
+    assert session.normalizer_step is None
+
+
+def test_empty_pipeline_yields_no_introspected_steps():
+    session = make_session()
+    assert session.relative_step is None
+    assert session.normalizer_step is None
+
+
+# ---------------------------------------------------------------------------
+# SessionRegistry
+# ---------------------------------------------------------------------------
+
+
+def test_registry_add_get_remove_len_snapshot():
+    registry = SessionRegistry()
+    assert len(registry) == 0
+    assert registry.get("missing") is None
+    assert registry.snapshot() == []
+
+    session_a = make_session("uuid-a")
+    session_b = make_session("uuid-b")
+    assert registry.add(session_a) is None
+    assert registry.add(session_b) is None
+
+    assert len(registry) == 2
+    assert registry.get("uuid-a") is session_a
+    assert registry.get("uuid-b") is session_b
+    assert set(registry.snapshot()) == {session_a, session_b}
+
+    removed = registry.remove("uuid-a")
+    assert removed is session_a
+    assert len(registry) == 1
+    assert registry.get("uuid-a") is None
+
+
+def test_registry_remove_missing_returns_none():
+    registry = SessionRegistry()
+    assert registry.remove("never-added") is None
+
+
+def test_registry_add_returns_displaced_same_uuid_session():
+    registry = SessionRegistry()
+    first = make_session("uuid-x")
+    second = make_session("uuid-x")
+
+    assert registry.add(first) is None
+    displaced = registry.add(second)
+
+    assert displaced is first
+    assert registry.get("uuid-x") is second
+    assert len(registry) == 1
+
+
+def test_registry_thread_safety_smoke():
+    registry = SessionRegistry()
+    errors: list[Exception] = []
+
+    def worker(prefix: str) -> None:
+        try:
+            for i in range(200):
+                session = make_session(f"{prefix}-{i}")
+                registry.add(session)
+                assert registry.get(session.client_uuid) is session
+                assert registry.remove(session.client_uuid) is session
+        except Exception as exc:  # noqa: BLE001 — surfaced to the main thread
+            errors.append(exc)
+
+    threads = [threading.Thread(target=worker, args=(p,)) for p in ("alpha", "beta")]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=10)
+        assert not t.is_alive()
+
+    assert errors == []
+    assert len(registry) == 0
+    assert registry.snapshot() == []

tests/policy_server/test_validation.py

@@ -0,0 +1,323 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for serving-mode classification and session-open validation.
+
+Uses tiny fake policy classes (deliberately NOT subclassing
+``PreTrainedPolicy``): classification keys off the ``name`` attribute and
+the presence of a ``predict_action_chunk`` override, never off the class
+hierarchy.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from lerobot.policy_server.schema import (
+    MIN_SUPPORTED_SCHEMA_VERSION,
+    SCHEMA_VERSION,
+    SessionOpenMsg,
+    StatusMsg,
+)
+from lerobot.policy_server.validation import (
+    PolicyClassification,
+    ServingClass,
+    classify_policy,
+    resolve_serving_mode,
+    validate_session_open,
+)
+from tests.policy_server.conftest import ACTION_NAMES, STATE_DIM, TASK, make_manifest
+
+# ---------------------------------------------------------------------------
+# Fake policy classes (classification only needs `name`, an optional
+# `predict_action_chunk` method, and `.config` for smolvla)
+# ---------------------------------------------------------------------------
+
+
+def _fake_policy(name: str, *, chunk_api: bool = True, n_obs_steps: int | None = None):
+    """Build a minimal fake policy instance with a class-level chunk method."""
+
+    namespace = {"name": name}
+    if chunk_api:
+        namespace["predict_action_chunk"] = lambda self, batch, **kwargs: None
+    cls = type(f"Fake_{name}_Policy", (), namespace)
+    policy = cls()
+    if n_obs_steps is not None:
+        policy.config = SimpleNamespace(n_obs_steps=n_obs_steps)
+    return policy
+
+
+# ---------------------------------------------------------------------------
+# classify_policy
+# ---------------------------------------------------------------------------
+
+
+def test_classify_act_is_shared_without_rtc():
+    classification = classify_policy(_fake_policy("act"))
+    assert classification.serving_class is ServingClass.SHARED
+    assert classification.supports_rtc is False
+    assert classification.needs_queue_population is False
+
+
+@pytest.mark.parametrize("name", ["pi0", "pi05"])
+def test_classify_pi_families_are_shared_with_rtc(name):
+    classification = classify_policy(_fake_policy(name))
+    assert classification.serving_class is ServingClass.SHARED
+    assert classification.supports_rtc is True
+    assert classification.needs_queue_population is False
+
+
+def test_classify_smolvla_single_obs_step_is_shared():
+    classification = classify_policy(_fake_policy("smolvla", n_obs_steps=1))
+    assert classification.serving_class is ServingClass.SHARED
+    assert classification.supports_rtc is True
+
+
+def test_classify_smolvla_with_history_is_exclusive():
+    classification = classify_policy(_fake_policy("smolvla", n_obs_steps=2))
+    assert classification.serving_class is ServingClass.EXCLUSIVE
+    assert classification.supports_rtc is True
+    assert classification.needs_queue_population is False
+
+
+def test_classify_diffusion_is_exclusive_with_queue_population():
+    classification = classify_policy(_fake_policy("diffusion"))
+    assert classification.serving_class is ServingClass.EXCLUSIVE
+    assert classification.supports_rtc is False
+    assert classification.needs_queue_population is True
+
+
+def test_classify_without_chunk_api_is_refused():
+    classification = classify_policy(_fake_policy("act", chunk_api=False))
+    assert classification.serving_class is ServingClass.REFUSED
+    assert classification.supports_rtc is False
+    assert "predict_action_chunk" in classification.reason
+
+
+def test_classify_unknown_name_with_chunk_api_is_exclusive():
+    classification = classify_policy(_fake_policy("totally_new_policy"))
+    assert classification.serving_class is ServingClass.EXCLUSIVE
+    assert classification.supports_rtc is False
+    assert classification.needs_queue_population is False
+    assert "verified" in classification.reason
+
+
+# ---------------------------------------------------------------------------
+# resolve_serving_mode
+# ---------------------------------------------------------------------------
+
+
+def _classification(serving_class: ServingClass, reason: str = "test") -> PolicyClassification:
+    return PolicyClassification(
+        serving_class, supports_rtc=False, needs_queue_population=False, reason=reason
+    )
+
+
+def test_resolve_auto_maps_shared_to_shared():
+    mode, max_sessions = resolve_serving_mode(
+        _classification(ServingClass.SHARED), make_manifest(serving_mode="auto", max_sessions=4)
+    )
+    assert mode == "shared"
+    assert max_sessions == 4
+
+
+def test_resolve_auto_maps_exclusive_to_exclusive():
+    mode, max_sessions = resolve_serving_mode(
+        _classification(ServingClass.EXCLUSIVE), make_manifest(serving_mode="auto", max_sessions=4)
+    )
+    assert mode == "exclusive"
+    assert max_sessions == 1
+
+
+def test_resolve_forced_shared_rejected_for_non_verified_policy():
+    with pytest.raises(ValueError, match="unsafe"):
+        resolve_serving_mode(_classification(ServingClass.EXCLUSIVE), make_manifest(serving_mode="shared"))
+
+
+def test_resolve_forced_shared_allowed_for_verified_policy():
+    mode, max_sessions = resolve_serving_mode(
+        _classification(ServingClass.SHARED), make_manifest(serving_mode="shared", max_sessions=4)
+    )
+    assert mode == "shared"
+    assert max_sessions == 4
+
+
+def test_resolve_forced_exclusive_allowed_for_shared_policy():
+    mode, _ = resolve_serving_mode(
+        _classification(ServingClass.SHARED), make_manifest(serving_mode="exclusive")
+    )
+    assert mode == "exclusive"
+
+
+def test_resolve_exclusive_forces_single_session():
+    mode, max_sessions = resolve_serving_mode(
+        _classification(ServingClass.SHARED), make_manifest(serving_mode="exclusive", max_sessions=4)
+    )
+    assert mode == "exclusive"
+    assert max_sessions == 1
+
+
+def test_resolve_refused_raises_with_reason():
+    with pytest.raises(ValueError, match="no chunk API here"):
+        resolve_serving_mode(
+            _classification(ServingClass.REFUSED, reason="no chunk API here"), make_manifest()
+        )
+
+
+# ---------------------------------------------------------------------------
+# validate_session_open
+# ---------------------------------------------------------------------------
+
+EXPECTED_CAMERAS = ["observation.images.front"]
+
+
+@pytest.fixture
+def capabilities() -> StatusMsg:
+    return StatusMsg(
+        model_repo="mock/model",
+        policy_type="mockchunk",
+        action_names=list(ACTION_NAMES),
+        expected_cameras=list(EXPECTED_CAMERAS),
+        state_dim=STATE_DIM,
+        chunk_size=20,
+        trained_fps=30.0,
+        supports_rtc=True,
+        min_schema_version=MIN_SUPPORTED_SCHEMA_VERSION,
+        max_schema_version=SCHEMA_VERSION,
+        max_sessions=4,
+    )
+
+
+def _open_msg(**overrides) -> SessionOpenMsg:
+    kwargs: dict = {
+        "client_uuid": "client-1",
+        "policy_type": "mockchunk",
+        "fps": 30.0,
+        "action_names": list(ACTION_NAMES),
+        "camera_names": list(EXPECTED_CAMERAS),
+        "state_dim": STATE_DIM,
+        "schema_version": SCHEMA_VERSION,
+        "task": TASK,
+    }
+    kwargs.update(overrides)
+    return SessionOpenMsg(**kwargs)
+
+
+def test_validate_happy_path(capabilities):
+    result = validate_session_open(_open_msg(), capabilities, make_manifest(), active_sessions=0)
+    assert result.ok
+    assert result.error == ""
+    assert result.warnings == []
+    assert result.rtc_downgraded is False
+
+
+def test_validate_action_name_order_is_a_hard_reject(capabilities):
+    # Same set of names, different order: chunk columns would map to the
+    # wrong motors, so this must be a hard reject.
+    result = validate_session_open(
+        _open_msg(action_names=list(reversed(ACTION_NAMES))),
+        capabilities,
+        make_manifest(),
+        active_sessions=0,
+    )
+    assert not result.ok
+    assert "action" in result.error
+    assert "mismatch" in result.error
+
+
+def test_validate_missing_camera_rejected(capabilities):
+    result = validate_session_open(
+        _open_msg(camera_names=[]), capabilities, make_manifest(), active_sessions=0
+    )
+    assert not result.ok
+    assert "observation.images.front" in result.error
+
+
+def test_validate_wrong_state_dim_rejected(capabilities):
+    result = validate_session_open(
+        _open_msg(state_dim=STATE_DIM + 1), capabilities, make_manifest(), active_sessions=0
+    )
+    assert not result.ok
+    assert "state dim" in result.error
+
+
+def test_validate_schema_version_out_of_range_rejected(capabilities):
+    result = validate_session_open(
+        _open_msg(schema_version=SCHEMA_VERSION + 99),
+        capabilities,
+        make_manifest(),
+        active_sessions=0,
+    )
+    assert not result.ok
+    assert "schema_version" in result.error
+
+
+def test_validate_at_capacity_rejected_with_load(capabilities):
+    result = validate_session_open(
+        _open_msg(), capabilities, make_manifest(), active_sessions=capabilities.max_sessions
+    )
+    assert not result.ok
+    assert "full" in result.error
+    assert f"{capabilities.max_sessions}/{capabilities.max_sessions}" in result.error
+
+
+def test_validate_pinned_task_rejects_other_task(capabilities):
+    result = validate_session_open(
+        _open_msg(task="another task"),
+        capabilities,
+        make_manifest(pin_task=True),
+        active_sessions=0,
+    )
+    assert not result.ok
+    assert "pinned" in result.error
+
+
+def test_validate_fps_mismatch_strict_rejects(capabilities):
+    result = validate_session_open(
+        _open_msg(fps=15.0), capabilities, make_manifest(strict_fps=True), active_sessions=0
+    )
+    assert not result.ok
+    assert "fps" in result.error
+
+
+def test_validate_fps_mismatch_lenient_warns_only(capabilities):
+    result = validate_session_open(
+        _open_msg(fps=15.0), capabilities, make_manifest(strict_fps=False), active_sessions=0
+    )
+    assert result.ok
+    assert len(result.warnings) == 1
+    assert "fps" in result.warnings[0]
+
+
+def test_validate_rtc_downgraded_when_unsupported(capabilities):
+    capabilities.supports_rtc = False
+    result = validate_session_open(
+        _open_msg(rtc_enabled=True), capabilities, make_manifest(), active_sessions=0
+    )
+    assert result.ok
+    assert result.rtc_downgraded is True
+    assert any("RTC" in warning for warning in result.warnings)
+
+
+def test_validate_empty_capability_action_names_skips_action_check(capabilities):
+    capabilities.action_names = []
+    result = validate_session_open(
+        _open_msg(action_names=["whatever.pos"]),
+        capabilities,
+        make_manifest(),
+        active_sessions=0,
+    )
+    assert result.ok