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

package/src/Projects/index.ts

@@ -0,0 +1,608 @@
+import ApiClient, { ApiError, ApiResponse } from '../ApiClient';
+import { SupernetType } from '../ApiClient/WebSocketClient/types';
+import {
+  ArtistCancelConfirmationEvent,
+  JobErrorEvent,
+  JobETAEvent,
+  JobInitiatingModelEvent,
+  JobLifecycleCompletedEvent,
+  JobProgressEvent,
+  JobResultEvent,
+  JobStartedEvent
+} from '../ApiClient/WebSocketClient/events';
+import TypedEventEmitter from '../lib/TypedEventEmitter';
+import getUUID from '../lib/getUUID';
+import { MediaInput, toFetchBody } from '../lib/mediaInput';
+import ErrorData from '../types/ErrorData';
+import createJobRequestMessage from './createJobRequestMessage';
+import {
+  AvailableModel,
+  ImageProjectParams,
+  Job,
+  JobStatus,
+  Project,
+  ProjectParams,
+  ProjectStatus,
+  VideoProjectParams
+} from './types';
+import { getMaxContextImages } from './ControlNet';
+import { RawProject, rawProjectToProject } from './restNormalize';
+
+interface ProjectsApiEvents {
+  /** Fired for any project state change (status, progress, jobs, etc.). */
+  changed: Project;
+  /** Progress tick on any project. */
+  progress: { project: Project; jobId: string };
+  /** Project reached terminal state 'completed' with all result URLs populated. */
+  completed: Project;
+  /** Project reached terminal state 'failed' or 'canceled'. */
+  failed: Project;
+}
+
+const TERMINAL_PROJECT_STATUSES: ProjectStatus[] = ['completed', 'failed', 'canceled'];
+const TERMINAL_JOB_STATUSES: JobStatus[] = ['completed', 'failed', 'canceled'];
+const GC_DELAY_MS = 30_000;
+
+class ProjectsApi extends TypedEventEmitter<ProjectsApiEvents> {
+  private _client: ApiClient;
+  private _projects = new Map<string, Project>();
+  private _subscribers = new Map<string, Set<(p: Project) => void>>();
+  private _gcTimers = new Map<string, ReturnType<typeof setTimeout>>();
+
+  constructor(client: ApiClient) {
+    super();
+    this._client = client;
+    this._wireSocketHandlers();
+  }
+
+  /** List models available on the given network. Defaults to the client's current network. */
+  async getAvailableModels(network?: SupernetType): Promise<AvailableModel[]> {
+    const net = network ?? this._client.socket.supernetType;
+    const res = await this._client.socketRest.get<ApiResponse<unknown>>(
+      `/api/v1/status/network/${net}/models`
+    );
+    // The endpoint returns either an array of models directly or wrapped in
+    // { models: [...] }. Accept either shape.
+    const data: any = res.data;
+    if (Array.isArray(data)) return data as AvailableModel[];
+    if (data && Array.isArray(data.models)) return data.models as AvailableModel[];
+    return [];
+  }
+
+  /**
+   * Create a generation project (image, video, or audio).
+   *
+   * Returns the initial Project snapshot immediately; subscribe via
+   * `subscribe(id, cb)` or `waitForCompletion(id)` to observe state changes as
+   * the server processes it.
+   *
+   * For image projects, `startingImage`, `contextImages[]`, and
+   * `controlNet.image` are uploaded first. For video projects, `referenceImage`,
+   * `referenceImageEnd`, `referenceAudio`, and `referenceVideo` are uploaded
+   * first. Reference uploads happen in parallel; the WS `jobRequest` is sent
+   * once all uploads complete.
+   */
+  async create(params: ProjectParams): Promise<Project> {
+    if (!this._client.socket.isConnected) {
+      throw new Error(
+        'Cannot create project: WebSocket is not connected. Authenticate first and wait for `connected` event.'
+      );
+    }
+    const id = getUUID();
+    const project: Project = {
+      id,
+      params,
+      status: 'pending',
+      progress: 0,
+      eta: null,
+      jobs: [],
+      resultUrls: [],
+      error: null,
+      createdAt: Date.now(),
+      updatedAt: Date.now()
+    };
+    this._projects.set(id, project);
+
+    const type = params.type ?? 'image';
+    if (type === 'image') {
+      await this._uploadImageRefs(id, params as ImageProjectParams);
+    } else if (type === 'video') {
+      await this._uploadVideoRefs(id, params as VideoProjectParams);
+    }
+
+    const message = createJobRequestMessage(id, params);
+    await this._client.socket.send('jobRequest', message);
+    this._notifyChanged(project);
+    return project;
+  }
+
+  private async _uploadImageRefs(projectId: string, params: ImageProjectParams) {
+    if (params.startingImage) {
+      await this._uploadMedia(projectId, 'startingImage', params.startingImage);
+    }
+    if (params.contextImages && params.contextImages.length) {
+      const max = getMaxContextImages(params.modelId);
+      if (params.contextImages.length > max) {
+        throw new Error(
+          `Model ${params.modelId} accepts at most ${max} context images (got ${params.contextImages.length})`
+        );
+      }
+      await Promise.all(
+        params.contextImages.map((img, idx) =>
+          this._uploadMedia(projectId, `contextImage${idx + 1}`, img)
+        )
+      );
+    }
+    if (params.controlNet?.image) {
+      await this._uploadMedia(projectId, 'cnImage', params.controlNet.image);
+    }
+  }
+
+  private async _uploadVideoRefs(projectId: string, params: VideoProjectParams) {
+    const uploads: Promise<unknown>[] = [];
+    if (params.referenceImage) {
+      uploads.push(this._uploadMedia(projectId, 'referenceImage', params.referenceImage));
+    }
+    if (params.referenceImageEnd) {
+      uploads.push(this._uploadMedia(projectId, 'referenceImageEnd', params.referenceImageEnd));
+    }
+    if (params.referenceAudio) {
+      uploads.push(this._uploadMedia(projectId, 'referenceAudio', params.referenceAudio));
+    }
+    if (params.referenceVideo) {
+      uploads.push(this._uploadMedia(projectId, 'referenceVideo', params.referenceVideo));
+    }
+    if (uploads.length) await Promise.all(uploads);
+  }
+
+  private async _uploadMedia(projectId: string, type: string, media: MediaInput): Promise<string> {
+    const imageId = getUUID();
+    const { body, contentType } = await toFetchBody(media);
+    const presignedUrl = await this._getUploadUrl({
+      imageId,
+      jobId: projectId,
+      type,
+      ...(contentType ? { contentType } : {})
+    });
+    const headers: Record<string, string> = {};
+    if (contentType) headers['Content-Type'] = contentType;
+    const res = await fetch(presignedUrl, {
+      method: 'PUT',
+      headers,
+      body
+    });
+    if (!res.ok) {
+      throw new ApiError(res.status, {
+        status: 'error',
+        errorCode: 0,
+        message: `Failed to upload ${type} (HTTP ${res.status})`
+      });
+    }
+    return imageId;
+  }
+
+  private async _getUploadUrl(params: {
+    imageId: string;
+    jobId: string;
+    type: string;
+    contentType?: string;
+  }): Promise<string> {
+    const res = await this._client.socketRest.get<ApiResponse<string | { url: string }>>(
+      '/v1/image/uploadUrl',
+      params
+    );
+    const d: any = res.data;
+    return typeof d === 'string' ? d : d.url;
+  }
+
+  /** Get the cached snapshot for a project, or undefined if unknown. */
+  get(projectId: string): Project | undefined {
+    return this._projects.get(projectId.toUpperCase());
+  }
+
+  /**
+   * Fetch a project's current state from REST and cache it.
+   *
+   * Use this on cold launch when you have a persisted project ID but no
+   * in-memory snapshot (e.g., the app was killed mid-generation). The server
+   * returns the latest known state; if the project is still in flight, the
+   * SDK will keep updating the snapshot via socket events.
+   *
+   * @param projectId Server project ID.
+   * @param paramsOverride Optional. The original ProjectParams used to create
+   * the project. The REST response only includes a subset of generation
+   * settings; pass the original params if you persisted them (e.g., in MMKV)
+   * so the returned snapshot's `params` field is fully populated.
+   */
+  async fetch(projectId: string, paramsOverride?: ProjectParams): Promise<Project> {
+    const id = projectId.toUpperCase();
+    const res = await this._client.rest.get<ApiResponse<RawProject>>(`/v1/projects/${id}`);
+    const project = rawProjectToProject(res.data, paramsOverride);
+    this._projects.set(project.id, project);
+    this._notifyChanged(project);
+    return project;
+  }
+
+  /** Subscribe to changes for a specific project. Returns an unsubscribe function. */
+  subscribe(projectId: string, listener: (project: Project) => void): () => void {
+    const id = projectId.toUpperCase();
+    let set = this._subscribers.get(id);
+    if (!set) {
+      set = new Set();
+      this._subscribers.set(id, set);
+    }
+    set.add(listener);
+    return () => {
+      const current = this._subscribers.get(id);
+      if (!current) return;
+      current.delete(listener);
+      if (current.size === 0) this._subscribers.delete(id);
+    };
+  }
+
+  /**
+   * Resolve when the project reaches a terminal state. On success, returns the
+   * downloadable result URLs (fetched from the server). On failure or cancel,
+   * rejects with an Error describing the reason.
+   */
+  async waitForCompletion(projectId: string): Promise<string[]> {
+    const id = projectId.toUpperCase();
+    const project = this._projects.get(id);
+    if (!project) {
+      throw new Error(`Project ${projectId} is not known to this client`);
+    }
+    if (project.status === 'completed') {
+      return project.resultUrls.length ? project.resultUrls : this._fetchResultUrls(project);
+    }
+    if (project.status === 'failed' || project.status === 'canceled') {
+      throw new Error(project.error?.message || `Project ${project.status}`);
+    }
+    return new Promise<string[]>((resolve, reject) => {
+      const unsubscribe = this.subscribe(id, (p) => {
+        if (p.status === 'completed') {
+          unsubscribe();
+          if (p.resultUrls.length) resolve(p.resultUrls);
+          else this._fetchResultUrls(p).then(resolve).catch(reject);
+        } else if (p.status === 'failed' || p.status === 'canceled') {
+          unsubscribe();
+          reject(new Error(p.error?.message || `Project ${p.status}`));
+        }
+      });
+    });
+  }
+
+  /** Request a download URL for a specific job's result image. */
+  async getResultUrl(projectId: string, jobId: string): Promise<string> {
+    const res = await this._client.socketRest.get<ApiResponse<string | { url: string }>>(
+      `/v1/image/downloadUrl`,
+      { projectId: projectId.toUpperCase(), jobId: jobId.toUpperCase(), type: 'complete' }
+    );
+    const d: any = res.data;
+    return typeof d === 'string' ? d : d.url;
+  }
+
+  /**
+   * Request cancellation of a project. Server confirms via
+   * `artistCancelConfirmation`; the project's status flips to 'canceled' if
+   * accepted, otherwise stays in its current state.
+   */
+  async cancel(projectId: string): Promise<void> {
+    const id = projectId.toUpperCase();
+    await this._client.socket.send('jobError', {
+      jobID: id,
+      error: 'artistCanceled',
+      error_message: 'Artist Cancelled Job',
+      isFromWorker: false
+    });
+  }
+
+  // ===========================================================================
+  // Internal: event wiring
+  // ===========================================================================
+
+  private _wireSocketHandlers() {
+    const s = this._client.socket;
+    s.on('jobStarted', (e) => this._onJobStarted(e));
+    s.on('jobInitiatingModel', (e) => this._onJobInitiatingModel(e));
+    s.on('jobLifecycleCompleted', (e) => this._onJobLifecycleCompleted(e));
+    s.on('jobProgress', (e) => this._onJobProgress(e));
+    s.on('jobResult', (e) => this._onJobResult(e));
+    s.on('jobError', (e) => this._onJobError(e));
+    s.on('jobETA', (e) => this._onJobETA(e));
+    s.on('artistCancelConfirmation', (e) => this._onCancelConfirmation(e));
+  }
+
+  private _onJobStarted(e: JobStartedEvent) {
+    const project = this._projects.get(e.projectId);
+    if (!project) return;
+    const next = this._withJobMutation(project, e.jobId, (job) => ({
+      ...job,
+      status: 'processing',
+      workerName: e.workerName
+    }));
+    this._setStatusIfPending(next, 'processing');
+    this._notifyChanged(next);
+  }
+
+  private _onJobInitiatingModel(e: JobInitiatingModelEvent) {
+    const project = this._projects.get(e.projectId);
+    if (!project) return;
+    const next = this._withJobMutation(project, e.jobId, (job) => ({
+      ...job,
+      status: 'initiatingModel'
+    }));
+    this._setStatusIfPending(next, 'processing');
+    this._notifyChanged(next);
+  }
+
+  private _onJobLifecycleCompleted(e: JobLifecycleCompletedEvent) {
+    // jobState.type === 'jobCompleted' is the lifecycle marker from server.
+    // The real result data arrives via jobResult; we just ensure the job is
+    // marked completed here if it wasn't already.
+    const project = this._projects.get(e.projectId);
+    if (!project) return;
+    const next = this._withJobMutation(project, e.jobId, (job) => ({
+      ...job,
+      status: job.status === 'failed' ? 'failed' : 'completed',
+      progress: 1
+    }));
+    this._maybeFinalizeProject(next);
+  }
+
+  private _onJobProgress(e: JobProgressEvent) {
+    const project = this._projects.get(e.projectId);
+    if (!project) return;
+    const next = this._withJobMutation(project, e.jobId, (job) => ({
+      ...job,
+      status: 'processing',
+      progress: e.progress,
+      step: e.step,
+      stepCount: e.stepCount
+    }));
+    this._setStatusIfPending(next, 'processing');
+    next.progress = this._computeOverallProgress(next.jobs);
+    this._notifyChanged(next);
+    this.emit('progress', { project: next, jobId: e.jobId });
+  }
+
+  private _onJobResult(e: JobResultEvent) {
+    const project = this._projects.get(e.projectId);
+    if (!project) return;
+    if (e.userCanceled) {
+      const error: ErrorData = {
+        code: 5004,
+        originalCode: 'artistCanceled',
+        message: 'Artist cancelled the job'
+      };
+      const next = this._withJobMutation(project, e.jobId, (job) => ({
+        ...job,
+        status: 'canceled',
+        error
+      }));
+      this._maybeFinalizeProject(next);
+      return;
+    }
+    if (e.triggeredNSFWFilter) {
+      const error: ErrorData = {
+        code: 5006,
+        originalCode: 'sensitiveContent',
+        message: 'NSFW filter triggered'
+      };
+      const next = this._withJobMutation(project, e.jobId, (job) => ({
+        ...job,
+        status: 'failed',
+        isNSFW: true,
+        seed: e.seed,
+        error
+      }));
+      this._maybeFinalizeProject(next);
+      return;
+    }
+    const next = this._withJobMutation(project, e.jobId, (job) => ({
+      ...job,
+      status: 'completed',
+      progress: 1,
+      seed: e.seed
+    }));
+    this._maybeFinalizeProject(next);
+  }
+
+  private _onJobError(e: JobErrorEvent) {
+    const project = this._projects.get(e.projectId);
+    if (!project) return;
+    const error: ErrorData = {
+      code: 5000,
+      originalCode: e.rawCode,
+      message: e.message
+    };
+    if (e.jobId) {
+      const next = this._withJobMutation(project, e.jobId, (job) => ({
+        ...job,
+        status: 'failed',
+        error
+      }));
+      this._maybeFinalizeProject(next);
+      return;
+    }
+    // Project-level error
+    const next: Project = {
+      ...project,
+      status: 'failed',
+      error,
+      updatedAt: Date.now()
+    };
+    this._projects.set(project.id, next);
+    this._notifyChanged(next);
+    this.emit('failed', next);
+    this._scheduleGC(next.id);
+  }
+
+  private _onJobETA(e: JobETAEvent) {
+    const project = this._projects.get(e.projectId);
+    if (!project) return;
+    const next: Project = { ...project, eta: e.etaSeconds, updatedAt: Date.now() };
+    this._projects.set(project.id, next);
+    this._notifyChanged(next);
+  }
+
+  private _onCancelConfirmation(e: ArtistCancelConfirmationEvent) {
+    const project = this._projects.get(e.projectId);
+    if (!project) return;
+    if (!e.didCancel) {
+      this._client.logger.warn(
+        `Cancel request for ${e.projectId} was rejected: ${e.errorMessage || 'unknown reason'}`
+      );
+      return;
+    }
+    const next: Project = {
+      ...project,
+      status: 'canceled',
+      error: { code: 5004, originalCode: 'artistCanceled', message: 'Artist canceled project' },
+      updatedAt: Date.now()
+    };
+    this._projects.set(project.id, next);
+    this._notifyChanged(next);
+    this.emit('failed', next);
+    this._scheduleGC(next.id);
+  }
+
+  // ===========================================================================
+  // Internal: state machinery
+  // ===========================================================================
+
+  private _withJobMutation(
+    project: Project,
+    jobId: string | null,
+    mutate: (job: Job) => Job
+  ): Project {
+    if (!jobId) {
+      // Some events (jobStarted before imgID is assigned) carry no jobId. We
+      // don't synthesize a job in that case; just bump the updatedAt.
+      const next = { ...project, updatedAt: Date.now() };
+      this._projects.set(project.id, next);
+      return next;
+    }
+    const existingIdx = project.jobs.findIndex((j) => j.id === jobId);
+    const existing: Job =
+      existingIdx >= 0 ? project.jobs[existingIdx]! : this._emptyJob(project.id, jobId);
+    const updated = mutate(existing);
+    const jobs =
+      existingIdx >= 0
+        ? project.jobs.map((j, i) => (i === existingIdx ? updated : j))
+        : [...project.jobs, updated];
+    const next: Project = { ...project, jobs, updatedAt: Date.now() };
+    this._projects.set(project.id, next);
+    return next;
+  }
+
+  private _emptyJob(projectId: string, jobId: string): Job {
+    return {
+      id: jobId,
+      projectId,
+      status: 'pending',
+      progress: 0,
+      step: 0,
+      stepCount: 0,
+      seed: null,
+      resultUrl: null,
+      previewUrl: null,
+      isNSFW: false,
+      workerName: null,
+      error: null
+    };
+  }
+
+  private _setStatusIfPending(project: Project, status: ProjectStatus) {
+    if (project.status !== 'pending') return;
+    const next: Project = { ...project, status, updatedAt: Date.now() };
+    this._projects.set(project.id, next);
+  }
+
+  private _computeOverallProgress(jobs: Job[]): number {
+    if (!jobs.length) return 0;
+    const sum = jobs.reduce((acc, j) => acc + (j.progress || 0), 0);
+    return sum / jobs.length;
+  }
+
+  private _maybeFinalizeProject(project: Project) {
+    const expected = Math.max(1, project.params.numberOfMedia);
+    const knownJobs = project.jobs.length;
+    const completedJobs = project.jobs.filter((j) => j.status === 'completed').length;
+    const allKnownTerminal =
+      knownJobs > 0 && project.jobs.every((j) => TERMINAL_JOB_STATUSES.includes(j.status));
+
+    if (knownJobs >= expected && allKnownTerminal) {
+      const anyCompleted = completedJobs > 0;
+      const next: Project = {
+        ...project,
+        status: anyCompleted ? 'completed' : 'failed',
+        progress: this._computeOverallProgress(project.jobs),
+        error: anyCompleted
+          ? null
+          : (project.jobs.find((j) => j.error)?.error ?? {
+              code: 5000,
+              message: 'All jobs failed'
+            }),
+        updatedAt: Date.now()
+      };
+      this._projects.set(project.id, next);
+      this._notifyChanged(next);
+      if (next.status === 'completed') {
+        this.emit('completed', next);
+      } else {
+        this.emit('failed', next);
+      }
+      this._scheduleGC(next.id);
+      return;
+    }
+
+    // Not terminal yet — just notify of the in-progress change.
+    const nextProgress = this._computeOverallProgress(project.jobs);
+    const next: Project = { ...project, progress: nextProgress, updatedAt: Date.now() };
+    this._projects.set(project.id, next);
+    this._notifyChanged(next);
+  }
+
+  private async _fetchResultUrls(project: Project): Promise<string[]> {
+    const urls: string[] = [];
+    for (const job of project.jobs) {
+      if (job.status !== 'completed') continue;
+      try {
+        const url = await this.getResultUrl(project.id, job.id);
+        urls.push(url);
+        // Update cached snapshot with the resolved URL
+        const cached = this._projects.get(project.id);
+        if (cached) {
+          const jobs = cached.jobs.map((j) => (j.id === job.id ? { ...j, resultUrl: url } : j));
+          const next: Project = { ...cached, jobs, resultUrls: urls.slice() };
+          this._projects.set(project.id, next);
+          this._notifyChanged(next);
+        }
+      } catch (e) {
+        this._client.logger.warn(`Failed to fetch result URL for job ${job.id}:`, e);
+      }
+    }
+    return urls;
+  }
+
+  private _notifyChanged(project: Project) {
+    const subs = this._subscribers.get(project.id);
+    if (subs) {
+      for (const fn of subs) fn(project);
+    }
+    this.emit('changed', project);
+  }
+
+  private _scheduleGC(projectId: string) {
+    const existing = this._gcTimers.get(projectId);
+    if (existing) clearTimeout(existing);
+    const timer = setTimeout(() => {
+      this._projects.delete(projectId);
+      this._subscribers.delete(projectId);
+      this._gcTimers.delete(projectId);
+    }, GC_DELAY_MS);
+    this._gcTimers.set(projectId, timer);
+  }
+}
+
+export default ProjectsApi;