@sogni-ai/sogni-client 4.0.0-alpha.11 → 4.0.0-alpha.13

package/src/ApiClient/WebSocketClient/BrowserWebSocketClient/WSCoordinator.ts

@@ -0,0 +1,425 @@
+import { Logger } from '../../../lib/DefaultLogger';
+import getUUID from '../../../lib/getUUID';
+import { AuthenticatedData, SocketEventMap, SocketEventName } from '../events';
+import { MessageType, SocketMessageMap } from '../messages';
+import { Balances } from '../../../Account/types';
+import {
+  ChannelMessage,
+  Heartbeat,
+  MessageEnvelope,
+  SocketEventReceived,
+  SendSocketMessage
+} from './types';
+
+interface StartingEvents {
+  authenticated?: AuthenticatedData;
+  balanceUpdate?: Balances;
+  swarmModels?: Record<string, number>;
+}
+
+interface WSCoordinatorCallbacks {
+  /**
+   * Invoked when authentication state changes (authenticated/unauthenticated).
+   * @param isAuthenticated - true if client is authenticated, false if not.
+   */
+  onAuthChanged: (isAuthenticated: boolean) => void;
+  /**
+   * Invoked when role changes (primary/secondary).
+   * @param isPrimary - true if client is primary, false if secondary.
+   */
+  onRoleChange: (isPrimary: boolean) => void;
+  /**
+   * Invoked when connection state must change (connected/disconnected).
+   * @param isConnected
+   */
+  onConnectionToggle: (isConnected: boolean) => void;
+  /**
+   * Invoked when secondary client receives a socket event from primary.
+   * @param message
+   */
+  onMessageFromPrimary: (message: SocketEventReceived | Heartbeat) => void;
+  /**
+   * Invoked when primary client receives a socket send request from secondary.
+   * @param message
+   */
+  onSendRequest: (message: SendSocketMessage) => Promise<void>;
+}
+
+class WSCoordinator {
+  private static readonly HEARTBEAT_INTERVAL = 2000;
+  private static readonly PRIMARY_TIMEOUT = 3000;
+  private static readonly CHANNEL_NAME = 'sogni-websocket-clients';
+  private static readonly ACK_TIMEOUT = 5000;
+
+  private id: string;
+  private callbacks: WSCoordinatorCallbacks;
+  private channel: BroadcastChannel;
+  private _isPrimary: boolean;
+  private _isConnected = false;
+  private lastPrimaryHeartbeat: number = 0;
+  private logger: Logger;
+  private heartbeatInterval: NodeJS.Timeout | null = null;
+  private primaryCheckInterval: NodeJS.Timeout | null = null;
+  private startingEvents: StartingEvents = {};
+  private ackCallbacks: Record<string, (error?: any) => void> = {};
+
+  constructor(callbacks: WSCoordinatorCallbacks, logger: Logger) {
+    this.id = getUUID();
+    this.logger = logger;
+    this.callbacks = callbacks;
+    this.channel = new BroadcastChannel(WSCoordinator.CHANNEL_NAME);
+    this.channel.onmessage = this.handleMessage.bind(this);
+    this._isPrimary = false;
+
+    // Listen for tab closing to gracefully release primary role
+    if (typeof window !== 'undefined') {
+      window.addEventListener('beforeunload', this.handleBeforeUnload.bind(this));
+    }
+  }
+
+  /**
+   * Initialize tab coordination and determine role
+   */
+  async initialize(): Promise<boolean> {
+    this.logger.info(`WSCoordinator ${this.id} initializing...`);
+
+    // Announce our presence
+    this.broadcast({
+      type: 'announce'
+    });
+
+    // Wait to see if there's an existing primary
+    await this.waitForPrimaryResponse();
+    if (!this._isPrimary) {
+      this.logger.info(`Client ${this.id} is secondary, primary exists`);
+      this.startPrimaryCheck();
+    } else {
+      this.logger.info(`Client ${this.id} becoming primary`);
+      this.becomePrimary();
+    }
+
+    return this._isPrimary;
+  }
+
+  connect() {
+    if (this._isPrimary) {
+      throw new Error('Primary should connect the socket directly.');
+    }
+    this.broadcast({
+      type: 'connection-toggle',
+      payload: { connected: true }
+    });
+  }
+
+  disconnect() {
+    if (this._isPrimary) {
+      throw new Error('Primary should disconnect socket directly.');
+    }
+    this.broadcast({
+      type: 'connection-toggle',
+      payload: { connected: false }
+    });
+  }
+
+  async changeAuthState(isAuthenticated: boolean) {
+    this.broadcast({
+      type: 'authentication',
+      payload: { authenticated: isAuthenticated }
+    });
+  }
+
+  /**
+   * Wait briefly to see if a primary tab responds
+   */
+  private waitForPrimaryResponse(): Promise<void> {
+    return new Promise((resolve) => {
+      const timeout = setTimeout(() => {
+        // No primary responded, we become primary
+        this._isPrimary = true;
+        resolve();
+      }, 500);
+
+      const messageHandler = (e: MessageEvent<MessageEnvelope>) => {
+        const envelope = e.data;
+        const message = envelope.payload;
+        if (message.type === 'primary-present' || message.type === 'primary-claim') {
+          // Primary exists
+          this._isPrimary = false;
+          this.lastPrimaryHeartbeat = envelope.timestamp;
+          clearTimeout(timeout);
+          this.channel.removeEventListener('message', messageHandler);
+          resolve();
+        }
+      };
+
+      this.channel.addEventListener('message', messageHandler);
+    });
+  }
+
+  /**
+   * Become the primary tab
+   */
+  private becomePrimary() {
+    this._isPrimary = true;
+    this.callbacks.onRoleChange(true);
+
+    // Broadcast that we're claiming primary role
+    this.broadcast({
+      type: 'primary-claim'
+    });
+
+    // Start sending heartbeats
+    this.startHeartbeat();
+  }
+
+  /**
+   * Release primary role (when closing)
+   */
+  private releasePrimary() {
+    if (this._isPrimary) {
+      this.broadcast({
+        type: 'primary-release'
+      });
+      this.stopHeartbeat();
+    }
+  }
+
+  /**
+   * Start sending heartbeat messages as primary
+   */
+  private startHeartbeat() {
+    if (this.heartbeatInterval)
+      throw new Error('Heartbeat interval already started. This should never happen.');
+    this.heartbeatInterval = setInterval(() => {
+      this.broadcast({
+        type: 'primary-present',
+        payload: { connected: this._isConnected }
+      });
+    }, WSCoordinator.HEARTBEAT_INTERVAL);
+  }
+
+  /**
+   * Stop sending heartbeat messages
+   */
+  private stopHeartbeat() {
+    if (this.heartbeatInterval) {
+      clearInterval(this.heartbeatInterval);
+      this.heartbeatInterval = null;
+    }
+  }
+
+  /**
+   * Start checking for primary heartbeat (as secondary)
+   */
+  private startPrimaryCheck() {
+    this.primaryCheckInterval = setInterval(() => {
+      const timeSinceLastHeartbeat = Date.now() - this.lastPrimaryHeartbeat;
+      if (timeSinceLastHeartbeat > WSCoordinator.PRIMARY_TIMEOUT) {
+        this.logger.warn(`Primary tab timeout, becoming primary`);
+        this.stopPrimaryCheck();
+        this.becomePrimary();
+      }
+    }, 1000);
+  }
+
+  /**
+   * Stop checking for primary heartbeat
+   */
+  private stopPrimaryCheck() {
+    if (this.primaryCheckInterval) {
+      clearInterval(this.primaryCheckInterval);
+      this.primaryCheckInterval = null;
+    }
+  }
+
+  private handleMessage(e: MessageEvent<MessageEnvelope>) {
+    const envelope = e.data;
+    const message = envelope.payload;
+    // If a message sent to specific recipient and not to us, ignore it
+    if (!!envelope.recipientId && envelope.recipientId !== this.id) {
+      return;
+    }
+    switch (message.type) {
+      case 'announce':
+        if (this._isPrimary) {
+          this.broadcast({
+            type: 'primary-present',
+            payload: { connected: this._isConnected }
+          });
+          // Re-broadcast starting events for the new tab
+          Object.entries(this.startingEvents).forEach(([eventType, payload]) => {
+            this.broadcast(
+              {
+                type: 'socket-event',
+                payload: { eventType: eventType as SocketEventName, payload }
+              },
+              envelope.senderId
+            );
+          });
+        }
+        return;
+      case 'primary-claim':
+      case 'primary-present':
+        this.lastPrimaryHeartbeat = envelope.timestamp;
+        if (this._isPrimary) {
+          this.logger.info(`Stepping down from primary, ${envelope.senderId} claimed it`);
+          this.stopHeartbeat();
+          this._isPrimary = false;
+          this.callbacks.onRoleChange(false);
+          this.startPrimaryCheck();
+        }
+        return;
+      case 'primary-release':
+        if (!this._isPrimary) {
+          // Wait random time (0 - 300ms) before claiming "primary" role, so not all tabs try to claim at the same time
+          setTimeout(
+            () => {
+              const timeSinceRelease = Date.now() - envelope.timestamp;
+              const timeSinceLastHeartbeat = Date.now() - this.lastPrimaryHeartbeat;
+              if (timeSinceLastHeartbeat > timeSinceRelease) {
+                this.logger.info(`Primary released, becoming primary`);
+                this.stopPrimaryCheck();
+                this.becomePrimary();
+              } else {
+                this.logger.info(`Another primary exists, do nothing`);
+              }
+            },
+            Math.round(Math.random() * 300)
+          );
+        }
+        return;
+      case 'authentication':
+        this.callbacks.onAuthChanged(message.payload.authenticated);
+        return;
+      case 'connection-toggle':
+        if (this._isPrimary) {
+          this.logger.info(
+            `Should ${message.payload.connected ? 'connect' : 'disconnect'} socket.`
+          );
+          this.callbacks.onConnectionToggle(message.payload.connected);
+        }
+        return;
+      case 'socket-event': {
+        if (!this._isPrimary) {
+          this.callbacks.onMessageFromPrimary(message);
+        }
+        return;
+      }
+      case 'socket-send': {
+        if (this._isPrimary) {
+          this.callbacks.onSendRequest(message).then(() => {
+            //Acknowledge the request
+            this.broadcast({
+              type: 'socket-ack',
+              payload: {
+                envelopeId: envelope.id
+              }
+            });
+          });
+        }
+        return;
+      }
+      case 'socket-ack': {
+        if (!this._isPrimary) {
+          if (this.ackCallbacks[message.payload.envelopeId]) {
+            this.ackCallbacks[message.payload.envelopeId]();
+          }
+        }
+      }
+    }
+  }
+
+  private broadcast(message: ChannelMessage, recipientId?: string): string {
+    const envelope: MessageEnvelope = {
+      id: getUUID(),
+      senderId: this.id,
+      timestamp: Date.now(),
+      payload: message
+    };
+    if (recipientId) {
+      envelope.recipientId = recipientId;
+    }
+    this.channel.postMessage(envelope);
+    return envelope.id;
+  }
+
+  /**
+   * Send a message to be transmitted over the socket (from secondary to primary)
+   */
+  sendToSocket<T extends MessageType = MessageType>(messageType: T, data: SocketMessageMap[T]) {
+    if (this._isPrimary) {
+      throw new Error('Primary tab should send directly');
+    }
+    this.logger.debug(`Sending socket message ${messageType}`, data);
+    const messageId = this.broadcast({
+      type: 'socket-send',
+      payload: { messageType, data }
+    });
+    return new Promise<void>((resolve, reject) => {
+      const ackTimeout = setTimeout(() => {
+        //If callback is not called within 5 seconds, call it with an error
+        if (this.ackCallbacks[messageId]) {
+          this.ackCallbacks[messageId](new Error('Message delivery timeout'));
+        }
+      }, WSCoordinator.ACK_TIMEOUT);
+      this.ackCallbacks[messageId] = (error?: any) => {
+        delete this.ackCallbacks[messageId];
+        clearTimeout(ackTimeout);
+        if (error) {
+          reject(error);
+        } else {
+          resolve();
+        }
+      };
+    });
+  }
+
+  /**
+   * Broadcast a socket event from primary to all secondaries
+   */
+  broadcastSocketEvent<E extends SocketEventName = SocketEventName>(
+    eventType: E,
+    payload: SocketEventMap[E]
+  ) {
+    if (!this._isPrimary) {
+      throw new Error('Only primary tab can broadcast socket events');
+    }
+    if (eventType === 'connected') {
+      this._isConnected = true;
+    } else if (eventType === 'disconnected') {
+      this._isConnected = false;
+    }
+    this.updateStartingState(eventType, payload);
+    this.logger.debug(`Broadcasting socket event ${eventType}`, payload);
+    this.broadcast({
+      type: 'socket-event',
+      payload: { eventType, payload }
+    });
+  }
+
+  private updateStartingState<E extends SocketEventName>(eventType: E, payload: SocketEventMap[E]) {
+    if (eventType === 'authenticated') {
+      this.startingEvents.authenticated = payload as AuthenticatedData;
+    } else if (eventType === 'balanceUpdate') {
+      this.startingEvents.balanceUpdate = payload as Balances;
+    } else if (eventType === 'swarmModels') {
+      this.startingEvents.swarmModels = payload as Record<string, number>;
+    }
+  }
+
+  isPrimary() {
+    return this._isPrimary;
+  }
+
+  /**
+   * Handle tab closing event
+   */
+  private handleBeforeUnload = () => {
+    if (this._isPrimary) {
+      this.logger.info(`Client ${this.id} closing, releasing primary role`);
+      this.releasePrimary();
+    }
+  };
+}
+
+export default WSCoordinator;

package/src/ApiClient/WebSocketClient/BrowserWebSocketClient/types.ts

@@ -0,0 +1,107 @@
+import { SocketEventMap, SocketEventName } from '../events';
+import { MessageType, SocketMessageMap } from '../messages';
+
+/**
+ * Primary tab to broadcast socket events to secondary tabs.
+ * @param eventType - The event type.
+ * @param payload - The event payload. See {@link SocketEventMap} for the list of available events.
+ */
+export interface SocketEventReceived<T extends SocketEventName = SocketEventName> {
+  type: 'socket-event';
+  payload: { eventType: T; payload: SocketEventMap[T] };
+}
+
+/**
+ * Sent by secondary tabs to the primary tab to send a message to WebSocket.
+ * @param messageType - The message type.
+ * @param data - The message payload. See {@link SocketMessageMap} for the list of available messages.
+ */
+export interface SendSocketMessage<T extends MessageType = MessageType> {
+  type: 'socket-send';
+  payload: { messageType: T; data: SocketMessageMap[T] };
+}
+
+/**
+ * Sent by the primary tab to acknowledge that a message was sent to WebSocket.
+ */
+export interface SocketMessageAck {
+  type: 'socket-ack';
+  payload: { envelopeId: string };
+}
+
+/**
+ * Sent by the primary tab to notify the secondary tabs that the primary tab is still alive.
+ * @param connected - true if the primary tab is connected to the server, false otherwise.
+ */
+export interface Heartbeat {
+  type: 'primary-present';
+  payload: { connected: boolean };
+}
+
+/**
+ * Used to tell the primary tab to connect/disconnect socket.
+ * @param connected - true to connect, false to disconnect.
+ */
+export interface ConnectionToggle {
+  type: 'connection-toggle';
+  payload: { connected: boolean };
+}
+
+/**
+ * Sent by tab when it is opened to notify other tabs. If another tab is present,
+ * it will respond with {@link Heartbeat}.
+ */
+export interface ClientAnnounce {
+  type: 'announce';
+}
+
+/**
+ * Sent by the client to let other know that it is claiming the primary role.
+ */
+export interface PrimaryClaim {
+  type: 'primary-claim';
+}
+
+/**
+ * Sent by the client to let other know that it is releasing the primary role.
+ * This usually happens when the tab is closed.
+ */
+export interface PrimaryRelease {
+  type: 'primary-release';
+}
+
+/**
+ * Sent by the tab where user has auth state changed.
+ * @param authenticated - true if the user is authenticated, false otherwise.
+ */
+export interface AuthenticationChange {
+  type: 'authentication';
+  payload: { authenticated: boolean };
+}
+
+export type ChannelMessage =
+  | AuthenticationChange
+  | SocketEventReceived
+  | SendSocketMessage
+  | SocketMessageAck
+  | ConnectionToggle
+  | Heartbeat
+  | ClientAnnounce
+  | PrimaryClaim
+  | PrimaryRelease;
+
+/**
+ * Envelope for messages sent between tabs.
+ * @param id - Unique message ID.
+ * @param senderId - ID of the tab that sent the message.
+ * @param recipientId - ID of the tab that should receive the message. If not specified, the message will be broadcasted to all tabs.
+ * @param timestamp - Timestamp of the message.
+ * @param payload - Message payload.
+ */
+export interface MessageEnvelope {
+  id: string;
+  senderId: string;
+  recipientId?: string;
+  timestamp: number;
+  payload: ChannelMessage;
+}

package/src/ApiClient/WebSocketClient/events.ts

@@ -148,3 +148,5 @@ export type SocketEventMap = {
 
   artistCancelConfirmation: ArtistCancelConfirmation;
 };
+
+export type SocketEventName = keyof SocketEventMap;

package/src/ApiClient/WebSocketClient/index.ts

@@ -1,7 +1,7 @@
 import { MessageType, SocketMessageMap } from './messages';
 import { SocketEventMap } from './events';
 import RestClient from '../../lib/RestClient';
-import { SupernetType } from './types';
+import { IWebSocketClient, SupernetType } from './types';
 import WebSocket, { CloseEvent, ErrorEvent, MessageEvent } from 'isomorphic-ws';
 import { base64Decode, base64Encode } from '../../lib/base64';
 import isNodejs from '../../lib/isNodejs';
@@ -13,7 +13,7 @@ const PROTOCOL_VERSION = '3.0.0';
 
 const PING_INTERVAL = 15000;
 
-class WebSocketClient extends RestClient<SocketEventMap> {
+class WebSocketClient extends RestClient<SocketEventMap> implements IWebSocketClient {
   appId: string;
   baseUrl: string;
   private socket: WebSocket | null = null;

package/src/ApiClient/WebSocketClient/types.ts

@@ -1 +1,17 @@
+import { MessageType, SocketMessageMap } from './messages';
+import RestClient from '../../lib/RestClient';
+import { SocketEventMap } from './events';
+
 export type SupernetType = 'relaxed' | 'fast';
+
+export interface IWebSocketClient extends RestClient<SocketEventMap> {
+  appId: string;
+  baseUrl: string;
+  isConnected: boolean;
+  supernetType: SupernetType;
+
+  connect(): Promise<void>;
+  disconnect(): void;
+  send<T extends MessageType>(messageType: T, data: SocketMessageMap[T]): Promise<void>;
+  switchNetwork(supernetType: SupernetType): Promise<SupernetType>;
+}

package/src/ApiClient/index.ts

@@ -3,12 +3,14 @@ import WebSocketClient from './WebSocketClient';
 import TypedEventEmitter from '../lib/TypedEventEmitter';
 import { ApiClientEvents } from './events';
 import { ServerConnectData, ServerDisconnectData } from './WebSocketClient/events';
-import { isNotRecoverable } from './WebSocketClient/ErrorCode';
+import { ErrorCode, isNotRecoverable } from './WebSocketClient/ErrorCode';
 import { JSONValue } from '../types/json';
-import { SupernetType } from './WebSocketClient/types';
+import { IWebSocketClient, SupernetType } from './WebSocketClient/types';
 import { Logger } from '../lib/DefaultLogger';
 import CookieAuthManager from '../lib/AuthManager/CookieAuthManager';
 import { AuthManager, TokenAuthManager } from '../lib/AuthManager';
+import isNodejs from '../lib/isNodejs';
+import BrowserWebSocketClient from './WebSocketClient/BrowserWebSocketClient/BrowserWebSocketClient';
 
 const WS_RECONNECT_ATTEMPTS = 5;
 
@@ -48,7 +50,7 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
   readonly appId: string;
   readonly logger: Logger;
   private _rest: RestClient;
-  private _socket: WebSocketClient;
+  private _socket: IWebSocketClient;
   private _auth: AuthManager;
   private _reconnectAttempts = WS_RECONNECT_ATTEMPTS;
   private _disableSocket: boolean = false;
@@ -68,7 +70,12 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     this._auth =
       authType === 'token' ? new TokenAuthManager(baseUrl, logger) : new CookieAuthManager(logger);
     this._rest = new RestClient(baseUrl, this._auth, logger);
-    this._socket = new WebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+    // Use coordinated WebSocket client in browser, regular in Node.js
+    if (this._auth instanceof TokenAuthManager || isNodejs) {
+      this._socket = new WebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+    } else {
+      this._socket = new BrowserWebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+    }
     this._disableSocket = disableSocket;
     this._auth.on('updated', this.handleAuthUpdated.bind(this));
     this._socket.on('connected', this.handleSocketConnect.bind(this));
@@ -83,7 +90,7 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     return this._auth;
   }
 
-  get socket(): WebSocketClient {
+  get socket(): IWebSocketClient {
     return this._socket;
   }
 
@@ -101,14 +108,26 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
   }
 
   handleSocketDisconnect(data: ServerDisconnectData) {
+    // If user is not authenticated, we don't need to reconnect
+    if (!this.auth.isAuthenticated) {
+      this.emit('disconnected', data);
+      return;
+    }
     if (!data.code || isNotRecoverable(data.code)) {
+      // If this is browser, another tab is probably claiming the connection, so we don't need to reconnect
+      if (
+        this._socket instanceof BrowserWebSocketClient &&
+        data.code === ErrorCode.SWITCH_CONNECTION
+      ) {
+        this.logger.debug('Switching network connection, not reconnecting');
+        return;
+      }
       this.auth.clear();
       this.emit('disconnected', data);
       this.logger.error('Not recoverable socket error', data);
       return;
     }
     if (this._reconnectAttempts <= 0) {
-      this.auth.clear();
       this.emit('disconnected', data);
       this._reconnectAttempts = WS_RECONNECT_ATTEMPTS;
       return;

package/src/Projects/index.ts

@@ -4,6 +4,7 @@ import {
   EnhancementStrength,
   EstimateRequest,
   ImageUrlParams,
+  CostEstimation,
   ProjectParams,
   SizePreset,
   SupportedModel
@@ -518,7 +519,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     guidance,
     sampler,
     contextImages
-  }: EstimateRequest) {
+  }: EstimateRequest): Promise<CostEstimation> {
     let apiVersion = 2;
     const pathParams = [
       tokenType || 'sogni',
@@ -553,7 +554,9 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     );
     return {
       token: r.quote.project.costInToken,
-      usd: r.quote.project.costInUSD
+      usd: r.quote.project.costInUSD,
+      spark: r.quote.project.costInSpark,
+      sogni: r.quote.project.costInSogni
     };
   }
 

package/src/Projects/types/EstimationResponse.ts

@@ -14,6 +14,8 @@ export interface Job {
   costInRenderSec: string;
   costInUSD: string;
   costInToken: string;
+  costInSpark: string;
+  costInSogni: string;
   calculatedStepCount?: number;
 }