MujocoPhysics.plugin.js

import {AViewerPluginSync} from '@kite3d/engine'
import {ModelView} from './view.js'

// One Worker per session: models, custom controllers and WASM memory stay isolated.
export class MujocoPhysicsPlugin extends AViewerPluginSync {
  static PluginType='MujocoPhysicsPlugin'
  constructor(){super();this.sessions=new Set()}
  get sessionCount(){return this.sessions.size}
  createSession(config){
    const session=new PhysicsSession(config,()=>this.sessions.delete(session))
    this.sessions.add(session)
    return session
  }
  async load(config){const session=this.createSession(config);try{await session.ready;return session}catch(error){session.close();throw error}}
  createView(description){return new ModelView(description)}
  onRemove(viewer){for(const session of [...this.sessions])session.close();super.onRemove(viewer)}
}
export default MujocoPhysicsPlugin
class PhysicsSession {
  constructor(config,onClose){
    this.onClose=onClose;this.pending=new Map();this.listeners=new Set();this.nextId=0;this.closed=false
    this.worker=new Worker(new URL('./session.worker.js',import.meta.url),{type:'module'})
    this.worker.onmessage=({data})=>{
      if(data.event){for(const listener of this.listeners)listener(data.event);return}
      const call=this.pending.get(data.id);if(!call)return
      this.pending.delete(data.id);clearTimeout(call.timer)
      if(data.error)call.reject(Error(data.error));else call.resolve(data.value)
    }
    this.worker.onerror=event=>{const error=Error(event.message||'MuJoCo worker failed');this.fail(error);this.close()}
    this.ready=this.request('load',config).then(description=>{this.description=description;return description})
    // Callers still receive the rejection; closing during loading must not create an unhandled rejection.
    this.ready.catch(()=>{})
  }
  request(method,args){
    if(this.closed)return Promise.reject(new DOMException('MuJoCo session closed','AbortError'))
    return new Promise((resolve,reject)=>{
      const id=++this.nextId,timer=setTimeout(()=>{this.pending.delete(id);reject(Error(`MuJoCo ${method} timed out`));this.close()},60000)
      this.pending.set(id,{resolve,reject,timer});this.worker.postMessage({id,method,args})
    })
  }
  onEvent(listener){this.listeners.add(listener);return()=>this.listeners.delete(listener)}
  command(command){if(!this.closed)this.ready.then(()=>{if(!this.closed)this.worker.postMessage({command})}).catch(error=>{if(!this.closed)for(const listener of this.listeners)listener({type:'error',error:String(error)})})}
  async step(action,frames=1){await this.ready;return this.request('step',{action,frames})}
  async reset(state={}){await this.ready;return this.request('reset',state)}
  async snapshot(){await this.ready;return this.request('snapshot')}
  async describe(){return this.ready}
  fail(error){for(const call of this.pending.values()){clearTimeout(call.timer);call.reject(error)}this.pending.clear()}
  close(){if(this.closed)return;this.closed=true;this.worker.terminate();this.worker=undefined;this.fail(new DOMException('MuJoCo session closed','AbortError'));this.listeners.clear();this.onClose()}
}