@sogni-ai/sogni-client 0.3.1 → 0.3.3

package/src/ApiClient/WebSocketClient/index.ts

@@ -0,0 +1,203 @@
+import { MessageType, SocketMessageMap } from './messages';
+import { SocketEventMap } from './events';
+import RestClient, { AuthData } from '../../lib/RestClient';
+import { SupernetType } from './types';
+import WebSocket, { CloseEvent, ErrorEvent, MessageEvent } from 'isomorphic-ws';
+import { base64Decode, base64Encode } from '../../lib/base64';
+import isNodejs from '../../lib/isNodejs';
+import Cookie from 'js-cookie';
+import { LIB_VERSION } from '../../version';
+import { Logger } from '../../lib/DefaultLogger';
+
+const PING_INTERVAL = 15000;
+
+class WebSocketClient extends RestClient<SocketEventMap> {
+  appId: string;
+  baseUrl: string;
+  private socket: WebSocket | null = null;
+  private _supernetType: SupernetType;
+  private _pingInterval: NodeJS.Timeout | null = null;
+
+  constructor(baseUrl: string, appId: string, supernetType: SupernetType, logger: Logger) {
+    super(baseUrl, logger);
+    this.appId = appId;
+    this.baseUrl = baseUrl;
+    this._supernetType = supernetType;
+  }
+
+  set auth(auth: AuthData | null) {
+    //In browser, set the cookie
+    if (!isNodejs) {
+      if (auth) {
+        Cookie.set('authorization', auth.token, {
+          domain: '.sogni.ai',
+          expires: 1
+        });
+      } else {
+        Cookie.remove('authorization', {
+          domain: '.sogni.ai'
+        });
+      }
+    }
+    this._auth = auth;
+  }
+
+  get supernetType(): SupernetType {
+    return this._supernetType;
+  }
+
+  get isConnected(): boolean {
+    return !!this.socket;
+  }
+
+  connect() {
+    if (this.socket) {
+      this.disconnect();
+    }
+    const userAgent = `Sogni/${LIB_VERSION} (sogni-client)`;
+    const url = new URL(this.baseUrl);
+    url.searchParams.set('appId', this.appId);
+    url.searchParams.set('clientName', userAgent);
+    url.searchParams.set('clientType', 'artist');
+    //At this point 'relaxed' does not work as expected, so we use 'fast' or empty
+    url.searchParams.set('forceWorkerId', this._supernetType === 'fast' ? 'fast' : '');
+    let params;
+    // In Node.js, ws package is used, so we need to set the auth header
+    if (isNodejs) {
+      params = {
+        headers: {
+          Authorization: this._auth?.token,
+          'User-Agent': userAgent
+        }
+      };
+    }
+    this.socket = new WebSocket(url.toString(), params);
+    this.socket.onerror = this.handleError.bind(this);
+    this.socket.onmessage = this.handleMessage.bind(this);
+    this.socket.onopen = this.handleOpen.bind(this);
+    this.socket.onclose = this.handleClose.bind(this);
+    this.startPing(this.socket);
+  }
+
+  disconnect() {
+    if (!this.socket) {
+      return;
+    }
+    const socket = this.socket;
+    this.socket = null;
+    socket.onerror = null;
+    socket.onmessage = null;
+    socket.onopen = null;
+    this.stopPing();
+    socket.close();
+  }
+
+  private startPing(socket: WebSocket) {
+    if (!isNodejs) {
+      return;
+    }
+    this._pingInterval = setInterval(() => {
+      socket.ping();
+    }, PING_INTERVAL);
+  }
+
+  private stopPing() {
+    if (this._pingInterval) {
+      clearInterval(this._pingInterval);
+      this._pingInterval = null;
+    }
+  }
+
+  switchNetwork(supernetType: SupernetType) {
+    this._supernetType = supernetType;
+    this.disconnect();
+    this.connect();
+  }
+
+  /**
+   * Ensure the WebSocket connection is open, waiting if necessary and throwing an error if it fails
+   * @private
+   */
+  private async waitForConnection(): Promise<void> {
+    if (!this.socket) {
+      throw new Error('WebSocket not connected');
+    }
+    if (this.socket.readyState === WebSocket.OPEN) {
+      return;
+    }
+    let attempts = 10;
+    while (this.socket?.readyState === WebSocket.CONNECTING) {
+      this._logger.info('Waiting for WebSocket connection...');
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+      attempts--;
+      if (attempts === 0) {
+        this.disconnect();
+        throw new Error('WebSocket connection timeout');
+      }
+    }
+    //@ts-expect-error State may change between checks
+    if (this.socket?.readyState !== WebSocket.OPEN) {
+      this.disconnect();
+      throw new Error('WebSocket connection failed');
+    }
+  }
+
+  private handleOpen() {
+    this.emit('connected', { network: this._supernetType });
+  }
+
+  private handleClose(e: CloseEvent) {
+    if (e.target === this.socket) {
+      this._logger.info('WebSocket disconnected, cleanup', e);
+      this.disconnect();
+      this.emit('disconnected', {
+        code: e.code,
+        reason: e.reason
+      });
+    }
+  }
+
+  private handleError(e: ErrorEvent) {
+    this._logger.error('WebSocket error:', e);
+  }
+
+  private handleMessage(e: MessageEvent) {
+    let dataPromise: Promise<string>;
+    // In Node.js, e.data is a Buffer, while in browser it's a Blob
+    if (isNodejs) {
+      dataPromise = Promise.resolve(e.data.toString());
+    } else {
+      const data = e.data as unknown as Blob;
+      dataPromise = data.text();
+    }
+    dataPromise
+      .then((str: string) => {
+        const data = JSON.parse(str);
+        let payload = null;
+        if (data.data) {
+          payload = JSON.parse(base64Decode(data.data));
+        }
+        // Convert jobID and imgID to uppercase for consistency
+        ['jobID', 'imgID'].forEach((idKey) => {
+          if (payload[idKey]) {
+            payload[idKey] = payload[idKey].toUpperCase();
+          }
+        });
+        this._logger.debug('WebSocket:', data.type, payload);
+        this.emit(data.type, payload);
+      })
+      .catch((err: any) => {
+        this._logger.error('Failed to parse WebSocket message:', err);
+      });
+  }
+
+  async send<T extends MessageType>(messageType: T, data: SocketMessageMap[T]) {
+    await this.waitForConnection();
+    this._logger.debug('WebSocket send:', messageType, data);
+    this.socket!.send(
+      JSON.stringify({ type: messageType, data: base64Encode(JSON.stringify(data)) })
+    );
+  }
+}
+
+export default WebSocketClient;

package/src/ApiClient/WebSocketClient/messages.ts

@@ -0,0 +1,7 @@
+import { JobRequestRaw } from '../../Projects/createJobRequestMessage';
+
+export interface SocketMessageMap {
+  jobRequest: JobRequestRaw;
+}
+
+export type MessageType = keyof SocketMessageMap;

package/src/ApiClient/WebSocketClient/types.ts

@@ -0,0 +1 @@
+export type SupernetType = 'relaxed' | 'fast';

package/src/ApiClient/events.ts

@@ -0,0 +1,20 @@
+import ErrorCode from './WebSocketClient/ErrorCode';
+import { SupernetType } from './WebSocketClient/types';
+
+export type ApiClientEvents = {
+  /**
+   * @event ApiClient#connected - The client has been connected to the server.
+   */
+  connected: {
+    network: SupernetType;
+  };
+  /**
+   * @event ApiClient#disconnected - The client has been disconnected by the server,
+   * either authentication is lost or the server is unreachable.This event is not triggered
+   * when the client manually disconnects.
+   */
+  disconnected: {
+    code: ErrorCode;
+    reason: string;
+  };
+};

package/src/ApiClient/index.ts

@@ -0,0 +1,124 @@
+import RestClient from '../lib/RestClient';
+import WebSocketClient from './WebSocketClient';
+import { jwtDecode } from 'jwt-decode';
+import TypedEventEmitter from '../lib/TypedEventEmitter';
+import { ApiClientEvents } from './events';
+import { ServerConnectData, ServerDisconnectData } from './WebSocketClient/events';
+import { isNotRecoverable } from './WebSocketClient/ErrorCode';
+import { JSONValue } from '../types/json';
+import { SupernetType } from './WebSocketClient/types';
+import { Logger } from '../lib/DefaultLogger';
+
+const WS_RECONNECT_ATTEMPTS = 5;
+
+export interface ApiReponse<D = JSONValue> {
+  status: 'success';
+  data: D;
+}
+
+/** @inline */
+export interface ApiErrorResponse {
+  status: 'error';
+  message: string;
+  errorCode: number;
+}
+
+export class ApiError extends Error {
+  status: number;
+  payload: ApiErrorResponse;
+  constructor(status: number, payload: ApiErrorResponse) {
+    super(payload.message);
+    this.status = status;
+    this.payload = payload;
+  }
+}
+
+/**
+ * @inline
+ */
+interface AuthData {
+  token: string;
+  walletAddress: string;
+  expiresAt: Date;
+}
+
+class ApiClient extends TypedEventEmitter<ApiClientEvents> {
+  readonly appId: string;
+  readonly logger: Logger;
+  private _rest: RestClient;
+  private _socket: WebSocketClient;
+  private _auth: AuthData | null = null;
+  private _reconnectAttempts = WS_RECONNECT_ATTEMPTS;
+
+  constructor(
+    baseUrl: string,
+    socketUrl: string,
+    appId: string,
+    networkType: SupernetType,
+    logger: Logger
+  ) {
+    super();
+    this.appId = appId;
+    this.logger = logger;
+    this._rest = new RestClient(baseUrl, logger);
+    this._socket = new WebSocketClient(socketUrl, appId, networkType, logger);
+    this._socket.on('connected', this.handleSocketConnect.bind(this));
+    this._socket.on('disconnected', this.handleSocketDisconnect.bind(this));
+  }
+
+  get isAuthenticated(): boolean {
+    return !!this._auth && this._auth.expiresAt > new Date();
+  }
+
+  get auth(): AuthData | null {
+    return this._auth && this._auth.expiresAt > new Date() ? this._auth : null;
+  }
+
+  get socket(): WebSocketClient {
+    return this._socket;
+  }
+
+  get rest(): RestClient {
+    return this._rest;
+  }
+
+  authenticate(token: string) {
+    const decoded = jwtDecode<{ addr: string; env: string; iat: number; exp: number }>(token);
+    this._auth = {
+      token,
+      walletAddress: decoded.addr,
+      expiresAt: new Date(decoded.exp * 1000)
+    };
+    this.rest.auth = { token };
+    this.socket.auth = { token };
+    this.socket.connect();
+  }
+
+  removeAuth() {
+    this._auth = null;
+    this.socket.disconnect();
+  }
+
+  handleSocketConnect({ network }: ServerConnectData) {
+    this._reconnectAttempts = WS_RECONNECT_ATTEMPTS;
+    this.emit('connected', { network });
+  }
+
+  handleSocketDisconnect(data: ServerDisconnectData) {
+    if (!data.code || isNotRecoverable(data.code)) {
+      this.removeAuth();
+      this.emit('disconnected', data);
+      this.logger.error('Not recoverable socket error', data);
+      return;
+    }
+    if (this._reconnectAttempts <= 0) {
+      this.emit('disconnected', data);
+      this._reconnectAttempts = WS_RECONNECT_ATTEMPTS;
+      return;
+    }
+    this._reconnectAttempts--;
+    setTimeout(() => this.socket.connect(), 1000);
+  }
+}
+
+export default ApiClient;

package/src/ApiGroup.ts

@@ -0,0 +1,25 @@
+import { AbstractProvider } from 'ethers';
+import ApiClient from './ApiClient';
+import EIP712Helper from './lib/EIP712Helper';
+import TypedEventEmitter, { EventMap } from './lib/TypedEventEmitter';
+
+export interface ApiConfig {
+  client: ApiClient;
+  provider: AbstractProvider;
+  eip712: EIP712Helper;
+}
+
+abstract class ApiGroup<E extends EventMap = {}> extends TypedEventEmitter<E> {
+  protected client: ApiClient;
+  protected provider: AbstractProvider;
+  protected eip712: EIP712Helper;
+
+  constructor(config: ApiConfig) {
+    super();
+    this.client = config.client;
+    this.provider = config.provider;
+    this.eip712 = config.eip712;
+  }
+}
+
+export default ApiGroup;

package/src/Projects/Job.ts

@@ -0,0 +1,124 @@
+import DataEntity, { EntityEvents } from '../lib/DataEntity';
+import ErrorData from '../types/ErrorData';
+
+export type JobStatus =
+  | 'pending'
+  | 'initiating'
+  | 'processing'
+  | 'completed'
+  | 'failed'
+  | 'canceled';
+
+/**
+ * @inline
+ */
+export interface JobData {
+  id: string;
+  status: JobStatus;
+  step: number;
+  stepCount: number;
+  seed?: number;
+  isNSFW?: boolean;
+  userCanceled?: boolean;
+  previewUrl?: string;
+  resultUrl?: string | null;
+  error?: ErrorData;
+}
+
+export interface JobEventMap extends EntityEvents {
+  progress: number;
+  completed: string;
+  failed: ErrorData;
+}
+
+class Job extends DataEntity<JobData, JobEventMap> {
+  constructor(data: JobData) {
+    super(data);
+    this.on('updated', this.handleUpdated.bind(this));
+  }
+
+  get id() {
+    return this.data.id;
+  }
+
+  /**
+   * Current status of the job.
+   */
+  get status() {
+    return this.data.status;
+  }
+
+  /**
+   * Progress of the job in percentage (0-100).
+   */
+  get progress() {
+    return Math.round((this.data.step / this.data.stepCount) * 100);
+  }
+
+  /**
+   * Current step of the job.
+   */
+  get step() {
+    return this.data.step;
+  }
+
+  /**
+   * Total number of steps that worker will perform.
+   */
+  get stepCount() {
+    return this.data.stepCount;
+  }
+
+  /**
+   * Seed used to generate the image. This property is only available when the job is completed.
+   */
+  get seed() {
+    return this.data.seed;
+  }
+
+  /**
+   * Last preview image URL generated by the worker.
+   */
+  get previewUrl() {
+    return this.data.previewUrl;
+  }
+
+  /**
+   * URL to the result image, could be null if the job was canceled or triggered NSFW filter while
+   * it was not disabled explicitly.
+   */
+  get resultUrl() {
+    return this.data.resultUrl;
+  }
+
+  get imageUrl() {
+    return this.data.resultUrl || this.data.previewUrl;
+  }
+
+  get error() {
+    return this.data.error;
+  }
+
+  /**
+   * Whether the image is NSFW or not. Only makes sense if job is completed.
+   * If NSFW filter is disabled, this property will always be false.
+   * If NSFW filter is enabled and the image is NSFW, image will not be available for download.
+   */
+  get isNSFW() {
+    return !!this.data.isNSFW;
+  }
+
+  private handleUpdated(keys: string[]) {
+    if (keys.includes('step') || keys.includes('stepCount')) {
+      this.emit('progress', this.progress);
+    }
+    if (keys.includes('status') && this.status === 'completed') {
+      this.emit('completed', this.resultUrl!);
+    }
+    if (keys.includes('status') && this.status === 'failed') {
+      this.emit('failed', this.data.error!);
+    }
+  }
+}
+
+export default Job;

package/src/Projects/Project.ts

@@ -0,0 +1,185 @@
+import Job, { JobData } from './Job';
+import DataEntity, { EntityEvents } from '../lib/DataEntity';
+import { ProjectParams } from './types';
+import cloneDeep from 'lodash/cloneDeep';
+import ErrorData from '../types/ErrorData';
+import getUUID from '../lib/getUUID';
+
+export type ProjectStatus = 'pending' | 'queued' | 'processing' | 'completed' | 'failed';
+
+/**
+ * @inline
+ */
+export interface ProjectData {
+  id: string;
+  startedAt: Date;
+  params: ProjectParams;
+  queuePosition: number;
+  status: ProjectStatus;
+  error?: ErrorData;
+}
+/** @inline */
+export interface SerializedProject extends ProjectData {
+  jobs: JobData[];
+}
+
+export interface ProjectEventMap extends EntityEvents {
+  progress: number;
+  completed: string[];
+  failed: ErrorData;
+  jobCompleted: Job;
+  jobFailed: Job;
+}
+
+class Project extends DataEntity<ProjectData, ProjectEventMap> {
+  private _jobs: Job[] = [];
+  private _lastEmitedProgress = -1;
+
+  constructor(data: ProjectParams) {
+    super({
+      id: getUUID(),
+      startedAt: new Date(),
+      params: data,
+      queuePosition: -1,
+      status: 'pending'
+    });
+
+    this.on('updated', this.handleUpdated.bind(this));
+  }
+
+  get id() {
+    return this.data.id;
+  }
+
+  get params() {
+    return this.data.params;
+  }
+
+  get status() {
+    return this.data.status;
+  }
+
+  get error() {
+    return this.data.error;
+  }
+
+  /**
+   * Progress of the project in percentage (0-100).
+   */
+  get progress() {
+    // Worker can reduce the number of steps in the job, so we need to calculate the progress based on the actual number of steps
+    const stepsPerJob = this.jobs.length ? this.jobs[0].stepCount : this.data.params.steps;
+    const jobCount = this.data.params.numberOfImages;
+    const stepsDone = this._jobs.reduce((acc, job) => acc + job.step, 0);
+    return Math.round((stepsDone / (stepsPerJob * jobCount)) * 100);
+  }
+
+  get queuePosition() {
+    return this.data.queuePosition;
+  }
+
+  /**
+   * List of jobs in the project. Note that jobs will be added to this list as
+   * workers start processing them. So initially this list will be empty.
+   * Subscribe to project `updated` event to get notified about any update, including new jobs.
+   * @example
+   * project.on('updated', (keys) => {
+   *  if (keys.includes('jobs')) {
+   *    // Project jobs have been updated
+   *  }
+   * });
+   */
+  get jobs() {
+    return this._jobs.slice(0);
+  }
+
+  /**
+   * List of result URLs for all completed jobs in the project.
+   */
+  get resultUrls() {
+    return this.jobs.map((job) => job.resultUrl).filter((r) => !!r) as string[];
+  }
+
+  /**
+   * Wait for the project to complete, then return the result URLs, or throw an error if the project fails.
+   * @returns Promise<string[]> - Promise that resolves to the list of result URLs
+   * @throws ErrorData
+   */
+  waitForCompletion(): Promise<string[]> {
+    if (this.status === 'completed') {
+      return Promise.resolve(this.resultUrls);
+    }
+    if (this.status === 'failed') {
+      return Promise.reject(this.error);
+    }
+
+    return new Promise((resolve, reject) => {
+      this.on('completed', (images) => {
+        resolve(images);
+      });
+      this.on('failed', (error) => {
+        reject(error);
+      });
+    });
+  }
+
+  /**
+   * Find a job by id
+   * @param id
+   */
+  job(id: string) {
+    return this._jobs.find((job) => job.id === id);
+  }
+
+  private handleUpdated(keys: string[]) {
+    const progress = this.progress;
+    if (progress !== this._lastEmitedProgress) {
+      this.emit('progress', progress);
+      this._lastEmitedProgress = progress;
+    }
+    if (keys.includes('status') || keys.includes('jobs')) {
+      const allJobsDone = this.jobs.every((job) =>
+        ['completed', 'failed', 'canceled'].includes(job.status)
+      );
+      if (this.data.status === 'completed' && allJobsDone) {
+        return this.emit('completed', this.resultUrls);
+      }
+      if (this.data.status === 'failed') {
+        this.emit('failed', this.data.error!);
+      }
+    }
+  }
+
+  /**
+   * This is internal method to add a job to the project. Do not call this directly.
+   * @internal
+   * @param data
+   */
+  _addJob(data: JobData) {
+    const job = new Job(data);
+    this._jobs.push(job);
+    job.on('updated', () => {
+      this.emit('updated', ['jobs']);
+    });
+    job.on('completed', () => {
+      this.emit('jobCompleted', job);
+    });
+    job.on('failed', () => {
+      this.emit('jobFailed', job);
+    });
+    return job;
+  }
+
+  /**
+   * Get full project data snapshot. Can be used to serialize the project and store it in a database.
+   */
+  toJSON(): SerializedProject {
+    const data = cloneDeep(this.data);
+    return {
+      ...data,
+      jobs: this._jobs.map((job) => job.toJSON())
+    };
+  }
+}
+
+export default Project;