Fix: Fix Image Generation Button Logic

ShenKSPZ · ShenKSPZ · commit 17aa4710d6b7 · 2025-05-03T20:17:31.000+08:00
diff --git a/src/components/chat/ChatMessageArea.tsx b/src/components/chat/ChatMessageArea.tsx
@@ -610,30 +610,30 @@ export const ChatMessageArea: React.FC<ChatMessageAreaProps> = ({
                       {t('chat.availableMcpServers')}
                     </div>
                     <div className="overflow-y-auto max-h-60">
-                      {Object.values(mcpServers).map((server) => (
-                        <div
-                          key={server.id}
-                          className={`flex items-center px-3 py-2 cursor-pointer rounded-md ${
-                            selectedMcpServers?.includes(server.id)
-                              ? 'image-generation-provider-selected'
-                              : 'image-generation-provider-item'
-                          }`}
-                          onClick={() => onToggleMcpServer?.(server.id)}
-                        >
-                          <ServerCog className="w-5 h-5 mr-2" />
-                          <div className="flex flex-col">
-                            <span className="text-sm font-medium">{server.name}</span>
-                            {server.isDefault && (
-                              <span className="text-xs text-gray-500">{t('mcpServer.default')}</span>
-                            )}
-                            {server.isImageGeneration && (
-                              <span className="text-xs text-gray-500">{t('mcpServer.imageGeneration')}</span>
-                            )}
+                      {Object.values(mcpServers)
+                        // Filter out image generation servers - they're handled by the ImageGenerationButton
+                        .filter(server => !server.isImageGeneration)
+                        .map((server) => (
+                          <div
+                            key={server.id}
+                            className={`flex items-center px-3 py-2 cursor-pointer rounded-md ${
+                              selectedMcpServers?.includes(server.id)
+                                ? 'image-generation-provider-selected'
+                                : 'image-generation-provider-item'
+                            }`}
+                            onClick={() => onToggleMcpServer?.(server.id)}
+                          >
+                            <ServerCog className="w-5 h-5 mr-2" />
+                            <div className="flex flex-col">
+                              <span className="text-sm font-medium">{server.name}</span>
+                              {server.isDefault && (
+                                <span className="text-xs text-gray-500">{t('mcpServer.default')}</span>
+                              )}
+                            </div>
                           </div>
-                        </div>
-                      ))}
+                        ))}
                       
-                      {Object.keys(mcpServers).length === 0 && (
+                      {Object.values(mcpServers).filter(server => !server.isImageGeneration).length === 0 && (
                         <div className="px-3 py-2 text-sm text-gray-500">
                           {t('chat.noMcpServersAvailable')}
                         </div>
diff --git a/src/components/chat/ImageGenerationButton.tsx b/src/components/chat/ImageGenerationButton.tsx
@@ -1,5 +1,5 @@
 import React, { useState, useRef, useEffect } from 'react';
-import { Image } from 'lucide-react';
+import { Image, ToggleLeft, ToggleRight } from 'lucide-react';
 import { SettingsService, SETTINGS_CHANGE_EVENT } from '../../services/settings-service';
 import { AIServiceCapability } from '../../types/capabilities';
 import ProviderIcon from '../ui/ProviderIcon';
@@ -25,10 +25,11 @@ const ImageGenerationButton: React.FC<ImageGenerationButtonProps> = ({
   const [providers, setProviders] = useState<ProviderModel[]>([]);
   const [selectedProvider, setSelectedProvider] = useState<string | null>(null);
   const [selectedModel, setSelectedModel] = useState<string | null>(null);
+  const [isEnabled, setIsEnabled] = useState(true);
   const popupRef = useRef<HTMLDivElement>(null);
   const buttonRef = useRef<HTMLButtonElement>(null);
 
-  // Load available image generation providers and models
+  // Load available image generation providers and models and settings
   useEffect(() => {
     const loadProviders = () => {
       const settingsService = SettingsService.getInstance();
@@ -38,6 +39,13 @@ const ImageGenerationButton: React.FC<ImageGenerationButtonProps> = ({
       const settings = settingsService.getSettings();
       const providerIds = Object.keys(settings.providers);
       
+      // Load current image generation enabled status
+      setIsEnabled(settings.imageGenerationEnabled !== false);
+      
+      // Load saved selected provider and model
+      const savedProvider = settings.imageGenerationProvider;
+      const savedModel = settings.imageGenerationModel;
+      
       for (const providerId of providerIds) {
         // Get the provider's settings
         const providerSettings = settingsService.getProviderSettings(providerId);
@@ -59,8 +67,11 @@ const ImageGenerationButton: React.FC<ImageGenerationButtonProps> = ({
 
       setProviders(availableProviders);
       
-      // Set default selected provider and model if available
-      if (availableProviders.length > 0) {
+      // Set saved or default selected provider and model
+      if (savedProvider && savedModel && availableProviders.some(p => p.providerName === savedProvider && p.modelId === savedModel)) {
+        setSelectedProvider(savedProvider);
+        setSelectedModel(savedModel);
+      } else if (availableProviders.length > 0) {
         setSelectedProvider(availableProviders[0].providerName);
         setSelectedModel(availableProviders[0].modelId);
       }
@@ -99,29 +110,57 @@ const ImageGenerationButton: React.FC<ImageGenerationButtonProps> = ({
     setIsPopupOpen(!isPopupOpen);
   };
 
-  const handleProviderModelSelect = (providerName: string, modelId: string) => {
+  const handleProviderModelSelect = async (providerName: string, modelId: string) => {
     setSelectedProvider(providerName);
     setSelectedModel(modelId);
-    setIsPopupOpen(false);
+    
+    // Save selected provider and model in settings
+    const settingsService = SettingsService.getInstance();
+    await settingsService.updateSettings({
+      imageGenerationProvider: providerName,
+      imageGenerationModel: modelId
+    });
     
     // If onImageGenerate is provided, call it with an empty prompt
     // The actual prompt will be filled in by the chat message
-    if (onImageGenerate) {
+    if (onImageGenerate && isEnabled) {
       onImageGenerate("", providerName, modelId);
     }
+    
+    setIsPopupOpen(false);
+  };
+  
+  const toggleImageGeneration = async () => {
+    const newEnabledState = !isEnabled;
+    setIsEnabled(newEnabledState);
+    
+    // Save the enabled state in settings
+    const settingsService = SettingsService.getInstance();
+    await settingsService.updateSettings({
+      imageGenerationEnabled: newEnabledState
+    });
   };
 
-  const isButtonEnabled = !disabled && providers.length > 0;
+  const isButtonEnabled = !disabled && providers.length > 0 && isEnabled;
+  const buttonClass = `flex items-center justify-center w-8 h-8 rounded-full focus:outline-none ${
+    isEnabled ? 'image-generation-button' : 'text-gray-400 bg-gray-100'
+  }`;
 
   return (
     <div className="relative">
       <button
         ref={buttonRef}
         type="button"
         onClick={togglePopup}
-        disabled={!isButtonEnabled}
-        className="flex items-center justify-center w-8 h-8 rounded-full image-generation-button focus:outline-none"
-        title={isButtonEnabled ? t('chat.generateImage') : t('chat.imageGenerationNotAvailable')}
+        disabled={!providers.length > 0 || disabled}
+        className={buttonClass}
+        title={
+          !providers.length > 0 
+            ? t('chat.imageGenerationNotAvailable') 
+            : isEnabled 
+              ? t('chat.generateImage') 
+              : t('chat.imageGenerationDisabled')
+        }
       >
         <Image size={20} />
       </button>
@@ -133,15 +172,28 @@ const ImageGenerationButton: React.FC<ImageGenerationButtonProps> = ({
           style={{ bottom: '100%', left: 0, minWidth: '220px' }}
         >
           <div className="p-2">
-            <div className="mb-2 text-sm font-medium text-gray-700">
-              {t('chat.selectImageProvider')}
+            <div className="flex items-center justify-between mb-3">
+              <div className="text-sm font-medium text-gray-700">
+                {t('chat.selectImageProvider')}
+              </div>
+              <button
+                onClick={toggleImageGeneration}
+                className="flex items-center text-sm focus:outline-none"
+                title={isEnabled ? t('chat.disableImageGeneration') : t('chat.enableImageGeneration')}
+              >
+                {isEnabled ? (
+                  <ToggleRight className="w-6 h-6 text-blue-500" />
+                ) : (
+                  <ToggleLeft className="w-6 h-6 text-gray-400" />
+                )}
+              </button>
             </div>
             <div className="overflow-y-auto max-h-60">
               {providers.map((provider) => (
                 <div
                   key={`${provider.providerName}-${provider.modelId}`}
                   className={`flex items-center px-3 py-2 cursor-pointer rounded-md ${
-                    selectedProvider === provider.providerName && selectedModel === provider.modelId
+                    isEnabled && selectedProvider === provider.providerName && selectedModel === provider.modelId
                       ? 'image-generation-provider-selected'
                       : 'image-generation-provider-item'
                   }`}
diff --git a/src/locales/en/translation.json b/src/locales/en/translation.json
@@ -33,6 +33,9 @@
     "pressShiftEnterToChangeLines": "Press Shift+Enter to change lines",
     "generateImage": "Generate Image",
     "imageGenerationNotAvailable": "Image generation not available",
+    "imageGenerationDisabled": "Image generation is disabled",
+    "enableImageGeneration": "Enable image generation",
+    "disableImageGeneration": "Disable image generation",
     "selectImageProvider": "Select Image Provider",
     "noImageProvidersAvailable": "No image providers available"
   },
diff --git a/src/locales/es/translation.json b/src/locales/es/translation.json
@@ -33,6 +33,9 @@
     "pressShiftEnterToChangeLines": "Presiona Shift+Enter para cambiar de línea",
     "generateImage": "Generar Imagen",
     "imageGenerationNotAvailable": "Generación de imágenes no disponible",
+    "imageGenerationDisabled": "La generación de imágenes está desactivada",
+    "enableImageGeneration": "Activar generación de imágenes",
+    "disableImageGeneration": "Desactivar generación de imágenes",
     "selectImageProvider": "Seleccionar Proveedor de Imágenes",
     "noImageProvidersAvailable": "No hay proveedores de imágenes disponibles"
   },
diff --git a/src/locales/ja/translation.json b/src/locales/ja/translation.json
@@ -33,6 +33,9 @@
     "pressShiftEnterToChangeLines": "Shift+Enterを押して行を変更",
     "generateImage": "画像を生成",
     "imageGenerationNotAvailable": "画像生成は利用できません",
+    "imageGenerationDisabled": "画像生成は無効になっています",
+    "enableImageGeneration": "画像生成を有効にする",
+    "disableImageGeneration": "画像生成を無効にする",
     "selectImageProvider": "画像プロバイダーを選択",
     "noImageProvidersAvailable": "利用可能な画像プロバイダーがありません"
   },
diff --git a/src/locales/ko/translation.json b/src/locales/ko/translation.json
@@ -33,6 +33,9 @@
     "pressShiftEnterToChangeLines": "줄을 바꾸려면 Shift+Enter를 누르세요",
     "generateImage": "이미지 생성",
     "imageGenerationNotAvailable": "이미지 생성을 사용할 수 없습니다",
+    "imageGenerationDisabled": "이미지 생성이 비활성화되었습니다",
+    "enableImageGeneration": "이미지 생성 활성화",
+    "disableImageGeneration": "이미지 생성 비활성화",
     "selectImageProvider": "이미지 제공자 선택",
     "noImageProvidersAvailable": "사용 가능한 이미지 제공자가 없습니다"
   },
diff --git a/src/locales/zh-CN/translation.json b/src/locales/zh-CN/translation.json
@@ -33,6 +33,9 @@
     "pressShiftEnterToChangeLines": "按 Shift+Enter 换行",
     "generateImage": "生成图片",
     "imageGenerationNotAvailable": "图片生成不可用",
+    "imageGenerationDisabled": "图片生成已禁用",
+    "enableImageGeneration": "启用图片生成",
+    "disableImageGeneration": "禁用图片生成",
     "selectImageProvider": "选择图片提供商",
     "noImageProvidersAvailable": "没有可用的图片提供商"
   },
diff --git a/src/locales/zh-TW/translation.json b/src/locales/zh-TW/translation.json
@@ -33,6 +33,9 @@
     "pressShiftEnterToChangeLines": "按 Shift+Enter 換行",
     "generateImage": "生成圖片",
     "imageGenerationNotAvailable": "圖片生成不可用",
+    "imageGenerationDisabled": "圖片生成已禁用",
+    "enableImageGeneration": "啟用圖片生成",
+    "disableImageGeneration": "禁用圖片生成",
     "selectImageProvider": "選擇圖片提供商",
     "noImageProvidersAvailable": "沒有可用的圖片提供商"
   },
diff --git a/src/types/settings.ts b/src/types/settings.ts
@@ -10,6 +10,10 @@ export interface UserSettings {
   useStreaming: boolean;
   webSearchEnabled: boolean;
   enableWebSearch_Preview: boolean;
+  // Image generation settings
+  imageGenerationEnabled?: boolean;
+  imageGenerationProvider?: string;
+  imageGenerationModel?: string;
   // General settings
   startWithSystem?: boolean;
   startupToTray?: boolean;
