From deecf9ae68eb4e0138f2ffe7354b8132336c065a Mon Sep 17 00:00:00 2001 From: cc Date: Tue, 20 Oct 2026 21:31:19 -0700 Subject: [PATCH] feat(transport): support CDR alongside existing message encodings --- dimos/core/coordination/test_blueprints.py | 8 +- dimos/core/test_dds_transport.py | 37 ++++++++ dimos/core/test_transport_factory.py | 39 +++++++- dimos/core/test_zenoh_transport.py | 16 +--- dimos/core/transport.py | 38 +++++--- dimos/core/transport_factory.py | 16 +++- dimos/protocol/pubsub/codec.py | 91 +++++++++++++++++++ dimos/protocol/pubsub/encoders.py | 83 ++++++++++++----- dimos/protocol/pubsub/impl/lcmpubsub.py | 31 +++++-- dimos/protocol/pubsub/impl/test_lcmpubsub.py | 27 ++++++ .../protocol/pubsub/impl/test_zenohpubsub.py | 39 ++++++++ .../pubsub/impl/webrtc/test_transport.py | 23 +++++ dimos/protocol/pubsub/impl/zenohpubsub.py | 25 +++-- dimos/protocol/pubsub/test_codec.py | 72 +++++++++++++++ dimos/protocol/pubsub/test_encoder.py | 65 ++++++++++++- docs/usage/blueprints.md | 4 +- 16 files changed, 530 insertions(+), 84 deletions(-) create mode 100644 dimos/core/test_dds_transport.py create mode 100644 dimos/protocol/pubsub/codec.py create mode 100644 dimos/protocol/pubsub/test_codec.py diff --git a/dimos/core/coordination/test_blueprints.py b/dimos/core/coordination/test_blueprints.py index 043e3bf0a4f..72e544c1e79 100644 --- a/dimos/core/coordination/test_blueprints.py +++ b/dimos/core/coordination/test_blueprints.py @@ -38,7 +38,7 @@ from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out -from dimos.core.transport import LCMTransport +from dimos.core.transport import pLCMTransport from dimos.spec.utils import Spec @@ -139,7 +139,7 @@ def test_config_rejects_unknown_fields_and_sections() -> None: def test_transports() -> None: - custom_transport = LCMTransport("/custom_topic", Data1) + custom_transport = pLCMTransport("/custom_topic") blueprint_set = autoconnect(ModuleA.blueprint(), ModuleB.blueprint()).transports( {("data1", Data1): custom_transport} ) @@ -414,12 +414,12 @@ def test_explicit_instance_name_sets_blueprint_identity() -> None: def test_namespace_prefixes_pinned_transports() -> None: blueprint = ( autoconnect(ModuleA.blueprint(), ModuleB.blueprint()) - .transports({("data1", Data1): LCMTransport("/custom_topic", Data1)}) + .transports({("data1", Data1): pLCMTransport("/custom_topic")}) .namespace("robot0") ) transport = blueprint.transport_map[("robot0/data1", Data1)] - assert transport.topic.pattern == "/robot0/custom_topic" + assert transport.topic == "/robot0/custom_topic" assert ("data1", Data1) not in blueprint.transport_map diff --git a/dimos/core/test_dds_transport.py b/dimos/core/test_dds_transport.py new file mode 100644 index 00000000000..14a5f2b6bd3 --- /dev/null +++ b/dimos/core/test_dds_transport.py @@ -0,0 +1,37 @@ +# Copyright 2026 Dimensional Inc. +# +# 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. + +from dataclasses import dataclass + +import pytest + +from dimos.core.transport import DDSTransport + +idl = pytest.importorskip("cyclonedds.idl") + + +@dataclass +class NativeSample(idl.IdlStruct): + value: int + + +def test_dds_transport_accepts_native_idl_type(): + transport = DDSTransport("native_sample", NativeSample) + try: + assert transport.topic.name == "native_sample" + assert transport.topic.data_type is NativeSample + sample = NativeSample(value=42) + assert transport.topic.data_type.deserialize(sample.serialize()) == sample + finally: + transport.stop() diff --git a/dimos/core/test_transport_factory.py b/dimos/core/test_transport_factory.py index 42504371c96..822f2282c79 100644 --- a/dimos/core/test_transport_factory.py +++ b/dimos/core/test_transport_factory.py @@ -18,7 +18,9 @@ from dimos.core.global_config import GlobalConfig from dimos.core.transport import ( + DDSTransport, LCMTransport, + ROSTransport, ZenohTransport, pLCMTransport, pZenohTransport, @@ -31,6 +33,8 @@ session_config, transport_topic, ) +from dimos.messages.sensor_msgs.msg import Image as CdrImage +from dimos.messages.std_msgs.msg import Int32 from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.sensor_msgs.Image import Image @@ -82,8 +86,9 @@ def test_make_transport_zenoh_pickled() -> None: assert t.topic == "dimos/human_input" -def test_default_zenoh_qos_high_rate_sensor_types_drop() -> None: - assert default_zenoh_qos("/camera/color", Image) == QOS_LATEST_WINS +@pytest.mark.parametrize("message_type", [Image, CdrImage]) +def test_default_zenoh_qos_high_rate_sensor_types_drop(message_type) -> None: + assert default_zenoh_qos("/camera/color", message_type) == QOS_LATEST_WINS def test_default_zenoh_qos_agent_channels_never_drop() -> None: @@ -155,3 +160,33 @@ def test_apply_transport_arg_bad_value_exits(argv: list[str]) -> None: apply_transport_arg(argv, g=g) assert exc.value.code == 2 assert g.transport == "lcm" # config left untouched + + +@pytest.mark.parametrize("config", [LCM, ZENOH]) +def test_factory_binds_generated_messages(config): + transport = make_transport("generated", Int32, g=config) + assert transport.topic.message.encoding == "cdr" + assert transport.topic.message.message_type is Int32 + + +@pytest.mark.parametrize("config", [LCM, ZENOH]) +@pytest.mark.parametrize("message_type", [dict, int, str, bool]) +def test_factory_preserves_pickle_for_python_types(config, message_type): + transport = make_transport("python", message_type, g=config) + expected = pLCMTransport if config.transport == "lcm" else pZenohTransport + assert isinstance(transport, expected) + + +@pytest.mark.parametrize("transport_type", [ROSTransport, DDSTransport]) +def test_generated_messages_require_ros_dds_type_adapters(transport_type): + with pytest.raises(TypeError, match="type adapter"): + transport_type("generated", Int32) + + +@pytest.mark.parametrize("config", [LCM, ZENOH]) +def test_factory_does_not_pickle_unregistered_cdr_types(config): + class Unregistered: + __msgtype__ = "test_msgs/msg/Unregistered" + + with pytest.raises(TypeError): + make_transport("invalid", Unregistered, g=config) diff --git a/dimos/core/test_zenoh_transport.py b/dimos/core/test_zenoh_transport.py index 31c1532292c..a141fa3030b 100644 --- a/dimos/core/test_zenoh_transport.py +++ b/dimos/core/test_zenoh_transport.py @@ -48,14 +48,6 @@ from dimos.protocol.service.zenohservice import ZenohSessionPool -class TypedMsg: - """A fake typed message with lcm_encode for testing.""" - - @staticmethod - def lcm_encode() -> bytes: - return b"" - - class UntypedMsg: """A message without lcm_encode. Triggers pickle transport.""" @@ -63,12 +55,12 @@ class UntypedMsg: class ProducerModule(Module): - typed_data: Out[TypedMsg] + typed_data: Out[Image] untyped_data: Out[UntypedMsg] class ConsumerModule(Module): - typed_data: In[TypedMsg] + typed_data: In[Image] untyped_data: In[UntypedMsg] @@ -147,7 +139,7 @@ def test_invalid_transport_is_rejected_on_update() -> None: def test_lcm_transport_returned_when_transport_is_lcm(use_lcm, blueprint) -> None: - transport = _get_transport_for(blueprint, "typed_data", TypedMsg) + transport = _get_transport_for(blueprint, "typed_data", Image) assert isinstance(transport, LCMTransport) @@ -157,7 +149,7 @@ def test_lcm_pickle_transport_returned_for_untyped_when_lcm(use_lcm, blueprint) def test_zenoh_transport_returned_when_transport_is_zenoh(use_zenoh, blueprint) -> None: - transport = _get_transport_for(blueprint, "typed_data", TypedMsg) + transport = _get_transport_for(blueprint, "typed_data", Image) assert isinstance(transport, ZenohTransport) diff --git a/dimos/core/transport.py b/dimos/core/transport.py index ffef78418dd..7f844669d87 100644 --- a/dimos/core/transport.py +++ b/dimos/core/transport.py @@ -29,6 +29,7 @@ from dimos.core.stream import Stream, Transport from dimos.msgs.protocol import DimosMsg +from dimos.protocol.pubsub.codec import describe_message from dimos.protocol.pubsub.impl.lcmpubsub import LCM, PickleLCM, Topic as LCMTopic from dimos.protocol.pubsub.impl.webrtc.providers.broker import BrokerConfig from dimos.protocol.pubsub.impl.webrtc.providers.spec import AudioProvider, ProviderConfig @@ -198,6 +199,11 @@ class ROSTransport(PubSubTransport[DimosMsg]): _ros: DimosROS | None = None def __init__(self, topic: str, msg_type: type[DimosMsg], **kwargs: Any) -> None: + # TODO: Revisit generated-message adapters; ROS also uses CDR encoding. + if hasattr(msg_type, "__msgtype__") and describe_message(msg_type).encoding == "cdr": + raise TypeError( + "Generated messages require a ROS type adapter; raw CDR is not a ROS message" + ) # rclpy (about 0.1 s with rospubsub): only loaded when a ROSTransport exists. from dimos.protocol.pubsub.impl.rospubsub import ROSTopic @@ -241,6 +247,11 @@ def stop(self) -> None: class DDSTransport(PubSubTransport[T]): def __init__(self, topic: str, type: type, **kwargs) -> None: # type: ignore[no-untyped-def] + # TODO: Revisit generated-message adapters; DDS also uses CDR encoding. + if hasattr(type, "__msgtype__") and describe_message(type).encoding == "cdr": + raise TypeError( + "Generated messages require a DDS type adapter; raw CDR is not a DDS type" + ) # cyclonedds: see DDS_AVAILABLE. from dimos.protocol.pubsub.impl.ddspubsub import DDS, Topic as DDSTopic @@ -280,7 +291,7 @@ def subscribe( return self.dds.subscribe(self.topic, lambda msg, topic: callback(msg)) -M = TypeVar("M", bound=DimosMsg) +M = TypeVar("M") def _rebuild_webrtc_transport( @@ -296,10 +307,9 @@ class WebRTCTransport(PubSubTransport[M]): also be used directly with an explicit ``config``. Two modes: * **Raw bytes** (``msg_type=None``): messages pass through as ``bytes``. - * **Typed LCM** (``msg_type=SomeMsg``): LCM-encoded on ``broadcast()``, - LCM-decoded on ``subscribe()`` (foreign types on the shared channel are skipped) — so multiple - transports sharing one multiplexed DataChannel each receive only - their own message type. + * **Typed messages**: bind a shared codec from ``msg_type``. Generated + messages use raw CDR; legacy messages keep their LCM wire format. + Peers must agree on the schema; remote schema validation is deferred. The transport itself holds no connection: the picklable ``config`` resolves to a per-process singleton provider on first use, so transports @@ -322,6 +332,7 @@ def __init__( ) -> None: super().__init__(topic) self._msg_type = msg_type + self._message = describe_message(msg_type) if msg_type is not None else None self._config = config or self._config_cls(**config_kwargs) self._pubsub: WebRTCPubSub | None = None # Guards first-use init: concurrent subscribe()/broadcast() must not @@ -337,7 +348,7 @@ def broadcast(self, _: Stream[M] | None, msg: M) -> None: if not self._started: self.start() assert self._pubsub is not None - data = msg.lcm_encode() if self._msg_type is not None else msg + data = self._message.codec.encode(msg) if self._message is not None else cast("bytes", msg) self._pubsub.publish(self.topic, data) # type: ignore[arg-type] def subscribe( @@ -347,18 +358,15 @@ def subscribe( self.start() assert self._pubsub is not None - if self._msg_type is not None: - msg_type = self._msg_type + if self._message is not None: + descriptor = self._message def _typed_cb(data: bytes, _topic: str) -> None: - # The channel is multiplexed (e.g. the browser sends Twists and - # Poses on cmd_unreliable); lcm_decode verifies the wire - # fingerprint and raises on other types — skip those. try: - msg = msg_type.lcm_decode(data) + msg = descriptor.codec.decode(data) except ValueError: return - callback(msg) # type: ignore[arg-type] + callback(msg) return self._pubsub.subscribe(self.topic, _typed_cb) return self._pubsub.subscribe(self.topic, lambda msg, _topic: callback(msg)) # type: ignore[arg-type] @@ -503,7 +511,7 @@ class CloudflareAudioTransport(WebRTCAudioTransport): class ZenohTransport(PubSubTransport[T]): - """Zenoh transport with LCM encoding for typed DimosMsg. + """Zenoh transport with a codec bound to the declared message type. Accepts either a plain topic string plus message type, or a full `ZenohTopic` carrying per-topic settings: `ZenohTransport(ZenohTopic("bla", @@ -546,7 +554,7 @@ def stop(self) -> None: def broadcast(self, _: Stream[T] | None, msg: T) -> None: if not self._started: self.start() - self.zenoh.publish(self.topic, cast("DimosMsg", msg)) + self.zenoh.publish(self.topic, msg) def subscribe( self, callback: Callable[[T], None], selfstream: Stream[T] | None = None diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index b86b271682a..5a3e88f352b 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -27,6 +27,7 @@ pLCMTransport, pZenohTransport, ) +from dimos.protocol.pubsub.codec import describe_message from dimos.protocol.pubsub.impl.zenohpubsub import ( QOS_LATEST_WINS, QOS_NEVER_DROP, @@ -91,7 +92,10 @@ def default_zenoh_qos_for(name: str, msg_name: str) -> ZenohQoS | None: def default_zenoh_qos(name: str, msg_type: type | None = None) -> ZenohQoS | None: """Default publisher QoS for a logical channel; None = zenoh defaults.""" - return default_zenoh_qos_for(name, getattr(msg_type, "msg_name", "")) + type_name = getattr(msg_type, "msg_name", "") + if msg_type is not None and hasattr(msg_type, "__msgtype__"): + type_name = describe_message(msg_type).name.replace("/msg/", ".") + return default_zenoh_qos_for(name, type_name) def make_transport( @@ -99,8 +103,8 @@ def make_transport( ) -> PubSubTransport[Any]: """Construct the active-backend pub/sub transport for a logical channel. - A pickled (self-describing) transport is used when no `msg_type` is given or - the type has no `lcm_encode`. Otherwise a typed transport is used. + Generated messages use CDR and legacy messages keep LCM encoding. + Other Python types and undeclared types preserve the existing pickle transport. A channel name alone doesn't fully define a topic: backends have per-topic settings (Zenoh publisher QoS, for one) that live on their Topic objects. @@ -110,7 +114,11 @@ def make_transport( transport map. LCM (UDP multicast) has no per-topic settings. """ - use_pickled = msg_type is None or getattr(msg_type, "lcm_encode", None) is None + use_pickled = msg_type is None or ( + not hasattr(msg_type, "__msgtype__") and getattr(msg_type, "lcm_encode", None) is None + ) + if msg_type is not None and not use_pickled: + describe_message(msg_type) topic = transport_topic(name, g) if g.transport == "zenoh": ztopic = ZenohTopic( diff --git a/dimos/protocol/pubsub/codec.py b/dimos/protocol/pubsub/codec.py new file mode 100644 index 00000000000..811b9450571 --- /dev/null +++ b/dimos/protocol/pubsub/codec.py @@ -0,0 +1,91 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# 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. + +"""Message codecs selected once from the declared message type.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Generic, Literal, Protocol, TypeVar, cast + +from dimos_cdrgen.registry import decode, encode, message_types + +from dimos.msgs.protocol import DimosMsg + +T = TypeVar("T") + + +class MessageCodec(Protocol[T]): + def encode(self, message: T) -> bytes: ... + + def decode(self, payload: bytes) -> T: ... + + +@dataclass(frozen=True) +class LCMCodec(Generic[T]): + message_type: type[T] + + def encode(self, message: T) -> bytes: + if not isinstance(message, self.message_type): + raise TypeError("Message does not match the declared type") + return cast("DimosMsg", message).lcm_encode() + + def decode(self, payload: bytes) -> T: + return cast("T", cast("type[DimosMsg]", self.message_type).lcm_decode(payload)) + + +@dataclass(frozen=True) +class CDRCodec(Generic[T]): + message_type: type[T] + + def encode(self, message: T) -> bytes: + if type(message) is not self.message_type: + raise TypeError("Message does not match the declared type") + return encode(message) + + def decode(self, payload: bytes) -> T: + if len(payload) < 4 or payload[:2] not in (b"\x00\x00", b"\x00\x01"): + raise ValueError("Expected standalone XCDR1 encapsulation") + return decode(payload, self.message_type) + + +@dataclass(frozen=True) +class MessageDescriptor(Generic[T]): + message_type: type[T] + encoding: Literal["lcm", "cdr"] + name: str + codec: MessageCodec[T] + + +def describe_message(message_type: type[T]) -> MessageDescriptor[T]: + legacy = all( + callable(getattr(message_type, method, None)) for method in ("lcm_encode", "lcm_decode") + ) + if legacy and not hasattr(message_type, "__msgtype__"): + name = getattr(message_type, "msg_name", None) + if not isinstance(name, str): + raise TypeError("Legacy message type must declare msg_name") + return MessageDescriptor(message_type, "lcm", name, LCMCodec(message_type)) + if not hasattr(message_type, "__msgtype__"): + raise TypeError(f"Unsupported message type: {message_type!r}") + names = [name for name, cls in message_types().items() if cls is message_type] + if (names and legacy) or len(names) > 1: + raise TypeError("Ambiguous message encoding") + if names: + return MessageDescriptor(message_type, "cdr", names[0], CDRCodec(message_type)) + if legacy and isinstance(getattr(message_type, "msg_name", None), str): + return MessageDescriptor( + message_type, "lcm", cast("Any", message_type).msg_name, LCMCodec(message_type) + ) + raise TypeError(f"Unsupported message type: {message_type!r}") diff --git a/dimos/protocol/pubsub/encoders.py b/dimos/protocol/pubsub/encoders.py index ff5c78a2f52..bfd2f217eff 100644 --- a/dimos/protocol/pubsub/encoders.py +++ b/dimos/protocol/pubsub/encoders.py @@ -18,9 +18,9 @@ from abc import ABC, abstractmethod import pickle -from typing import TYPE_CHECKING, Generic, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast -from dimos.msgs.protocol import DimosMsg +from dimos.protocol.pubsub.codec import MessageDescriptor if TYPE_CHECKING: from collections.abc import Callable @@ -33,8 +33,6 @@ class DecodingError(Exception): """Raised by decode() to skip a message without calling the callback.""" - pass - class PubSubEncoderMixin(Generic[TopicT, MsgT, EncodingT], ABC): """Mixin that encodes messages before publishing and decodes them after receiving. @@ -62,6 +60,9 @@ def encode(self, msg: MsgT, topic: TopicT) -> EncodingT: ... @abstractmethod def decode(self, msg: EncodingT, topic: TopicT) -> MsgT: ... + def decode_received(self, payload: EncodingT, declared: TopicT, received: TopicT) -> MsgT: + return self.decode(payload, received) + def publish(self, topic: TopicT, message: MsgT) -> None: """Encode the message and publish it.""" encoded_message = self.encode(message, topic) @@ -74,9 +75,11 @@ def subscribe( ) -> Callable[[], None]: """Subscribe with automatic decoding.""" + declared_topic = topic + def wrapper_cb(encoded_data: EncodingT, topic: TopicT) -> None: try: - decoded_message = self.decode(encoded_data, topic) + decoded_message = self.decode_received(encoded_data, declared_topic, topic) except DecodingError: return callback(decoded_message, topic) @@ -95,32 +98,64 @@ def decode(self, msg: bytes, _: TopicT) -> MsgT: class LCMTopicProto(Protocol): - """Protocol for topics usable with LCM encoders.""" + """Compatibility interface for legacy encoders such as JPEG.""" + + topic: str + lcm_type: type[Any] | None - topic: str # At decode time, always concrete string - lcm_type: type[DimosMsg] | None +class MessageTopic(Protocol): + @property + def message(self) -> MessageDescriptor[Any] | None: ... -class LCMEncoderMixin(PubSubEncoderMixin[LCMTopicProto, DimosMsg, bytes]): - """Encoder mixin for DimosMsg using LCM binary encoding.""" + +class MessageEncoderMixin(PubSubEncoderMixin[MessageTopic, Any, bytes]): + """Apply the topic's bound codec without changing transport behavior.""" def subscribe( - self, topic: LCMTopicProto, callback: Callable[[DimosMsg, LCMTopicProto], None] + self, topic: MessageTopic, callback: Callable[[Any, MessageTopic], None] ) -> Callable[[], None]: - # Import the type's heavy decode dependencies now, on the subscriber's - # thread. Left to the first decode, the import would run on the LCM - # handler thread and block it for seconds, dropping messages meanwhile. - warmup = getattr(topic.lcm_type, "lcm_warmup", None) - if warmup is not None: - warmup() + descriptor = topic.message + if descriptor is not None and descriptor.encoding == "lcm": + warmup = getattr(descriptor.message_type, "lcm_warmup", None) + if warmup is not None: + warmup() + return super().subscribe(topic, callback) - def encode(self, msg: DimosMsg | bytes, _: LCMTopicProto) -> bytes: - if isinstance(msg, bytes): + def decode_received( + self, payload: bytes, declared: MessageTopic, received: MessageTopic + ) -> Any: + descriptor = declared.message + actual = received.message + if descriptor is not None and descriptor.encoding == "cdr": + if ( + actual is None + or actual.encoding != "cdr" + or actual.message_type is not descriptor.message_type + ): + raise DecodingError("Message metadata does not match the declared type") + return self.decode(payload, declared) + if descriptor is not None and actual is not None and actual.encoding != descriptor.encoding: + raise DecodingError("Message encoding does not match the declared type") + return self.decode(payload, received) + + def encode(self, msg: Any, topic: MessageTopic) -> bytes: + descriptor = topic.message + if isinstance(msg, bytes) and (descriptor is None or descriptor.encoding == "lcm"): return msg - return msg.lcm_encode() + if descriptor is None: + raise TypeError("Typed publication requires a declared message type") + return descriptor.codec.encode(msg) + + def decode(self, msg: bytes, topic: MessageTopic) -> Any: + if topic.message is None: + raise DecodingError("Typed subscription requires a declared message type") + try: + return topic.message.codec.decode(msg) + except ValueError as error: + raise DecodingError(str(error)) from error + - def decode(self, msg: bytes, topic: LCMTopicProto) -> DimosMsg: - if topic.lcm_type is None: - raise DecodingError(f"Cannot decode: topic {topic.topic!r} has no lcm_type") - return topic.lcm_type.lcm_decode(msg) +# Keep existing imports working while callers migrate to the generic name. +LCMEncoderMixin = MessageEncoderMixin diff --git a/dimos/protocol/pubsub/impl/lcmpubsub.py b/dimos/protocol/pubsub/impl/lcmpubsub.py index f47280f35cb..fdeaca2fc28 100644 --- a/dimos/protocol/pubsub/impl/lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/lcmpubsub.py @@ -15,14 +15,15 @@ from __future__ import annotations from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field import re import threading from typing import Any from dimos.msgs.protocol import DimosMsg +from dimos.protocol.pubsub.codec import MessageDescriptor, describe_message from dimos.protocol.pubsub.encoders import ( - LCMEncoderMixin, + MessageEncoderMixin, PickleEncoderMixin, ) from dimos.protocol.pubsub.patterns import Glob @@ -36,7 +37,21 @@ @dataclass class Topic: topic: str | re.Pattern[str] | Glob - lcm_type: type[DimosMsg] | None = None + lcm_type: type[Any] | None = None + message_type: type[Any] | None = field(default=None, kw_only=True) + message: MessageDescriptor[Any] | None = field(init=False, default=None, repr=False) + + def __post_init__(self) -> None: + if ( + self.lcm_type is not None + and self.message_type is not None + and self.lcm_type is not self.message_type + ): + raise TypeError("Conflicting declared message types") + declared = self.message_type if self.message_type is not None else self.lcm_type + self.message = describe_message(declared) if declared is not None else None + self.message_type = declared + self.lcm_type = declared @property def is_pattern(self) -> bool: @@ -51,9 +66,10 @@ def pattern(self) -> str: return self.topic def __str__(self) -> str: - if self.lcm_type is None: + if self.message is None or self.message.encoding == "cdr": return self.pattern - return f"{self.pattern}#{self.lcm_type.msg_name}" + assert self.message is not None + return f"{self.pattern}#{self.message.name}" @staticmethod def from_channel_str(channel: str, default_lcm_type: type[DimosMsg] | None = None) -> Topic: @@ -64,6 +80,9 @@ def from_channel_str(channel: str, default_lcm_type: type[DimosMsg] | None = Non """ from dimos.msgs.helpers import resolve_msg_type + if default_lcm_type is not None and describe_message(default_lcm_type).encoding == "cdr": + return Topic(channel, message_type=default_lcm_type) + if "#" not in channel: return Topic(topic=channel, lcm_type=default_lcm_type) @@ -149,7 +168,7 @@ def unsubscribe() -> None: class LCM( # type: ignore[misc] - LCMEncoderMixin, + MessageEncoderMixin, LCMPubSubBase, ): ... diff --git a/dimos/protocol/pubsub/impl/test_lcmpubsub.py b/dimos/protocol/pubsub/impl/test_lcmpubsub.py index 4ab7030b107..9290065dcf6 100644 --- a/dimos/protocol/pubsub/impl/test_lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/test_lcmpubsub.py @@ -18,6 +18,8 @@ import pytest +from dimos.core.transport import LCMTransport +from dimos.messages.std_msgs.msg import Int32 from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -184,3 +186,28 @@ def test_lcm_geometry_msgs_autopickle_pubsub(test_message: Any, pickle_lcm: Pick assert isinstance(received_topic, Topic) assert received_topic == topic + + +def test_generated_message_transport_and_unsubscribe(): + transport = LCMTransport("codec-check", Int32, url="memq://") + received = [] + delivered = threading.Event() + + def receive(value): + received.append(value) + delivered.set() + + try: + unsubscribe = transport.subscribe(receive) + LCMPubSubBase.publish(transport.lcm, transport.topic, b"broken") + transport.broadcast(None, Int32(42)) + assert delivered.wait(3) + assert received == [Int32(42)] + unsubscribe() + delivered.clear() + transport.subscribe(lambda value: delivered.set()) + transport.broadcast(None, Int32(43)) + assert delivered.wait(3) + assert received == [Int32(42)] + finally: + transport.stop() diff --git a/dimos/protocol/pubsub/impl/test_zenohpubsub.py b/dimos/protocol/pubsub/impl/test_zenohpubsub.py index d1c6d03f1e5..cc43ff71d35 100644 --- a/dimos/protocol/pubsub/impl/test_zenohpubsub.py +++ b/dimos/protocol/pubsub/impl/test_zenohpubsub.py @@ -27,6 +27,8 @@ import zenoh +from dimos.core.transport import ZenohTransport +from dimos.messages.std_msgs.msg import Int32 from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.sensor_msgs.Image import Image from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic @@ -493,3 +495,40 @@ def test_round_trip_untyped(self) -> None: reconstructed = _key_expr_to_topic(key) assert reconstructed.topic == original.topic assert reconstructed.lcm_type is None + + +def test_generated_message_transport_qos_and_unsubscribe(): + pool = ZenohSessionPool() + qos = ZenohQoS(reliability="reliable", congestion_control="block") + transport = ZenohTransport( + Topic("codec-check", message_type=Int32, qos=qos), + session_pool=pool, + connect=[], + listen=[LOOPBACK_LISTEN], + scouting=False, + multicast=False, + gossip=False, + ) + received = [] + delivered = threading.Event() + + def receive(value): + received.append(value) + delivered.set() + + try: + unsubscribe = transport.subscribe(receive) + assert transport.publish_qos == qos.to_wire() + ZenohPubSubBase.publish(transport.zenoh, transport.topic, b"broken") + transport.broadcast(None, Int32(42)) + assert delivered.wait(3) + assert received == [Int32(42)] + unsubscribe() + delivered.clear() + transport.subscribe(lambda value: delivered.set()) + transport.broadcast(None, Int32(43)) + assert delivered.wait(3) + assert received == [Int32(42)] + finally: + transport.stop() + pool.close_all() diff --git a/dimos/protocol/pubsub/impl/webrtc/test_transport.py b/dimos/protocol/pubsub/impl/webrtc/test_transport.py index d2a81dd5103..5f1529f0295 100644 --- a/dimos/protocol/pubsub/impl/webrtc/test_transport.py +++ b/dimos/protocol/pubsub/impl/webrtc/test_transport.py @@ -28,7 +28,9 @@ from dimos.core.coordination.blueprints import Blueprint from dimos.core.coordination.module_coordinator import _materialize_transports from dimos.core.transport import WebRTCTransport +from dimos.messages.std_msgs.msg import Int32 from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos.protocol.pubsub.codec import describe_message from dimos.protocol.pubsub.impl.webrtc.providers import spec from dimos.protocol.pubsub.impl.webrtc.providers.broker import BrokerConfig from dimos.protocol.pubsub.impl.webrtc.providers.spec import ( @@ -512,3 +514,24 @@ def test_subscribe_all_fires_once_per_message() -> None: ps.subscribe("u", lambda data, t: None) ps.publish("u", b"y") assert seen == [(b"x", "t"), (b"y", "u")] + + +def test_cdr_transport_raw_payload_errors_and_unsubscribe(): + config = MockConfig(name="cdr-raw") + transport = MockTransport("messages", Int32, config=config) + received = [] + wire = [] + try: + config.provider().subscribe("messages", lambda data, topic: wire.append(data)) + unsubscribe = transport.subscribe(received.append) + config.provider().publish("messages", b"\x00\x01\x00\x00") + assert received == [] + transport.broadcast(None, Int32(7)) + assert received == [Int32(7)] + assert wire[-1] == describe_message(Int32).codec.encode(Int32(7)) + unsubscribe() + transport.broadcast(None, Int32(8)) + assert received == [Int32(7)] + finally: + transport.stop() + config.provider().stop() diff --git a/dimos/protocol/pubsub/impl/zenohpubsub.py b/dimos/protocol/pubsub/impl/zenohpubsub.py index 802fc57a7a3..0ce32074b6a 100644 --- a/dimos/protocol/pubsub/impl/zenohpubsub.py +++ b/dimos/protocol/pubsub/impl/zenohpubsub.py @@ -18,12 +18,13 @@ from dataclasses import dataclass from functools import lru_cache import threading -from typing import Any, Literal +from typing import Any, Literal, cast import zenoh from dimos.msgs.helpers import resolve_msg_type -from dimos.protocol.pubsub.encoders import LCMEncoderMixin, PickleEncoderMixin +from dimos.protocol.pubsub.codec import describe_message +from dimos.protocol.pubsub.encoders import MessageEncoderMixin, PickleEncoderMixin from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic from dimos.protocol.pubsub.spec import AllPubSub from dimos.protocol.service.zenohservice import ZenohService @@ -91,8 +92,8 @@ def key_expr(self) -> str: Topic("dimos/cmd_vel", Twist) -> "dimos/cmd_vel/geometry_msgs.Twist" Topic("dimos/data") -> "dimos/data" """ - if self.lcm_type is not None: - return f"{self.pattern}/{self.lcm_type.msg_name}" + if self.message is not None and self.message.encoding == "lcm": + return f"{self.pattern}/{self.message.name}" return self.pattern @@ -118,6 +119,8 @@ def _key_expr_to_topic(key_expr: str, default_lcm_type: type | None = None) -> T "dimos/data" -> Topic("dimos/data", default_lcm_type) "dimos/data/unknown.Foo" -> Topic("dimos/data/unknown.Foo", default_lcm_type) """ + if default_lcm_type is not None and describe_message(default_lcm_type).encoding == "cdr": + return Topic(key_expr, message_type=default_lcm_type) # Try to resolve the last segment as a message type parts = key_expr.rsplit("/", 1) if len(parts) == 2: @@ -195,7 +198,7 @@ def publish(self, topic: Topic, message: bytes) -> None: publisher = self._get_publisher(key_expr, qos) publisher.put(message) except Exception: - logger.error(f"Error publishing to {key_expr}", exc_info=True) + logger.exception(f"Error publishing to {key_expr}") def subscribe( self, topic: Topic, callback: Callable[[bytes, Topic], None] @@ -207,7 +210,7 @@ def on_sample(sample: zenoh.Sample) -> None: try: data = sample.payload.to_bytes() except Exception: - logger.error(f"Error reading payload from {key_expr}", exc_info=True) + logger.exception(f"Error reading payload from {key_expr}") return # Concrete subscriptions only ever receive their own key, so the # subscribed topic can be passed through without re-parsing. @@ -215,7 +218,7 @@ def on_sample(sample: zenoh.Sample) -> None: if sample_key == key_expr: recv_topic = topic else: - recv_topic = _key_expr_to_topic(sample_key, topic.lcm_type) + recv_topic = _key_expr_to_topic(sample_key, cast("Hashable", topic.lcm_type)) callback(data, recv_topic) sub = self.session.declare_subscriber(key_expr, on_sample) @@ -279,7 +282,7 @@ def drain() -> None: try: callback(msg, topic) except Exception: - logger.error("Error in subscribe_all callback", exc_info=True) + logger.exception("Error in subscribe_all callback") thread = threading.Thread(target=drain, name="zenoh-subscribe-all", daemon=True) @@ -329,18 +332,14 @@ def stop(self) -> None: class Zenoh( # type: ignore[misc] - LCMEncoderMixin, + MessageEncoderMixin, ZenohPubSubBase, ): """Zenoh pub/sub with LCM encoding for typed DimosMsg.""" - ... - class PickleZenoh( PickleEncoderMixin, # type: ignore[type-arg] ZenohPubSubBase, ): """Zenoh pub/sub with pickle encoding for arbitrary Python objects.""" - - ... diff --git a/dimos/protocol/pubsub/test_codec.py b/dimos/protocol/pubsub/test_codec.py new file mode 100644 index 00000000000..712c789b589 --- /dev/null +++ b/dimos/protocol/pubsub/test_codec.py @@ -0,0 +1,72 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# 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. + +"""Codec contracts and channel metadata, without network access.""" + +import pickle + +import pytest + +from dimos.messages.geometry_msgs.msg import Vector3 as CDRVector3 +from dimos.messages.std_msgs.msg import Int32, String +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.protocol.pubsub.codec import describe_message +from dimos.protocol.pubsub.impl.lcmpubsub import Topic +from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic, _key_expr_to_topic + + +@pytest.mark.parametrize( + "payload", [b"", b"bad!", b"\x00\x01\x00\x00", b"\x00\x03\x00\x00" + b"\x00" * 4] +) +def test_cdr_rejects_invalid_payload_without_fallback(payload): + with pytest.raises(ValueError): + describe_message(Int32).codec.decode(payload) + + +@pytest.mark.parametrize( + "declared,message", + [(Int32, String("wrong")), (Vector3, CDRVector3(1, 2, 3)), (CDRVector3, Vector3(1, 2, 3))], +) +def test_codec_rejects_wrong_declared_type(declared, message): + with pytest.raises(TypeError, match="declared type"): + describe_message(declared).codec.encode(message) + + +def test_unsupported_and_ambiguous_types_fail_at_binding(monkeypatch): + with pytest.raises(TypeError, match="Unsupported"): + Topic("test", message_type=dict) + monkeypatch.setattr(Int32, "lcm_encode", lambda self: b"", raising=False) + monkeypatch.setattr(Int32, "lcm_decode", lambda data: None, raising=False) + with pytest.raises(TypeError, match="Ambiguous"): + Topic("test", message_type=Int32) + + +@pytest.mark.parametrize("topic_type", [Topic, ZenohTopic]) +def test_generic_and_legacy_constructor_names_are_compatible(topic_type): + topic = topic_type("test", message_type=Int32) + assert topic == topic_type("test", Int32) + assert pickle.loads(pickle.dumps(topic)) == topic + with pytest.raises(TypeError, match="Conflicting"): + topic_type("test", Vector3, message_type=Int32) + + +def test_legacy_channel_names_and_plain_cdr_topics(): + assert str(Topic("test", Vector3)) == "test#geometry_msgs.Vector3" + assert ZenohTopic("test", Vector3).key_expr == "test/geometry_msgs.Vector3" + cdr = Topic("test", message_type=Int32) + assert str(cdr) == "test" + assert Topic.from_channel_str(str(cdr), Int32) == cdr + zenoh = ZenohTopic("test", message_type=Int32) + assert zenoh.key_expr == "test" + assert _key_expr_to_topic(zenoh.key_expr, Int32) == zenoh diff --git a/dimos/protocol/pubsub/test_encoder.py b/dimos/protocol/pubsub/test_encoder.py index dec1e429726..31ca775a5c8 100644 --- a/dimos/protocol/pubsub/test_encoder.py +++ b/dimos/protocol/pubsub/test_encoder.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python3 - # Copyright 2025-2026 Dimensional Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,9 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import dataclass import json from typing import Any +from dimos_cdrgen.registry import encode as encode_cdr +import pytest + +from dimos.messages.std_msgs.msg import Int32, String +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.protocol.pubsub.codec import MessageDescriptor, describe_message +from dimos.protocol.pubsub.encoders import MessageEncoderMixin from dimos.protocol.pubsub.impl.memory import Memory, MemoryWithJSONEncoder @@ -169,3 +175,58 @@ def callback(message: Any, topic: str) -> None: # Verify it's actually JSON decoded_raw = json.loads(raw_message.decode("utf-8")) assert decoded_raw == original_message + + +@dataclass(frozen=True) +class _TypedTopic: + name: str + message: MessageDescriptor + + +class _EncodedMemory(MessageEncoderMixin, Memory): + pass + + +@pytest.mark.parametrize("message", [Vector3(1, 2, 3), Int32(42), String("hello")]) +def test_message_encoder_publish_subscribe_and_unsubscribe(message): + pubsub = _EncodedMemory() + topic = _TypedTopic("messages", describe_message(type(message))) + received = [] + wire = [] + remove_raw = Memory.subscribe(pubsub, topic, lambda data, topic: wire.append(data)) + unsubscribe = pubsub.subscribe(topic, lambda value, topic: received.append(value)) + try: + pubsub.publish(topic, message) + expected = message.lcm_encode() if isinstance(message, Vector3) else encode_cdr(message) + assert wire == [expected] + assert received == [message] + assert type(received[0]) is type(message) + unsubscribe() + pubsub.publish(topic, message) + assert received == [message] + assert wire == [expected, expected] + finally: + unsubscribe() + remove_raw() + + +def test_message_encoder_keeps_legacy_and_cdr_subscriptions_separate(): + pubsub = _EncodedMemory() + legacy = _TypedTopic("legacy", describe_message(Vector3)) + cdr = _TypedTopic("cdr", describe_message(Int32)) + received_legacy = [] + received_cdr = [] + stop_legacy = pubsub.subscribe(legacy, lambda value, topic: received_legacy.append(value)) + stop_cdr = pubsub.subscribe(cdr, lambda value, topic: received_cdr.append(value)) + try: + pubsub.publish(legacy, Vector3(1, 2, 3)) + Memory.publish(pubsub, cdr, b"broken") + assert received_cdr == [] + with pytest.raises(TypeError, match="declared type"): + pubsub.publish(cdr, Vector3(1, 2, 3)) + pubsub.publish(cdr, Int32(7)) + assert received_legacy == [Vector3(1, 2, 3)] + assert received_cdr == [Int32(7)] + finally: + stop_legacy() + stop_cdr() diff --git a/docs/usage/blueprints.md b/docs/usage/blueprints.md index 5420afc204b..dbaacccf50a 100644 --- a/docs/usage/blueprints.md +++ b/docs/usage/blueprints.md @@ -153,7 +153,7 @@ from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out, In -from dimos.msgs.sensor_msgs import Image +from dimos.msgs.sensor_msgs.Image import Image class ModuleA(Module): image: Out[Image] @@ -214,7 +214,7 @@ from dimos.core.coordination.blueprints import autoconnect from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out, In -from dimos.msgs.sensor_msgs import Image +from dimos.msgs.sensor_msgs.Image import Image class ConnectionModule(Module): color_image: Out[Image] # Outputs on 'color_image'