@mochabug/adapt-web
Version:
The client library to execute automations, without effort, in a browser environment
536 lines • 20 kB
JavaScript
import { fromBinary, toJson } from "@bufbuild/protobuf";
import { timestampDate } from "@bufbuild/protobuf/wkt";
import WebSocket from "isomorphic-ws";
import { parse as parseUuid } from "uuid";
import { SessionSchema, Status, UrlSchema, WebsocketMessageSchema, } from "./genproto/mochabugapis/adapt/automations/v1/automations_pb.js";
import { getConfig } from "./index.js";
// Connection states
export var ConnectionState;
(function (ConnectionState) {
ConnectionState["DISCONNECTED"] = "disconnected";
ConnectionState["CONNECTING"] = "connecting";
ConnectionState["CONNECTED"] = "connected";
ConnectionState["CLOSING"] = "closing";
})(ConnectionState || (ConnectionState = {}));
// Logger class for consistent logging
class Logger {
constructor(debug) {
this.debug = debug;
}
log(level, message, data) {
if (!this.debug && level === "debug")
return;
const timestamp = new Date().toISOString();
const prefix = `[PubSub ${timestamp}]`;
const fullMessage = `${prefix} ${message}`;
if (data !== undefined) {
console[level](fullMessage, data);
}
else {
console[level](fullMessage);
}
}
}
// ACK tracker to prevent duplicate message processing
class AckTracker {
constructor(logger) {
this.logger = logger;
this.ackedMessages = new Map();
this.MESSAGE_TTL = 60000; // 60 seconds
this.cleanupTimer = null;
}
start() {
// Cleanup old messages every 30 seconds
this.cleanupTimer = setInterval(() => this.cleanup(), 30000);
}
stop() {
if (this.cleanupTimer) {
clearInterval(this.cleanupTimer);
this.cleanupTimer = null;
}
this.ackedMessages.clear();
}
isProcessed(messageId) {
return this.ackedMessages.has(messageId);
}
markProcessed(messageId) {
this.ackedMessages.set(messageId, Date.now());
}
cleanup() {
const now = Date.now();
const deadline = now - this.MESSAGE_TTL;
let removed = 0;
for (const [messageId, timestamp] of this.ackedMessages.entries()) {
if (timestamp < deadline) {
this.ackedMessages.delete(messageId);
removed++;
}
}
if (removed > 0) {
this.logger.log("debug", `Cleaned up ${removed} old ACKed messages`);
}
}
}
// Reconnection manager with exponential backoff
class ReconnectionManager {
constructor(logger) {
this.logger = logger;
this.attempts = 0;
this.timer = null;
this.FAST_INITIAL_DELAY = 50; // 50ms for first 10 attempts
this.FAST_ATTEMPTS_THRESHOLD = 10; // First 10 attempts are fast
this.SLOW_INITIAL_DELAY = 1000; // 1 second after fast attempts
this.MAX_DELAY = 30000; // 30 seconds
this.MAX_ATTEMPTS = -1; // -1 = infinite
}
reset() {
this.attempts = 0;
this.cancel();
}
cancel() {
if (this.timer) {
clearTimeout(this.timer);
this.timer = null;
}
}
scheduleReconnect(callback, immediate = false) {
if (this.MAX_ATTEMPTS !== -1 && this.attempts >= this.MAX_ATTEMPTS) {
this.logger.log("error", "Max reconnection attempts reached");
return false;
}
this.attempts++;
// Calculate delay with fast exponential backoff for first 10 attempts
let delay;
if (immediate) {
// Immediate reconnection for server restart
delay = 0;
this.logger.log("info", "Scheduling immediate reconnect (server restart)");
}
else if (this.attempts <= this.FAST_ATTEMPTS_THRESHOLD) {
// Fast reconnection: 50 => 100 => 200 => 400 => 800 => 1600 => 3200 => 6400 => 12800 => 25600
delay = this.FAST_INITIAL_DELAY * Math.pow(2, this.attempts - 1);
}
else {
// After fast attempts, use slower backoff starting from SLOW_INITIAL_DELAY
const slowAttempt = this.attempts - this.FAST_ATTEMPTS_THRESHOLD;
delay = this.SLOW_INITIAL_DELAY * Math.pow(2, slowAttempt - 1);
}
// Cap at MAX_DELAY (except for immediate reconnects)
if (!immediate) {
delay = Math.min(delay, this.MAX_DELAY);
}
if (!immediate) {
this.logger.log("info", `Scheduling reconnect attempt ${this.attempts} in ${delay}ms`);
}
this.timer = setTimeout(() => {
this.timer = null;
callback();
}, delay);
return true;
}
}
// WebSocket connection manager
class WebSocketManager {
constructor(logger) {
this.logger = logger;
this.ws = null;
this.closePromise = null;
this.closeResolve = null;
}
connect(url, onOpen, onMessage, onClose, onError) {
if (this.ws) {
throw new Error("WebSocket already exists");
}
this.logger.log("info", `Connecting to: ${url}`);
// Create a promise that resolves when the WebSocket closes
this.closePromise = new Promise((resolve) => {
this.closeResolve = resolve;
});
this.ws = new WebSocket(url);
this.ws.binaryType = "arraybuffer";
this.ws.onopen = () => {
this.logger.log("info", "✓ WebSocket connected");
onOpen();
};
this.ws.onmessage = (event) => {
if (event.data instanceof ArrayBuffer) {
onMessage(event.data);
}
else if (event.data instanceof Blob) {
// Handle Blob data (some browsers may return Blob instead of ArrayBuffer)
event.data
.arrayBuffer()
.then(onMessage)
.catch((error) => {
this.logger.log("error", "Failed to convert Blob to ArrayBuffer:", error);
});
}
else {
this.logger.log("warn", "Received non-binary message, ignoring");
}
};
this.ws.onerror = (error) => {
this.logger.log("error", "WebSocket error:", error);
onError(error);
};
this.ws.onclose = (event) => {
this.logger.log("info", `WebSocket closed: ${event.code} - ${event.reason}`);
onClose(event.code, event.reason);
this.ws = null;
// Resolve the close promise
this.closeResolve();
this.closeResolve = null;
this.closePromise = null;
};
}
sendAck(messageId) {
if (!this.isConnected()) {
this.logger.log("warn", `Cannot send ACK (not connected): ${messageId}`);
return false;
}
try {
// Parse UUID string to bytes (16 bytes)
const uuid = new Uint8Array(parseUuid(messageId));
this.ws.send(uuid);
this.logger.log("debug", `Sent ACK: ${messageId}`);
return true;
}
catch (error) {
this.logger.log("error", `Failed to send ACK for ${messageId}:`, error);
return false;
}
}
async close(code = 1000 /* CloseCode.NORMAL */, reason = "Client disconnecting") {
if (!this.ws) {
return;
}
if (this.ws.readyState === WebSocket.OPEN ||
this.ws.readyState === WebSocket.CONNECTING) {
this.ws.close(code, reason);
}
await this.closePromise;
}
isConnected() {
return this.ws !== null && this.ws.readyState === WebSocket.OPEN;
}
}
// Message processor for handling incoming WebSocket messages
class MessageProcessor {
constructor(logger, handlers) {
this.logger = logger;
this.handlers = handlers;
this.sessionEnded = false;
}
process(data) {
try {
const message = fromBinary(WebsocketMessageSchema, new Uint8Array(data));
this.logger.log("debug", `Processing message: ${message.id}`);
// Process based on message type
switch (message.message.case) {
case "output":
this.processOutput(message.message.value);
break;
case "session":
this.processSession(message.message.value);
break;
case "url":
this.processUrl(message.message.value);
break;
default:
this.logger.log("warn", `Unknown message type: ${message.message.case}`);
}
return { messageId: message.id, sessionEnded: this.sessionEnded };
}
catch (error) {
this.logger.log("error", "Failed to process message:", error);
return { messageId: null, sessionEnded: this.sessionEnded };
}
}
hasSessionEnded() {
return this.sessionEnded;
}
processOutput(output) {
if (!this.handlers.onOutput)
return;
try {
this.handlers.onOutput({
vertex: output.vertex,
fork: output.fork,
data: output.data,
created: output.created ? timestampDate(output.created) : undefined,
});
}
catch (error) {
this.logger.log("error", "Error in output handler:", error);
}
}
processSession(session) {
// Check for terminal status
switch (session.status) {
case Status.STOPPED:
case Status.COMPLETED:
case Status.ERRORED:
case Status.EXPIRED:
this.sessionEnded = true;
this.logger.log("info", `Session ended with status: ${Status[session.status]}`);
break;
}
if (!this.handlers.onSession)
return;
try {
this.handlers.onSession(toJson(SessionSchema, session));
}
catch (error) {
this.logger.log("error", "Error in session handler:", error);
}
}
processUrl(url) {
if (!this.handlers.onUrl)
return;
try {
this.handlers.onUrl(toJson(UrlSchema, url));
}
catch (error) {
this.logger.log("error", "Error in URL handler:", error);
}
}
}
// Main PubsubClient class
export class PubsubClient {
constructor(debug = false) {
this.state = ConnectionState.DISCONNECTED;
this.sessionInfo = null;
this.messageProcessor = null;
this.shouldReconnect = true;
this.sessionComplete = false;
this.logger = new Logger(debug);
this.ackTracker = new AckTracker(this.logger);
this.reconnectManager = new ReconnectionManager(this.logger);
this.wsManager = new WebSocketManager(this.logger);
}
async subscribe(opts) {
this.logger.log("debug", "Subscribe called", {
automation: opts.id,
hasToken: !!opts.sessionToken,
});
// Validate input
if (!opts.sessionToken || !opts.id) {
throw new Error("Session token and automation are required");
}
// Check if already connected to same session
if (this.isSameSession(opts)) {
this.logger.log("debug", "Already connected to same session, updating handlers");
this.updateHandlers(opts);
return;
}
// Check state
if (this.state === ConnectionState.CONNECTING) {
throw new Error("Already connecting, please wait");
}
// Disconnect if connected to different session
if (this.state === ConnectionState.CONNECTED) {
await this.unsubscribe();
}
// Store configuration
this.sessionInfo = {
token: opts.sessionToken,
id: opts.id,
};
// Create message processor with handlers
this.messageProcessor = new MessageProcessor(this.logger, {
onOutput: opts.onOutput,
onSession: opts.onSession,
onUrl: opts.onUrl,
});
// Reset state
this.shouldReconnect = true;
this.sessionComplete = false;
this.reconnectManager.reset();
// Start ACK tracker
this.ackTracker.start();
// Connect
await this.connect();
}
async unsubscribe() {
if (this.state === ConnectionState.DISCONNECTED ||
this.state === ConnectionState.CLOSING) {
return;
}
this.logger.log("info", "Unsubscribing...");
// Update state
this.state = ConnectionState.CLOSING;
this.shouldReconnect = false;
// Cancel any pending reconnects
this.reconnectManager.cancel();
// Close WebSocket and wait for it to fully close
await this.wsManager.close();
// Stop ACK tracker
this.ackTracker.stop();
// Clear state
this.sessionInfo = null;
this.messageProcessor = null;
this.state = ConnectionState.DISCONNECTED;
this.logger.log("info", "Unsubscribed successfully");
}
isConnected() {
return this.state === ConnectionState.CONNECTED;
}
getConnectionState() {
return this.state;
}
getSessionInfo() {
if (!this.sessionInfo) {
throw new Error("No active session");
}
return this.sessionInfo;
}
async connect() {
if (!this.sessionInfo) {
throw new Error("No session configuration");
}
this.state = ConnectionState.CONNECTING;
const config = getConfig();
const wsUrl = `${config.wsBaseUrl}/${this.sessionInfo.id}/ws?token=${encodeURIComponent(this.sessionInfo.token)}`;
try {
this.wsManager.connect(wsUrl, () => this.handleOpen(), (data) => this.handleMessage(data), (code, reason) => this.handleClose(code, reason), (error) => this.handleError(error));
}
catch (error) {
this.state = ConnectionState.DISCONNECTED;
throw error;
}
}
handleOpen() {
this.state = ConnectionState.CONNECTED;
this.reconnectManager.reset();
}
handleMessage(data) {
if (!this.messageProcessor) {
this.logger.log("error", "No message processor configured");
return;
}
// Process the message
const result = this.messageProcessor.process(data);
if (!result.messageId) {
return; // Failed to parse
}
// Update session complete flag if session ended
if (result.sessionEnded && !this.sessionComplete) {
this.logger.log("info", "Session has ended (detected from message)");
this.sessionComplete = true;
}
// Check for duplicate
if (this.ackTracker.isProcessed(result.messageId)) {
this.logger.log("debug", `Duplicate message: ${result.messageId}`);
}
else {
this.ackTracker.markProcessed(result.messageId);
}
// ACK the message (skip for session-ending messages as connection will close)
if (!result.sessionEnded) {
this.wsManager.sendAck(result.messageId);
}
}
handleClose(code, reason) {
this.state = ConnectionState.DISCONNECTED;
// Analyze close code
const reconnectStrategy = this.analyzeCloseCode(code, reason);
// Schedule reconnect if appropriate
if (reconnectStrategy.shouldReconnect &&
this.shouldReconnect &&
!this.sessionComplete) {
this.scheduleReconnect(reconnectStrategy.immediate);
}
else {
// Clean shutdown (fire and forget - cleanup is async but we don't need to wait)
this.cleanup().catch((error) => {
this.logger.log("error", "Error during cleanup:", error);
});
}
}
handleError(error) {
this.logger.log("error", "WebSocket error:", error);
}
analyzeCloseCode(code, reason) {
switch (code) {
case 1000 /* CloseCode.NORMAL */:
// Normal closure - session complete or client requested
this.logger.log("info", "Normal closure");
this.sessionComplete = true;
return { shouldReconnect: false, immediate: false };
case 1001 /* CloseCode.GOING_AWAY */:
// Server shutting down
this.logger.log("warn", "Server going away");
return { shouldReconnect: true, immediate: false };
case 1012 /* CloseCode.SERVER_RESTART */:
// Server restart - reconnect immediately (no backoff)
this.logger.log("warn", "Server restarting - will reconnect immediately");
return { shouldReconnect: true, immediate: true };
case 1013 /* CloseCode.TRY_AGAIN_LATER */:
// Service temporarily unavailable - reconnect with backoff
this.logger.log("warn", "Service temporarily unavailable, will retry");
return { shouldReconnect: true, immediate: false };
case 1002 /* CloseCode.PROTOCOL_ERROR */:
case 1003 /* CloseCode.UNSUPPORTED_DATA */:
case 1008 /* CloseCode.POLICY_VIOLATION */:
// Critical errors - don't reconnect
this.logger.log("error", `Critical error (${code}): ${reason}`);
return { shouldReconnect: false, immediate: false };
case 1006 /* CloseCode.ABNORMAL */:
// Connection lost - try to reconnect
this.logger.log("warn", "Connection lost (abnormal closure)");
return { shouldReconnect: true, immediate: false };
default:
// Unknown code - try to reconnect
this.logger.log("warn", `Unknown close code ${code}: ${reason}`);
return { shouldReconnect: true, immediate: false };
}
}
scheduleReconnect(immediate = false) {
const scheduled = this.reconnectManager.scheduleReconnect(() => {
// Double-check state to avoid race conditions
if (this.shouldReconnect &&
this.state === ConnectionState.DISCONNECTED &&
!this.sessionComplete) {
this.connect().catch((error) => {
this.logger.log("error", "Reconnection failed:", error);
// Try again with normal backoff (not immediate)
this.scheduleReconnect(false);
});
}
}, immediate);
if (!scheduled) {
// Max attempts reached (fire and forget cleanup)
this.cleanup().catch((error) => {
this.logger.log("error", "Error during cleanup:", error);
});
}
}
async cleanup() {
this.shouldReconnect = false;
this.reconnectManager.cancel();
this.ackTracker.stop();
await this.wsManager.close();
this.sessionInfo = null;
this.messageProcessor = null;
this.state = ConnectionState.DISCONNECTED;
}
isSameSession(opts) {
if (!this.sessionInfo || this.state !== ConnectionState.CONNECTED) {
return false;
}
return (this.sessionInfo.token === opts.sessionToken &&
this.sessionInfo.id === opts.id);
}
updateHandlers(opts) {
if (!this.messageProcessor) {
return;
}
// Update handlers in message processor
this.messageProcessor = new MessageProcessor(this.logger, {
onOutput: opts.onOutput,
onSession: opts.onSession,
onUrl: opts.onUrl,
});
}
}
//# sourceMappingURL=pubsub-client.js.map