Skip to content

feat: add websockets for dotbot commands #201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aabadie merged 3 commits into from
Jan 23, 2026
Merged

feat: add websockets for dotbot commands #201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
171dbd5
feat: add websockets
WilliamTakeshi Jan 21, 2026
63c4bfb
test: add case for ws
WilliamTakeshi Jan 23, 2026
f3501cd
test: case for invalid messages on WS
WilliamTakeshi Jan 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 42 additions & 20 deletions dotbot/examples/charging_station.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,12 @@
DotBotMoveRawCommandModel,
DotBotRgbLedCommandModel,
DotBotWaypoints,
WSRgbLed,
WSWaypoints,
)
from dotbot.protocol import ApplicationType
from dotbot.rest import RestClient
from dotbot.rest import RestClient, rest_client
from dotbot.websocket import DotBotWsClient

THRESHOLD = 30 # Acceptable distance error to consider a waypoint reached
DT = 0.05 # Control loop period (seconds)
Expand All @@ -40,16 +43,18 @@

async def queue_robots(
client: RestClient,
ws: DotBotWsClient,
dotbots: List[DotBotModel],
params: OrcaParams,
) -> None:
sorted_bots = order_bots(dotbots, QUEUE_HEAD_X, QUEUE_HEAD_Y)
goals = assign_queue_goals(sorted_bots, QUEUE_HEAD_X, QUEUE_HEAD_Y, QUEUE_SPACING)
await send_to_goal(client, goals, params)
await send_to_goal(client, ws, goals, params)


async def charge_robots(
client: RestClient,
ws: DotBotWsClient,
params: OrcaParams,
) -> None:
dotbots = await client.fetch_active_dotbots()
Expand All @@ -76,7 +81,7 @@ async def charge_robots(
"x": PARK_X,
"y": PARK_Y + parked_count * PARK_SPACING,
}
await send_to_goal(client, goals, params)
await send_to_goal(client, ws, goals, params)

if len(remaining) == 0:
break
Expand Down Expand Up @@ -123,6 +128,7 @@ async def disengage_from_charger(client: RestClient, dotbot: DotBotModel):

async def send_to_goal(
client: RestClient,
ws: DotBotWsClient,
goals: Dict[str, dict],
params: OrcaParams,
) -> None:
Expand Down Expand Up @@ -178,11 +184,15 @@ async def send_to_goal(
)
],
)
await client.send_waypoint_command(
address=agent.id,
application=ApplicationType.DotBot,
command=waypoints,
await ws.send(
WSWaypoints(
cmd="waypoints",
address=agent.id,
application=ApplicationType.DotBot,
data=waypoints,
)
)

await asyncio.sleep(DT)
return None

Expand Down Expand Up @@ -299,22 +309,34 @@ async def main() -> None:
url = os.getenv("DOTBOT_CONTROLLER_URL", "localhost")
port = os.getenv("DOTBOT_CONTROLLER_PORT", "8000")
use_https = os.getenv("DOTBOT_CONTROLLER_USE_HTTPS", False)
client = RestClient(url, port, use_https)

dotbots = await client.fetch_active_dotbots()
async with rest_client(url, port, use_https) as client:
dotbots = await client.fetch_active_dotbots()

# Cosmetic: all bots are red
for dotbot in dotbots:
await client.send_rgb_led_command(
address=dotbot.address,
command=DotBotRgbLedCommandModel(red=255, green=0, blue=0),
)
ws = DotBotWsClient(url, port)
await ws.connect()
try:
# Cosmetic: all bots are red
for dotbot in dotbots:
await ws.send(
WSRgbLed(
cmd="rgb_led",
address=dotbot.address,
application=ApplicationType.DotBot,
data=DotBotRgbLedCommandModel(
red=255,
green=0,
blue=0,
),
)
)

# Phase 1: initial queue
await queue_robots(client, dotbots, params)
# Phase 1: initial queue
await queue_robots(client, ws, dotbots, params)

# Phase 2: charging loop
await charge_robots(client, params)
# Phase 2: charging loop
await charge_robots(client, ws, params)
finally:
await ws.close()

return None

Expand Down
30 changes: 29 additions & 1 deletion dotbot/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
# pylint: disable=too-few-public-methods,no-name-in-module

from enum import IntEnum
from typing import Any, List, Optional, Union
from typing import Any, List, Literal, Optional, Union

from pydantic import BaseModel

Expand Down Expand Up @@ -172,3 +172,31 @@ class DotBotModel(BaseModel):
position_history: List[Union[DotBotLH2Position, DotBotGPSPosition]] = []
calibrated: bool = False
battery: float = 0.0 # Voltage in Volts


class WSBase(BaseModel):
cmd: str
address: str
application: ApplicationType


class WSRgbLed(WSBase):
cmd: Literal["rgb_led"]
data: DotBotRgbLedCommandModel


class WSMoveRaw(WSBase):
cmd: Literal["move_raw"]
data: DotBotMoveRawCommandModel


class WSWaypoints(WSBase):
cmd: Literal["waypoints"]
data: DotBotWaypoints


WSMessage = Union[
WSRgbLed,
WSMoveRaw,
WSWaypoints,
]
68 changes: 68 additions & 0 deletions dotbot/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from fastapi.staticfiles import StaticFiles
from pydantic import TypeAdapter, ValidationError
from starlette.middleware.base import BaseHTTPMiddleware

from dotbot import pydotbot_version
Expand All @@ -25,6 +26,10 @@
DotBotQueryModel,
DotBotRgbLedCommandModel,
DotBotWaypoints,
WSMessage,
WSMoveRaw,
WSRgbLed,
WSWaypoints,
)
from dotbot.protocol import (
ApplicationType,
Expand All @@ -40,6 +45,8 @@
"PYDOTBOT_FRONTEND_BASE_URL", "https://dotbots.github.io/PyDotBot"
)

ws_adapter = TypeAdapter(WSMessage)


class ReverseProxyMiddleware(BaseHTTPMiddleware):

Expand Down Expand Up @@ -98,6 +105,10 @@ async def dotbots_move_raw(
if address not in api.controller.dotbots:
raise HTTPException(status_code=404, detail="No matching dotbot found")

_dotbots_move_raw(address=address, command=command)


def _dotbots_move_raw(address: str, command: DotBotMoveRawCommandModel):
payload = PayloadCommandMoveRaw(
left_x=command.left_x,
left_y=command.left_y,
Expand All @@ -120,6 +131,10 @@ async def dotbots_rgb_led(
if address not in api.controller.dotbots:
raise HTTPException(status_code=404, detail="No matching dotbot found")

_dotbots_rgb_led(address=address, command=command)


def _dotbots_rgb_led(address: str, command: DotBotRgbLedCommandModel):
payload = PayloadCommandRgbLed(
red=command.red, green=command.green, blue=command.blue
)
Expand All @@ -141,6 +156,16 @@ async def dotbots_waypoints(
if address not in api.controller.dotbots:
raise HTTPException(status_code=404, detail="No matching dotbot found")

await _dotbots_waypoints(
address=address, application=application, waypoints=waypoints
)


async def _dotbots_waypoints(
address: str,
application: int,
waypoints: DotBotWaypoints,
):
waypoints_list = waypoints.waypoints
if application == ApplicationType.SailBot.value:
if api.controller.dotbots[address].gps_position is not None:
Expand Down Expand Up @@ -236,6 +261,49 @@ async def websocket_endpoint(websocket: WebSocket):
api.controller.websockets.remove(websocket)


@api.websocket("/controller/ws/dotbots")
async def ws_dotbots(websocket: WebSocket):
Copy link
Contributor

@aabadie aabadie Jan 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be possible to add a test function for this new endpoint? server.py is 100% covered by test (there are not so many modules in this case, so let's not break that one ;) )

Copy link
Contributor Author

@WilliamTakeshi WilliamTakeshi Jan 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good call! Testcases added on:

63c4bfb (this PR)
f3501cd (this PR)

await websocket.accept()
try:
while True:
raw = await websocket.receive_json()

try:
msg = ws_adapter.validate_python(raw)
except ValidationError as e:
await websocket.send_json(
{
"error": "invalid_message",
"details": e.errors(),
}
)
continue

if msg.address not in api.controller.dotbots:
# ignore messages where address doesn't exist
continue

if isinstance(msg, WSRgbLed):
_dotbots_rgb_led(
address=msg.address,
command=msg.data,
)
elif isinstance(msg, WSMoveRaw):
_dotbots_move_raw(
address=msg.address,
command=msg.data,
)
elif isinstance(msg, WSWaypoints):
await _dotbots_waypoints(
address=msg.address,
application=msg.application,
waypoints=msg.data,
)

except WebSocketDisconnect:
LOGGER.debug("WebSocket client disconnected")


# Mount static files after all routes are defined
FRONTEND_DIR = os.path.join(os.path.dirname(__file__), "frontend", "build")
api.mount("/PyDotBot", StaticFiles(directory=FRONTEND_DIR, html=True), name="PyDotBot")
59 changes: 57 additions & 2 deletions dotbot/tests/test_experiment_charging_station.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
DotBotRgbLedCommandModel,
DotBotStatus,
DotBotWaypoints,
WSMessage,
)
from dotbot.protocol import ApplicationType

Expand Down Expand Up @@ -130,6 +131,56 @@ async def send_rgb_led_command(
self.rgb_commands.append((address, command))


class FakeDotBotWsClient:
"""
Fake WebSocket client for testing control logic.

- Accepts typed WSMessage objects
- Dispatches to FakeRestClient logic
- Records all WS messages for assertions
"""

def __init__(self, rest_client: FakeRestClient):
self.rest = rest_client
self.sent_messages: list[WSMessage] = []
self.connected = False

async def connect(self):
self.connected = True

async def close(self):
self.connected = False

async def send(self, msg: WSMessage):
if not self.connected:
raise RuntimeError("FakeDotBotWsClient is not connected")

self.sent_messages.append(msg)

if msg.cmd == "rgb_led":
await self.rest.send_rgb_led_command(
address=msg.address,
command=msg.data,
)

elif msg.cmd == "move_raw":
await self.rest.send_move_raw_command(
address=msg.address,
application=msg.application,
command=msg.data,
)

elif msg.cmd == "waypoints":
await self.rest.send_waypoint_command(
address=msg.address,
application=msg.application,
command=msg.data,
)

else:
raise ValueError(f"Unknown WS command: {msg.cmd}")


def fake_bot(address: str, x: float, y: float) -> DotBotModel:
return DotBotModel(
address=address,
Expand All @@ -151,9 +202,11 @@ async def test_queue_robots_converges_to_queue_positions(_):
]

client = FakeRestClient(bots)
ws = FakeDotBotWsClient(client)
await ws.connect()
params = OrcaParams(time_horizon=5 * DT, time_step=DT)

await queue_robots(client, bots, params)
await queue_robots(client, ws, bots, params)

# Bots should be ordered A, B, C along the queue
expected = {
Expand Down Expand Up @@ -184,9 +237,11 @@ async def test_charge_robots_moves_all_bots_to_parking(_):
]

client = FakeRestClient(bots)
ws = FakeDotBotWsClient(client)
await ws.connect()
params = OrcaParams(time_horizon=5 * DT, time_step=DT)

await charge_robots(client, params)
await charge_robots(client, ws, params)

# --- Assertions: all bots parked ---
# Bots should be ordered A, B, C along the park slots
Expand Down
Loading