Export netimgui files in MuJoCo Copybara configuration for use with the upcoming Web Viewer

PiperOrigin-RevId: 954769526
Change-Id: I4cdee083082f7d55c2babd5e664ec0283a2aaf63

python/mujoco/experimental/netimgui/google/NetImgui_NetworkWASM.cpp

@@ -0,0 +1,324 @@
+// 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.
+
+// WASM networking backend for NetImgui.
+//
+// Implements the NetImgui::Internal::Network interface using Emscripten
+// WebSockets. This file is compiled only under the Emscripten toolchain and
+// provides the browser-side transport layer for netimgui draw data.
+#include "NetImgui_Shared.h"
+
+#ifndef __EMSCRIPTEN__
+#error "This file must be compiled with emscripten."
+#endif
+
+#include <emscripten/console.h>
+#include <emscripten/websocket.h>
+
+#include <algorithm>
+#include <atomic>
+#include <cstring>
+#include <mutex>
+#include <vector>
+
+#include "NetImgui_CmdPackets.h"
+#include "google/logging.h"
+#include "google/network_status.h"
+
+namespace NetImgui {
+namespace Internal {
+namespace Network {
+
+struct SocketInfo {
+  EMSCRIPTEN_WEBSOCKET_T mSocket = 0;
+  std::atomic<bool> mConnected{false};
+  std::atomic<bool> mError{false};
+  std::atomic<bool> mClosed{false};
+  std::atomic<int> mCloseCode{0};
+
+  std::vector<uint8_t> mBuffer;
+  std::mutex mBufferMutex;
+  int mSendSizeMax =
+      1024 * 1024;  // Interface compatibility with other backends
+};
+
+// --- WebSocket event callbacks ---
+
+static EM_BOOL OnWebSocketOpen(int /*event_type*/,
+                               const EmscriptenWebSocketOpenEvent* /*event*/,
+                               void* user_data) {
+  auto* socket = static_cast<SocketInfo*>(user_data);
+  if (socket) {
+    socket->mConnected = true;
+  }
+  return EM_TRUE;
+}
+
+static EM_BOOL OnWebSocketMessage(int /*event_type*/,
+                                  const EmscriptenWebSocketMessageEvent* event,
+                                  void* user_data) {
+  auto* socket = static_cast<SocketInfo*>(user_data);
+  if (socket && !event->isText) {
+    std::lock_guard<std::mutex> lock(socket->mBufferMutex);
+    size_t old_size = socket->mBuffer.size();
+    socket->mBuffer.insert(socket->mBuffer.end(), event->data,
+                           event->data + event->numBytes);
+    static int msg_count = 0;
+    ++msg_count;
+    VLOG(1, "onmessage #%d: %d bytes, buffer: %zu -> %zu", msg_count,
+         event->numBytes, old_size, socket->mBuffer.size());
+  } else if (socket && event->isText) {
+    VLOG(1, "onmessage: TEXT frame (%d bytes), IGNORED", event->numBytes);
+  }
+  return EM_TRUE;
+}
+
+static EM_BOOL OnWebSocketClose(int /*event_type*/,
+                                const EmscriptenWebSocketCloseEvent* event,
+                                void* user_data) {
+  auto* socket = static_cast<SocketInfo*>(user_data);
+  if (socket) {
+    socket->mClosed = true;
+    socket->mCloseCode = event->code;
+  }
+  return EM_TRUE;
+}
+
+static EM_BOOL OnWebSocketError(int /*event_type*/,
+                                const EmscriptenWebSocketErrorEvent* /*event*/,
+                                void* user_data) {
+  auto* socket = static_cast<SocketInfo*>(user_data);
+  if (socket) {
+    socket->mError = true;
+  }
+  return EM_TRUE;
+}
+
+// --- Network interface implementation ---
+
+bool Startup() { return emscripten_websocket_is_supported(); }
+
+void Shutdown() {}
+
+SocketInfo* Connect(const char* server_host, uint32_t /*server_port*/) {
+  if (!emscripten_websocket_is_supported()) return nullptr;
+
+  SocketInfo* socket_info = netImguiNew<SocketInfo>();
+
+  EmscriptenWebSocketCreateAttributes attr;
+  emscripten_websocket_init_create_attributes(&attr);
+
+  // The web viewer always passes a complete WebSocket URL (e.g.
+  // "ws://host:8080/ui", built from the page origin by WsUrl), so it is used
+  // verbatim and server_port is ignored — a browser reaches the viewer's
+  // paths on the page's own shared port, not a dedicated NetImgui port.
+  const std::string url(server_host);
+  attr.url = url.c_str();
+  attr.createOnMainThread = EM_TRUE;
+  LOG(Info, "Connecting WebSocket to: %s", url.c_str());
+
+  socket_info->mSocket = emscripten_websocket_new(&attr);
+  if (socket_info->mSocket <= 0) {
+    netImguiDelete(socket_info);
+    return nullptr;
+  }
+
+  emscripten_websocket_set_onopen_callback(socket_info->mSocket, socket_info,
+                                           OnWebSocketOpen);
+  emscripten_websocket_set_onmessage_callback(socket_info->mSocket, socket_info,
+                                              OnWebSocketMessage);
+  emscripten_websocket_set_onclose_callback(socket_info->mSocket, socket_info,
+                                            OnWebSocketClose);
+  emscripten_websocket_set_onerror_callback(socket_info->mSocket, socket_info,
+                                            OnWebSocketError);
+
+  return socket_info;
+}
+
+// Abandoned sockets are not freed immediately: with -pthread, websocket
+// events are queued across threads, so an already-queued close/error event
+// can still dereference the SocketInfo after emscripten_websocket_delete().
+// When the freed block was recycled for the next socket, such a late event
+// stamped a stale mClosed flag onto a healthy connection, which the
+// reconnect logic then tore down — a self-sustaining reconnect loop. Keep
+// abandoned sockets in a small ring and free them several disconnects
+// later, when any queued events are long gone.
+static SocketInfo* s_socket_graveyard[8] = {};
+static int s_socket_graveyard_idx = 0;
+
+void Disconnect(SocketInfo* client_socket) {
+  if (!client_socket) return;
+  client_socket->mClosed = true;
+  // Detach this socket from future events; queued events may already hold
+  // the pointer (the graveyard above covers those).
+  emscripten_websocket_set_onopen_callback(client_socket->mSocket, nullptr,
+                                           OnWebSocketOpen);
+  emscripten_websocket_set_onmessage_callback(client_socket->mSocket, nullptr,
+                                              OnWebSocketMessage);
+  emscripten_websocket_set_onclose_callback(client_socket->mSocket, nullptr,
+                                            OnWebSocketClose);
+  emscripten_websocket_set_onerror_callback(client_socket->mSocket, nullptr,
+                                            OnWebSocketError);
+  emscripten_websocket_close(client_socket->mSocket, 1000,
+                             "Normal Disconnection");
+  emscripten_websocket_delete(client_socket->mSocket);
+  {
+    // Release the receive buffer now; only the flags must stay valid.
+    std::lock_guard<std::mutex> lock(client_socket->mBufferMutex);
+    client_socket->mBuffer.clear();
+    client_socket->mBuffer.shrink_to_fit();
+  }
+  if (s_socket_graveyard[s_socket_graveyard_idx]) {
+    netImguiDelete(s_socket_graveyard[s_socket_graveyard_idx]);
+  }
+  s_socket_graveyard[s_socket_graveyard_idx] = client_socket;
+  s_socket_graveyard_idx = (s_socket_graveyard_idx + 1) % 8;
+}
+
+bool DataReceivePending(SocketInfo* client_socket) {
+  if (!client_socket) return false;
+
+  if (client_socket->mError || client_socket->mClosed) {
+    // Connection is dead — flush any buffered data so we stop processing
+    // stale commands that arrived before the close.
+    std::lock_guard<std::mutex> lock(client_socket->mBufferMutex);
+    if (!client_socket->mBuffer.empty()) {
+      LOG(Warning, "Connection closed/error. Discarding %zu buffered bytes.",
+          client_socket->mBuffer.size());
+      client_socket->mBuffer.clear();
+    }
+    return false;
+  }
+
+  std::lock_guard<std::mutex> lock(client_socket->mBufferMutex);
+  return !client_socket->mBuffer.empty();
+}
+
+void DataReceive(SocketInfo* client_socket, PendingCom& pending_rcv) {
+  if (!client_socket || !pending_rcv.pCommand) {
+    pending_rcv.bError = true;
+    return;
+  }
+
+  if (!client_socket->mConnected) {
+    pending_rcv.bError = false;  // Not ready yet, caller will retry.
+    return;
+  }
+
+  // The size field comes off the wire; a value smaller than what has already
+  // been read (e.g. a command header claiming < 8 bytes, from a corrupted or
+  // desynced stream) would underflow the subtraction below into a huge
+  // size_t and memcpy past the destination command buffer.
+  if (pending_rcv.pCommand->mSize < pending_rcv.SizeCurrent) {
+    LOG(Error, "DataReceive: wire size %u < %zu bytes already read; stream "
+        "is corrupt",
+        pending_rcv.pCommand->mSize,
+        static_cast<size_t>(pending_rcv.SizeCurrent));
+    pending_rcv.bError = true;
+    return;
+  }
+
+  size_t bytes_to_read = pending_rcv.pCommand->mSize - pending_rcv.SizeCurrent;
+  if (bytes_to_read == 0) return;
+
+  std::lock_guard<std::mutex> lock(client_socket->mBufferMutex);
+
+  VLOG(1, "DataReceive: want=%zu, have=%zu, cmd_size=%u, progress=%zu",
+       bytes_to_read, client_socket->mBuffer.size(),
+       pending_rcv.pCommand->mSize, pending_rcv.SizeCurrent);
+
+  if (client_socket->mBuffer.empty()) {
+    if (client_socket->mError || client_socket->mClosed) {
+      pending_rcv.bError = true;
+    }
+    return;
+  }
+
+  size_t bytes_to_consume =
+      std::min(bytes_to_read, client_socket->mBuffer.size());
+  if (bytes_to_consume > 0) {
+    memcpy(reinterpret_cast<uint8_t*>(pending_rcv.pCommand) +
+               pending_rcv.SizeCurrent,
+           client_socket->mBuffer.data(), bytes_to_consume);
+    client_socket->mBuffer.erase(
+        client_socket->mBuffer.begin(),
+        client_socket->mBuffer.begin() + bytes_to_consume);
+    pending_rcv.SizeCurrent += bytes_to_consume;
+    pending_rcv.bError = false;
+  }
+}
+
+void DataSend(SocketInfo* client_socket, PendingCom& pending_send) {
+  if (!client_socket || client_socket->mClosed || client_socket->mError ||
+      !pending_send.pCommand) {
+    pending_send.bError = true;
+    return;
+  }
+
+  if (!client_socket->mConnected) {
+    pending_send.bError = false;  // Not ready yet, caller will retry.
+    return;
+  }
+
+  size_t bytes_remaining =
+      pending_send.pCommand->mSize - pending_send.SizeCurrent;
+  if (bytes_remaining == 0) return;
+
+  EMSCRIPTEN_RESULT result = emscripten_websocket_send_binary(
+      client_socket->mSocket,
+      reinterpret_cast<uint8_t*>(pending_send.pCommand) +
+          pending_send.SizeCurrent,
+      bytes_remaining);
+
+  if (result == EMSCRIPTEN_RESULT_SUCCESS) {
+    pending_send.SizeCurrent += bytes_remaining;
+    pending_send.bError = false;
+  } else {
+    pending_send.bError = true;
+  }
+}
+
+SocketInfo* ListenStart(uint32_t /*listen_port*/) {
+  return nullptr;  // Browsers cannot open listening ports.
+}
+
+SocketInfo* ListenConnect(SocketInfo* /*listen_socket*/) { return nullptr; }
+
+int GetCloseCode(SocketInfo* client_socket) {
+  return client_socket ? client_socket->mCloseCode.load() : 0;
+}
+
+ReadyState GetReadyState(SocketInfo* client_socket) {
+  if (!client_socket) return ReadyState::kDisconnected;
+  if (client_socket->mError) return ReadyState::kError;
+  if (client_socket->mClosed) return ReadyState::kClosed;
+
+  uint16_t ready_state = 0;
+  emscripten_websocket_get_ready_state(client_socket->mSocket, &ready_state);
+  switch (ready_state) {
+    case 0:
+      return ReadyState::kConnecting;
+    case 1:
+      return ReadyState::kOpen;
+    case 2:
+      return ReadyState::kClosing;
+    case 3:
+      return ReadyState::kClosed;
+  }
+  return ReadyState::kError;
+}
+
+}  // namespace Network
+}  // namespace Internal
+}  // namespace NetImgui

python/mujoco/experimental/netimgui/google/logging.h

@@ -0,0 +1,100 @@
+// 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.
+
+#ifndef THIRD_PARTY_NETIMGUI_GOOGLE_LOGGING_H_
+#define THIRD_PARTY_NETIMGUI_GOOGLE_LOGGING_H_
+
+#include <time.h>
+
+#include <chrono>
+#include <cstdarg>
+#include <cstdio>
+#include <cstring>
+
+#if defined(__EMSCRIPTEN__)
+#include <emscripten/console.h>
+#endif
+
+namespace NetImgui {
+
+enum class LogSeverity { kInfo, kWarning, kError };
+
+// Set to 1 or higher to enable VLOG messages at runtime.
+constexpr int kLogVerbosity = 0;
+
+__attribute__((format(printf, 4, 5))) inline void NetImGuiLog(
+    LogSeverity severity, const char* file, int line, const char* fmt, ...) {
+  char buf[1024];
+  const char* basename = strrchr(file, '/');
+  basename = basename ? basename + 1 : file;
+
+  auto now_tp = std::chrono::system_clock::now();
+  time_t now = std::chrono::system_clock::to_time_t(now_tp);
+  auto now_usec = std::chrono::duration_cast<std::chrono::microseconds>(
+                      now_tp.time_since_epoch())
+                      .count() %
+                  1000000;
+  struct tm tm_info;
+  localtime_r(&now, &tm_info);
+
+  const char severity_char = severity == LogSeverity::kError     ? 'E'
+                             : severity == LogSeverity::kWarning ? 'W'
+                                                                 : 'I';
+
+  int prefix_len = snprintf(
+      buf, sizeof(buf), "%c%02d%02d %02d:%02d:%02d.%06d %s:%d] ", severity_char,
+      tm_info.tm_mon + 1, tm_info.tm_mday, tm_info.tm_hour, tm_info.tm_min,
+      tm_info.tm_sec, static_cast<int>(now_usec), basename, line);
+  if (prefix_len < 0 || prefix_len >= static_cast<int>(sizeof(buf))) {
+    prefix_len = 0;  // overwrite prefix on error or truncation
+  }
+
+  va_list args;
+  va_start(args, fmt);
+  vsnprintf(buf + prefix_len, sizeof(buf) - prefix_len, fmt, args);
+  va_end(args);
+
+  if (severity == LogSeverity::kInfo) {
+#if defined(__EMSCRIPTEN__)
+    emscripten_out(buf);
+#else
+    fputs(buf, stdout);
+    fputc('\n', stdout);
+    fflush(stdout);
+#endif
+  } else {
+#if defined(__EMSCRIPTEN__)
+    emscripten_err(buf);
+#else
+    fputs(buf, stderr);
+    fputc('\n', stderr);
+    fflush(stderr);
+#endif
+  }
+}
+
+}  // namespace NetImgui
+
+#define LOG(severity, fmt, ...)                                           \
+  ::NetImgui::NetImGuiLog(::NetImgui::LogSeverity::k##severity, __FILE__, \
+                          __LINE__, fmt, ##__VA_ARGS__)
+
+#define VLOG(level, fmt, ...)                                           \
+  do {                                                                  \
+    if (::NetImgui::kLogVerbosity >= (level))                           \
+      ::NetImgui::NetImGuiLog(::NetImgui::LogSeverity::kInfo, __FILE__, \
+                              __LINE__, fmt, ##__VA_ARGS__);            \
+  } while (0)
+
+#endif  // THIRD_PARTY_NETIMGUI_GOOGLE_LOGGING_H_

python/mujoco/experimental/netimgui/google/network_status.h

@@ -0,0 +1,75 @@
+// 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.
+
+// Connection-state query for the WASM WebSocket network backend.
+//
+// Stock NetImgui has no such query and does not need one: its Connect()
+// implementations block until the connection is established, so a non-null
+// SocketInfo is always usable and later failures surface as DataSend /
+// DataReceive errors inside the client thread. In the browser,
+// emscripten_websocket_new() returns a socket handle immediately while the
+// connection completes (or fails) asynchronously, and no client state
+// machine is watching it. Callers poll this state to hold back traffic
+// until the socket is actually open and to detect closure for reconnecting.
+
+#ifndef THIRD_PARTY_NETIMGUI_GOOGLE_NETWORK_STATUS_H_
+#define THIRD_PARTY_NETIMGUI_GOOGLE_NETWORK_STATUS_H_
+
+namespace NetImgui {
+namespace Internal {
+namespace Network {
+
+struct SocketInfo;
+
+enum class ReadyState {
+  kDisconnected,  // Null socket.
+  kConnecting,
+  kOpen,
+  kClosing,
+  kClosed,
+  kError,
+};
+
+// Implemented in NetImgui_NetworkWASM.cpp (Emscripten builds only).
+ReadyState GetReadyState(SocketInfo* client_socket);
+
+// The WebSocket close code once the socket has closed, else 0. Lets callers
+// distinguish a deliberate server-side rejection (e.g. 4001 = driver slot
+// taken) from an ordinary drop. Implemented in NetImgui_NetworkWASM.cpp.
+int GetCloseCode(SocketInfo* client_socket);
+
+// Human-readable state name, for status overlays and logs.
+inline const char* ReadyStateName(ReadyState state) {
+  switch (state) {
+    case ReadyState::kDisconnected:
+      return "Disconnected";
+    case ReadyState::kConnecting:
+      return "Connecting";
+    case ReadyState::kOpen:
+      return "Open";
+    case ReadyState::kClosing:
+      return "Closing";
+    case ReadyState::kClosed:
+      return "Closed";
+    case ReadyState::kError:
+      return "Error";
+  }
+  return "Unknown";
+}
+
+}  // namespace Network
+}  // namespace Internal
+}  // namespace NetImgui
+
+#endif  // THIRD_PARTY_NETIMGUI_GOOGLE_NETWORK_STATUS_H_