Mujoco_WASM/web_platform/src/robot/ExternalControlClient.ts

import type { PhysicsAdapter } from '../simulation/PhysicsAdapter';
import { RobotError, type RobotIdentity } from './types';
import { record, exactKeys, integer, identifier } from './validation';

export interface BridgeStatus {
  phase: 'disconnected' | 'connecting' | 'ready';
  error?: string;
}
type Backend = Pick<
  PhysicsAdapter,
  | 'describeRobot'
  | 'robotObservation'
  | 'externalControlStatus'
  | 'claimExternalControlLease'
  | 'sendRobotAction'
  | 'stopExternalControl'
  | 'reset'
>;

/** Bounded numeric transport only. Never owns physics or a render/UI snapshot loop. */
export class ExternalControlClient {
  private socket?: WebSocket;
  private timer?: ReturnType<typeof setInterval>;
  private handshakeTimer?: ReturnType<typeof setTimeout>;
  private sessionId?: string;
  private lease?: RobotIdentity;
  private leaseGeneration?: number;
  private rejectConnect?: (error: Error) => void;
  private statusValue: BridgeStatus = { phase: 'disconnected' };
  constructor(
    private readonly backend: () => Backend | undefined,
    private readonly onStatus: (status: BridgeStatus) => void = () => {},
  ) {}
  status(): BridgeStatus {
    return { ...this.statusValue };
  }
  private statusChanged(value: BridgeStatus): void {
    this.statusValue = value;
    this.onStatus(value);
  }
  private state() {
    const backend = this.backend(),
      observation = backend?.robotObservation(),
      control = backend?.externalControlStatus();
    if (
      !backend ||
      !observation ||
      !control ||
      (this.sessionId && observation.sessionId !== this.sessionId)
    )
      throw new RobotError('DISCONNECTED', '模型会话已变化，请重新连接');
    return {
      type: 'state',
      observation,
      enabled: control.enabled,
      authorizationGeneration: control.authorizationGeneration,
    };
  }
  async connect(endpoint: string, token: string): Promise<void> {
    this.disconnect('重新连接');
    let url: URL;
    try {
      url = new URL(endpoint);
    } catch {
      return Promise.reject(new RobotError('INVALID_MESSAGE', '桥接地址无效'));
    }
    if (
      url.protocol !== 'http:' ||
      !['127.0.0.1', 'localhost'].includes(url.hostname) ||
      url.username ||
      url.password ||
      url.search ||
      url.hash ||
      url.pathname !== '/'
    )
      return Promise.reject(
        new RobotError('UNAUTHORIZED', 'V1 仅允许本机 http://127.0.0.1:port 或 localhost'),
      );
    if (!/^[\x21-\x7e]{16,4096}$/.test(token.trim()))
      return Promise.reject(
        new RobotError('UNAUTHORIZED', '控制 token 需16–4096位可打印 ASCII，不能含空白'),
      );
    const descriptor = this.backend()?.describeRobot(),
      initial = this.state();
    if (!descriptor) return Promise.reject(new RobotError('UNSUPPORTED', '请先应用机器人 profile'));
    this.sessionId = initial.observation.sessionId;
    url.protocol = 'ws:';
    url.pathname = '/ws/control/v1';
    const socket = new WebSocket(url);
    this.socket = socket;
    this.statusChanged({ phase: 'connecting' });
    return new Promise<void>((resolve, reject) => {
      this.rejectConnect = reject;
      this.handshakeTimer = setTimeout(() => this.disconnect('桥接认证超时'), 5000);
      socket.onopen = () => {
        if (this.socket === socket)
          this.send({ type: 'auth', protocolVersion: 1, token: token.trim() });
      };
      socket.onerror = () => {
        if (this.socket === socket) this.disconnect('无法连接本机控制桥接');
      };
      socket.onclose = () => {
        if (this.socket === socket) this.disconnect('桥接已断开，请重新授权');
      };
      socket.onmessage = (event) => {
        if (this.socket !== socket) return;
        void this.message(event.data, socket, descriptor, resolve).catch((error: unknown) => {
          if (this.socket === socket)
            this.disconnect(error instanceof Error ? error.message : '桥接协议错误');
        });
      };
    });
  }
  private send(message: unknown): void {
    const json = JSON.stringify(message),
      socket = this.socket;
    if (
      !socket ||
      socket.readyState !== WebSocket.OPEN ||
      socket.bufferedAmount + new TextEncoder().encode(json).length > 65536
    )
      throw new RobotError('DISCONNECTED', '桥接发送缓冲区已满或已关闭');
    socket.send(json);
  }
  private async message(
    raw: unknown,
    socket: WebSocket,
    descriptor: unknown,
    ready: () => void,
  ): Promise<void> {
    if (typeof raw !== 'string' || new TextEncoder().encode(raw).length > 65536)
      throw new RobotError('INVALID_MESSAGE', '无效桥接帧');
    const message = record(JSON.parse(raw));
    if (message.type === 'authenticated') {
      exactKeys(message, ['type']);
      this.send({ ...this.state(), type: 'register', descriptor });
    } else if (message.type === 'ready') {
      exactKeys(message, ['type']);
      clearTimeout(this.handshakeTimer);
      this.rejectConnect = undefined;
      this.statusChanged({ phase: 'ready' });
      if (this.timer) clearInterval(this.timer);
      this.timer = setInterval(() => this.sync(), 1000 / 30);
      ready();
    } else if (message.type === 'error') {
      const error = record(message.error);
      throw new RobotError(
        'DISCONNECTED',
        typeof error.message === 'string' ? error.message : '桥接错误',
      );
    } else if (message.type === 'stop') {
      exactKeys(message, ['type', 'reason', 'sessionId', 'authorizationGeneration']);
      const current = this.state();
      if (
        message.sessionId === current.observation.sessionId &&
        message.authorizationGeneration === current.authorizationGeneration
      ) {
        this.backend()?.stopExternalControl(
          typeof message.reason === 'string' ? message.reason : '服务端撤销控制',
        );
        this.lease = undefined;
        this.sync();
      }
    } else if (message.type === 'request') {
      exactKeys(message, ['type', 'id', 'op', 'payload']);
      const id = identifier(message.id);
      let value: unknown,
        ok = true;
      try {
        value = await this.request(message.op, message.payload);
      } catch (error: unknown) {
        ok = false;
        value = {
          code: error instanceof RobotError ? error.code : 'DISCONNECTED',
          message: error instanceof Error ? error.message.slice(0, 1024) : '控制请求失败',
        };
      }
      // A delayed physics acknowledgement must never land on a replacement socket.
      if (this.socket === socket) this.send({ type: 'result', id, ok, value });
    } else throw new RobotError('INVALID_MESSAGE', '未知桥接消息类型');
  }
  private async request(op: unknown, payload: unknown): Promise<unknown> {
    if (this.statusValue.phase !== 'ready') throw new RobotError('UNAUTHORIZED', '桥接握手未完成');
    const current = this.state(),
      backend = this.backend()!;
    const data = record(payload);
    if (op === 'claim') {
      exactKeys(data, [
        'sessionId',
        'modelEpoch',
        'leaseId',
        'modelFingerprint',
        'authorizationGeneration',
      ]);
      identifier(data.sessionId);
      identifier(data.leaseId);
      integer(data.modelEpoch, 'modelEpoch');
      if (
        data.sessionId !== current.observation.sessionId ||
        data.modelEpoch !== current.observation.modelEpoch ||
        data.authorizationGeneration !== current.authorizationGeneration ||
        data.modelFingerprint !== backend.describeRobot()?.modelFingerprint
      )
        throw new RobotError('STALE', '模型或授权代次已变化');
      this.lease = backend.claimExternalControlLease(data.leaseId as string);
      this.leaseGeneration = current.authorizationGeneration;
      return this.lease;
    }
    if (
      !this.lease ||
      !current.enabled ||
      !backend.externalControlStatus()?.connected ||
      current.authorizationGeneration !== this.leaseGeneration ||
      current.observation.modelEpoch !== this.lease.modelEpoch ||
      data.leaseId !== this.lease.leaseId ||
      data.sessionId !== this.lease.sessionId ||
      data.modelEpoch !== this.lease.modelEpoch
    )
      throw new RobotError('STALE', '控制租约失效');
    if (op === 'action') return backend.sendRobotAction(data);
    exactKeys(data, ['sessionId', 'modelEpoch', 'leaseId']);
    if (op === 'release') {
      backend.stopExternalControl('Python 控制者已断开');
      this.lease = undefined;
      return { released: true };
    }
    if (op === 'reset') {
      backend.reset();
      this.lease = undefined;
      return backend.robotObservation();
    }
    throw new RobotError('UNSUPPORTED', '不支持的机器人操作');
  }
  /** Also called after explicit UI authorization; never resends a stale observation as fresh. */
  sync(): void {
    if (this.statusValue.phase !== 'ready') return;
    try {
      this.send(this.state());
    } catch (error: unknown) {
      this.disconnect(error instanceof Error ? error.message : '桥接状态异常');
    }
  }
  disconnect(reason = '桥接连接已关闭'): void {
    clearInterval(this.timer);
    clearTimeout(this.handshakeTimer);
    this.timer = undefined;
    const socket = this.socket;
    this.socket = undefined;
    this.rejectConnect?.(new RobotError('DISCONNECTED', reason));
    this.rejectConnect = undefined;
    socket?.close();
    try {
      if (this.sessionId && this.backend()?.robotObservation()?.sessionId === this.sessionId)
        this.backend()?.stopExternalControl(reason);
    } catch {
      /* Session teardown already revoked controls before deleting WASM handles. */
    }
    this.sessionId = undefined;
    this.lease = undefined;
    this.statusChanged({ phase: 'disconnected', error: reason });
  }
}