@sogni-ai/sogni-client 0.4.2 → 0.5.0-alpha.0

package/src/ApiClient/WebSocketClient/index.ts

@@ -1,13 +1,13 @@
 import { MessageType, SocketMessageMap } from './messages';
 import { SocketEventMap } from './events';
-import RestClient, { AuthData } from '../../lib/RestClient';
+import RestClient 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';
+import AuthManager from '../../lib/AuthManager';
 
 const PING_INTERVAL = 15000;
 
@@ -18,34 +18,23 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   private _supernetType: SupernetType;
   private _pingInterval: NodeJS.Timeout | null = null;
 
-  constructor(baseUrl: string, appId: string, supernetType: SupernetType, logger: Logger) {
+  constructor(
+    baseUrl: string,
+    auth: AuthManager,
+    appId: string,
+    supernetType: SupernetType,
+    logger: Logger
+  ) {
     const _baseUrl = new URL(baseUrl);
     if (_baseUrl.protocol === 'wss:') {
       _baseUrl.protocol = 'https:';
     }
-    super(_baseUrl.toString(), logger);
+    super(_baseUrl.toString(), auth, logger);
     this.appId = appId;
     this.baseUrl = _baseUrl.toString();
     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;
   }
@@ -54,7 +43,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     return !!this.socket;
   }
 
-  connect() {
+  async connect() {
     if (this.socket) {
       this.disconnect();
     }
@@ -71,7 +60,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     if (isNodejs) {
       params = {
         headers: {
-          Authorization: this._auth?.token,
+          Authorization: await this.auth.getToken(),
           'User-Agent': userAgent
         }
       };
@@ -114,7 +103,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   }
 
   switchNetwork(supernetType: SupernetType): Promise<SupernetType> {
-    return new Promise<SupernetType>(async (resolve, reject) => {
+    return new Promise<SupernetType>(async (resolve) => {
       this.once('changeNetwork', ({ network }) => {
         this._supernetType = network;
         resolve(network);

package/src/ApiClient/index.ts

@@ -1,6 +1,5 @@
 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';
@@ -8,6 +7,7 @@ import { isNotRecoverable } from './WebSocketClient/ErrorCode';
 import { JSONValue } from '../types/json';
 import { SupernetType } from './WebSocketClient/types';
 import { Logger } from '../lib/DefaultLogger';
+import AuthManager, { Tokens } from '../lib/AuthManager';
 
 const WS_RECONNECT_ATTEMPTS = 5;
 
@@ -33,21 +33,12 @@ export class ApiError extends Error {
   }
 }
 
-/**
- * @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 _auth: AuthManager;
   private _reconnectAttempts = WS_RECONNECT_ATTEMPTS;
 
   constructor(
@@ -60,18 +51,21 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     super();
     this.appId = appId;
     this.logger = logger;
-    this._rest = new RestClient(baseUrl, logger);
-    this._socket = new WebSocketClient(socketUrl, appId, networkType, logger);
+    this._auth = new AuthManager(baseUrl, logger);
+    this._rest = new RestClient(baseUrl, this._auth, logger);
+    this._socket = new WebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+
+    this._auth.on('refreshFailed', this.handleRefreshFailed.bind(this));
     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();
+    return this.auth.isAuthenticated;
   }
 
-  get auth(): AuthData | null {
-    return this._auth && this._auth.expiresAt > new Date() ? this._auth : null;
+  get auth(): AuthManager {
+    return this._auth;
   }
 
   get socket(): WebSocketClient {
@@ -82,20 +76,13 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     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();
+  async authenticate(tokens: Tokens) {
+    await this.auth.setTokens(tokens);
+    await this.socket.connect();
   }
 
   removeAuth() {
-    this._auth = null;
+    this.auth.clear();
     this.socket.disconnect();
   }
 
@@ -112,6 +99,7 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
       return;
     }
     if (this._reconnectAttempts <= 0) {
+      this.removeAuth();
       this.emit('disconnected', data);
       this._reconnectAttempts = WS_RECONNECT_ATTEMPTS;
       return;
@@ -119,6 +107,10 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     this._reconnectAttempts--;
     setTimeout(() => this.socket.connect(), 1000);
   }
+
+  handleRefreshFailed() {
+    this.removeAuth();
+  }
 }
 
 export default ApiClient;

package/src/Projects/index.ts

@@ -72,7 +72,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     }, {});
     this._availableModels = Object.entries(data).map(([id, workerCount]) => ({
       id,
-      name: modelIndex[id].modelShortName || id.replace(/-/g, ' '),
+      name: modelIndex[id]?.modelShortName || id.replace(/-/g, ' '),
       workerCount
     }));
     this.emit('availableModels', this._availableModels);

package/src/lib/AuthManager.ts

@@ -0,0 +1,172 @@
+import { decodeToken, decodeRefreshToken } from './utils';
+import { ApiError, ApiErrorResponse } from '../ApiClient';
+import { Logger } from './DefaultLogger';
+import isNodejs from './isNodejs';
+import Cookie from 'js-cookie';
+import TypedEventEmitter from './TypedEventEmitter';
+
+/**
+ * Tokens object, containing the token and refresh token
+ * @typedef {Object} Tokens
+ * @property {string} [token] - The JWT token. Optonal, if missing it will be retrieved from the server
+ * @property {string} refreshToken - The refresh token
+ */
+export interface Tokens {
+  token?: string;
+  refreshToken: string;
+}
+
+export type AuthUpdatedEvent =
+  | { token: string; refreshToken: string; walletAddress: string }
+  | { token: null; refreshToken: null; walletAddress: null };
+
+interface AuthManagerEvents {
+  updated: AuthUpdatedEvent;
+  refreshFailed: ApiErrorResponse;
+}
+
+class AuthManager extends TypedEventEmitter<AuthManagerEvents> {
+  private _token?: string;
+  private _tokenExpiresAt: Date = new Date(0);
+  private _refreshToken?: string;
+  private _refreshTokenExpiresAt: Date = new Date(0);
+  private _logger: Logger;
+  private _baseUrl: string;
+  private _walletAddress?: string;
+  private _renewTokenPromise?: Promise<string>;
+
+  constructor(baseUrl: string, logger: Logger) {
+    super();
+    this._logger = logger;
+    this._baseUrl = baseUrl;
+  }
+
+  get refreshToken() {
+    return this._refreshToken;
+  }
+
+  get walletAddress() {
+    return this._walletAddress;
+  }
+
+  get isAuthenticated() {
+    return !!this._refreshToken && this._refreshTokenExpiresAt > new Date();
+  }
+
+  private _updateCookies() {
+    if (isNodejs) {
+      return;
+    }
+    const token = this._token;
+    if (token) {
+      Cookie.set('authorization', token, {
+        domain: '.sogni.ai',
+        expires: 1
+      });
+    } else {
+      Cookie.remove('authorization', {
+        domain: '.sogni.ai'
+      });
+    }
+  }
+
+  async setTokens({ refreshToken, token }: { refreshToken: string; token?: string }) {
+    if (token) {
+      this._updateTokens({ token, refreshToken });
+      return;
+    }
+    this._refreshToken = refreshToken;
+    const { expiresAt: refreshExpiresAt } = decodeRefreshToken(refreshToken);
+    this._refreshTokenExpiresAt = refreshExpiresAt;
+    await this.renewToken();
+  }
+
+  async getToken(): Promise<string | undefined> {
+    //If there is a token, and it is not expired, return it
+    if (this._token && this._tokenExpiresAt > new Date()) {
+      return this._token;
+    }
+    //If there is no refresh token, return undefined, to make unauthorized requests
+    if (!this._refreshToken) {
+      return;
+    }
+    //If there is a refresh token, try to renew the token
+    return this.renewToken();
+  }
+
+  async renewToken(): Promise<string> {
+    if (this._renewTokenPromise) {
+      return this._renewTokenPromise;
+    }
+    this._renewTokenPromise = this._renewToken();
+    this._renewTokenPromise.finally(() => {
+      this._renewTokenPromise = undefined;
+    });
+    return this._renewTokenPromise;
+  }
+
+  clear() {
+    // Prevent duplicate events
+    if (!this._token && !this._refreshToken) {
+      return;
+    }
+    this._refreshToken = undefined;
+    this._refreshTokenExpiresAt = new Date(0);
+    this._token = undefined;
+    this._tokenExpiresAt = new Date(0);
+    this._walletAddress = undefined;
+    this.emit('updated', { token: null, refreshToken: null, walletAddress: null });
+  }
+
+  private _updateTokens({ token, refreshToken }: { token: string; refreshToken: string }) {
+    // Prevent duplicate events
+    if (this._token === token && this._refreshToken === refreshToken) {
+      return;
+    }
+    this._token = token;
+    const { expiresAt, walletAddress } = decodeToken(token);
+    this._walletAddress = walletAddress;
+    this._tokenExpiresAt = expiresAt;
+    this._refreshToken = refreshToken;
+    const { expiresAt: refreshExpiresAt } = decodeRefreshToken(refreshToken);
+    this._refreshTokenExpiresAt = refreshExpiresAt;
+    this._updateCookies();
+    this.emit('updated', { token, refreshToken, walletAddress });
+  }
+
+  private async _renewToken(): Promise<string> {
+    if (this._refreshTokenExpiresAt < new Date()) {
+      throw new Error('Refresh token expired');
+    }
+    const url = new URL('/v1/account/refresh-token', this._baseUrl).toString();
+    const response = await fetch(url, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json'
+      },
+      body: JSON.stringify({ refreshToken: this._refreshToken })
+    });
+    let responseData: any;
+    try {
+      responseData = await response.json();
+    } catch (e) {
+      this.emit('refreshFailed', {
+        status: 'error',
+        errorCode: 0,
+        message: 'Failed to parse response'
+      });
+      this._logger.error('Failed to parse response:', e);
+      throw new Error('Failed to parse response');
+    }
+    if (!response.ok) {
+      this.emit('refreshFailed', responseData);
+      this.clear();
+      throw new ApiError(response.status, responseData as ApiErrorResponse);
+    }
+    const { token, refreshToken } = responseData.data;
+    this._updateTokens({ token, refreshToken });
+    return this._token!;
+  }
+}
+
+export default AuthManager;

package/src/lib/RestClient.ts

@@ -2,30 +2,24 @@ import { ApiError, ApiErrorResponse } from '../ApiClient';
 import TypedEventEmitter, { EventMap } from './TypedEventEmitter';
 import { JSONValue } from '../types/json';
 import { Logger } from './DefaultLogger';
-
-export interface AuthData {
-  token: string;
-}
+import AuthManager from './AuthManager';
 
 class RestClient<E extends EventMap = never> extends TypedEventEmitter<E> {
   readonly baseUrl: string;
-  protected _auth: AuthData | null = null;
+  protected _auth: AuthManager;
   protected _logger: Logger;
 
-  constructor(baseUrl: string, logger: Logger) {
+  constructor(baseUrl: string, auth: AuthManager, logger: Logger) {
     super();
     this.baseUrl = baseUrl;
+    this._auth = auth;
     this._logger = logger;
   }
 
-  get auth(): AuthData | null {
+  get auth(): AuthManager {
     return this._auth;
   }
 
-  set auth(auth: AuthData | null) {
-    this._auth = auth;
-  }
-
   private formatUrl(relativeUrl: string, query: Record<string, string> = {}): string {
     const url = new URL(relativeUrl, this.baseUrl);
     Object.keys(query).forEach((key) => {
@@ -34,12 +28,13 @@ class RestClient<E extends EventMap = never> extends TypedEventEmitter<E> {
     return url.toString();
   }
 
-  private request<T = JSONValue>(url: string, options: RequestInit = {}): Promise<T> {
+  private async request<T = JSONValue>(url: string, options: RequestInit = {}): Promise<T> {
+    const token = await this.auth.getToken();
     return fetch(url, {
       ...options,
       headers: {
         ...(options.headers || {}),
-        ...(this.auth ? { Authorization: this.auth.token } : {})
+        ...(token ? { Authorization: token } : {})
       }
     }).then((r) => this.processResponse(r) as T);
   }

package/src/lib/utils.ts

@@ -0,0 +1,17 @@
+import { jwtDecode } from 'jwt-decode';
+
+export function decodeToken(token: string) {
+  const data = jwtDecode<{ addr: string; env: string; iat: number; exp: number }>(token);
+  return {
+    walletAddress: data.addr,
+    expiresAt: new Date(data.exp * 1000)
+  };
+}
+
+export function decodeRefreshToken(token: string) {
+  const data = jwtDecode<{ type: string; env: string; iat: number; exp: number }>(token);
+  return {
+    env: data.env,
+    expiresAt: new Date(data.exp * 1000)
+  };
+}

package/src/version.ts

@@ -1 +1 @@
-export const LIB_VERSION = "0.4.2";
+export const LIB_VERSION = "0.5.0-alpha.0";