Repair Stream

dimos/agents/skills/gps_nav_skill.py

@@ -84,7 +84,7 @@ def set_gps_travel_points(self, *points: dict[str, float]) -> str:
         logger.info(f"Set travel points: {new_points}")
 
         if self.gps_goal._transport is not None:
-            self.gps_goal.publish(new_points)
+            self.gps_goal.publish(new_points)  # type: ignore[arg-type]
 
         if self._set_gps_travel_goal_points:
             self._set_gps_travel_goal_points(new_points)

dimos/core/stream.py

@@ -82,14 +82,16 @@ class State(enum.Enum):
 
 class Transport(Resource, ObservableMixin[T]):
     # used by local Output
-    def broadcast(self, selfstream: Out[T], value: T) -> None: ...
+    def broadcast(self, selfstream: Out[T], value: T) -> None:
+        raise NotImplementedError
+
+    # used by local Input
+    def subscribe(self, callback: Callable[[T], Any], selfstream: Stream[T]) -> Callable[[], None]:
+        raise NotImplementedError
 
     def publish(self, msg: T) -> None:
         self.broadcast(None, msg)  # type: ignore[arg-type]
 
-    # used by local Input
-    def subscribe(self, selfstream: In[T], callback: Callable[[T], any]) -> None: ...  # type: ignore[valid-type]
-
 
 class Stream(Generic[T]):
     _transport: Transport | None  # type: ignore[type-arg]
@@ -139,9 +141,11 @@ def __str__(self) -> str:
 
 class Out(Stream[T], ObservableMixin[T]):
     _transport: Transport  # type: ignore[type-arg]
+    _subscribers: list[Callable[[T], Any]]
 
     def __init__(self, *argv, **kwargs) -> None:  # type: ignore[no-untyped-def]
         super().__init__(*argv, **kwargs)
+        self._subscribers = []
 
     @property
     def transport(self) -> Transport[T]:
@@ -168,22 +172,19 @@ def __reduce__(self):  # type: ignore[no-untyped-def]
             ),
         )
 
-    def publish(self, msg) -> None:  # type: ignore[no-untyped-def]
-        if not hasattr(self, "_transport") or self._transport is None:
-            logger.warning(f"Trying to publish on Out {self} without a transport")
-            return
-        self._transport.broadcast(self, msg)
+    def publish(self, msg: T) -> None:
+        if hasattr(self, "_transport") and self._transport is not None:
+            self._transport.broadcast(self, msg)
+        for cb in self._subscribers:
+            cb(msg)
 
-    def subscribe(self, cb) -> Callable[[], None]:  # type: ignore[no-untyped-def]
-        """Subscribe to this output stream.
+    def subscribe(self, cb: Callable[[T], Any]) -> Callable[[], None]:
+        self._subscribers.append(cb)
 
-        Args:
-            cb: Callback function to receive messages
+        def unsubscribe() -> None:
+            self._subscribers.remove(cb)
 
-        Returns:
-            Unsubscribe function
-        """
-        return self.transport.subscribe(cb, self)  # type: ignore[arg-type, func-returns-value, no-any-return]
+        return unsubscribe
 
 
 class RemoteStream(Stream[T]):
@@ -206,8 +207,8 @@ class RemoteOut(RemoteStream[T]):
     def connect(self, other: RemoteIn[T]):  # type: ignore[no-untyped-def]
         return other.connect(self)
 
-    def subscribe(self, cb) -> Callable[[], None]:  # type: ignore[no-untyped-def]
-        return self.transport.subscribe(cb, self)  # type: ignore[arg-type, func-returns-value, no-any-return]
+    def subscribe(self, cb: Callable[[T], Any]) -> Callable[[], None]:
+        return self.transport.subscribe(cb, self)
 
 
 # representation of Input
@@ -249,8 +250,8 @@ def state(self) -> State:
         return State.UNBOUND if self.owner is None else State.READY
 
     # returns unsubscribe function
-    def subscribe(self, cb) -> Callable[[], None]:  # type: ignore[no-untyped-def]
-        return self.transport.subscribe(cb, self)  # type: ignore[arg-type, func-returns-value, no-any-return]
+    def subscribe(self, cb: Callable[[T], Any]) -> Callable[[], None]:
+        return self.transport.subscribe(cb, self)
 
 
 # representation of input outside of module

dimos/core/transport.py

@@ -25,7 +25,7 @@
     TypeVar,
 )
 
-from dimos.core.stream import In, Transport
+from dimos.core.stream import In, Out, Stream, Transport
 from dimos.protocol.pubsub.jpeg_shm import JpegSharedMemory
 from dimos.protocol.pubsub.lcmpubsub import LCM, JpegLCM, PickleLCM, Topic as LCMTopic
 from dimos.protocol.pubsub.shmpubsub import PickleSharedMemory, SharedMemory
@@ -60,18 +60,20 @@ def __init__(self, topic: str, **kwargs) -> None:  # type: ignore[no-untyped-def
     def __reduce__(self):  # type: ignore[no-untyped-def]
         return (pLCMTransport, (self.topic,))
 
-    def broadcast(self, _, msg) -> None:  # type: ignore[no-untyped-def]
+    def broadcast(self, _: Out[T] | None, msg: T) -> None:
         if not self._started:
             self.lcm.start()
             self._started = True
 
         self.lcm.publish(self.topic, msg)
 
-    def subscribe(self, callback: Callable[[T], None], selfstream: In[T] = None) -> None:  # type: ignore[assignment, override]
+    def subscribe(
+        self, callback: Callable[[T], Any], selfstream: Stream[T] | None = None
+    ) -> Callable[[], None]:
         if not self._started:
             self.lcm.start()
             self._started = True
-        return self.lcm.subscribe(self.topic, lambda msg, topic: callback(msg))  # type: ignore[return-value]
+        return self.lcm.subscribe(self.topic, lambda msg, topic: callback(msg))
 
     def start(self) -> None: ...