LibreChat/client/src/components/Chat/Messages/Content/Parts/OpenAIImageGen/OpenAIImageGen.tsx

import { useState, useEffect, useRef, useCallback } from 'react';
import { PixelCard } from '@librechat/client';
import type { TAttachment, TFile, TAttachmentMetadata } from 'librechat-data-provider';
import { ToolIcon, isError } from '~/components/Chat/Messages/Content/ToolOutput';
import Image from '~/components/Chat/Messages/Content/Image';
import { useProgress, useLocalize } from '~/hooks';
import ProgressText from './ProgressText';
import { AGENT_STYLE_TOOLS } from '.';
import { scaleImage } from '~/utils';

function computeCancelled(
  isSubmitting: boolean | undefined,
  initialProgress: number,
  hasError: boolean,
): boolean {
  if (isSubmitting !== undefined) {
    return (!isSubmitting && initialProgress < 1) || hasError;
  }
  // Legacy path: in-progress (0 < progress < 1) is never cancelled
  // because legacy image gen lacks a submitting signal.
  if (initialProgress < 1 && initialProgress > 0) {
    return false;
  }
  return hasError;
}

export default function OpenAIImageGen({
  initialProgress = 0.1,
  isSubmitting,
  toolName,
  args: _args = '',
  output,
  attachments,
}: {
  initialProgress: number;
  isSubmitting?: boolean;
  toolName?: string;
  args: string | Record<string, unknown>;
  output?: string | null;
  attachments?: TAttachment[];
}) {
  const localize = useLocalize();
  const isAgentStyle = toolName != null && AGENT_STYLE_TOOLS.has(toolName);
  const [agentProgress, setAgentProgress] = useState(initialProgress);
  const legacyProgress = useProgress(isAgentStyle ? 1 : initialProgress);
  const progress = isAgentStyle ? agentProgress : legacyProgress;
  const intervalRef = useRef<NodeJS.Timeout | null>(null);

  const hasError = typeof output === 'string' && isError(output);

  /**
   * Determines if the image generation was cancelled.
   * - Agent path (isSubmitting defined): cancelled if not submitting + incomplete, or on error.
   * - Legacy path (isSubmitting undefined): in-progress (0 < progress < 1) is never cancelled
   *   because legacy image gen lacks a submitting signal — only errors cancel.
   */
  const cancelled = computeCancelled(isSubmitting, initialProgress, hasError);

  let width: number | undefined;
  let height: number | undefined;
  let quality: 'low' | 'medium' | 'high' = 'high';

  let parsedArgs: Record<string, unknown> = {};
  try {
    parsedArgs = typeof _args === 'string' ? JSON.parse(_args) : _args;
  } catch {
    parsedArgs = {};
  }

  try {
    const argsObj = parsedArgs;

    if (argsObj && typeof argsObj.size === 'string') {
      const [w, h] = argsObj.size.split('x').map((v: string) => parseInt(v, 10));
      if (!isNaN(w) && !isNaN(h)) {
        width = w;
        height = h;
      }
    } else if (argsObj && (typeof argsObj.size !== 'string' || !argsObj.size)) {
      width = undefined;
      height = undefined;
    }

    if (argsObj && typeof argsObj.quality === 'string') {
      const q = argsObj.quality.toLowerCase();
      if (q === 'low' || q === 'medium' || q === 'high') {
        quality = q;
      }
    }
  } catch {
    width = undefined;
    height = undefined;
  }

  const attachment = attachments?.[0];
  const {
    width: imageWidth,
    height: imageHeight,
    filepath = null,
    filename = '',
  } = (attachment as TFile & TAttachmentMetadata) || {};

  let origWidth = width ?? imageWidth;
  let origHeight = height ?? imageHeight;

  if (origWidth === undefined || origHeight === undefined) {
    origWidth = 1024;
    origHeight = 1024;
  }

  const [dimensions, setDimensions] = useState({ width: 'auto', height: 'auto' });
  const containerRef = useRef<HTMLDivElement>(null);

  const updateDimensions = useCallback(() => {
    if (origWidth && origHeight && containerRef.current) {
      const scaled = scaleImage({
        originalWidth: origWidth,
        originalHeight: origHeight,
        containerRef,
      });
      setDimensions(scaled);
    }
  }, [origWidth, origHeight]);

  useEffect(() => {
    if (!isAgentStyle) {
      return;
    }

    if (isSubmitting) {
      setAgentProgress(initialProgress);

      if (intervalRef.current) {
        clearInterval(intervalRef.current);
      }

      let baseDuration = 20000;
      if (quality === 'low') {
        baseDuration = 10000;
      } else if (quality === 'high') {
        baseDuration = 50000;
      }
      const jitter = Math.floor(baseDuration * 0.3);
      const totalDuration = Math.floor(Math.random() * jitter) + baseDuration;
      const updateInterval = 200;
      const totalSteps = totalDuration / updateInterval;
      let currentStep = 0;

      intervalRef.current = setInterval(() => {
        currentStep++;

        if (currentStep >= totalSteps) {
          clearInterval(intervalRef.current as NodeJS.Timeout);
          setAgentProgress(0.9);
        } else {
          const progressRatio = currentStep / totalSteps;
          let mapRatio: number;
          if (progressRatio < 0.8) {
            mapRatio = Math.pow(progressRatio, 1.1);
          } else {
            const sub = (progressRatio - 0.8) / 0.2;
            mapRatio = 0.8 + (1 - Math.pow(1 - sub, 2)) * 0.2;
          }
          const scaledProgress = 0.1 + mapRatio * 0.8;

          setAgentProgress(scaledProgress);
        }
      }, updateInterval);
    }

    return () => {
      if (intervalRef.current) {
        clearInterval(intervalRef.current);
      }
    };
  }, [isSubmitting, initialProgress, quality, isAgentStyle]);

  useEffect(() => {
    if (!isAgentStyle) {
      return;
    }

    if (initialProgress >= 1 || cancelled) {
      setAgentProgress(initialProgress);
      if (intervalRef.current) {
        clearInterval(intervalRef.current);
      }
    }
  }, [initialProgress, cancelled, isAgentStyle]);

  useEffect(() => {
    updateDimensions();

    const resizeObserver = new ResizeObserver(() => {
      updateDimensions();
    });

    if (containerRef.current) {
      resizeObserver.observe(containerRef.current);
    }

    return () => {
      resizeObserver.disconnect();
    };
  }, [updateDimensions]);

  const isInProgress = progress < 1 && !cancelled;

  return (
    <>
      <span className="sr-only" aria-live="polite" aria-atomic="true">
        {(() => {
          if (progress < 1 && !cancelled) {
            return '';
          }
          if (cancelled && hasError) {
            return localize('com_ui_image_gen_failed');
          }
          if (cancelled) {
            return localize('com_ui_cancelled');
          }
          return localize('com_ui_image_created');
        })()}
      </span>
      <div className="relative my-1 flex h-5 shrink-0 items-center gap-2">
        <ToolIcon type="image_gen" isAnimating={isInProgress} />
        <ProgressText progress={progress} error={cancelled} toolName={toolName} />
      </div>
      {isAgentStyle && (
        <div className="relative mb-2 flex w-full justify-start">
          <div ref={containerRef} className="w-full max-w-lg">
            {dimensions.width !== 'auto' && progress < 1 && (
              <PixelCard
                variant="default"
                progress={progress}
                randomness={0.6}
                width={dimensions.width}
                height={dimensions.height}
              />
            )}
            <Image
              altText={filename}
              imagePath={filepath ?? ''}
              width={Number(dimensions.width?.split('px')[0])}
              height={Number(dimensions.height?.split('px')[0])}
              args={parsedArgs}
            />
          </div>
        </div>
      )}
    </>
  );
}