@sogni-ai/expo-client 1.0.0-alpha.10

package/src/ApiClient/WebSocketClient/events.ts

@@ -0,0 +1,168 @@
+import { SupernetType } from './types';
+import { JobErrorReason } from '../../Projects/types';
+import { LLMModelInfo, ToolCallDelta, TokenUsage } from '../../Chat/types';
+
+export interface BalanceData {
+  settled: string;
+  credit: string;
+  debit: string;
+  net: string;
+}
+
+export interface SparkBalanceData extends BalanceData {
+  premiumCredit: string;
+}
+
+export interface Balances {
+  sogni: BalanceData;
+  spark: SparkBalanceData;
+}
+
+export interface AuthenticatedEvent {
+  username: string;
+  walletAddress: string;
+  balance: Balances;
+  isMainnet: boolean;
+  accountWasMigrated: boolean;
+  hasUnclaimedAirdrop: boolean;
+  firstLoginAfterMigration: boolean;
+}
+
+export interface BalanceUpdateEvent {
+  balance: Balances;
+}
+
+export interface ChangeNetworkEvent {
+  network: SupernetType;
+}
+
+export interface ConnectedEvent {
+  network: SupernetType;
+}
+
+export interface DisconnectedEvent {
+  code?: number;
+  reason?: string;
+}
+
+export interface ToastMessageEvent {
+  level: 'info' | 'success' | 'warning' | 'error';
+  message: string;
+  stickyID?: string;
+}
+
+/* === Project / job lifecycle (Phase 3). All IDs are normalized; on the wire,
+ *     `jobID` is the project id and `imgID` is the job id. The rest of the SDK
+ *     and consumers only see `projectId` and `jobId`. === */
+
+export interface JobStartedEvent {
+  projectId: string;
+  jobId: string | null;
+  workerName: string | null;
+}
+
+export interface JobInitiatingModelEvent {
+  projectId: string;
+  jobId: string | null;
+}
+
+export interface JobLifecycleCompletedEvent {
+  projectId: string;
+  jobId: string | null;
+}
+
+export interface JobProgressEvent {
+  projectId: string;
+  jobId: string;
+  /** Derived as step/stepCount in [0,1]. */
+  progress: number;
+  step: number;
+  stepCount: number;
+}
+
+export interface JobResultEvent {
+  projectId: string;
+  jobId: string;
+  seed: number;
+  performedStepCount: number;
+  /** True if the artist cancelled this job. The SDK turns these into JobFailed. */
+  userCanceled: boolean;
+  /** True if NSFW filter triggered. The SDK marks the job result as withheld. */
+  triggeredNSFWFilter: boolean;
+}
+
+export interface JobErrorEvent {
+  projectId: string;
+  jobId: string | null;
+  reason: JobErrorReason;
+  /** Original server-side error string, preserved for forward compat. */
+  rawCode: string;
+  message: string;
+  source: 'worker' | 'server';
+}
+
+export interface JobETAEvent {
+  projectId: string;
+  jobId: string;
+  etaSeconds: number;
+}
+
+export interface ArtistCancelConfirmationEvent {
+  projectId: string;
+  didCancel: boolean;
+  errorMessage?: string;
+}
+
+/* === Chat / LLM lifecycle (Phase 4b). Chat jobs have a single jobId (no
+ *     project/job split like image projects). The wire field stays jobID
+ *     and we just rename to jobId for camelCase consistency. === */
+
+export interface ChatTokensEvent {
+  jobId: string;
+  content: string;
+  role?: string;
+  finishReason?: string | null;
+  usage?: TokenUsage;
+  tool_calls?: ToolCallDelta[];
+}
+
+export interface ChatResultEvent {
+  jobId: string;
+  usage: TokenUsage;
+  timeTaken: number;
+  workerName?: string;
+}
+
+export interface ChatErrorEvent {
+  jobId: string;
+  rawCode: string;
+  message: string;
+  workerName?: string;
+}
+
+export interface SwarmLLMModelsEvent {
+  /** Map of model ID -> info. Worker counts are 0 if no workers serve the model. */
+  models: Record<string, LLMModelInfo>;
+}
+
+/** Public event map for WebSocketClient. */
+export interface SocketEventMap {
+  connected: ConnectedEvent;
+  disconnected: DisconnectedEvent;
+  authenticated: AuthenticatedEvent;
+  balanceUpdate: BalanceUpdateEvent;
+  changeNetwork: ChangeNetworkEvent;
+  toastMessage: ToastMessageEvent;
+  jobStarted: JobStartedEvent;
+  jobInitiatingModel: JobInitiatingModelEvent;
+  jobLifecycleCompleted: JobLifecycleCompletedEvent;
+  jobProgress: JobProgressEvent;
+  jobResult: JobResultEvent;
+  jobError: JobErrorEvent;
+  jobETA: JobETAEvent;
+  artistCancelConfirmation: ArtistCancelConfirmationEvent;
+  chatTokens: ChatTokensEvent;
+  chatResult: ChatResultEvent;
+  chatError: ChatErrorEvent;
+  swarmLLMModels: SwarmLLMModelsEvent;
+}

package/src/ApiClient/WebSocketClient/index.ts

@@ -0,0 +1,345 @@
+import TypedEventEmitter from '../../lib/TypedEventEmitter';
+import { base64Decode, base64Encode } from '../../lib/base64';
+import { Logger } from '../../lib/DefaultLogger';
+import { AuthManager } from '../../lib/AuthManager';
+import { LIB_VERSION } from '../../version';
+import { ConnectionStatus, SupernetType } from './types';
+import { SocketEventMap } from './events';
+import {
+  normalizeArtistCancelConfirmation,
+  normalizeAuthenticated,
+  normalizeBalanceUpdate,
+  normalizeChangeNetwork,
+  normalizeChatError,
+  normalizeChatResult,
+  normalizeChatTokens,
+  normalizeJobError,
+  normalizeJobETA,
+  normalizeJobProgress,
+  normalizeJobResult,
+  normalizeJobState,
+  normalizeSwarmLLMModels,
+  normalizeToastMessage
+} from './normalize';
+import {
+  RawArtistCancelConfirmation,
+  RawAuthenticated,
+  RawBalanceUpdate,
+  RawChangeNetwork,
+  RawJobError,
+  RawJobETA,
+  RawJobProgress,
+  RawJobResult,
+  RawJobState,
+  RawJobTokens,
+  RawLLMJobError,
+  RawLLMJobResult,
+  RawSwarmLLMModels,
+  RawToastMessage
+} from './wireTypes';
+
+const PROTOCOL_VERSION = '3.0.0';
+const RECONNECT_DELAY_MS = 1000;
+const MAX_RECONNECT_ATTEMPTS = 5;
+
+/** RN WebSocket constructor accepts a third options argument with `headers`. */
+type WebSocketCtorOptions = { headers?: Record<string, string> } | undefined;
+type WebSocketCtor = new (
+  url: string,
+  protocols?: string | string[] | null,
+  options?: WebSocketCtorOptions
+) => WebSocket;
+
+/** Minimal shape of the RN AppState module — we lazy-require it. */
+interface AppStateModule {
+  currentState?: string;
+  addEventListener(type: 'change', listener: (state: string) => void): { remove(): void };
+}
+
+function loadAppState(logger: Logger): AppStateModule | null {
+  try {
+    // eslint-disable-next-line @typescript-eslint/no-var-requires
+    const rn = require('react-native');
+    return rn?.AppState ?? null;
+  } catch (e) {
+    logger.debug('react-native not available, skipping AppState integration:', e);
+    return null;
+  }
+}
+
+class WebSocketClient extends TypedEventEmitter<SocketEventMap> {
+  readonly appId: string;
+  readonly appSource?: string;
+  readonly baseUrl: string;
+  protected _logger: Logger;
+  private _auth: AuthManager;
+  private _supernetType: SupernetType;
+  private socket: WebSocket | null = null;
+  private _status: ConnectionStatus = 'disconnected';
+  private _reconnectAttempts = 0;
+  private _wantConnected = false;
+  private _appStateSubscription: { remove(): void } | null = null;
+
+  constructor(
+    baseUrl: string,
+    auth: AuthManager,
+    appId: string,
+    supernetType: SupernetType,
+    logger: Logger,
+    appSource?: string
+  ) {
+    super();
+    this.baseUrl = normalizeBaseUrl(baseUrl);
+    this._auth = auth;
+    this.appId = appId;
+    this.appSource = appSource?.trim() || undefined;
+    this._supernetType = supernetType;
+    this._logger = logger;
+    this._subscribeAppState();
+  }
+
+  /**
+   * Subscribe to RN AppState 'change' events. When the app returns to the
+   * foreground, if we *want* to be connected but aren't, reconnect (and reset
+   * the backoff counter so the user immediately sees a real attempt). When
+   * the app backgrounds we do nothing — iOS will probably suspend the socket
+   * itself, and we'll reconnect on next foreground transition.
+   */
+  private _subscribeAppState(): void {
+    const AppState = loadAppState(this._logger);
+    if (!AppState) return;
+    this._appStateSubscription = AppState.addEventListener('change', (state) => {
+      if (state !== 'active') return;
+      if (!this._wantConnected || this._status === 'connected') return;
+      this._logger.debug('App returned to foreground — reconnecting WebSocket');
+      this._reconnectAttempts = 0;
+      void this.connect();
+    });
+  }
+
+  get supernetType(): SupernetType {
+    return this._supernetType;
+  }
+
+  get isConnected(): boolean {
+    return this._status === 'connected';
+  }
+
+  get status(): ConnectionStatus {
+    return this._status;
+  }
+
+  async connect(): Promise<void> {
+    this._wantConnected = true;
+    if (this.socket) {
+      this._cleanupSocket();
+    }
+    this._status = 'connecting';
+
+    const userAgent = `Sogni/${PROTOCOL_VERSION} (sogni-expo-client) ${LIB_VERSION}`;
+    const url = new URL(this.baseUrl);
+    url.protocol = url.protocol === 'https:' || url.protocol === 'wss:' ? 'wss:' : 'ws:';
+    url.searchParams.set('appId', this.appId);
+    if (this.appSource) url.searchParams.set('appSource', this.appSource);
+    url.searchParams.set('clientName', userAgent);
+    url.searchParams.set('clientType', 'artist');
+    url.searchParams.set('forceWorkerId', this._supernetType === 'fast' ? 'fast' : '');
+
+    const auth = await this._auth.socketAuth();
+    if (auth.queryParams) {
+      for (const [k, v] of Object.entries(auth.queryParams)) {
+        url.searchParams.set(k, v);
+      }
+    }
+
+    const options: WebSocketCtorOptions = auth.headers ? { headers: auth.headers } : undefined;
+    const Ctor = (globalThis as unknown as { WebSocket: WebSocketCtor }).WebSocket;
+    if (!Ctor) {
+      throw new Error('No global WebSocket available — are you running outside of React Native?');
+    }
+    const socket = new Ctor(url.toString(), undefined, options);
+    this.socket = socket;
+    socket.onopen = () => this._handleOpen();
+    socket.onmessage = (e) => this._handleMessage(e);
+    socket.onclose = (e) => this._handleClose(e);
+    socket.onerror = (e) => this._handleError(e);
+  }
+
+  /** Tear down completely — closes the socket AND unsubscribes from AppState. */
+  dispose(): void {
+    this.disconnect();
+    if (this._appStateSubscription) {
+      this._appStateSubscription.remove();
+      this._appStateSubscription = null;
+    }
+    this.removeAllListeners();
+  }
+
+  disconnect(): void {
+    this._wantConnected = false;
+    if (!this.socket) return;
+    const socket = this.socket;
+    this._cleanupSocket();
+    try {
+      socket.close(1000, 'Client disconnected');
+    } catch (e) {
+      this._logger.warn('Error closing socket:', e);
+    }
+    this._status = 'disconnected';
+    this.emit('disconnected', { code: 1000, reason: 'Client disconnected' });
+  }
+
+  async switchNetwork(target: SupernetType): Promise<SupernetType> {
+    return new Promise<SupernetType>((resolve) => {
+      this.once('changeNetwork', ({ network }) => {
+        this._supernetType = network;
+        resolve(network);
+      });
+      void this.send('changeNetwork', target);
+    });
+  }
+
+  async send(messageType: string, data: unknown): Promise<void> {
+    if (!this.socket || this._status !== 'connected') {
+      throw new Error(`Cannot send ${messageType}: socket not connected`);
+    }
+    const envelope = JSON.stringify({
+      type: messageType,
+      data: base64Encode(JSON.stringify(data))
+    });
+    this._logger.debug('WebSocket send:', messageType, data);
+    this.socket.send(envelope);
+  }
+
+  private _handleOpen() {
+    this._status = 'connected';
+    this._reconnectAttempts = 0;
+    this.emit('connected', { network: this._supernetType });
+  }
+
+  private _handleClose(e: { code?: number; reason?: string }) {
+    const wasConnected = this._status === 'connected';
+    this._status = 'disconnected';
+    this._cleanupSocket();
+    this.emit('disconnected', { code: e.code, reason: e.reason });
+    if (this._wantConnected && wasConnected && e.code !== 1000) {
+      this._scheduleReconnect();
+    }
+  }
+
+  private _handleError(e: unknown) {
+    this._logger.error('WebSocket error:', e);
+  }
+
+  private _scheduleReconnect() {
+    if (this._reconnectAttempts >= MAX_RECONNECT_ATTEMPTS) {
+      this._logger.warn('WebSocket reconnect attempts exhausted');
+      return;
+    }
+    this._reconnectAttempts++;
+    setTimeout(() => {
+      if (this._wantConnected) {
+        void this.connect();
+      }
+    }, RECONNECT_DELAY_MS);
+  }
+
+  private _cleanupSocket() {
+    if (!this.socket) return;
+    const s = this.socket;
+    s.onopen = null;
+    s.onmessage = null;
+    s.onclose = null;
+    s.onerror = null;
+    this.socket = null;
+  }
+
+  private _handleMessage(e: MessageEvent) {
+    Promise.resolve(typeof e.data === 'string' ? e.data : (e.data as Blob).text())
+      .then((str: string) => {
+        const envelope = JSON.parse(str) as { type: string; data?: string };
+        const payload: any = envelope.data ? JSON.parse(base64Decode(envelope.data)) : null;
+        this._logger.debug('WebSocket recv:', envelope.type, payload);
+        this._dispatch(envelope.type, payload);
+      })
+      .catch((err) => this._logger.error('Failed to parse WebSocket message:', err));
+  }
+
+  private _dispatch(type: string, payload: any) {
+    switch (type) {
+      case 'authenticated':
+        this.emit('authenticated', normalizeAuthenticated(payload as RawAuthenticated));
+        break;
+      case 'balanceUpdate':
+        this.emit('balanceUpdate', normalizeBalanceUpdate(payload as RawBalanceUpdate));
+        break;
+      case 'changeNetwork':
+        this.emit('changeNetwork', normalizeChangeNetwork(payload as RawChangeNetwork));
+        break;
+      case 'toastMessage':
+        this.emit('toastMessage', normalizeToastMessage(payload as RawToastMessage));
+        break;
+      case 'jobState': {
+        const parsed = normalizeJobState(payload as RawJobState);
+        if (parsed) this.emit(parsed.kind, parsed.event as any);
+        break;
+      }
+      case 'jobProgress': {
+        const ev = normalizeJobProgress(payload as RawJobProgress);
+        if (ev) this.emit('jobProgress', ev);
+        break;
+      }
+      case 'jobResult': {
+        const ev = normalizeJobResult(payload as RawJobResult);
+        if (ev) this.emit('jobResult', ev);
+        break;
+      }
+      case 'jobError': {
+        const ev = normalizeJobError(payload as RawJobError);
+        if (ev) this.emit('jobError', ev);
+        break;
+      }
+      case 'jobETA': {
+        const ev = normalizeJobETA(payload as RawJobETA);
+        if (ev) this.emit('jobETA', ev);
+        break;
+      }
+      case 'artistCancelConfirmation': {
+        const ev = normalizeArtistCancelConfirmation(payload as RawArtistCancelConfirmation);
+        if (ev) this.emit('artistCancelConfirmation', ev);
+        break;
+      }
+      case 'jobTokens': {
+        const ev = normalizeChatTokens(payload as RawJobTokens);
+        if (ev) this.emit('chatTokens', ev);
+        break;
+      }
+      case 'llmJobResult': {
+        const ev = normalizeChatResult(payload as RawLLMJobResult);
+        if (ev) this.emit('chatResult', ev);
+        break;
+      }
+      case 'llmJobError': {
+        const ev = normalizeChatError(payload as RawLLMJobError);
+        if (ev) this.emit('chatError', ev);
+        break;
+      }
+      case 'swarmLLMModels': {
+        this.emit('swarmLLMModels', normalizeSwarmLLMModels(payload as RawSwarmLLMModels));
+        break;
+      }
+      default:
+        this._logger.debug('Unhandled WebSocket message type:', type);
+        break;
+    }
+  }
+}
+
+function normalizeBaseUrl(url: string): string {
+  const u = new URL(url);
+  if (u.protocol === 'http:' || u.protocol === 'ws:') u.protocol = 'http:';
+  else u.protocol = 'https:';
+  return u.toString();
+}
+
+export default WebSocketClient;