woolball-client
Version:
Client-side library for Woolball enabling secure browser resource sharing for distributed AI task processing
336 lines (335 loc) • 12.8 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
const utils_1 = require("../utils");
const web_worker_1 = __importDefault(require("web-worker"));
const TaskAvailability_1 = require("./TaskAvailability");
const browser_1 = require("../utils/browser");
class Woolball {
constructor(id, url = 'ws://localhost:9003/ws', options = {}) {
this.wsConnection = null;
this.eventListeners = new Map();
this.options = options;
// Definir ambiente padrão como 'browser' se não for especificado
if (!this.options.environment) {
this.options.environment = 'browser';
}
if (['browser', 'extension'].includes(this.options.environment)) {
(0, utils_1.verifyBrowserCompatibility)();
}
this.clientId = id;
this.wsUrl = url;
this.eventListeners.set('started', new Set());
this.eventListeners.set('success', new Set());
this.eventListeners.set('error', new Set());
this.eventListeners.set('node_count', new Set());
this.eventListeners.set('tps_validation', new Set());
this.workerTypes = new Map();
this.activeWorkers = new Set();
// Initialize worker types based on task configurations
Object.keys(TaskAvailability_1.TASK_CONFIGS).forEach((taskType) => {
const handler = (0, TaskAvailability_1.getTaskHandler)(taskType, this.getCurrentEnvironment());
if (handler) {
this.workerTypes.set(taskType, handler);
}
});
}
async start() {
if (this.wsConnection) {
console.warn('WebSocket connection already exists');
return;
}
// Calculate TPS before establishing WebSocket connection (for all environments)
let tps = 0;
try {
console.log('Calculating TPS benchmark...');
tps = await (0, browser_1.calculateTPS)();
console.log(`TPS benchmark completed: ${tps}`);
}
catch (error) {
console.warn('TPS calculation failed, using 0:', error);
tps = 0;
}
// Validate minimum TPS requirement
const minimumTPS = 30;
if (tps < minimumTPS) {
// Emit error event for insufficient TPS
this.emitEvent('error', {
id: this.clientId,
type: 'tps_validation',
status: 'error',
tpsValue: tps,
});
throw new Error();
}
// Emit success event for TPS validation
this.emitEvent('tps_validation', {
id: this.clientId,
type: 'tps_validation',
status: 'success',
tpsValue: tps
});
console.log(`TPS validation passed: ${tps} >= ${minimumTPS}`);
this.connectWebSocket(this.wsUrl, tps);
}
destroy() {
if (this.wsConnection) {
this.wsConnection.close();
this.wsConnection = null;
}
this.activeWorkers.forEach(worker => {
worker.terminate();
});
this.activeWorkers.clear();
this.eventListeners.clear();
this.eventListeners.set('started', new Set());
this.eventListeners.set('success', new Set());
this.eventListeners.set('error', new Set());
this.eventListeners.set('node_count', new Set());
this.eventListeners.set('tps_validation', new Set());
this.workerTypes.clear();
}
/**
* Establishes WebSocket connection and sets up message handlers
*/
connectWebSocket(url, tps = 0) {
const wsUrl = `${url}/${this.clientId}?tps=${tps}`;
this.wsConnection = new WebSocket(wsUrl);
this.wsConnection.onopen = () => {
console.log(`WebSocket connection established with TPS: ${tps}`);
};
this.wsConnection.onmessage = (event) => {
if (event.data === 'ping') {
return;
}
if (event.data.startsWith('node_count:')) {
const nodeCountStr = event.data.split(':')[1];
const nodeCount = parseInt(nodeCountStr, 10);
if (!isNaN(nodeCount)) {
this.emitEvent('node_count', {
id: '',
type: 'node_count',
status: 'node_count',
nodeCount: nodeCount
});
}
return;
}
try {
this.handleWebSocketMessage(JSON.parse(event.data));
}
catch (parseError) {
console.error('Failed to parse WebSocket message:', parseError);
console.error('Raw message:', event.data);
}
};
this.wsConnection.onerror = (error) => {
console.error('WebSocket error:', error);
};
this.wsConnection.onclose = (event) => {
console.log(`WebSocket connection closed: ${event.code} ${event.reason}`);
};
}
/**
* Handles incoming WebSocket messages
*/
async handleWebSocketMessage(message) {
const { Id, Key, Value } = message;
if (!Id || !Key || !Value) {
console.error('Invalid message format:', message);
return;
}
try {
this.emitEvent('started', {
id: Id,
type: Key,
status: 'started'
});
const response = await this.processEvent(Key, Value);
if (response.error) {
const errorData = {
id: Id,
type: Key,
status: 'error',
};
console.error(`Error processing ${Key}:`, response.error);
this.emitEvent('error', errorData);
this.sendWebSocketMessage({
type: 'ERROR',
data: {
requestId: Id,
error: response.error,
}
});
return;
}
this.emitEvent('success', {
id: Id,
type: Key,
status: 'success',
});
this.sendWebSocketMessage({
type: 'PROCESS_RESULT',
data: {
requestId: Id,
response
}
});
}
catch (error) {
console.error('Error handling WebSocket message:', error);
this.sendWebSocketMessage({
type: 'ERROR',
data: {
requestId: Id,
error: error instanceof Error ? error.message : 'Unknown error',
}
});
}
}
/**
* Sends a message to the WebSocket server
*/
sendWebSocketMessage(message) {
if (this.wsConnection && this.wsConnection.readyState === WebSocket.OPEN) {
this.wsConnection.send(JSON.stringify(message));
return true;
}
return false;
}
setWorkerSource(type, source) {
this.workerTypes.set(type, source);
}
createWorker(type) {
const workerSource = this.workerTypes.get(type);
if (!workerSource) {
throw new Error(`Worker type not found: ${type}`);
}
try {
const blob = new Blob([workerSource], { type: 'application/javascript' });
const workerUrl = URL.createObjectURL(blob);
const worker = new web_worker_1.default(workerUrl);
URL.revokeObjectURL(workerUrl);
this.activeWorkers.add(worker);
return worker;
}
catch (error) {
console.error('Error creating worker:', error);
throw new Error(`Failed to create worker: ${error.message || 'Unknown error'}`);
}
}
/**
* Gets the current environment based on options or detection
* @returns The current environment
*/
getCurrentEnvironment() {
if (this.options.environment) {
return this.options.environment;
}
else if (typeof window !== 'undefined') {
return 'browser';
}
else if (typeof process !== 'undefined' && process.versions && process.versions.node) {
return 'node';
}
else {
// Fallback, though this should not happen in normal circumstances
console.warn('Could not determine environment, defaulting to node');
return 'node';
}
}
terminateWorker(worker) {
worker.terminate();
if (worker._blobUrl) {
URL.revokeObjectURL(worker._blobUrl);
delete worker._blobUrl;
}
this.activeWorkers.delete(worker);
}
async processEvent(type, value) {
// Convert string boolean values to actual booleans
for (const key in value) {
if (value[key] === 'true') {
value[key] = true;
}
if (value[key] === 'false') {
value[key] = false;
}
}
const currentEnvironment = this.getCurrentEnvironment();
const executionType = (0, TaskAvailability_1.getTaskExecutionType)(type, currentEnvironment);
if (!executionType) {
return { error: `Task type '${type}' is not supported in ${currentEnvironment} environment` };
}
// Handle tasks based on their execution type
switch (executionType) {
case 'browser':
try {
const handler = (0, TaskAvailability_1.getTaskHandler)(type, currentEnvironment);
const result = await handler(value);
return result;
}
catch (processorError) {
console.error(`[Browser] Error in ${type} processor:`, processorError);
const errorMessage = processorError instanceof Error ? processorError.message : String(processorError);
return { error: errorMessage };
}
case 'node_worker': {
// Validate provider type for Node.js (keep existing validation)
if (value.provider && value.provider !== 'transformers') {
throw new Error(`Unsupported provider for Node.js: ${value.provider}. Only 'transformers' is supported.`);
}
const { processWithoutNodeWorker } = await import('./node-worker.js');
return processWithoutNodeWorker(type, value);
}
case 'worker':
const worker = this.createWorker(type);
return new Promise((resolve, reject) => {
const messageHandler = (e) => {
worker.removeEventListener('message', messageHandler);
worker.removeEventListener('error', errorHandler);
this.terminateWorker(worker);
if (e.data.error) {
resolve({ error: e.data.error });
}
else {
resolve(e.data);
}
};
const errorHandler = (e) => {
worker.removeEventListener('message', messageHandler);
worker.removeEventListener('error', errorHandler);
this.terminateWorker(worker);
const errorMessage = e.message || 'Unknown worker error';
resolve({ error: errorMessage });
};
worker.addEventListener('message', messageHandler);
worker.addEventListener('error', errorHandler);
worker.postMessage(value);
});
default:
return { error: `Unknown execution type: ${executionType}` };
}
}
on(status, listener) {
const listeners = this.eventListeners.get(status);
if (listeners) {
listeners.add(listener);
}
}
off(status, listener) {
const listeners = this.eventListeners.get(status);
if (listeners) {
listeners.delete(listener);
}
}
emitEvent(status, data) {
const listeners = this.eventListeners.get(status);
if (listeners) {
listeners.forEach(listener => listener(data));
}
}
}
exports.default = Woolball;