cinny/src/app/state/upload.ts

import { atom, useAtom } from 'jotai';
import { atomFamily } from 'jotai/utils';
import { MatrixClient, UploadResponse, UploadProgress, MatrixError } from 'matrix-js-sdk';
import { useCallback, useRef } from 'react';
import { useThrottle } from '../hooks/useThrottle';
import { uploadContent, TUploadContent } from '../utils/matrix';
import { uploadProgressAtom } from './uploadProgress';

export enum UploadStatus {
  Idle = 'idle',
  Loading = 'loading',
  Success = 'success',
  Error = 'error',
}

export type UploadIdle = {
  file: TUploadContent;
  status: UploadStatus.Idle;
};

export type UploadLoading = {
  file: TUploadContent;
  status: UploadStatus.Loading;
  promise: Promise<UploadResponse>;
  progress: UploadProgress;
};

export type UploadSuccess = {
  file: TUploadContent;
  status: UploadStatus.Success;
  mxc: string;
};

export type UploadError = {
  file: TUploadContent;
  status: UploadStatus.Error;
  error: MatrixError;
};

export type Upload = UploadIdle | UploadLoading | UploadSuccess | UploadError;

export type UploadAtomAction =
  | {
      promise: Promise<UploadResponse>;
    }
  | {
      progress: UploadProgress;
    }
  | {
      mxc: string;
    }
  | {
      error: MatrixError;
    };

export const createUploadAtom = (file: TUploadContent) => {
  const baseUploadAtom = atom<Upload>({
    file,
    status: UploadStatus.Idle,
  });
  return atom<Upload, [UploadAtomAction], undefined>(
    (get) => get(baseUploadAtom),
    (get, set, update) => {
      const uploadState = get(baseUploadAtom);
      // [cinny-desktop #10] Mirror into the app-wide progress map for the
      // desktop taskbar bar. Cancelled uploads just disappear (no red bar).
      const track = (entry: { loaded: number; total: number; failed: boolean } | undefined) =>
        set(uploadProgressAtom, (prev) => {
          const next = new Map(prev);
          if (entry) next.set(file, { ...entry, updatedAt: Date.now() });
          else next.delete(file);
          return next;
        });
      if ('promise' in update) {
        track({ loaded: 0, total: file.size, failed: false });
        set(baseUploadAtom, {
          status: UploadStatus.Loading,
          file,
          promise: update.promise,
          progress: { loaded: 0, total: file.size },
        });
        return;
      }
      if ('progress' in update && uploadState.status === UploadStatus.Loading) {
        track({
          loaded: update.progress.loaded,
          total: update.progress.total || file.size,
          failed: false,
        });
        set(baseUploadAtom, {
          ...uploadState,
          progress: update.progress,
        });
        return;
      }
      if ('mxc' in update) {
        track(undefined);
        set(baseUploadAtom, {
          status: UploadStatus.Success,
          file,
          mxc: update.mxc,
        });
        return;
      }
      if ('error' in update) {
        const cancelled =
          update.error?.name === 'AbortError' || /abort/i.test(update.error?.message ?? '');
        const loaded =
          uploadState.status === UploadStatus.Loading ? uploadState.progress.loaded : 0;
        track(cancelled ? undefined : { loaded, total: file.size, failed: true });
        set(baseUploadAtom, {
          status: UploadStatus.Error,
          file,
          error: update.error,
        });
      }
    },
  );
};
export type TUploadAtom = ReturnType<typeof createUploadAtom>;

export const useBindUploadAtom = (
  mx: MatrixClient,
  uploadAtom: TUploadAtom,
  hideFilename?: boolean,
) => {
  const [upload, setUpload] = useAtom(uploadAtom);
  const { file } = upload;

  const handleProgress = useThrottle(
    useCallback((progress: UploadProgress) => setUpload({ progress }), [setUpload]),
    { immediate: true, wait: 200 },
  );

  const abortRef = useRef<AbortController | undefined>(undefined);

  const startUpload = useCallback(() => {
    const controller = new AbortController();
    abortRef.current = controller;
    return uploadContent(mx, file, {
      hideFilename,
      signal: controller.signal,
      onPromise: (promise: Promise<UploadResponse>) => setUpload({ promise }),
      onProgress: handleProgress,
      onSuccess: (mxc) => setUpload({ mxc }),
      onError: (error) => setUpload({ error }),
    });
  }, [mx, file, hideFilename, setUpload, handleProgress]);

  const cancelUpload = useCallback(async () => {
    // Abort the retry loop first (covers a cancel during the back-off sleep, when
    // there is no in-flight request), then abort any in-flight upload request.
    abortRef.current?.abort();
    if (upload.status === UploadStatus.Loading) {
      await mx.cancelUpload(upload.promise);
    }
  }, [mx, upload]);

  return {
    upload,
    startUpload,
    cancelUpload,
  };
};

export const createUploadAtomFamily = () =>
  atomFamily<TUploadContent, TUploadAtom>(createUploadAtom);
export type TUploadAtomFamily = ReturnType<typeof createUploadAtomFamily>;

export const createUploadFamilyObserverAtom = (
  uploadFamily: TUploadAtomFamily,
  uploads: TUploadContent[],
) => atom<Upload[]>((get) => uploads.map((upload) => get(uploadFamily(upload))));
export type TUploadFamilyObserverAtom = ReturnType<typeof createUploadFamilyObserverAtom>;