@@ -9,9 +9,8 @@ import { getReadablePath } from "../../utils/path"
99import { isPathOutsideWorkspace } from "../../utils/pathUtils"
1010import { EXPERIMENT_IDS , experiments } from "../../shared/experiments"
1111import { OpenRouterHandler } from "../../api/providers/openrouter"
12-
13- // Hardcoded list of image generation models for now
14- const IMAGE_GENERATION_MODELS = [ "google/gemini-2.5-flash-image-preview" , "google/gemini-2.5-flash-image-preview:free" ]
12+ import { GeminiHandler } from "../../api/providers/gemini"
13+ import { ImageGenerationProvider , getImageGenerationModelsForProvider } from "@roo-code/types"
1514
1615export async function generateImageTool (
1716 cline : Task ,
@@ -128,25 +127,60 @@ export async function generateImageTool(
128127 // Check if file is write-protected
129128 const isWriteProtected = cline . rooProtectedController ?. isWriteProtected ( relPath ) || false
130129
131- // Get OpenRouter API key from global settings (experimental image generation )
132- const openRouterApiKey = state ?. openRouterImageApiKey
130+ // Get the selected provider from settings (default to openrouter )
131+ const selectedProvider = ( state ?. imageGenerationProvider || "openrouter" ) as ImageGenerationProvider
133132
134- if ( ! openRouterApiKey ) {
135- await cline . say (
136- "error" ,
137- "OpenRouter API key is required for image generation. Please configure it in the Image Generation experimental settings." ,
138- )
139- pushToolResult (
140- formatResponse . toolError (
133+ // Get selected model from settings based on provider
134+ let selectedModel : string
135+ let apiKey : string | undefined
136+
137+ if ( selectedProvider === "openrouter" ) {
138+ apiKey = state ?. openRouterImageApiKey
139+ if ( ! apiKey ) {
140+ await cline . say (
141+ "error" ,
141142 "OpenRouter API key is required for image generation. Please configure it in the Image Generation experimental settings." ,
142- ) ,
143- )
143+ )
144+ pushToolResult (
145+ formatResponse . toolError (
146+ "OpenRouter API key is required for image generation. Please configure it in the Image Generation experimental settings." ,
147+ ) ,
148+ )
149+ return
150+ }
151+ // Get selected model or use default for OpenRouter
152+ const models = getImageGenerationModelsForProvider ( "openrouter" )
153+ selectedModel =
154+ state ?. openRouterImageGenerationSelectedModel ||
155+ ( models [ 0 ] ?. modelId ?? "google/gemini-2.5-flash-image-preview" )
156+ } else if ( selectedProvider === "gemini" ) {
157+ // For Gemini, we can use the existing Gemini API key from the provider settings
158+ // Check for a dedicated image generation API key first, then fall back to the provider's API key
159+ apiKey =
160+ state ?. geminiImageApiKey ||
161+ ( state ?. apiConfiguration ?. apiProvider === "gemini" ? state ?. apiConfiguration ?. geminiApiKey : undefined )
162+ if ( ! apiKey ) {
163+ await cline . say (
164+ "error" ,
165+ "Gemini API key is required for image generation. Please configure it in the Image Generation experimental settings or in the Gemini provider settings." ,
166+ )
167+ pushToolResult (
168+ formatResponse . toolError (
169+ "Gemini API key is required for image generation. Please configure it in the Image Generation experimental settings or in the Gemini provider settings." ,
170+ ) ,
171+ )
172+ return
173+ }
174+ // Get selected model or use default for Gemini
175+ const models = getImageGenerationModelsForProvider ( "gemini" )
176+ selectedModel =
177+ state ?. geminiImageGenerationSelectedModel || ( models [ 0 ] ?. modelId ?? "gemini-2.5-flash-image-preview" )
178+ } else {
179+ await cline . say ( "error" , `Unsupported image generation provider: ${ selectedProvider } ` )
180+ pushToolResult ( formatResponse . toolError ( `Unsupported image generation provider: ${ selectedProvider } ` ) )
144181 return
145182 }
146183
147- // Get selected model from settings or use default
148- const selectedModel = state ?. openRouterImageGenerationSelectedModel || IMAGE_GENERATION_MODELS [ 0 ]
149-
150184 // Determine if the path is outside the workspace
151185 const fullPath = path . resolve ( cline . cwd , removeClosingTag ( "path" , relPath ) )
152186 const isOutsideWorkspace = isPathOutsideWorkspace ( fullPath )
@@ -176,16 +210,28 @@ export async function generateImageTool(
176210 return
177211 }
178212
179- // Create a temporary OpenRouter handler with minimal options
180- const openRouterHandler = new OpenRouterHandler ( { } as any )
181-
182- // Call the generateImage method with the explicit API key and optional input image
183- const result = await openRouterHandler . generateImage (
184- prompt ,
185- selectedModel ,
186- openRouterApiKey ,
187- inputImageData ,
188- )
213+ // Generate image based on provider
214+ let result
215+
216+ if ( selectedProvider === "openrouter" ) {
217+ // Create a temporary OpenRouter handler with minimal options
218+ const openRouterHandler = new OpenRouterHandler ( { } as any )
219+
220+ // Call the generateImage method with the explicit API key and optional input image
221+ result = await openRouterHandler . generateImage ( prompt , selectedModel , apiKey ! , inputImageData )
222+ } else if ( selectedProvider === "gemini" ) {
223+ // Create a temporary Gemini handler with minimal options
224+ const geminiHandler = new GeminiHandler ( { geminiApiKey : apiKey } as any )
225+
226+ // Call the generateImage method with the optional input image
227+ result = await geminiHandler . generateImage ( prompt , selectedModel , apiKey , inputImageData )
228+ } else {
229+ // This should not happen due to earlier check, but for type safety
230+ result = {
231+ success : false ,
232+ error : `Unsupported provider: ${ selectedProvider } ` ,
233+ }
234+ }
189235
190236 if ( ! result . success ) {
191237 await cline . say ( "error" , result . error || "Failed to generate image" )
0 commit comments