@auxot/worker-cli 0.1.2 → 0.1.3

package/dist/websocket.js

@@ -1,433 +0,0 @@
-/**
- * WebSocket Connection Manager
- *
- * Handles connection to Auxot WebSocket server with:
- * - Hello handshake
- * - Heartbeat keepalive
- * - Job message handling
- * - Automatic reconnection with exponential backoff
- */
-import WebSocket from 'ws';
-import { logServerToClient, logClientToServer } from './debug.js';
-import { validatePolicy } from './policy-validator.js';
-const HEARTBEAT_INTERVAL = 5000; // 5 seconds
-const INITIAL_RETRY_DELAY = 1000; // 1 second
-const MAX_RETRY_DELAY = 60000; // 60 seconds
-const RETRY_MULTIPLIER = 2; // Exponential backoff multiplier
-export class WebSocketConnection {
-    ws = null;
-    heartbeatTimer = null;
-    reconnectTimer = null;
-    gpuKey;
-    capabilities;
-    onJobCallback = null;
-    onCancelCallback = null;
-    onPolicyCallback = null;
-    onConfigAckCallback = null;
-    wsUrl = '';
-    retryDelay = INITIAL_RETRY_DELAY;
-    isConnected = false;
-    shouldReconnect = true;
-    isReconnecting = false;
-    policy = null;
-    constructor(gpuKey, capabilities) {
-        this.gpuKey = gpuKey;
-        this.capabilities = capabilities;
-    }
-    /**
-     * Connect to WebSocket server and send hello message
-     */
-    async connect(wsUrl) {
-        this.wsUrl = wsUrl;
-        this.shouldReconnect = true;
-        return this.attemptConnection();
-    }
-    /**
-     * Attempt to establish WebSocket connection
-     */
-    async attemptConnection() {
-        return new Promise((resolve, reject) => {
-            if (!this.wsUrl) {
-                reject(new Error('No WebSocket URL configured'));
-                return;
-            }
-            // Only log if not already reconnecting (first attempt or successful reconnection)
-            if (!this.isReconnecting) {
-                console.log(`Connecting to ${this.wsUrl}...`);
-            }
-            try {
-                this.ws = new WebSocket(this.wsUrl);
-            }
-            catch (error) {
-                this.scheduleReconnect();
-                reject(error);
-                return;
-            }
-            // Set a connection timeout
-            const connectionTimeout = setTimeout(() => {
-                if (!this.isConnected) {
-                    this.ws?.close();
-                    this.scheduleReconnect();
-                    reject(new Error('Connection timeout'));
-                }
-            }, 10000);
-            this.ws.on('open', () => {
-                // Send hello message (server will assign GPU ID)
-                const helloMsg = {
-                    type: 'hello',
-                    gpu_key: this.gpuKey,
-                    capabilities: this.capabilities,
-                };
-                this.send(helloMsg);
-            });
-            this.ws.on('message', async (data) => {
-                try {
-                    const message = JSON.parse(data.toString());
-                    // Debug log (skip heartbeat_ack to reduce noise)
-                    if (message.type !== 'heartbeat_ack') {
-                        logServerToClient(message);
-                    }
-                    if (message.type === 'hello_ack') {
-                        clearTimeout(connectionTimeout);
-                        if (message.success) {
-                            // Store policy from server
-                            if (!message.policy) {
-                                const errorMsg = 'Server did not send policy in hello_ack';
-                                console.error(`✗ ${errorMsg}`);
-                                this.shouldReconnect = false;
-                                this.ws?.close();
-                                reject(new Error(errorMsg));
-                                return;
-                            }
-                            this.policy = message.policy;
-                            // If policy callback is registered, defer validation (will validate via config message after spawning llama.cpp)
-                            if (this.onPolicyCallback) {
-                                // Spawn llama.cpp and discover capabilities (validation will happen via config message)
-                                try {
-                                    await this.onPolicyCallback(message.policy);
-                                }
-                                catch (error) {
-                                    console.error('[Policy Callback] Error:', error);
-                                    this.shouldReconnect = false;
-                                    this.ws?.close();
-                                    reject(error);
-                                    return;
-                                }
-                                // Validation deferred - will happen via config message
-                                // Connection is established but validation happens later
-                                console.log('✓ Successfully authenticated with server');
-                                console.log(`  Policy: ${message.policy.model_name} (${message.policy.quantization})`);
-                                console.log('  Spawning llama.cpp process...');
-                                console.log('  (Capabilities validation will happen via config message)');
-                            }
-                            else {
-                                // No policy callback - validate immediately (legacy flow)
-                                const validation = await validatePolicy(this.capabilities, message.policy);
-                                // Log warnings if any
-                                if (validation.warnings && validation.warnings.length > 0) {
-                                    console.warn('⚠ Policy validation warnings:');
-                                    validation.warnings.forEach((warning) => {
-                                        console.warn(`  - ${warning}`);
-                                    });
-                                }
-                                if (!validation.valid) {
-                                    console.error('✗ Policy validation failed:');
-                                    console.error('  Expected:');
-                                    console.error(`    Model: ${message.policy.model_name}`);
-                                    console.error(`    Quantization: ${message.policy.quantization}`);
-                                    console.error(`    Context Size: ${message.policy.context_size}`);
-                                    console.error(`    Capabilities: ${message.policy.capabilities.join(', ')}`);
-                                    if (message.policy.parameters) {
-                                        console.error(`    Parameters: ${message.policy.parameters}`);
-                                    }
-                                    if (message.policy.family) {
-                                        console.error(`    Family: ${message.policy.family}`);
-                                    }
-                                    console.error('  Discovered:');
-                                    console.error(`    Model: ${this.capabilities.model}`);
-                                    console.error(`    Context Size: ${this.capabilities.ctx_size}`);
-                                    if (this.capabilities.parameters) {
-                                        console.error(`    Parameters: ${this.capabilities.parameters}`);
-                                    }
-                                    console.error('  Errors:');
-                                    validation.errors.forEach((error) => {
-                                        console.error(`    - ${error}`);
-                                    });
-                                    console.error('\nPlease configure your llama.cpp server to match the policy requirements.');
-                                    this.shouldReconnect = false;
-                                    this.ws?.close();
-                                    reject(new Error(`Policy validation failed: ${validation.errors.join('; ')}`));
-                                    return;
-                                }
-                                console.log('✓ Successfully authenticated with server');
-                                console.log(`  Policy: ${message.policy.model_name} (${message.policy.quantization})`);
-                                console.log('✓ Capabilities validated against policy');
-                            }
-                            this.isConnected = true;
-                            this.isReconnecting = false;
-                            this.retryDelay = INITIAL_RETRY_DELAY; // Reset retry delay on success
-                            this.startHeartbeat();
-                            resolve();
-                        }
-                        else {
-                            // Check if there's a reconnect_in_seconds field (duplicate connection)
-                            const errorMessage = message.error || 'Authentication failed';
-                            if (message.reconnect_in_seconds || message.reconnectInSeconds) {
-                                const reconnectIn = message.reconnect_in_seconds || message.reconnectInSeconds;
-                                console.error(`✗ Worker UUID already connected! Are you trying to connect the same worker ID twice?`);
-                                console.error(`   The existing connection's presence key expires in ${reconnectIn} seconds.`);
-                                console.error(`   Waiting ${reconnectIn} seconds and retrying automatically...`);
-                                // Close connection and schedule reconnect (do NOT reject - let retry handle it)
-                                this.ws?.close();
-                                // Wait for the specified time before reconnecting
-                                this.retryDelay = Math.max(reconnectIn * 1000, INITIAL_RETRY_DELAY);
-                                this.scheduleReconnect();
-                                // Resolve instead of reject - this allows the app to continue with auto-retry
-                                resolve();
-                            }
-                            else {
-                                // Authentication errors (invalid key, key not found, etc.) are fatal - stop retrying
-                                console.error('✗ Authentication failed:', errorMessage);
-                                console.error('   Fatal error: Stopping retry attempts. Please check your GPU key and try again.');
-                                // Stop retrying for authentication errors
-                                this.shouldReconnect = false;
-                                this.ws?.close();
-                                // Don't schedule reconnect - this is a fatal error
-                                reject(new Error(errorMessage));
-                            }
-                        }
-                    }
-                    else if (message.type === 'config_ack') {
-                        // Config acknowledged (validation happened on server)
-                        if (this.onConfigAckCallback) {
-                            this.onConfigAckCallback(message.success, message.error);
-                        }
-                        else {
-                            // Fallback: log and handle errors
-                            if (message.success) {
-                                console.log('✓ Capabilities validated by server');
-                            }
-                            else {
-                                console.error(`✗ Config validation failed: ${message.error}`);
-                                this.shouldReconnect = false;
-                                this.ws?.close();
-                            }
-                        }
-                    }
-                    else if (message.type === 'heartbeat_ack') {
-                        // Heartbeat acknowledged (silent)
-                    }
-                    else if (message.type === 'cancel') {
-                        // Cancel message
-                        console.log(`Received cancel request for job ${message.job_id}`);
-                        if (this.onCancelCallback) {
-                            this.onCancelCallback(message);
-                        }
-                        else {
-                            console.warn('No cancel callback registered!');
-                        }
-                    }
-                    else if (message.type === 'job') {
-                        // Job received
-                        console.log(`Received job ${message.job_id}`);
-                        if (this.onJobCallback) {
-                            this.onJobCallback(message).catch((err) => {
-                                console.error('Error processing job:', err);
-                                this.sendError(message.job_id, err.message);
-                            });
-                        }
-                        else {
-                            console.warn('No job callback registered!');
-                        }
-                    }
-                }
-                catch (error) {
-                    console.error('Error parsing message:', error);
-                }
-            });
-            this.ws.on('error', (error) => {
-                clearTimeout(connectionTimeout);
-                // Silently handle connection errors during reconnection
-            });
-            this.ws.on('close', (code, reason) => {
-                clearTimeout(connectionTimeout);
-                this.isConnected = false;
-                this.stopHeartbeat();
-                if (this.shouldReconnect) {
-                    if (!this.isReconnecting) {
-                        console.log('WebSocket disconnected, will continue to retry...');
-                        this.isReconnecting = true;
-                    }
-                    this.scheduleReconnect();
-                }
-                else {
-                    console.log('WebSocket disconnected');
-                }
-            });
-        });
-    }
-    /**
-     * Schedule reconnection with exponential backoff
-     */
-    scheduleReconnect() {
-        if (this.reconnectTimer || !this.shouldReconnect) {
-            return;
-        }
-        this.reconnectTimer = setTimeout(async () => {
-            this.reconnectTimer = null;
-            try {
-                await this.attemptConnection();
-            }
-            catch (error) {
-                // Connection failed, exponential backoff will continue (silently)
-                this.retryDelay = Math.min(this.retryDelay * RETRY_MULTIPLIER, MAX_RETRY_DELAY);
-            }
-        }, this.retryDelay);
-    }
-    /**
-     * Start sending heartbeat messages
-     */
-    startHeartbeat() {
-        this.stopHeartbeat(); // Clear any existing timer
-        this.heartbeatTimer = setInterval(() => {
-            if (this.isConnected) {
-                const heartbeatMsg = {
-                    type: 'heartbeat',
-                    // Server identifies worker by WebSocket connection, not GPU ID
-                };
-                this.send(heartbeatMsg);
-            }
-        }, HEARTBEAT_INTERVAL);
-    }
-    /**
-     * Stop heartbeat timer
-     */
-    stopHeartbeat() {
-        if (this.heartbeatTimer) {
-            clearInterval(this.heartbeatTimer);
-            this.heartbeatTimer = null;
-        }
-    }
-    /**
-     * Send message to server
-     */
-    send(message) {
-        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
-            const jsonString = JSON.stringify(message);
-            // Debug log (skip heartbeats to reduce noise)
-            if (message.type !== 'heartbeat') {
-                logClientToServer(message);
-            }
-            this.ws.send(jsonString);
-        }
-    }
-    /**
-     * Send token to server
-     */
-    sendToken(jobId, token) {
-        this.send({
-            type: 'token',
-            job_id: jobId,
-            token,
-        });
-    }
-    /**
-     * Send completion to server with metadata
-     */
-    sendComplete(jobId, fullResponse, durationMs, inputTokens, outputTokens, toolCalls) {
-        this.send({
-            type: 'complete',
-            job_id: jobId,
-            full_response: fullResponse,
-            duration_ms: durationMs,
-            input_tokens: inputTokens,
-            output_tokens: outputTokens,
-            tool_calls: toolCalls,
-        });
-    }
-    /**
-     * Send error to server
-     */
-    sendError(jobId, error) {
-        this.send({
-            type: 'error',
-            job_id: jobId,
-            error,
-        });
-    }
-    /**
-     * Register callback for job messages
-     */
-    onJob(callback) {
-        this.onJobCallback = callback;
-    }
-    /**
-     * Register callback for cancel messages
-     */
-    onCancel(callback) {
-        this.onCancelCallback = callback;
-    }
-    /**
-     * Register callback for policy (called when policy is received in hello_ack)
-     * Used to spawn llama.cpp process before validation
-     */
-    onPolicy(callback) {
-        this.onPolicyCallback = callback;
-    }
-    /**
-     * Register callback for config_ack (called when server validates our config)
-     */
-    onConfigAck(callback) {
-        this.onConfigAckCallback = callback;
-    }
-    /**
-     * Update capabilities (called after spawning llama.cpp)
-     */
-    updateCapabilities(capabilities) {
-        this.capabilities = capabilities;
-    }
-    /**
-     * Send config message to server (called after discovering capabilities)
-     */
-    sendConfig(capabilities) {
-        this.send({
-            type: 'config',
-            capabilities,
-        });
-    }
-    /**
-     * Check if connected
-     */
-    get connected() {
-        return this.isConnected;
-    }
-    /**
-     * Get policy received from server
-     */
-    getPolicy() {
-        return this.policy;
-    }
-    /**
-     * Get current capabilities
-     */
-    getCapabilities() {
-        return this.capabilities;
-    }
-    /**
-     * Close connection
-     */
-    close() {
-        this.shouldReconnect = false;
-        this.stopHeartbeat();
-        if (this.reconnectTimer) {
-            clearTimeout(this.reconnectTimer);
-            this.reconnectTimer = null;
-        }
-        if (this.ws) {
-            this.ws.close();
-            this.ws = null;
-        }
-        this.isConnected = false;
-        this.policy = null; // Clear policy on disconnect
-    }
-}