Skip to content

Commit c3d84d2

Browse files
daniel-lxsdaniel-lxs
authored
refactor: flatten image generation settings structure (#7536)
1 parent 01458f1 commit c3d84d2

File tree

17 files changed

+260
-116
lines changed

17 files changed

+260
-116
lines changed

‎packages/types/src/global-settings.ts‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@ export const globalSettingsSchema = z.object({
4242
customInstructions: z.string().optional(),
4343
taskHistory: z.array(historyItemSchema).optional(),
4444

45+
// Image generation settings (experimental) - flattened for simplicity
46+
openRouterImageApiKey: z.string().optional(),
47+
openRouterImageGenerationSelectedModel: z.string().optional(),
48+
4549
condensingApiConfigId: z.string().optional(),
4650
customCondensingPrompt: z.string().optional(),
4751

@@ -200,11 +204,24 @@ export const SECRET_STATE_KEYS = [
200204
"featherlessApiKey",
201205
"ioIntelligenceApiKey",
202206
"vercelAiGatewayApiKey",
203-
] as const satisfies readonly (keyof ProviderSettings)[]
204-
export type SecretState = Pick<ProviderSettings, (typeof SECRET_STATE_KEYS)[number]>
207+
] as const
208+
209+
// Global secrets that are part of GlobalSettings (not ProviderSettings)
210+
export const GLOBAL_SECRET_KEYS = [
211+
"openRouterImageApiKey", // For image generation
212+
] as const
213+
214+
// Type for the actual secret storage keys
215+
type ProviderSecretKey = (typeof SECRET_STATE_KEYS)[number]
216+
type GlobalSecretKey = (typeof GLOBAL_SECRET_KEYS)[number]
217+
218+
// Type representing all secrets that can be stored
219+
export type SecretState = Pick<ProviderSettings, Extract<ProviderSecretKey, keyof ProviderSettings>> & {
220+
[K in GlobalSecretKey]?: string
221+
}
205222

206223
export const isSecretStateKey = (key: string): key is Keys<SecretState> =>
207-
SECRET_STATE_KEYS.includes(key as Keys<SecretState>)
224+
SECRET_STATE_KEYS.includes(key as ProviderSecretKey) || GLOBAL_SECRET_KEYS.includes(key as GlobalSecretKey)
208225

209226
/**
210227
* GlobalState
@@ -213,7 +230,7 @@ export const isSecretStateKey = (key: string): key is Keys<SecretState> =>
213230
export type GlobalState = Omit<RooCodeSettings, Keys<SecretState>>
214231

215232
export const GLOBAL_STATE_KEYS = [...GLOBAL_SETTINGS_KEYS, ...PROVIDER_SETTINGS_KEYS].filter(
216-
(key: Keys<RooCodeSettings>) => !SECRET_STATE_KEYS.includes(key as Keys<SecretState>),
233+
(key: Keys<RooCodeSettings>) => !isSecretStateKey(key),
217234
) as Keys<GlobalState>[]
218235

219236
export const isGlobalStateKey = (key: string): key is Keys<GlobalState> =>

‎packages/types/src/provider-settings.ts‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -142,13 +142,6 @@ const openRouterSchema = baseProviderSettingsSchema.extend({
142142
openRouterBaseUrl: z.string().optional(),
143143
openRouterSpecificProvider: z.string().optional(),
144144
openRouterUseMiddleOutTransform: z.boolean().optional(),
145-
// Image generation settings (experimental)
146-
openRouterImageGenerationSettings: z
147-
.object({
148-
openRouterApiKey: z.string().optional(),
149-
selectedModel: z.string().optional(),
150-
})
151-
.optional(),
152145
})
153146

154147
const bedrockSchema = apiModelIdProviderModelSchema.extend({

‎src/core/config/ContextProxy.ts‎

Lines changed: 102 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
GLOBAL_SETTINGS_KEYS,
77
SECRET_STATE_KEYS,
88
GLOBAL_STATE_KEYS,
9+
GLOBAL_SECRET_KEYS,
910
type ProviderSettings,
1011
type GlobalSettings,
1112
type SecretState,
@@ -61,19 +62,77 @@ export class ContextProxy {
6162
}
6263
}
6364

64-
const promises = SECRET_STATE_KEYS.map(async (key) => {
65-
try {
66-
this.secretCache[key] = await this.originalContext.secrets.get(key)
67-
} catch (error) {
68-
logger.error(`Error loading secret ${key}: ${error instanceof Error ? error.message : String(error)}`)
69-
}
70-
})
65+
const promises = [
66+
...SECRET_STATE_KEYS.map(async (key) => {
67+
try {
68+
this.secretCache[key] = await this.originalContext.secrets.get(key)
69+
} catch (error) {
70+
logger.error(
71+
`Error loading secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
72+
)
73+
}
74+
}),
75+
...GLOBAL_SECRET_KEYS.map(async (key) => {
76+
try {
77+
this.secretCache[key] = await this.originalContext.secrets.get(key)
78+
} catch (error) {
79+
logger.error(
80+
`Error loading global secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
81+
)
82+
}
83+
}),
84+
]
7185

7286
await Promise.all(promises)
7387

88+
// Migration: Check for old nested image generation settings and migrate them
89+
await this.migrateImageGenerationSettings()
90+
7491
this._isInitialized = true
7592
}
7693

94+
/**
95+
* Migrates old nested openRouterImageGenerationSettings to the new flattened structure
96+
*/
97+
private async migrateImageGenerationSettings() {
98+
try {
99+
// Check if there's an old nested structure
100+
const oldNestedSettings = this.originalContext.globalState.get<any>("openRouterImageGenerationSettings")
101+
102+
if (oldNestedSettings && typeof oldNestedSettings === "object") {
103+
logger.info("Migrating old nested image generation settings to flattened structure")
104+
105+
// Migrate the API key if it exists and we don't already have one
106+
if (oldNestedSettings.openRouterApiKey && !this.secretCache.openRouterImageApiKey) {
107+
await this.originalContext.secrets.store(
108+
"openRouterImageApiKey",
109+
oldNestedSettings.openRouterApiKey,
110+
)
111+
this.secretCache.openRouterImageApiKey = oldNestedSettings.openRouterApiKey
112+
logger.info("Migrated openRouterImageApiKey to secrets")
113+
}
114+
115+
// Migrate the selected model if it exists and we don't already have one
116+
if (oldNestedSettings.selectedModel && !this.stateCache.openRouterImageGenerationSelectedModel) {
117+
await this.originalContext.globalState.update(
118+
"openRouterImageGenerationSelectedModel",
119+
oldNestedSettings.selectedModel,
120+
)
121+
this.stateCache.openRouterImageGenerationSelectedModel = oldNestedSettings.selectedModel
122+
logger.info("Migrated openRouterImageGenerationSelectedModel to global state")
123+
}
124+
125+
// Clean up the old nested structure
126+
await this.originalContext.globalState.update("openRouterImageGenerationSettings", undefined)
127+
logger.info("Removed old nested openRouterImageGenerationSettings")
128+
}
129+
} catch (error) {
130+
logger.error(
131+
`Error during image generation settings migration: ${error instanceof Error ? error.message : String(error)}`,
132+
)
133+
}
134+
}
135+
77136
public get extensionUri() {
78137
return this.originalContext.extensionUri
79138
}
@@ -152,20 +211,34 @@ export class ContextProxy {
152211
* This is useful when you need to ensure the cache has the latest values
153212
*/
154213
async refreshSecrets(): Promise<void> {
155-
const promises = SECRET_STATE_KEYS.map(async (key) => {
156-
try {
157-
this.secretCache[key] = await this.originalContext.secrets.get(key)
158-
} catch (error) {
159-
logger.error(
160-
`Error refreshing secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
161-
)
162-
}
163-
})
214+
const promises = [
215+
...SECRET_STATE_KEYS.map(async (key) => {
216+
try {
217+
this.secretCache[key] = await this.originalContext.secrets.get(key)
218+
} catch (error) {
219+
logger.error(
220+
`Error refreshing secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
221+
)
222+
}
223+
}),
224+
...GLOBAL_SECRET_KEYS.map(async (key) => {
225+
try {
226+
this.secretCache[key] = await this.originalContext.secrets.get(key)
227+
} catch (error) {
228+
logger.error(
229+
`Error refreshing global secret ${key}: ${error instanceof Error ? error.message : String(error)}`,
230+
)
231+
}
232+
}),
233+
]
164234
await Promise.all(promises)
165235
}
166236

167237
private getAllSecretState(): SecretState {
168-
return Object.fromEntries(SECRET_STATE_KEYS.map((key) => [key, this.getSecret(key)]))
238+
return Object.fromEntries([
239+
...SECRET_STATE_KEYS.map((key) => [key, this.getSecret(key as SecretStateKey)]),
240+
...GLOBAL_SECRET_KEYS.map((key) => [key, this.getSecret(key as SecretStateKey)]),
241+
])
169242
}
170243

171244
/**
@@ -232,18 +305,24 @@ export class ContextProxy {
232305
* RooCodeSettings
233306
*/
234307

235-
public setValue<K extends RooCodeSettingsKey>(key: K, value: RooCodeSettings[K]) {
236-
return isSecretStateKey(key) ? this.storeSecret(key, value as string) : this.updateGlobalState(key, value)
308+
public async setValue<K extends RooCodeSettingsKey>(key: K, value: RooCodeSettings[K]) {
309+
return isSecretStateKey(key)
310+
? this.storeSecret(key as SecretStateKey, value as string)
311+
: this.updateGlobalState(key as GlobalStateKey, value)
237312
}
238313

239314
public getValue<K extends RooCodeSettingsKey>(key: K): RooCodeSettings[K] {
240315
return isSecretStateKey(key)
241-
? (this.getSecret(key) as RooCodeSettings[K])
242-
: (this.getGlobalState(key) as RooCodeSettings[K])
316+
? (this.getSecret(key as SecretStateKey) as RooCodeSettings[K])
317+
: (this.getGlobalState(key as GlobalStateKey) as RooCodeSettings[K])
243318
}
244319

245320
public getValues(): RooCodeSettings {
246-
return { ...this.getAllGlobalState(), ...this.getAllSecretState() }
321+
const globalState = this.getAllGlobalState()
322+
const secretState = this.getAllSecretState()
323+
324+
// Simply merge all states - no nested secrets to handle
325+
return { ...globalState, ...secretState }
247326
}
248327

249328
public async setValues(values: RooCodeSettings) {
@@ -285,6 +364,7 @@ export class ContextProxy {
285364
await Promise.all([
286365
...GLOBAL_STATE_KEYS.map((key) => this.originalContext.globalState.update(key, undefined)),
287366
...SECRET_STATE_KEYS.map((key) => this.originalContext.secrets.delete(key)),
367+
...GLOBAL_SECRET_KEYS.map((key) => this.originalContext.secrets.delete(key)),
288368
])
289369

290370
await this.initialize()

‎src/core/config/__tests__/ContextProxy.spec.ts‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import * as vscode from "vscode"
44

5-
import { GLOBAL_STATE_KEYS, SECRET_STATE_KEYS } from "@roo-code/types"
5+
import { GLOBAL_STATE_KEYS, SECRET_STATE_KEYS, GLOBAL_SECRET_KEYS } from "@roo-code/types"
66

77
import { ContextProxy } from "../ContextProxy"
88

@@ -70,17 +70,23 @@ describe("ContextProxy", () => {
7070

7171
describe("constructor", () => {
7272
it("should initialize state cache with all global state keys", () => {
73-
expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length)
73+
// +1 for the migration check of old nested settings
74+
expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length + 1)
7475
for (const key of GLOBAL_STATE_KEYS) {
7576
expect(mockGlobalState.get).toHaveBeenCalledWith(key)
7677
}
78+
// Also check for migration call
79+
expect(mockGlobalState.get).toHaveBeenCalledWith("openRouterImageGenerationSettings")
7780
})
7881

7982
it("should initialize secret cache with all secret keys", () => {
80-
expect(mockSecrets.get).toHaveBeenCalledTimes(SECRET_STATE_KEYS.length)
83+
expect(mockSecrets.get).toHaveBeenCalledTimes(SECRET_STATE_KEYS.length + GLOBAL_SECRET_KEYS.length)
8184
for (const key of SECRET_STATE_KEYS) {
8285
expect(mockSecrets.get).toHaveBeenCalledWith(key)
8386
}
87+
for (const key of GLOBAL_SECRET_KEYS) {
88+
expect(mockSecrets.get).toHaveBeenCalledWith(key)
89+
}
8490
})
8591
})
8692

@@ -93,8 +99,8 @@ describe("ContextProxy", () => {
9399
const result = proxy.getGlobalState("apiProvider")
94100
expect(result).toBe("deepseek")
95101

96-
// Original context should be called once during updateGlobalState
97-
expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length) // Only from initialization
102+
// Original context should be called once during updateGlobalState (+1 for migration check)
103+
expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length + 1) // From initialization + migration check
98104
})
99105

100106
it("should handle default values correctly", async () => {
@@ -403,9 +409,12 @@ describe("ContextProxy", () => {
403409
for (const key of SECRET_STATE_KEYS) {
404410
expect(mockSecrets.delete).toHaveBeenCalledWith(key)
405411
}
412+
for (const key of GLOBAL_SECRET_KEYS) {
413+
expect(mockSecrets.delete).toHaveBeenCalledWith(key)
414+
}
406415

407416
// Total calls should equal the number of secret keys
408-
expect(mockSecrets.delete).toHaveBeenCalledTimes(SECRET_STATE_KEYS.length)
417+
expect(mockSecrets.delete).toHaveBeenCalledTimes(SECRET_STATE_KEYS.length + GLOBAL_SECRET_KEYS.length)
409418
})
410419

411420
it("should reinitialize caches after reset", async () => {

‎src/core/tools/__tests__/generateImageTool.test.ts‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,12 +46,8 @@ describe("generateImageTool", () => {
4646
experiments: {
4747
[EXPERIMENT_IDS.IMAGE_GENERATION]: true,
4848
},
49-
apiConfiguration: {
50-
openRouterImageGenerationSettings: {
51-
openRouterApiKey: "test-api-key",
52-
selectedModel: "google/gemini-2.5-flash-image-preview",
53-
},
54-
},
49+
openRouterImageApiKey: "test-api-key",
50+
openRouterImageGenerationSelectedModel: "google/gemini-2.5-flash-image-preview",
5551
}),
5652
}),
5753
},

‎src/core/tools/generateImageTool.ts‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -129,10 +129,8 @@ export async function generateImageTool(
129129
// Check if file is write-protected
130130
const isWriteProtected = cline.rooProtectedController?.isWriteProtected(relPath) || false
131131

132-
// Get OpenRouter API key from experimental settings ONLY (no fallback to profile)
133-
const apiConfiguration = state?.apiConfiguration
134-
const imageGenerationSettings = apiConfiguration?.openRouterImageGenerationSettings
135-
const openRouterApiKey = imageGenerationSettings?.openRouterApiKey
132+
// Get OpenRouter API key from global settings (experimental image generation)
133+
const openRouterApiKey = state?.openRouterImageApiKey
136134

137135
if (!openRouterApiKey) {
138136
await cline.say(
@@ -148,7 +146,7 @@ export async function generateImageTool(
148146
}
149147

150148
// Get selected model from settings or use default
151-
const selectedModel = imageGenerationSettings?.selectedModel || IMAGE_GENERATION_MODELS[0]
149+
const selectedModel = state?.openRouterImageGenerationSelectedModel || IMAGE_GENERATION_MODELS[0]
152150

153151
// Determine if the path is outside the workspace
154152
const fullPath = path.resolve(cline.cwd, removeClosingTag("path", relPath))

‎src/core/webview/ClineProvider.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1762,6 +1762,8 @@ export class ClineProvider
17621762
maxDiagnosticMessages,
17631763
includeTaskHistoryInEnhance,
17641764
remoteControlEnabled,
1765+
openRouterImageApiKey,
1766+
openRouterImageGenerationSelectedModel,
17651767
} = await this.getState()
17661768

17671769
const telemetryKey = process.env.POSTHOG_API_KEY
@@ -1893,6 +1895,8 @@ export class ClineProvider
18931895
maxDiagnosticMessages: maxDiagnosticMessages ?? 50,
18941896
includeTaskHistoryInEnhance: includeTaskHistoryInEnhance ?? true,
18951897
remoteControlEnabled,
1898+
openRouterImageApiKey,
1899+
openRouterImageGenerationSelectedModel,
18961900
}
18971901
}
18981902

@@ -2092,6 +2096,9 @@ export class ClineProvider
20922096
return false
20932097
}
20942098
})(),
2099+
// Add image generation settings
2100+
openRouterImageApiKey: stateValues.openRouterImageApiKey,
2101+
openRouterImageGenerationSelectedModel: stateValues.openRouterImageGenerationSelectedModel,
20952102
}
20962103
}
20972104

‎src/core/webview/__tests__/ClineProvider.spec.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -546,6 +546,8 @@ describe("ClineProvider", () => {
546546
profileThresholds: {},
547547
hasOpenedModeSelector: false,
548548
diagnosticsEnabled: true,
549+
openRouterImageApiKey: undefined,
550+
openRouterImageGenerationSelectedModel: undefined,
549551
}
550552

551553
const message: ExtensionMessage = {

‎src/core/webview/webviewMessageHandler.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1313,6 +1313,14 @@ export const webviewMessageHandler = async (
13131313
await updateGlobalState("language", message.text as Language)
13141314
await provider.postStateToWebview()
13151315
break
1316+
case "openRouterImageApiKey":
1317+
await provider.contextProxy.setValue("openRouterImageApiKey", message.text)
1318+
await provider.postStateToWebview()
1319+
break
1320+
case "openRouterImageGenerationSelectedModel":
1321+
await provider.contextProxy.setValue("openRouterImageGenerationSelectedModel", message.text)
1322+
await provider.postStateToWebview()
1323+
break
13161324
case "showRooIgnoredFiles":
13171325
await updateGlobalState("showRooIgnoredFiles", message.bool ?? false)
13181326
await provider.postStateToWebview()

0 commit comments

Comments
 (0)