MuJoCo Web Viewer: add web client containing code that runs in the browser

PiperOrigin-RevId: 956698568 Change-Id: Ia4bebcb25b488d255994018da03e9115187b890c
2026-07-30 13:16:41 -07:00
parent 8b78378868
commit 4dd70d2367
11 changed files with 3043 additions and 9 deletions
@@ -0,0 +1,249 @@
+// Copyright 2026 DeepMind Technologies Limited
+//
+// 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
+//
+//     https://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.
+
+// The session: this page's relationship with the Python-side viewer.
+//
+// Owns the /state WebSocket (simulation payloads in, control messages out),
+// the session wire protocol (roster, grants, heartbeats, acks), the role
+// state machine (claiming -> controlling / spectating), and the
+// model-change / page-reload / close-code policies. Everything outside the
+// session goes through the Callbacks interface: applying a payload to the
+// scene, and driving the remote UI stream on role transitions. The session
+// also implements SessionActions, so the role window's intents land here
+// directly.
+
+#ifndef MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_SESSION_H_
+#define MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_SESSION_H_
+
+#include <emscripten/websocket.h>
+
+#include <cstdint>
+#include <optional>
+#include <string>
+
+#include "state_payload.h"
+
+namespace mujoco::studio {
+
+// Our custom WebSocket close codes in range 4xxx. The client shows a notice
+// and retries slowly.
+constexpr int kWsCloseControllerTaken = 4001;  // /ui: another browser controls.
+constexpr int kWsCloseSessionFull = 4002;      // /state: spectator limit hit.
+constexpr int kWsCloseInactive = 4003;         // /state: hidden tab kicked.
+constexpr int kWsCloseNotController = 4004;    // /drop: only controller may
+                                               //   load models.
+
+// The page's role in the collaborative session. Every page starts by
+// claiming the controller slot; the claim either succeeds (kControlling)
+// or the page settles into spectating. A control grant puts a spectator
+// back into kClaiming while it reconnects to /ui.
+enum class SessionRole {
+  kClaiming = 0,  // /ui claim in flight; the role is not yet resolved.
+  kControlling,   // This page holds the open /ui connection.
+  kSpectating,    // Another page controls; scene + local role window only.
+};
+
+// Read-only snapshot of the session, passed to the local UI each frame.
+struct SessionView {
+  SessionRole role = SessionRole::kClaiming;
+  int viewers = 0;
+  int queue_pos = 0;  // 1-based position in the control queue; 0 = unqueued.
+  int queue_len = 0;
+  int max_spectators = 0;
+  uint64_t gui_bytes_per_sec = 0;
+  uint64_t sim_bytes_per_sec = 0;
+  bool have_remote_frame = false;
+  int camera_mode = 0;  // [SpectatorCamMode].
+};
+
+// User intent reported by the role window. Session implements this; the
+// interface is the complete list of effects the local UI can cause.
+class SessionActions {
+ public:
+  virtual ~SessionActions() = default;
+  virtual void RequestControl() = 0;
+  virtual void LeaveQueue() = 0;
+  virtual void StealControl() = 0;
+  virtual void ReleaseControl() = 0;
+  virtual void SetCameraMode(int mode) = 0;      // mode is [SpectatorCamMode].
+  virtual void SetMaxSpectators(int count) = 0;  // Already clamped by the UI.
+};
+
+// The roster: the server's membership broadcast, sent as a text frame on
+// /state whenever the session changes (a viewer joins or leaves, queues
+// for control, or control moves). It tells this page how many viewers are
+// connected, which role the server currently assigns it, and where it
+// stands in the control queue.
+struct Roster {
+  int viewers = 0;
+  bool spectator = false;  // The server's view: true = not the controller.
+  int queue_pos = 0;  // 1-based position in the control queue; 0 = unqueued.
+  int queue_len = 0;
+  int max_spectators = 8;  // Runtime spectator limit.
+};
+
+// Parses a roster line; returns false when text is not a roster.
+bool ParseRoster(const char* text, Roster* roster);
+
+// The remote UI stream's connection state, reported to the role state
+// machine by the app once per frame.
+enum class RemoteUiState {
+  kNoSocket = 0,   // No connection attempt exists.
+  kConnecting,     // In flight (or closing); the machine waits.
+  kOpen,           // The claim succeeded: this page controls.
+  kClosedOrError,  // Rejected or dropped; the machine retries or settles.
+};
+
+class Session : public SessionActions {
+ public:
+  // Everything the session needs from the rest of the application.
+  class Callbacks {
+   public:
+    virtual ~Callbacks() = default;
+    // Payloads are dropped until this returns true (model loaded).
+    virtual bool ReadyForPayload() = 0;
+    // Applies a parsed payload to the application.
+    virtual void OnPayload(const StatePayloadView& view) = 0;
+    // Role transition: claim the controller slot.
+    virtual void ConnectRemoteUi() = 0;
+    // Role transition: drop the stream when spectating
+    virtual void ShutdownRemoteUi() = 0;
+    // Spectator camera mode change.
+    virtual void SetCameraMode(int mode) = 0;
+  };
+
+  explicit Session(Callbacks& callbacks) : callbacks_(callbacks) {}
+
+  void Connect(const std::string& url);
+
+  // Records the CRC32 of the model this page actually loaded.
+  // Must match zlib.crc32 (used to compute model_crc32 in web_viewer.py).
+  void SetModelCrc32(uint32_t crc) { model_crc32_ = crc; }
+
+  // True while a connect attempt exists; used to pace reconnects.
+  // emscripten_websocket_new returns a handle immediately, so this is NOT the
+  // same as Connected().
+  bool HasSocket() const { return socket_ != 0; }
+
+  // True only while the WebSocket is actually open.
+  bool Connected() const { return connected_; }
+
+  // True once a payload with a new model has scheduled a page reload; all
+  // traffic is dropped from then on.
+  bool ReloadPending() const { return reload_pending_; }
+
+  // The close code from the server deliberately ending this connection (codes
+  // 4000-4999, e.g. kWsCloseSessionFull), else 0. Such conditions are transient
+  // (a slot frees up, the user returns to the tab), so the page shows a notice
+  // and retries slowly; the code clears when a connection opens again.
+  int ServerCloseCode() const { return server_close_code_; }
+
+  // Returns the bytes received since the last call and resets the counter.
+  uint64_t ConsumeByteCount() {
+    uint64_t bytes = bytes_accum_;
+    bytes_accum_ = 0;
+    return bytes;
+  }
+
+  // Wall-clock seconds of the last received message, or 0 before the first one.
+  // Payloads stream at ~60Hz while the Python side is alive, so staleness here
+  // means the server is gone, even if the socket still looks open (a suspended
+  // process keeps its sockets established).
+  double LastMessageTime() const { return last_message_time_; }
+
+  SessionRole Role() const { return role_; }
+
+  // Fills the role and roster fields of the view.
+  void FillView(SessionView* view) const;
+
+  // Periodic session upkeep (the ~30s liveness heartbeat); call once per frame.
+  void Update();
+
+  // Feeds the role state machine the remote UI stream's connection state; call
+  // once per frame. Owns claim retry pacing, promotion to kControlling when a
+  // claim opens, instant settling when an open stream closes with
+  // kWsCloseControllerTaken (ousted by Steal Control), and the retries then
+  // settle rule for rejected claims.
+  void HandleRemoteUiState(RemoteUiState state, int close_code);
+
+  // Parses one WebSocket message and applies the model-change/reload policy.
+  void HandleMessage(const uint8_t* data, uint32_t num_bytes);
+
+  // SessionActions (used to implement the role window UI).
+  void RequestControl() override;
+  void LeaveQueue() override;
+  void StealControl() override;
+  void ReleaseControl() override;
+  void SetCameraMode(int mode) override;
+  void SetMaxSpectators(int count) override;
+
+ private:
+  // Detaches callbacks and frees socket_ (if any), resetting to disconnected.
+  void CloseSocket();
+
+  // Sends a session message (control requests, acks, activity reports) to the
+  // server as a text frame. Dropped silently while not connected.
+  void SendText(const char* text);
+
+  // Updates the role and mirrors it into JS (Module.isSpectator), which gates
+  // controller-only page behavior (model drag-and-drop upload).
+  void SetRole(SessionRole role);
+
+  // Routes a session text frame: roster updates, control grants.
+  void OnSessionText(const char* text);
+
+  static EM_BOOL OnWsMessage(int event_type,
+                             const EmscriptenWebSocketMessageEvent* event,
+                             void* user_data);
+  static EM_BOOL OnWsOpen(int event_type,
+                          const EmscriptenWebSocketOpenEvent* event,
+                          void* user_data);
+  static EM_BOOL OnWsError(int event_type,
+                           const EmscriptenWebSocketErrorEvent* event,
+                           void* user_data);
+  static EM_BOOL OnWsClose(int event_type,
+                           const EmscriptenWebSocketCloseEvent* event,
+                           void* user_data);
+
+  Callbacks& callbacks_;
+
+  EMSCRIPTEN_WEBSOCKET_T socket_ = 0;
+  bool connected_ = false;
+
+  // CRC32 of the model this page loaded. When the payload's crc changes, the
+  // Python side has swapped models so reload the page; this refetches
+  // /model.mjb and reconnects everything.
+  std::optional<uint32_t> model_crc32_;
+  bool reload_pending_ = false;
+
+  int server_close_code_ = 0;
+
+  uint64_t bytes_accum_ = 0;
+  double last_message_time_ = 0;
+
+  // Role state machine and roster.
+  SessionRole role_ = SessionRole::kClaiming;
+  Roster roster_;
+
+  // Consecutive rejected /ui claims; the page eventually stops claiming and
+  // settles into spectating.
+  int ui_reject_count_ = 0;
+  RemoteUiState remote_ui_state_ = RemoteUiState::kNoSocket;
+  double last_ui_retry_time_ = 0;
+  double last_heartbeat_time_ = 0;
+};
+
+}  // namespace mujoco::studio
+
+#endif  // MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_SESSION_H_