Browse Source

refactor: flatten image generation settings structure (#7536)

Daniel 4 months ago
parent
commit
c3d84d295d

+ 21 - 4
packages/types/src/global-settings.ts

@@ -42,6 +42,10 @@ export const globalSettingsSchema = z.object({
 	customInstructions: z.string().optional(),
 	taskHistory: z.array(historyItemSchema).optional(),
 
+	// Image generation settings (experimental) - flattened for simplicity
+	openRouterImageApiKey: z.string().optional(),
+	openRouterImageGenerationSelectedModel: z.string().optional(),
+
 	condensingApiConfigId: z.string().optional(),
 	customCondensingPrompt: z.string().optional(),
 
@@ -200,11 +204,24 @@ export const SECRET_STATE_KEYS = [
 	"featherlessApiKey",
 	"ioIntelligenceApiKey",
 	"vercelAiGatewayApiKey",
-] as const satisfies readonly (keyof ProviderSettings)[]
-export type SecretState = Pick<ProviderSettings, (typeof SECRET_STATE_KEYS)[number]>
+] as const
+
+// Global secrets that are part of GlobalSettings (not ProviderSettings)
+export const GLOBAL_SECRET_KEYS = [
+	"openRouterImageApiKey", // For image generation
+] as const
+
+// Type for the actual secret storage keys
+type ProviderSecretKey = (typeof SECRET_STATE_KEYS)[number]
+type GlobalSecretKey = (typeof GLOBAL_SECRET_KEYS)[number]
+
+// Type representing all secrets that can be stored
+export type SecretState = Pick<ProviderSettings, Extract<ProviderSecretKey, keyof ProviderSettings>> & {
+	[K in GlobalSecretKey]?: string
+}
 
 export const isSecretStateKey = (key: string): key is Keys<SecretState> =>
-	SECRET_STATE_KEYS.includes(key as Keys<SecretState>)
+	SECRET_STATE_KEYS.includes(key as ProviderSecretKey) || GLOBAL_SECRET_KEYS.includes(key as GlobalSecretKey)
 
 /**
  * GlobalState
@@ -213,7 +230,7 @@ export const isSecretStateKey = (key: string): key is Keys<SecretState> =>
 export type GlobalState = Omit<RooCodeSettings, Keys<SecretState>>
 
 export const GLOBAL_STATE_KEYS = [...GLOBAL_SETTINGS_KEYS, ...PROVIDER_SETTINGS_KEYS].filter(
-	(key: Keys<RooCodeSettings>) => !SECRET_STATE_KEYS.includes(key as Keys<SecretState>),
+	(key: Keys<RooCodeSettings>) => !isSecretStateKey(key),
 ) as Keys<GlobalState>[]
 
 export const isGlobalStateKey = (key: string): key is Keys<GlobalState> =>

+ 0 - 7
packages/types/src/provider-settings.ts

@@ -142,13 +142,6 @@ const openRouterSchema = baseProviderSettingsSchema.extend({
 	openRouterBaseUrl: z.string().optional(),
 	openRouterSpecificProvider: z.string().optional(),
 	openRouterUseMiddleOutTransform: z.boolean().optional(),
-	// Image generation settings (experimental)
-	openRouterImageGenerationSettings: z
-		.object({
-			openRouterApiKey: z.string().optional(),
-			selectedModel: z.string().optional(),
-		})
-		.optional(),
 })
 
 const bedrockSchema = apiModelIdProviderModelSchema.extend({

+ 102 - 22
src/core/config/ContextProxy.ts

@@ -6,6 +6,7 @@ import {
 	GLOBAL_SETTINGS_KEYS,
 	SECRET_STATE_KEYS,
 	GLOBAL_STATE_KEYS,
+	GLOBAL_SECRET_KEYS,
 	type ProviderSettings,
 	type GlobalSettings,
 	type SecretState,
@@ -61,19 +62,77 @@ export class ContextProxy {
 			}
 		}
 
-		const promises = SECRET_STATE_KEYS.map(async (key) => {
-			try {
-				this.secretCache[key] = await this.originalContext.secrets.get(key)
-			} catch (error) {
-				logger.error(`Error loading secret ${key}: ${error instanceof Error ? error.message : String(error)}`)
-			}
-		})
+		const promises = [
+			...SECRET_STATE_KEYS.map(async (key) => {
+				try {
+					this.secretCache[key] = await this.originalContext.secrets.get(key)
+				} catch (error) {
+					logger.error(
+						`Error loading secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
+					)
+				}
+			}),
+			...GLOBAL_SECRET_KEYS.map(async (key) => {
+				try {
+					this.secretCache[key] = await this.originalContext.secrets.get(key)
+				} catch (error) {
+					logger.error(
+						`Error loading global secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
+					)
+				}
+			}),
+		]
 
 		await Promise.all(promises)
 
+		// Migration: Check for old nested image generation settings and migrate them
+		await this.migrateImageGenerationSettings()
+
 		this._isInitialized = true
 	}
 
+	/**
+	 * Migrates old nested openRouterImageGenerationSettings to the new flattened structure
+	 */
+	private async migrateImageGenerationSettings() {
+		try {
+			// Check if there's an old nested structure
+			const oldNestedSettings = this.originalContext.globalState.get<any>("openRouterImageGenerationSettings")
+
+			if (oldNestedSettings && typeof oldNestedSettings === "object") {
+				logger.info("Migrating old nested image generation settings to flattened structure")
+
+				// Migrate the API key if it exists and we don't already have one
+				if (oldNestedSettings.openRouterApiKey && !this.secretCache.openRouterImageApiKey) {
+					await this.originalContext.secrets.store(
+						"openRouterImageApiKey",
+						oldNestedSettings.openRouterApiKey,
+					)
+					this.secretCache.openRouterImageApiKey = oldNestedSettings.openRouterApiKey
+					logger.info("Migrated openRouterImageApiKey to secrets")
+				}
+
+				// Migrate the selected model if it exists and we don't already have one
+				if (oldNestedSettings.selectedModel && !this.stateCache.openRouterImageGenerationSelectedModel) {
+					await this.originalContext.globalState.update(
+						"openRouterImageGenerationSelectedModel",
+						oldNestedSettings.selectedModel,
+					)
+					this.stateCache.openRouterImageGenerationSelectedModel = oldNestedSettings.selectedModel
+					logger.info("Migrated openRouterImageGenerationSelectedModel to global state")
+				}
+
+				// Clean up the old nested structure
+				await this.originalContext.globalState.update("openRouterImageGenerationSettings", undefined)
+				logger.info("Removed old nested openRouterImageGenerationSettings")
+			}
+		} catch (error) {
+			logger.error(
+				`Error during image generation settings migration: ${error instanceof Error ? error.message : String(error)}`,
+			)
+		}
+	}
+
 	public get extensionUri() {
 		return this.originalContext.extensionUri
 	}
@@ -152,20 +211,34 @@ export class ContextProxy {
 	 * This is useful when you need to ensure the cache has the latest values
 	 */
 	async refreshSecrets(): Promise<void> {
-		const promises = SECRET_STATE_KEYS.map(async (key) => {
-			try {
-				this.secretCache[key] = await this.originalContext.secrets.get(key)
-			} catch (error) {
-				logger.error(
-					`Error refreshing secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
-				)
-			}
-		})
+		const promises = [
+			...SECRET_STATE_KEYS.map(async (key) => {
+				try {
+					this.secretCache[key] = await this.originalContext.secrets.get(key)
+				} catch (error) {
+					logger.error(
+						`Error refreshing secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
+					)
+				}
+			}),
+			...GLOBAL_SECRET_KEYS.map(async (key) => {
+				try {
+					this.secretCache[key] = await this.originalContext.secrets.get(key)
+				} catch (error) {
+					logger.error(
+						`Error refreshing global secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
+					)
+				}
+			}),
+		]
 		await Promise.all(promises)
 	}
 
 	private getAllSecretState(): SecretState {
-		return Object.fromEntries(SECRET_STATE_KEYS.map((key) => [key, this.getSecret(key)]))
+		return Object.fromEntries([
+			...SECRET_STATE_KEYS.map((key) => [key, this.getSecret(key as SecretStateKey)]),
+			...GLOBAL_SECRET_KEYS.map((key) => [key, this.getSecret(key as SecretStateKey)]),
+		])
 	}
 
 	/**
@@ -232,18 +305,24 @@ export class ContextProxy {
 	 * RooCodeSettings
 	 */
 
-	public setValue<K extends RooCodeSettingsKey>(key: K, value: RooCodeSettings[K]) {
-		return isSecretStateKey(key) ? this.storeSecret(key, value as string) : this.updateGlobalState(key, value)
+	public async setValue<K extends RooCodeSettingsKey>(key: K, value: RooCodeSettings[K]) {
+		return isSecretStateKey(key)
+			? this.storeSecret(key as SecretStateKey, value as string)
+			: this.updateGlobalState(key as GlobalStateKey, value)
 	}
 
 	public getValue<K extends RooCodeSettingsKey>(key: K): RooCodeSettings[K] {
 		return isSecretStateKey(key)
-			? (this.getSecret(key) as RooCodeSettings[K])
-			: (this.getGlobalState(key) as RooCodeSettings[K])
+			? (this.getSecret(key as SecretStateKey) as RooCodeSettings[K])
+			: (this.getGlobalState(key as GlobalStateKey) as RooCodeSettings[K])
 	}
 
 	public getValues(): RooCodeSettings {
-		return { ...this.getAllGlobalState(), ...this.getAllSecretState() }
+		const globalState = this.getAllGlobalState()
+		const secretState = this.getAllSecretState()
+
+		// Simply merge all states - no nested secrets to handle
+		return { ...globalState, ...secretState }
 	}
 
 	public async setValues(values: RooCodeSettings) {
@@ -285,6 +364,7 @@ export class ContextProxy {
 		await Promise.all([
 			...GLOBAL_STATE_KEYS.map((key) => this.originalContext.globalState.update(key, undefined)),
 			...SECRET_STATE_KEYS.map((key) => this.originalContext.secrets.delete(key)),
+			...GLOBAL_SECRET_KEYS.map((key) => this.originalContext.secrets.delete(key)),
 		])
 
 		await this.initialize()

+ 15 - 6
src/core/config/__tests__/ContextProxy.spec.ts

@@ -2,7 +2,7 @@
 
 import * as vscode from "vscode"
 
-import { GLOBAL_STATE_KEYS, SECRET_STATE_KEYS } from "@roo-code/types"
+import { GLOBAL_STATE_KEYS, SECRET_STATE_KEYS, GLOBAL_SECRET_KEYS } from "@roo-code/types"
 
 import { ContextProxy } from "../ContextProxy"
 
@@ -70,17 +70,23 @@ describe("ContextProxy", () => {
 
 	describe("constructor", () => {
 		it("should initialize state cache with all global state keys", () => {
-			expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length)
+			// +1 for the migration check of old nested settings
+			expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length + 1)
 			for (const key of GLOBAL_STATE_KEYS) {
 				expect(mockGlobalState.get).toHaveBeenCalledWith(key)
 			}
+			// Also check for migration call
+			expect(mockGlobalState.get).toHaveBeenCalledWith("openRouterImageGenerationSettings")
 		})
 
 		it("should initialize secret cache with all secret keys", () => {
-			expect(mockSecrets.get).toHaveBeenCalledTimes(SECRET_STATE_KEYS.length)
+			expect(mockSecrets.get).toHaveBeenCalledTimes(SECRET_STATE_KEYS.length + GLOBAL_SECRET_KEYS.length)
 			for (const key of SECRET_STATE_KEYS) {
 				expect(mockSecrets.get).toHaveBeenCalledWith(key)
 			}
+			for (const key of GLOBAL_SECRET_KEYS) {
+				expect(mockSecrets.get).toHaveBeenCalledWith(key)
+			}
 		})
 	})
 
@@ -93,8 +99,8 @@ describe("ContextProxy", () => {
 			const result = proxy.getGlobalState("apiProvider")
 			expect(result).toBe("deepseek")
 
-			// Original context should be called once during updateGlobalState
-			expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length) // Only from initialization
+			// Original context should be called once during updateGlobalState (+1 for migration check)
+			expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length + 1) // From initialization + migration check
 		})
 
 		it("should handle default values correctly", async () => {
@@ -403,9 +409,12 @@ describe("ContextProxy", () => {
 			for (const key of SECRET_STATE_KEYS) {
 				expect(mockSecrets.delete).toHaveBeenCalledWith(key)
 			}
+			for (const key of GLOBAL_SECRET_KEYS) {
+				expect(mockSecrets.delete).toHaveBeenCalledWith(key)
+			}
 
 			// Total calls should equal the number of secret keys
-			expect(mockSecrets.delete).toHaveBeenCalledTimes(SECRET_STATE_KEYS.length)
+			expect(mockSecrets.delete).toHaveBeenCalledTimes(SECRET_STATE_KEYS.length + GLOBAL_SECRET_KEYS.length)
 		})
 
 		it("should reinitialize caches after reset", async () => {

+ 2 - 6
src/core/tools/__tests__/generateImageTool.test.ts

@@ -46,12 +46,8 @@ describe("generateImageTool", () => {
 						experiments: {
 							[EXPERIMENT_IDS.IMAGE_GENERATION]: true,
 						},
-						apiConfiguration: {
-							openRouterImageGenerationSettings: {
-								openRouterApiKey: "test-api-key",
-								selectedModel: "google/gemini-2.5-flash-image-preview",
-							},
-						},
+						openRouterImageApiKey: "test-api-key",
+						openRouterImageGenerationSelectedModel: "google/gemini-2.5-flash-image-preview",
 					}),
 				}),
 			},

+ 3 - 5
src/core/tools/generateImageTool.ts

@@ -129,10 +129,8 @@ export async function generateImageTool(
 	// Check if file is write-protected
 	const isWriteProtected = cline.rooProtectedController?.isWriteProtected(relPath) || false
 
-	// Get OpenRouter API key from experimental settings ONLY (no fallback to profile)
-	const apiConfiguration = state?.apiConfiguration
-	const imageGenerationSettings = apiConfiguration?.openRouterImageGenerationSettings
-	const openRouterApiKey = imageGenerationSettings?.openRouterApiKey
+	// Get OpenRouter API key from global settings (experimental image generation)
+	const openRouterApiKey = state?.openRouterImageApiKey
 
 	if (!openRouterApiKey) {
 		await cline.say(
@@ -148,7 +146,7 @@ export async function generateImageTool(
 	}
 
 	// Get selected model from settings or use default
-	const selectedModel = imageGenerationSettings?.selectedModel || IMAGE_GENERATION_MODELS[0]
+	const selectedModel = state?.openRouterImageGenerationSelectedModel || IMAGE_GENERATION_MODELS[0]
 
 	// Determine if the path is outside the workspace
 	const fullPath = path.resolve(cline.cwd, removeClosingTag("path", relPath))

+ 7 - 0
src/core/webview/ClineProvider.ts

@@ -1762,6 +1762,8 @@ export class ClineProvider
 			maxDiagnosticMessages,
 			includeTaskHistoryInEnhance,
 			remoteControlEnabled,
+			openRouterImageApiKey,
+			openRouterImageGenerationSelectedModel,
 		} = await this.getState()
 
 		const telemetryKey = process.env.POSTHOG_API_KEY
@@ -1893,6 +1895,8 @@ export class ClineProvider
 			maxDiagnosticMessages: maxDiagnosticMessages ?? 50,
 			includeTaskHistoryInEnhance: includeTaskHistoryInEnhance ?? true,
 			remoteControlEnabled,
+			openRouterImageApiKey,
+			openRouterImageGenerationSelectedModel,
 		}
 	}
 
@@ -2092,6 +2096,9 @@ export class ClineProvider
 					return false
 				}
 			})(),
+			// Add image generation settings
+			openRouterImageApiKey: stateValues.openRouterImageApiKey,
+			openRouterImageGenerationSelectedModel: stateValues.openRouterImageGenerationSelectedModel,
 		}
 	}
 

+ 2 - 0
src/core/webview/__tests__/ClineProvider.spec.ts

@@ -546,6 +546,8 @@ describe("ClineProvider", () => {
 			profileThresholds: {},
 			hasOpenedModeSelector: false,
 			diagnosticsEnabled: true,
+			openRouterImageApiKey: undefined,
+			openRouterImageGenerationSelectedModel: undefined,
 		}
 
 		const message: ExtensionMessage = {

+ 8 - 0
src/core/webview/webviewMessageHandler.ts

@@ -1313,6 +1313,14 @@ export const webviewMessageHandler = async (
 			await updateGlobalState("language", message.text as Language)
 			await provider.postStateToWebview()
 			break
+		case "openRouterImageApiKey":
+			await provider.contextProxy.setValue("openRouterImageApiKey", message.text)
+			await provider.postStateToWebview()
+			break
+		case "openRouterImageGenerationSelectedModel":
+			await provider.contextProxy.setValue("openRouterImageGenerationSelectedModel", message.text)
+			await provider.postStateToWebview()
+			break
 		case "showRooIgnoredFiles":
 			await updateGlobalState("showRooIgnoredFiles", message.bool ?? false)
 			await provider.postStateToWebview()

+ 2 - 0
src/shared/ExtensionMessage.ts

@@ -273,6 +273,7 @@ export type ExtensionState = Pick<
 	| "includeDiagnosticMessages"
 	| "maxDiagnosticMessages"
 	| "remoteControlEnabled"
+	| "openRouterImageGenerationSelectedModel"
 > & {
 	version: string
 	clineMessages: ClineMessage[]
@@ -326,6 +327,7 @@ export type ExtensionState = Pick<
 	marketplaceInstalledMetadata?: { project: Record<string, any>; global: Record<string, any> }
 	profileThresholds: Record<string, number>
 	hasOpenedModeSelector: boolean
+	openRouterImageApiKey?: string
 }
 
 export interface ClineSayTool {

+ 4 - 0
src/shared/WebviewMessage.ts

@@ -212,6 +212,9 @@ export interface WebviewMessage {
 		| "createCommand"
 		| "insertTextIntoTextarea"
 		| "showMdmAuthRequiredNotification"
+		| "imageGenerationSettings"
+		| "openRouterImageApiKey"
+		| "openRouterImageGenerationSelectedModel"
 	text?: string
 	editedMessageContent?: string
 	tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "account"
@@ -248,6 +251,7 @@ export interface WebviewMessage {
 	messageTs?: number
 	historyPreviewCollapsed?: boolean
 	filters?: { type?: string; search?: string; tags?: string[] }
+	settings?: any
 	url?: string // For openExternal
 	mpItem?: MarketplaceItem
 	mpInstallOptions?: InstallMarketplaceItemOptions

+ 4 - 2
src/shared/checkExistApiConfig.ts

@@ -1,4 +1,4 @@
-import { SECRET_STATE_KEYS, ProviderSettings } from "@roo-code/types"
+import { SECRET_STATE_KEYS, GLOBAL_SECRET_KEYS, ProviderSettings } from "@roo-code/types"
 
 export function checkExistKey(config: ProviderSettings | undefined) {
 	if (!config) {
@@ -14,7 +14,9 @@ export function checkExistKey(config: ProviderSettings | undefined) {
 	}
 
 	// Check all secret keys from the centralized SECRET_STATE_KEYS array.
-	const hasSecretKey = SECRET_STATE_KEYS.some((key) => config[key] !== undefined)
+	// Filter out keys that are not part of ProviderSettings (global secrets are stored separately)
+	const providerSecretKeys = SECRET_STATE_KEYS.filter((key) => !GLOBAL_SECRET_KEYS.includes(key as any))
+	const hasSecretKey = providerSecretKeys.some((key) => config[key as keyof ProviderSettings] !== undefined)
 
 	// Check additional non-secret configuration properties
 	const hasOtherConfig = [

+ 17 - 3
webview-ui/src/components/settings/ExperimentalSettings.tsx

@@ -19,6 +19,10 @@ type ExperimentalSettingsProps = HTMLAttributes<HTMLDivElement> & {
 	setExperimentEnabled: SetExperimentEnabled
 	apiConfiguration?: any
 	setApiConfigurationField?: any
+	openRouterImageApiKey?: string
+	openRouterImageGenerationSelectedModel?: string
+	setOpenRouterImageApiKey?: (apiKey: string) => void
+	setImageGenerationSelectedModel?: (model: string) => void
 }
 
 export const ExperimentalSettings = ({
@@ -26,6 +30,10 @@ export const ExperimentalSettings = ({
 	setExperimentEnabled,
 	apiConfiguration,
 	setApiConfigurationField,
+	openRouterImageApiKey,
+	openRouterImageGenerationSelectedModel,
+	setOpenRouterImageApiKey,
+	setImageGenerationSelectedModel,
 	className,
 	...props
 }: ExperimentalSettingsProps) => {
@@ -56,7 +64,11 @@ export const ExperimentalSettings = ({
 								/>
 							)
 						}
-						if (config[0] === "IMAGE_GENERATION" && apiConfiguration && setApiConfigurationField) {
+						if (
+							config[0] === "IMAGE_GENERATION" &&
+							setOpenRouterImageApiKey &&
+							setImageGenerationSelectedModel
+						) {
 							return (
 								<ImageGenerationSettings
 									key={config[0]}
@@ -64,8 +76,10 @@ export const ExperimentalSettings = ({
 									onChange={(enabled) =>
 										setExperimentEnabled(EXPERIMENT_IDS.IMAGE_GENERATION, enabled)
 									}
-									apiConfiguration={apiConfiguration}
-									setApiConfigurationField={setApiConfigurationField}
+									openRouterImageApiKey={openRouterImageApiKey}
+									openRouterImageGenerationSelectedModel={openRouterImageGenerationSelectedModel}
+									setOpenRouterImageApiKey={setOpenRouterImageApiKey}
+									setImageGenerationSelectedModel={setImageGenerationSelectedModel}
 								/>
 							)
 						}

+ 20 - 32
webview-ui/src/components/settings/ImageGenerationSettings.tsx

@@ -1,17 +1,14 @@
 import React, { useState, useEffect } from "react"
 import { VSCodeCheckbox, VSCodeTextField, VSCodeDropdown, VSCodeOption } from "@vscode/webview-ui-toolkit/react"
 import { useAppTranslation } from "@/i18n/TranslationContext"
-import type { ProviderSettings } from "@roo-code/types"
 
 interface ImageGenerationSettingsProps {
 	enabled: boolean
 	onChange: (enabled: boolean) => void
-	apiConfiguration: ProviderSettings
-	setApiConfigurationField: <K extends keyof ProviderSettings>(
-		field: K,
-		value: ProviderSettings[K],
-		isUserAction?: boolean,
-	) => void
+	openRouterImageApiKey?: string
+	openRouterImageGenerationSelectedModel?: string
+	setOpenRouterImageApiKey: (apiKey: string) => void
+	setImageGenerationSelectedModel: (model: string) => void
 }
 
 // Hardcoded list of image generation models
@@ -24,43 +21,34 @@ const IMAGE_GENERATION_MODELS = [
 export const ImageGenerationSettings = ({
 	enabled,
 	onChange,
-	apiConfiguration,
-	setApiConfigurationField,
+	openRouterImageApiKey,
+	openRouterImageGenerationSelectedModel,
+	setOpenRouterImageApiKey,
+	setImageGenerationSelectedModel,
 }: ImageGenerationSettingsProps) => {
 	const { t } = useAppTranslation()
 
-	// Get image generation settings from apiConfiguration
-	const imageGenerationSettings = apiConfiguration?.openRouterImageGenerationSettings || {}
-	const [openRouterApiKey, setOpenRouterApiKey] = useState(imageGenerationSettings.openRouterApiKey || "")
+	const [apiKey, setApiKey] = useState(openRouterImageApiKey || "")
 	const [selectedModel, setSelectedModel] = useState(
-		imageGenerationSettings.selectedModel || IMAGE_GENERATION_MODELS[0].value,
+		openRouterImageGenerationSelectedModel || IMAGE_GENERATION_MODELS[0].value,
 	)
 
-	// Update local state when apiConfiguration changes (e.g., when switching profiles)
+	// Update local state when props change (e.g., when switching profiles)
 	useEffect(() => {
-		setOpenRouterApiKey(imageGenerationSettings.openRouterApiKey || "")
-		setSelectedModel(imageGenerationSettings.selectedModel || IMAGE_GENERATION_MODELS[0].value)
-	}, [imageGenerationSettings.openRouterApiKey, imageGenerationSettings.selectedModel])
-
-	// Helper function to update settings
-	const updateSettings = (newApiKey: string, newModel: string) => {
-		const newSettings = {
-			openRouterApiKey: newApiKey,
-			selectedModel: newModel,
-		}
-		setApiConfigurationField("openRouterImageGenerationSettings", newSettings, true)
-	}
+		setApiKey(openRouterImageApiKey || "")
+		setSelectedModel(openRouterImageGenerationSelectedModel || IMAGE_GENERATION_MODELS[0].value)
+	}, [openRouterImageApiKey, openRouterImageGenerationSelectedModel])
 
 	// Handle API key changes
 	const handleApiKeyChange = (value: string) => {
-		setOpenRouterApiKey(value)
-		updateSettings(value, selectedModel)
+		setApiKey(value)
+		setOpenRouterImageApiKey(value)
 	}
 
 	// Handle model selection changes
 	const handleModelChange = (value: string) => {
 		setSelectedModel(value)
-		updateSettings(openRouterApiKey, value)
+		setImageGenerationSelectedModel(value)
 	}
 
 	return (
@@ -84,7 +72,7 @@ export const ImageGenerationSettings = ({
 							{t("settings:experimental.IMAGE_GENERATION.openRouterApiKeyLabel")}
 						</label>
 						<VSCodeTextField
-							value={openRouterApiKey}
+							value={apiKey}
 							onInput={(e: any) => handleApiKeyChange(e.target.value)}
 							placeholder={t("settings:experimental.IMAGE_GENERATION.openRouterApiKeyPlaceholder")}
 							className="w-full"
@@ -123,13 +111,13 @@ export const ImageGenerationSettings = ({
 					</div>
 
 					{/* Status Message */}
-					{enabled && !openRouterApiKey && (
+					{enabled && !apiKey && (
 						<div className="p-2 bg-vscode-editorWarning-background text-vscode-editorWarning-foreground rounded text-sm">
 							{t("settings:experimental.IMAGE_GENERATION.warningMissingKey")}
 						</div>
 					)}
 
-					{enabled && openRouterApiKey && (
+					{enabled && apiKey && (
 						<div className="p-2 bg-vscode-editorInfo-background text-vscode-editorInfo-foreground rounded text-sm">
 							{t("settings:experimental.IMAGE_GENERATION.successConfigured")}
 						</div>

+ 27 - 0
webview-ui/src/components/settings/SettingsView.tsx

@@ -181,6 +181,8 @@ const SettingsView = forwardRef<SettingsViewRef, SettingsViewProps>(({ onDone, t
 		includeDiagnosticMessages,
 		maxDiagnosticMessages,
 		includeTaskHistoryInEnhance,
+		openRouterImageApiKey,
+		openRouterImageGenerationSelectedModel,
 	} = cachedState
 
 	const apiConfiguration = useMemo(() => cachedState.apiConfiguration ?? {}, [cachedState.apiConfiguration])
@@ -260,6 +262,20 @@ const SettingsView = forwardRef<SettingsViewRef, SettingsViewProps>(({ onDone, t
 		})
 	}, [])
 
+	const setOpenRouterImageApiKey = useCallback((apiKey: string) => {
+		setCachedState((prevState) => {
+			setChangeDetected(true)
+			return { ...prevState, openRouterImageApiKey: apiKey }
+		})
+	}, [])
+
+	const setImageGenerationSelectedModel = useCallback((model: string) => {
+		setCachedState((prevState) => {
+			setChangeDetected(true)
+			return { ...prevState, openRouterImageGenerationSelectedModel: model }
+		})
+	}, [])
+
 	const setCustomSupportPromptsField = useCallback((prompts: Record<string, string | undefined>) => {
 		setCachedState((prevState) => {
 			if (JSON.stringify(prevState.customSupportPrompts) === JSON.stringify(prompts)) {
@@ -343,6 +359,11 @@ const SettingsView = forwardRef<SettingsViewRef, SettingsViewProps>(({ onDone, t
 			vscode.postMessage({ type: "upsertApiConfiguration", text: currentApiConfigName, apiConfiguration })
 			vscode.postMessage({ type: "telemetrySetting", text: telemetrySetting })
 			vscode.postMessage({ type: "profileThresholds", values: profileThresholds })
+			vscode.postMessage({ type: "openRouterImageApiKey", text: openRouterImageApiKey })
+			vscode.postMessage({
+				type: "openRouterImageGenerationSelectedModel",
+				text: openRouterImageGenerationSelectedModel,
+			})
 			setChangeDetected(false)
 		}
 	}
@@ -723,6 +744,12 @@ const SettingsView = forwardRef<SettingsViewRef, SettingsViewProps>(({ onDone, t
 							experiments={experiments}
 							apiConfiguration={apiConfiguration}
 							setApiConfigurationField={setApiConfigurationField}
+							openRouterImageApiKey={openRouterImageApiKey as string | undefined}
+							openRouterImageGenerationSelectedModel={
+								openRouterImageGenerationSelectedModel as string | undefined
+							}
+							setOpenRouterImageApiKey={setOpenRouterImageApiKey}
+							setImageGenerationSelectedModel={setImageGenerationSelectedModel}
 						/>
 					)}
 

+ 24 - 29
webview-ui/src/components/settings/__tests__/ImageGenerationSettings.spec.tsx

@@ -1,7 +1,5 @@
 import { render, fireEvent } from "@testing-library/react"
 
-import type { ProviderSettings } from "@roo-code/types"
-
 import { ImageGenerationSettings } from "../ImageGenerationSettings"
 
 // Mock the translation context
@@ -12,14 +10,17 @@ vi.mock("@/i18n/TranslationContext", () => ({
 }))
 
 describe("ImageGenerationSettings", () => {
-	const mockSetApiConfigurationField = vi.fn()
+	const mockSetOpenRouterImageApiKey = vi.fn()
+	const mockSetImageGenerationSelectedModel = vi.fn()
 	const mockOnChange = vi.fn()
 
 	const defaultProps = {
 		enabled: false,
 		onChange: mockOnChange,
-		apiConfiguration: {} as ProviderSettings,
-		setApiConfigurationField: mockSetApiConfigurationField,
+		openRouterImageApiKey: undefined,
+		openRouterImageGenerationSelectedModel: undefined,
+		setOpenRouterImageApiKey: mockSetOpenRouterImageApiKey,
+		setImageGenerationSelectedModel: mockSetImageGenerationSelectedModel,
 	}
 
 	beforeEach(() => {
@@ -27,30 +28,31 @@ describe("ImageGenerationSettings", () => {
 	})
 
 	describe("Initial Mount Behavior", () => {
-		it("should not call setApiConfigurationField on initial mount with empty configuration", () => {
+		it("should not call setter functions on initial mount with empty configuration", () => {
 			render(<ImageGenerationSettings {...defaultProps} />)
 
-			// Should NOT call setApiConfigurationField on initial mount to prevent dirty state
-			expect(mockSetApiConfigurationField).not.toHaveBeenCalled()
+			// Should NOT call setter functions on initial mount to prevent dirty state
+			expect(mockSetOpenRouterImageApiKey).not.toHaveBeenCalled()
+			expect(mockSetImageGenerationSelectedModel).not.toHaveBeenCalled()
 		})
 
-		it("should not call setApiConfigurationField on initial mount with existing configuration", () => {
-			const apiConfiguration = {
-				openRouterImageGenerationSettings: {
-					openRouterApiKey: "existing-key",
-					selectedModel: "google/gemini-2.5-flash-image-preview:free",
-				},
-			} as ProviderSettings
-
-			render(<ImageGenerationSettings {...defaultProps} apiConfiguration={apiConfiguration} />)
+		it("should not call setter functions on initial mount with existing configuration", () => {
+			render(
+				<ImageGenerationSettings
+					{...defaultProps}
+					openRouterImageApiKey="existing-key"
+					openRouterImageGenerationSelectedModel="google/gemini-2.5-flash-image-preview:free"
+				/>,
+			)
 
-			// Should NOT call setApiConfigurationField on initial mount to prevent dirty state
-			expect(mockSetApiConfigurationField).not.toHaveBeenCalled()
+			// Should NOT call setter functions on initial mount to prevent dirty state
+			expect(mockSetOpenRouterImageApiKey).not.toHaveBeenCalled()
+			expect(mockSetImageGenerationSelectedModel).not.toHaveBeenCalled()
 		})
 	})
 
 	describe("User Interaction Behavior", () => {
-		it("should call setApiConfigurationField when user changes API key", async () => {
+		it("should call setimageGenerationSettings when user changes API key", async () => {
 			const { getByPlaceholderText } = render(<ImageGenerationSettings {...defaultProps} enabled={true} />)
 
 			const apiKeyInput = getByPlaceholderText(
@@ -60,15 +62,8 @@ describe("ImageGenerationSettings", () => {
 			// Simulate user typing
 			fireEvent.input(apiKeyInput, { target: { value: "new-api-key" } })
 
-			// Should call setApiConfigurationField with isUserAction=true
-			expect(mockSetApiConfigurationField).toHaveBeenCalledWith(
-				"openRouterImageGenerationSettings",
-				{
-					openRouterApiKey: "new-api-key",
-					selectedModel: "google/gemini-2.5-flash-image-preview",
-				},
-				true, // This should be true for user actions
-			)
+			// Should call setimageGenerationSettings
+			expect(defaultProps.setOpenRouterImageApiKey).toHaveBeenCalledWith("new-api-key")
 		})
 
 		// Note: Testing VSCode dropdown components is complex due to their custom nature

+ 2 - 0
webview-ui/src/context/ExtensionStateContext.tsx

@@ -252,6 +252,8 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 		alwaysAllowUpdateTodoList: true,
 		includeDiagnosticMessages: true,
 		maxDiagnosticMessages: 50,
+		openRouterImageApiKey: "",
+		openRouterImageGenerationSelectedModel: "",
 	})
 
 	const [didHydrateState, setDidHydrateState] = useState(false)