uploadRunner.ts8.6 KBView on GitHub
/**
 * Upload pump — drives the upload state machine for a single attempt.
 *
 * Validates the file, hashes it, calls trpc.files.requestUpload, PUTs to S3
 * via XHR (so we get progress events), then calls completeUpload. On any
 * failure it stamps the upload as `failed` and surfaces the error.
 *
 * Always reads fresh state via `useCedarStore.getState()` because Zustand+Immer
 * replaces the state ref after each set() — captured snapshots go stale.
 */

import { useCedarStore } from '@/modules/store';

import type { UploadAttempt } from '../store/documentsSlice';
import { sha256OfFile, validateUpload } from './allowlist';

/**
 * Side-channel map of in-flight XHR handles, keyed by opId. Storing XHR in
 * the Zustand store would freeze the object via immer and prevent abort(),
 * so we keep them outside the store.
 */
const inFlightXhrs = new Map<string, XMLHttpRequest>();

export function abortUpload(opId: string): void {
  const xhr = inFlightXhrs.get(opId);
  if (xhr) {
    try {
      xhr.abort();
    } catch {
      /* ignore */
    }
    inFlightXhrs.delete(opId);
  }
}

interface CompleteUploadResult {
  metadata?: {
    extraction?: { status?: 'pending' | 'done' | 'failed' | 'skipped' };
  } | null;
}

interface FileUploadTrpcClient {
  files: {
    requestUpload: { mutate: (input: unknown) => Promise<RequestUploadResult> };
    completeUpload: { mutate: (input: unknown) => Promise<CompleteUploadResult> };
    cancelUpload: { mutate: (input: unknown) => Promise<unknown> };
  };
}

interface UploadRunnerArgs {
  /** Resolved op id in the slice. */
  opId: string;
  /** The user-selected File object. */
  file: File;
  /** Vanilla tRPC client (NOT the React hooks variant). */
  trpc: FileUploadTrpcClient;
}

interface RequestUploadResult {
  documentId: string;
  presignedUrl: string;
  headers: Record<string, string>;
  expiresAt: string;
  storageKey=[redacted];
  path: string;
}

export async function runUpload(args: UploadRunnerArgs): Promise<void> {
  const { opId, file, trpc } = args;

  // Snapshot the upload now — fields like scope/parentId don't change after
  // attachToChat/optimisticCreateAttachment. Status is read fresh below.
  const initial = useCedarStore.getState().uploads[opId];
  if (!initial) return;
  if (initial.status !== 'validating') return;

  const setStatus = (patch: Partial<UploadAttempt>) =>
    useCedarStore.getState().setUploadStatus(opId, patch);

  // 1. Validate
  const validation = validateUpload({
    filename: file.name,
    mimeType: file.type || guessMimeType(file.name),
    sizeBytes: file.size,
  });
  if (!validation.ok) {
    setStatus({ status: 'failed', error: { code: validation.code, message: validation.message } });
    return;
  }

  // 2. Hash
  let sha256: string;
  try {
    sha256 = await sha256OfFile(file);
  } catch {
    setStatus({ status: 'failed', error: { code: 'hash', message: 'Could not hash file' } });
    return;
  }
  setStatus({ status: 'requesting', sha256 });

  // 3. Request presigned URL
  let presign: RequestUploadResult;
  try {
    presign = await trpc.files.requestUpload.mutate({
      scope:
        initial.scope.type === 'chat_thread'
          ? { type: 'chat_thread', id: initial.scope.id }
          : { type: initial.scope.type, id: initial.scope.id },
      parentId: initial.parentId,
      filename: file.name,
      mimeType: validation.mimeType,
      sizeBytes: file.size,
      sha256,
      // The client calls a KB item 'kb_item'; the server's DOCUMENT_TYPE calls it 'document'.
      // 'custom' has no server equivalent, so it is simply not forwarded.
      ...(initial.documentType === 'kb_item'
        ? { documentType: 'document' as const }
        : initial.documentType === 'attachment'
          ? { documentType: 'attachment' as const }
          : {}),
      ...(initial.chatThreadId ? { chatThreadId: initial.chatThreadId } : {}),
    });
  } catch (e) {
    setStatus({ status: 'failed', error: { code: 'requesting', message: errorMessage(e) } });
    return;
  }
  useCedarStore.getState().setUploadDocumentId(opId, presign.documentId);
  setStatus({ status: 'uploading', progress: 0, documentPath: presign.path });

  // 4. PUT via XHR (for progress)
  try {
    await xhrPut({
      url: presign.presignedUrl,
      headers: presign.headers,
      file,
      onProgress: (progress) => setStatus({ progress }),
      onXhrCreated: (xhr) => inFlightXhrs.set(opId, xhr),
      isAborted: () => useCedarStore.getState().uploads[opId]?.status === 'cancelled',
    });
    inFlightXhrs.delete(opId);
  } catch (e) {
    inFlightXhrs.delete(opId);
    if (errorMessage(e) === 'aborted') {
      setStatus({ status: 'cancelled' });
      void trpc.files.cancelUpload.mutate({ documentId: presign.documentId }).catch(() => {});
      return;
    }
    setStatus({ status: 'failed', error: { code: 'uploading', message: errorMessage(e) } });
    void trpc.files.cancelUpload.mutate({ documentId: presign.documentId }).catch(() => {});
    return;
  }

  // 5. Complete
  setStatus({ status: 'completing' });
  let completeResult: CompleteUploadResult;
  try {
    completeResult = await trpc.files.completeUpload.mutate({
      documentId: presign.documentId,
    });
  } catch (e) {
    setStatus({ status: 'failed', error: { code: 'completing', message: errorMessage(e) } });
    return;
  }

  // 6. If the server kicked off background extraction, surface that to the UI
  // so the badge shows "Extracting" instead of jumping straight to done. We
  // can't reliably poll across all surfaces, so flip to 'done' after a short
  // ceiling — extraction doesn't gate the file's usability for the user
  // (read/download work immediately; only search/grep waits for text).
  const extractionStatus = completeResult?.metadata?.extraction?.status;
  if (extractionStatus === 'pending') {
    setStatus({ status: 'extracting', progress: 1 });
    setTimeout(() => {
      const current = useCedarStore.getState().uploads[opId];
      if (current?.status === 'extracting') {
        useCedarStore.getState().setUploadStatus(opId, { status: 'done' });
      }
    }, 5000);
    return;
  }

  setStatus({ status: 'done', progress: 1 });
}

function errorMessage(e: unknown): string {
  if (e instanceof Error) return e.message;
  if (typeof e === 'string') return e;
  return 'Unknown error';
}

function guessMimeType(filename: string): string {
  const lower = filename.toLowerCase();
  if (lower.endsWith('.md') || lower.endsWith('.markdown')) return 'text/markdown';
  if (lower.endsWith('.pdf')) return 'application/pdf';
  if (lower.endsWith('.docx'))
    return 'application/vnd.openxmlformats-officedocument.wordprocessingml.document';
  if (lower.endsWith('.doc')) return 'application/msword';
  if (lower.endsWith('.pptx'))
    return 'application/vnd.openxmlformats-officedocument.presentationml.presentation';
  if (lower.endsWith('.ppt')) return 'application/vnd.ms-powerpoint';
  if (lower.endsWith('.png')) return 'image/png';
  if (lower.endsWith('.jpg') || lower.endsWith('.jpeg')) return 'image/jpeg';
  if (lower.endsWith('.gif')) return 'image/gif';
  if (lower.endsWith('.webp')) return 'image/webp';
  if (lower.endsWith('.svg')) return 'image/svg+xml';
  return 'application/octet-stream';
}

interface XhrPutArgs {
  url: string;
  headers: Record<string, string>;
  file: File;
  onProgress: (progress: number) => void;
  onXhrCreated: (xhr: XMLHttpRequest) => void;
  isAborted: () => boolean;
}

function xhrPut(args: XhrPutArgs): Promise<void> {
  return new Promise((resolve, reject) => {
    const xhr = new XMLHttpRequest();
    args.onXhrCreated(xhr);
    xhr.open('PUT', args.url, true);
    Object.entries(args.headers).forEach(([k, v]) => xhr.setRequestHeader(k, v));
    xhr.upload.onprogress = (e) => {
      if (!e.lengthComputable) return;
      args.onProgress(e.loaded / e.total);
    };
    xhr.onload = () => {
      if (xhr.status >= 200 && xhr.status < 300) {
        resolve();
        return;
      }
      const body = xhr.responseText?.slice(0, 500) ?? '';
      console.error('[S3 PUT] non-2xx', { status: xhr.status, body, url: args.url });
      reject(new Error(`S3 PUT failed: ${xhr.status}${body ? ` — ${body}` : ''}`));
    };
    xhr.onerror = () => {
      // CORS rejection, DNS failure, or other network error — XHR cannot
      // surface a status. Log the URL so we can pull it from the network tab.
      console.error('[S3 PUT] network error (likely CORS or DNS)', { url: args.url });
      reject(new Error('Network error during S3 PUT (likely CORS or DNS — see console)'));
    };
    xhr.onabort = () => reject(new Error('aborted'));
    xhr.send(args.file);
    if (args.isAborted()) {
      try {
        xhr.abort();
      } catch {
        /* ignore */
      }
    }
  });
}