@sogni-ai/sogni-client 4.0.0-alpha.5 → 4.0.0-alpha.50

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

@@ -0,0 +1,237 @@
+import { IWebSocketClient, SupernetType } from '../types';
+import { AuthManager, TokenAuthManager } from '../../../lib/AuthManager';
+import { Logger } from '../../../lib/DefaultLogger';
+import WebSocketClient from '../index';
+import RestClient from '../../../lib/RestClient';
+import { SocketEventMap } from '../events';
+import { MessageType, SocketMessageMap } from '../messages';
+import ChannelCoordinator from './ChannelCoordinator';
+
+interface SocketSend<T extends MessageType = MessageType> {
+  type: 'socket-send';
+  payload: { type: T; data: SocketMessageMap[T] };
+}
+
+interface SocketConnect {
+  type: 'connect';
+}
+
+interface SocketDisconnect {
+  type: 'disconnect';
+}
+
+interface SwitchNetwork {
+  type: 'switchNetwork';
+  payload: SupernetType;
+}
+
+type Message = SocketConnect | SocketDisconnect | SocketSend | SwitchNetwork;
+
+interface EventNotification<T extends keyof SocketEventMap = keyof SocketEventMap> {
+  type: 'socket-event';
+  payload: { type: T; data: SocketEventMap[T] };
+}
+
+interface AuthStateChanged {
+  type: 'auth-state-changed';
+  payload: boolean;
+}
+
+type Notification = EventNotification | AuthStateChanged;
+
+type EventInterceptor<T extends keyof SocketEventMap = keyof SocketEventMap> = (
+  eventType: T,
+  payload: SocketEventMap[T]
+) => void;
+
+class WrappedClient extends WebSocketClient {
+  private interceptor: EventInterceptor | undefined = undefined;
+  intercept(interceptor: EventInterceptor) {
+    this.interceptor = interceptor;
+  }
+  protected emit<T extends keyof SocketEventMap>(event: T, data: SocketEventMap[T]) {
+    super.emit(event, data);
+    if (this.interceptor) {
+      this.interceptor(event, data);
+    }
+  }
+}
+
+class BrowserWebSocketClient extends RestClient<SocketEventMap> implements IWebSocketClient {
+  appId: string;
+  baseUrl: string;
+  private socketClient: WrappedClient;
+  private coordinator: ChannelCoordinator<Message, Notification>;
+  private _isConnected = false;
+  private _supernetType: SupernetType;
+
+  constructor(
+    baseUrl: string,
+    auth: AuthManager,
+    appId: string,
+    supernetType: SupernetType,
+    logger: Logger
+  ) {
+    const socketClient = new WrappedClient(baseUrl, auth, appId, supernetType, logger);
+    super(socketClient.baseUrl, auth, logger);
+    this.socketClient = socketClient;
+    this.appId = appId;
+    this.baseUrl = socketClient.baseUrl;
+    this._supernetType = supernetType;
+    this.coordinator = new ChannelCoordinator({
+      callbacks: {
+        onRoleChange: this.handleRoleChange.bind(this),
+        onMessage: this.handleMessage.bind(this),
+        onNotification: this.handleNotification.bind(this)
+      },
+      logger
+    });
+    this.auth.on('updated', this.handleAuthUpdated.bind(this));
+    this.socketClient.intercept(this.handleSocketEvent.bind(this));
+  }
+
+  get isConnected() {
+    return this.coordinator.isPrimary ? this.socketClient.isConnected : this._isConnected;
+  }
+
+  get supernetType() {
+    return this.coordinator.isPrimary ? this.socketClient.supernetType : this._supernetType;
+  }
+
+  async connect(): Promise<void> {
+    await this.coordinator.isReady();
+    if (this.coordinator.isPrimary) {
+      await this.socketClient.connect();
+    } else {
+      return this.coordinator.sendMessage({
+        type: 'connect'
+      });
+    }
+  }
+
+  async disconnect() {
+    await this.coordinator.isReady();
+    if (this.coordinator.isPrimary) {
+      this.socketClient.disconnect();
+    } else {
+      this.coordinator.sendMessage({
+        type: 'disconnect'
+      });
+    }
+  }
+
+  async switchNetwork(supernetType: SupernetType): Promise<SupernetType> {
+    await this.coordinator.isReady();
+    if (this.coordinator.isPrimary) {
+      return this.socketClient.switchNetwork(supernetType);
+    }
+    await this.coordinator.sendMessage({
+      type: 'switchNetwork',
+      payload: supernetType
+    });
+    this._supernetType = supernetType;
+    return supernetType;
+  }
+
+  async send<T extends MessageType>(messageType: T, data: SocketMessageMap[T]): Promise<void> {
+    await this.coordinator.isReady();
+    if (this.coordinator.isPrimary) {
+      if (!this.socketClient.isConnected) {
+        await this.socketClient.connect();
+      }
+      return this.socketClient.send(messageType, data);
+    }
+    return this.coordinator.sendMessage({
+      type: 'socket-send',
+      payload: { type: messageType, data }
+    });
+  }
+
+  private async handleMessage(message: Message) {
+    this._logger.debug('Received control message', message);
+    switch (message.type) {
+      case 'socket-send': {
+        if (!this.socketClient.isConnected) {
+          await this.socketClient.connect();
+        }
+        return this.socketClient.send(message.payload.type, message.payload.data);
+      }
+      case 'connect': {
+        if (!this.socketClient.isConnected) {
+          await this.socketClient.connect();
+        }
+        return;
+      }
+      case 'disconnect': {
+        if (this.socketClient.isConnected) {
+          this.socketClient.disconnect();
+        }
+        return;
+      }
+      case 'switchNetwork': {
+        await this.switchNetwork(message.payload);
+        return;
+      }
+      default: {
+        this._logger.error('Received unknown message type:', message);
+      }
+    }
+  }
+
+  private async handleNotification(notification: Notification) {
+    this._logger.debug('Received notification', notification.type, notification.payload);
+    switch (notification.type) {
+      case 'socket-event': {
+        this.emit(notification.payload.type, notification.payload.data);
+        return;
+      }
+      case 'auth-state-changed': {
+        this.handleAuthChanged(notification.payload);
+        return;
+      }
+      default: {
+        this._logger.error('Received unknown notification type:', notification);
+      }
+    }
+  }
+
+  private handleAuthChanged(isAuthenticated: boolean) {
+    if (this.auth instanceof TokenAuthManager) {
+      throw new Error('TokenAuthManager is not supported in multi client mode');
+    }
+    if (this.auth.isAuthenticated !== isAuthenticated) {
+      if (isAuthenticated) {
+        this.auth.authenticate();
+      } else {
+        this.auth.clear();
+      }
+    }
+  }
+
+  private handleSocketEvent(eventType: keyof SocketEventMap, payload: any) {
+    if (this.coordinator.isPrimary) {
+      this.coordinator.notify({
+        type: 'socket-event',
+        payload: { type: eventType, data: payload }
+      });
+      this.emit(eventType, payload);
+    }
+  }
+
+  private handleAuthUpdated(isAuthenticated: boolean) {
+    this.coordinator.notify({
+      type: 'auth-state-changed',
+      payload: isAuthenticated
+    });
+  }
+
+  private handleRoleChange(isPrimary: boolean) {
+    if (isPrimary && !this.socketClient.isConnected && this.isConnected) {
+      this.socketClient.connect();
+    } else if (!isPrimary && this.socketClient.isConnected) {
+      this.socketClient.disconnect();
+    }
+  }
+}
+
+export default BrowserWebSocketClient;

package/src/ApiClient/WebSocketClient/events.ts

@@ -48,6 +48,12 @@ export type JobProgressData = {
   stepCount: number;
 };
 
+export type JobETAData = {
+  jobID: string;
+  imgID?: string;
+  etaSeconds: number;
+};
+
 export type JobResultData = {
   jobID: string;
   imgID: string;
@@ -55,6 +61,7 @@ export type JobResultData = {
   lastSeed: string;
   userCanceled: boolean;
   triggeredNSFWFilter: boolean;
+  resultUrl?: string;
 };
 
 export type JobStateData =
@@ -121,6 +128,11 @@ export type SocketEventMap = {
    * @event WebSocketClient#jobProgress - Job progress update
    */
   jobProgress: JobProgressData;
+  /**
+   * @event WebSocketClient#jobETA - Job ETA update (sent every second during inference by ComfyUI workers)
+   * Note: Only available for ComfyUI-based workers during video generation
+   */
+  jobETA: JobETAData;
   /**
    * @event WebSocketClient#jobResult - Job result received
    */
@@ -148,3 +160,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;
@@ -86,7 +86,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     socket.onmessage = null;
     socket.onopen = null;
     this.stopPing();
-    socket.close();
+    socket.close(1000, 'Client disconnected');
   }
 
   private startPing(socket: WebSocket) {
@@ -148,9 +148,16 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   }
 
   private handleClose(e: CloseEvent) {
-    if (e.target === this.socket) {
+    const socket = e.target;
+    socket.onerror = null;
+    socket.onmessage = null;
+    socket.onopen = null;
+    if (socket === this.socket || !this.socket) {
       this._logger.info('WebSocket disconnected, cleanup', e);
-      this.disconnect();
+      if (socket === this.socket) {
+        this.stopPing();
+        this.socket = null;
+      }
       this.emit('disconnected', {
         code: e.code,
         reason: e.reason
@@ -193,6 +200,9 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   }
 
   async send<T extends MessageType>(messageType: T, data: SocketMessageMap[T]) {
+    if (!this.isConnected) {
+      await this.connect();
+    }
     await this.waitForConnection();
     this._logger.debug('WebSocket send:', messageType, data);
     this.socket!.send(

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';
 
 const WS_RECONNECT_ATTEMPTS = 5;
 
@@ -42,13 +44,14 @@ export interface ApiClientOptions {
   logger: Logger;
   authType: 'token' | 'cookies';
   disableSocket?: boolean;
+  multiInstance?: boolean;
 }
 
 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;
@@ -60,7 +63,8 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     networkType,
     authType,
     logger,
-    disableSocket = false
+    disableSocket = false,
+    multiInstance = false
   }: ApiClientOptions) {
     super();
     this.appId = appId;
@@ -68,7 +72,13 @@ 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);
+    const supportMultiInstance = !isNodejs && this._auth instanceof CookieAuthManager;
+    if (supportMultiInstance && multiInstance) {
+      // Use coordinated WebSocket client to share single connection between tabs
+      this._socket = new BrowserWebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+    } else {
+      this._socket = new WebSocketClient(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 +93,7 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     return this._auth;
   }
 
-  get socket(): WebSocketClient {
+  get socket(): IWebSocketClient {
     return this._socket;
   }
 
@@ -101,14 +111,26 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
   }
 
   handleSocketDisconnect(data: ServerDisconnectData) {
+    // If user is not authenticated, we don't need to reconnect
+    if (!this.auth.isAuthenticated || data.code === 1000) {
+      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;
@@ -122,7 +144,7 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
       if (this.socket.isConnected) {
         this.socket.disconnect();
       }
-    } else if (!this._disableSocket) {
+    } else if (!this._disableSocket && !this.socket.isConnected) {
       this.socket.connect();
     }
   }

package/src/Projects/Job.ts

@@ -9,6 +9,7 @@ import Project from './Project';
 import { SupernetType } from '../ApiClient/WebSocketClient/types';
 import { getEnhacementStrength } from './utils';
 import { TokenType } from '../types/token';
+import { has } from 'lodash';
 
 export const enhancementDefaults = {
   network: 'fast' as SupernetType,
@@ -19,7 +20,7 @@ export const enhancementDefaults = {
   startingImageStrength: 0.5,
   steps: 5,
   guidance: 1,
-  numberOfImages: 1,
+  numberOfMedia: 1,
   numberOfPreviews: 0
 };
 
@@ -61,6 +62,17 @@ export interface JobData {
   positivePrompt?: string;
   negativePrompt?: string;
   jobIndex?: number;
+  /**
+   * Estimated time remaining in seconds (for long-running jobs like video generation).
+   * Updated by ComfyUI workers during inference.
+   * @deprecated Use `eta` instead.
+   */
+  etaSeconds?: number;
+  /**
+   * Estimate completion time of the job (for long-running jobs like video generation).
+   * Updated by ComfyUI workers during inference.
+   */
+  eta?: Date;
 }
 
 export interface JobEventMap extends EntityEvents {
@@ -178,10 +190,21 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.error;
   }
 
-  get hasResultImage() {
+  /**
+   * Whether this job has a result media file available for download.
+   * Returns true if completed and not NSFW filtered.
+   */
+  get hasResultMedia() {
     return this.status === 'completed' && !this.isNSFW;
   }
 
+  /**
+   * Whether this job produces video output (based on the model used)
+   */
+  get type(): 'image' | 'video' {
+    return this._api.isVideoModelId(this._project.params.modelId) ? 'video' : 'image';
+  }
+
   get enhancedImage() {
     if (!this._enhancementProject) {
       return null;
@@ -199,17 +222,27 @@ class Job extends DataEntity<JobData, JobEventMap> {
 
   /**
    * Get the result URL of the job. This method will make a request to the API to get signed URL.
-   * IMPORTANT: URL expires after 30 minutes, so make sure to download the image as soon as possible.
+   * IMPORTANT: URL expires after 30 minutes, so make sure to download the result as soon as possible.
+   * For video jobs, this returns a video URL. For image jobs, this returns an image URL.
    */
   async getResultUrl(): Promise<string> {
     if (this.data.status !== 'completed') {
       throw new Error('Job is not completed yet');
     }
-    const url = await this._api.downloadUrl({
-      jobId: this.projectId,
-      imageId: this.id,
-      type: 'complete'
-    });
+    let url: string;
+    if (this.type === 'video') {
+      url = await this._api.mediaDownloadUrl({
+        jobId: this.projectId,
+        id: this.id,
+        type: 'complete'
+      });
+    } else {
+      url = await this._api.downloadUrl({
+        jobId: this.projectId,
+        imageId: this.id,
+        type: 'complete'
+      });
+    }
     this._update({ resultUrl: url });
     return url;
   }
@@ -230,6 +263,26 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.workerName;
   }
 
+  /**
+   * Estimated time remaining in seconds for long-running jobs (e.g., video generation).
+   * Only available for ComfyUI-based workers during inference.
+   * Returns undefined if no ETA has been received.
+   * @deprecated Use `timeLeft` instead.
+   */
+  get etaSeconds() {
+    return this.data.etaSeconds;
+  }
+
+  /**
+   * Estimate completion time of the job.
+   * Only available for ComfyUI-based workers during inference.
+   * Is useful when data is persisted
+   * Returns undefined if no ETA has been received.
+   */
+  get eta() {
+    return this.data.eta;
+  }
+
   /**
    * Syncs the job data with the data received from the REST API.
    * @internal
@@ -247,11 +300,19 @@ class Job extends DataEntity<JobData, JobEventMap> {
     }
     if (!this.data.resultUrl && delta.status === 'completed' && !data.triggeredNSFWFilter) {
       try {
-        delta.resultUrl = await this._api.downloadUrl({
-          jobId: this.projectId,
-          imageId: this.id,
-          type: 'complete'
-        });
+        if (this.type === 'video') {
+          delta.resultUrl = await this._api.mediaDownloadUrl({
+            jobId: this.projectId,
+            id: this.id,
+            type: 'complete'
+          });
+        } else {
+          delta.resultUrl = await this._api.downloadUrl({
+            jobId: this.projectId,
+            imageId: this.id,
+            type: 'complete'
+          });
+        }
       } catch (error) {
         this._logger.error(error);
       }
@@ -259,6 +320,21 @@ class Job extends DataEntity<JobData, JobEventMap> {
     this._update(delta);
   }
 
+  /**
+   * Updates the job data with the provided delta.
+   * @internal
+   * @param delta
+   */
+  _update(delta: Partial<JobData>) {
+    if (has(delta, 'eta')) {
+      // Keeping etaSeconds for backwards compatibility
+      if (delta.eta) {
+        delta.etaSeconds = Math.round((delta.eta.getTime() - Date.now()) / 1000);
+      }
+    }
+    super._update(delta);
+  }
+
   private handleUpdated(keys: string[]) {
     if (keys.includes('step') || keys.includes('stepCount')) {
       this.emit('progress', this.progress);
@@ -276,8 +352,8 @@ class Job extends DataEntity<JobData, JobEventMap> {
   }
 
   async getResultData() {
-    if (!this.hasResultImage) {
-      throw new Error('No result image available');
+    if (!this.hasResultMedia) {
+      throw new Error('No result media available');
     }
     const url = await this.getResultUrl();
     const response = await fetch(url);
@@ -297,6 +373,10 @@ class Job extends DataEntity<JobData, JobEventMap> {
     strength: EnhancementStrength,
     overrides: { positivePrompt?: string; stylePrompt?: string; tokenType?: TokenType } = {}
   ) {
+    const parentProjectParams = this._project.params;
+    if (parentProjectParams.type !== 'image') {
+      throw new Error('Enhancement is only available for images');
+    }
     if (this.status !== 'completed') {
       throw new Error('Job is not completed yet');
     }
@@ -309,6 +389,7 @@ class Job extends DataEntity<JobData, JobEventMap> {
     }
     const imageData = await this.getResultData();
     const project = await this._api.create({
+      type: 'image',
       ...enhancementDefaults,
       positivePrompt: overrides.positivePrompt || this._project.params.positivePrompt,
       stylePrompt: overrides.stylePrompt || this._project.params.stylePrompt,
@@ -316,7 +397,7 @@ class Job extends DataEntity<JobData, JobEventMap> {
       seed: this.seed || this._project.params.seed,
       startingImage: imageData,
       startingImageStrength: 1 - getEnhacementStrength(strength),
-      sizePreset: this._project.params.sizePreset
+      sizePreset: parentProjectParams.sizePreset
     });
     this._enhancementProject = project;
     this._enhancementProject.on('updated', this.handleEnhancementUpdate);