diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index a0dd59198e1..74d71ed63e0 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -795,6 +795,7 @@ import type { Config } from '@qwen-code/qwen-code-core'; import type { LoadedSettings } from '../config/settings.js'; import type { CliArgs } from '../config/config.js'; import { + AuthType, SessionEndReason, MCPServerConfig, SessionService, @@ -4021,6 +4022,14 @@ describe('QwenAgent MCP SSE/HTTP support', () => { baseUrl: 'https://api.example.com', isRuntimeModel: false, }, + { + id: 'qwen-image-2.0', + label: 'Qwen Image 2.0', + authType: 'qwen', + baseUrl: 'https://api.example.com', + imageOnly: true, + isRuntimeModel: false, + }, ]), getToolRegistry: vi .fn() @@ -4064,6 +4073,26 @@ describe('QwenAgent MCP SSE/HTTP support', () => { expect( preflight.cells.find((c) => c.kind === 'tool_registry')?.status, ).toBe('ok'); + expect(preflight.cells.find((c) => c.kind === 'providers')).toMatchObject({ + status: 'ok', + detail: { count: 1, providers: ['qwen'] }, + }); + + vi.mocked(mockConfig.getAllConfiguredModels).mockReturnValue([ + { + id: 'qwen-image-2.0', + label: 'Qwen Image 2.0', + authType: AuthType.QWEN_OAUTH, + imageOnly: true, + }, + ]); + const imageOnlyPreflight = (await agent.extMethod( + SERVE_STATUS_EXT_METHODS.workspacePreflight, + {}, + )) as { cells: Array<{ kind: string; status: string }> }; + expect( + imageOnlyPreflight.cells.find((c) => c.kind === 'providers')?.status, + ).toBe('error'); mockConnectionState.resolve(); await agentPromise; diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index fb353894abe..40b9ac636f8 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -5855,7 +5855,9 @@ class QwenAgent implements Agent { private buildProvidersPreflightCell(config: Config): ServePreflightCell { try { - const models = config.getAllConfiguredModels(); + const models = config + .getAllConfiguredModels() + .filter((model) => !model.imageOnly); const authType = config.getAuthType?.(); if (models.length === 0) { // `authType` set but zero models = the next `POST /session` will diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index d78fcce1c77..2a817d8602d 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -1327,6 +1327,24 @@ describe('loadCliConfig', () => { ); }); + it('should propagate the image model selection', async () => { + process.argv = ['node', 'script.js']; + const argv = await parseArguments(); + + await loadCliConfig( + { + imageModel: 'openai:qwen-image-2.0\0https://images.example.com/api/v1', + }, + argv, + ); + + expect(mockConfigConstructorParams).toHaveBeenCalledWith( + expect.objectContaining({ + imageModel: 'openai:qwen-image-2.0\0https://images.example.com/api/v1', + }), + ); + }); + it('places session-injected (ACP/IDE) MCP servers at the top precedence tier', async () => { process.argv = ['node', 'script.js']; const argv = await parseArguments(); diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 98fe9284444..b4dcbab5022 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -2277,6 +2277,7 @@ export async function loadCliConfig( webSearch: bareMode || safeMode ? undefined : resolveWebSearchSettings(settings), visionModel: settings.visionModel || undefined, + imageModel: settings.imageModel || undefined, visionBridgeTimeoutMs: settings.visionBridgeTimeoutMs, modelFallbacks: resolveModelFallbacks( argv.fallbackModel, diff --git a/packages/cli/src/config/settingsSchema.test.ts b/packages/cli/src/config/settingsSchema.test.ts index 916841d9d1a..acc70778c2f 100644 --- a/packages/cli/src/config/settingsSchema.test.ts +++ b/packages/cli/src/config/settingsSchema.test.ts @@ -35,6 +35,7 @@ describe('SettingsSchema', () => { 'advanced', 'plansDirectory', 'voiceModel', + 'imageModel', ]; expectedSettings.forEach((setting) => { @@ -179,6 +180,16 @@ describe('SettingsSchema', () => { expect(voiceModel.showInDialog).toBe(false); }); + it('should define the image model setting', () => { + const imageModel = getSettingsSchema().imageModel; + + expect(imageModel.type).toBe('string'); + expect(imageModel.category).toBe('Model'); + expect(imageModel.default).toBe(''); + expect(imageModel.requiresRestart).toBe(false); + expect(imageModel.showInDialog).toBe(false); + }); + it('should define the built-in Explore model setting', () => { const exploreModel = getSettingsSchema().agents.properties.builtin.properties.exploreModel; diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 110f25df167..800a7170ee4 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -1309,6 +1309,17 @@ const SETTINGS_SCHEMA = { showInDialog: true, }, + imageModel: { + type: 'string', + label: 'Image Model', + category: 'Model', + requiresRestart: false, + default: '', + description: + 'Model used by the built-in image_gen tool. Set with /model --image. The selected model must be marked imageOnly in modelProviders.', + showInDialog: false, + }, + visionBridgeTimeoutMs: { type: 'integer', label: 'Vision Bridge Timeout (ms)', diff --git a/packages/cli/src/i18n/locales/en.js b/packages/cli/src/i18n/locales/en.js index bd1238cf5de..4e7c40213a2 100644 --- a/packages/cli/src/i18n/locales/en.js +++ b/packages/cli/src/i18n/locales/en.js @@ -219,6 +219,7 @@ export default { 'toolDisplayName.ExitWorktree': 'toolDisplayName.ExitWorktree', 'toolDisplayName.Workflow': 'toolDisplayName.Workflow', 'toolDisplayName.ReadMcpResource': 'toolDisplayName.ReadMcpResource', + 'toolDisplayName.ImageGen': 'toolDisplayName.ImageGen', // ============================================================================ // Help / UI Components // ============================================================================ @@ -1490,6 +1491,8 @@ export default { 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, [model-id] to switch immediately).', 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).': 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).', + 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --image for the image generation model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).': + 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --image for the image generation model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).', "Inline one-shot override isn't supported in this mode — run '/model {{model}}' first, then send your prompt.": "Inline one-shot override isn't supported in this mode — run '/model {{model}}' first, then send your prompt.", "Inline one-shot override can't switch providers. '{{model}}' belongs to a different provider — run '/model {{model}}' first, then send your prompt.": @@ -1503,17 +1506,23 @@ export default { 'Set the model for voice transcription', 'Set the image-capable model used to transcribe images for a text-only main model': 'Set the image-capable model used to transcribe images for a text-only main model', + 'Set the model used to generate images': + 'Set the model used to generate images', 'Persist the model selection to the project settings (workspace scope)': 'Persist the model selection to the project settings (workspace scope)', 'Persist the model selection to the user settings (global scope)': 'Persist the model selection to the user settings (global scope)', 'Select Fast Model': 'Select Fast Model', 'Select Vision Model': 'Select Vision Model', + 'Select Image Model': 'Select Image Model', 'Select Voice Model': 'Select Voice Model', 'Vision Model': 'Vision Model', + 'Image Model': 'Image Model', 'Voice Model': 'Voice Model', 'Selected voice model is unavailable.': 'Selected voice model is unavailable.', + 'Selected image model is unavailable.': + 'Selected image model is unavailable.', "Voice model '{{model}}' is configured more than once. Remove duplicate model ids before selecting it for voice transcription.": "Voice model '{{model}}' is configured more than once. Remove duplicate model ids before selecting it for voice transcription.", 'Voice dictation: {{status}} (mode: {{mode}}, {{modelText}}).': @@ -1761,8 +1770,16 @@ export default { 'Current voice model: {{voiceModel}}\nUse "/model --voice " to set voice model.', 'Current vision model: {{visionModel}}\nUse "/model --vision " to set the vision bridge model.': 'Current vision model: {{visionModel}}\nUse "/model --vision " to set the vision bridge model.', + 'Current image model: {{imageModel}}\nUse "/model --image " to set the image generation model.': + 'Current image model: {{imageModel}}\nUse "/model --image " to set the image generation model.', "Voice model '{{modelName}}' is ambiguous. Configure a unique model id before using /model --voice.": "Voice model '{{modelName}}' is ambiguous. Configure a unique model id before using /model --voice.", + "Image model '{{modelName}}' matches multiple configured endpoints. Run /model --image without an argument and choose the exact endpoint.": + "Image model '{{modelName}}' matches multiple configured endpoints. Run /model --image without an argument and choose the exact endpoint.", + "Image model '{{modelName}}' must declare a valid HTTPS baseUrl and credential environment variable.": + "Image model '{{modelName}}' must declare a valid HTTPS baseUrl and credential environment variable.", + "'{{model}}' must declare a valid HTTPS baseUrl and credential environment variable.": + "'{{model}}' must declare a valid HTTPS baseUrl and credential environment variable.", none: 'none', unknown: 'unknown', // ============================================================================ diff --git a/packages/cli/src/i18n/locales/zh-TW.js b/packages/cli/src/i18n/locales/zh-TW.js index f7915e16a5c..8053dca514e 100644 --- a/packages/cli/src/i18n/locales/zh-TW.js +++ b/packages/cli/src/i18n/locales/zh-TW.js @@ -210,6 +210,7 @@ export default { 'toolDisplayName.ExitWorktree': '退出 Worktree', 'toolDisplayName.Workflow': '工作流程', 'toolDisplayName.ReadMcpResource': '讀取 MCP 資源', + 'toolDisplayName.ImageGen': '圖像生成', '↑ to manage attachments': '↑ 管理附件', '← → select, Delete to remove, ↓ to exit': '← → 選擇，Delete 刪除，↓ 退出', @@ -1316,6 +1317,8 @@ export default { '切換此會話的模型（--fast 可設置建議模型，--voice 可設置語音轉寫模型，[model-id] 可立即切換）', 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).': '切換此會話的模型（--fast 建議模型，--voice 語音轉寫模型，--vision 視覺橋接模型，--project 持久化到專案設定，--global 持久化到使用者設定，[model-id] 立即切換，或用 [model-id] [prompt] 在另一個模型上執行一次性提示；內聯提示按原文發送，不展開 @file）', + 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --image for the image generation model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).': + '切換此會話的模型（--fast 建議模型，--voice 語音轉寫模型，--vision 視覺橋接模型，--image 圖像生成模型，--project 持久化到專案設定，--global 持久化到使用者設定，[model-id] 立即切換，或用 [model-id] [prompt] 在另一個模型上執行一次性提示；內聯提示按原文發送，不展開 @file）', "Inline one-shot override isn't supported in this mode — run '/model {{model}}' first, then send your prompt.": "此模式不支援內聯一次性覆寫——請先執行 '/model {{model}}'，再發送你的提示。", "Inline one-shot override can't switch providers. '{{model}}' belongs to a different provider — run '/model {{model}}' first, then send your prompt.": @@ -1328,16 +1331,20 @@ export default { 'Set the model for voice transcription': '設定語音轉寫模型', 'Set the image-capable model used to transcribe images for a text-only main model': '設定用於為純文字主模型轉寫圖像的圖像能力模型', + 'Set the model used to generate images': '設置用於生成圖像的模型', 'Persist the model selection to the project settings (workspace scope)': '將模型選擇持久化到專案設定（工作區）', 'Persist the model selection to the user settings (global scope)': '將模型選擇持久化到使用者設定（全域）', 'Select Fast Model': '選擇快速模型', 'Select Vision Model': '選擇視覺模型', + 'Select Image Model': '選擇圖像模型', 'Select Voice Model': '選擇語音模型', 'Vision Model': '視覺模型', + 'Image Model': '圖像模型', 'Voice Model': '語音模型', 'Selected voice model is unavailable.': '所選語音模型不可用。', + 'Selected image model is unavailable.': '所選圖像模型不可用。', "Voice model '{{model}}' is configured more than once. Remove duplicate model ids before selecting it for voice transcription.": "語音模型 '{{model}}' 被配置了多次。請先移除重複的模型 ID，再將其選為語音轉寫模型。", 'Voice dictation: {{status}} (mode: {{mode}}, {{modelText}}).': @@ -1543,8 +1550,16 @@ export default { '當前語音模型：{{voiceModel}}\n使用 "/model --voice " 設置語音模型。', 'Current vision model: {{visionModel}}\nUse "/model --vision " to set the vision bridge model.': '當前視覺模型：{{visionModel}}\n使用 "/model --vision " 設置視覺橋接模型。', + 'Current image model: {{imageModel}}\nUse "/model --image " to set the image generation model.': + '當前圖像模型：{{imageModel}}\n使用 "/model --image " 設置圖像生成模型。', "Voice model '{{modelName}}' is ambiguous. Configure a unique model id before using /model --voice.": "語音模型 '{{modelName}}' 不唯一。請先配置唯一的模型 ID，再使用 /model --voice。", + "Image model '{{modelName}}' matches multiple configured endpoints. Run /model --image without an argument and choose the exact endpoint.": + "圖像模型 '{{modelName}}' 匹配了多個已配置的端點。請執行 /model --image（不帶參數）並選擇確切的端點。", + "Image model '{{modelName}}' must declare a valid HTTPS baseUrl and credential environment variable.": + "圖像模型 '{{modelName}}' 必須宣告有效的 HTTPS baseUrl 和憑據環境變數。", + "'{{model}}' must declare a valid HTTPS baseUrl and credential environment variable.": + "'{{model}}' 必須宣告有效的 HTTPS baseUrl 和憑據環境變數。", none: '無', unknown: '未知', 'Manage folder trust settings': '管理檔案夾信任設置', diff --git a/packages/cli/src/i18n/locales/zh.js b/packages/cli/src/i18n/locales/zh.js index b97be1aa15a..5e43b0c4ea7 100644 --- a/packages/cli/src/i18n/locales/zh.js +++ b/packages/cli/src/i18n/locales/zh.js @@ -211,6 +211,7 @@ export default { 'toolDisplayName.ExitWorktree': '退出 Worktree', 'toolDisplayName.Workflow': '工作流', 'toolDisplayName.ReadMcpResource': '读取 MCP 资源', + 'toolDisplayName.ImageGen': '图像生成', // ============================================================================ // Help / UI Components // ============================================================================ @@ -1430,6 +1431,8 @@ export default { '切换此会话的模型（--fast 可设置建议模型，--voice 可设置语音转写模型，[model-id] 可立即切换）', 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).': '切换此会话的模型（--fast 建议模型，--voice 语音转写模型，--vision 视觉桥接模型，--project 持久化到项目设置，--global 持久化到用户设置，[model-id] 立即切换，或用 [model-id] [prompt] 在另一个模型上运行一次性提示；内联提示按原文发送，不展开 @file）', + 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --image for the image generation model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).': + '切换此会话的模型（--fast 建议模型，--voice 语音转写模型，--vision 视觉桥接模型，--image 图像生成模型，--project 持久化到项目设置，--global 持久化到用户设置，[model-id] 立即切换，或用 [model-id] [prompt] 在另一个模型上运行一次性提示；内联提示按原文发送，不展开 @file）', "Inline one-shot override isn't supported in this mode — run '/model {{model}}' first, then send your prompt.": "此模式不支持内联一次性覆盖——请先运行 '/model {{model}}'，再发送你的提示。", "Inline one-shot override can't switch providers. '{{model}}' belongs to a different provider — run '/model {{model}}' first, then send your prompt.": @@ -1442,16 +1445,20 @@ export default { 'Set the model for voice transcription': '设置语音转写模型', 'Set the image-capable model used to transcribe images for a text-only main model': '设置用于为纯文本主模型转写图像的图像能力模型', + 'Set the model used to generate images': '设置用于生成图像的模型', 'Persist the model selection to the project settings (workspace scope)': '将模型选择持久化到项目设置（工作区）', 'Persist the model selection to the user settings (global scope)': '将模型选择持久化到用户设置（全局）', 'Select Fast Model': '选择快速模型', 'Select Vision Model': '选择视觉模型', + 'Select Image Model': '选择图像模型', 'Select Voice Model': '选择语音模型', 'Vision Model': '视觉模型', + 'Image Model': '图像模型', 'Voice Model': '语音模型', 'Selected voice model is unavailable.': '所选语音模型不可用。', + 'Selected image model is unavailable.': '所选图像模型不可用。', "Voice model '{{model}}' is configured more than once. Remove duplicate model ids before selecting it for voice transcription.": "语音模型 '{{model}}' 被配置了多次。请先移除重复的模型 ID，再将其选为语音转写模型。", 'Voice dictation: {{status}} (mode: {{mode}}, {{modelText}}).': @@ -1685,8 +1692,16 @@ export default { '当前语音模型：{{voiceModel}}\n使用 "/model --voice " 设置语音模型。', 'Current vision model: {{visionModel}}\nUse "/model --vision " to set the vision bridge model.': '当前视觉模型：{{visionModel}}\n使用 "/model --vision " 设置视觉桥接模型。', + 'Current image model: {{imageModel}}\nUse "/model --image " to set the image generation model.': + '当前图像模型：{{imageModel}}\n使用 "/model --image " 设置图像生成模型。', "Voice model '{{modelName}}' is ambiguous. Configure a unique model id before using /model --voice.": "语音模型 '{{modelName}}' 不唯一。请先配置唯一的模型 ID，再使用 /model --voice。", + "Image model '{{modelName}}' matches multiple configured endpoints. Run /model --image without an argument and choose the exact endpoint.": + "图像模型 '{{modelName}}' 匹配了多个已配置的端点。请运行 /model --image（不带参数）并选择确切的端点。", + "Image model '{{modelName}}' must declare a valid HTTPS baseUrl and credential environment variable.": + "图像模型 '{{modelName}}' 必须声明有效的 HTTPS baseUrl 和凭据环境变量。", + "'{{model}}' must declare a valid HTTPS baseUrl and credential environment variable.": + "'{{model}}' 必须声明有效的 HTTPS baseUrl 和凭据环境变量。", none: '无', unknown: '未知', // ============================================================================ diff --git a/packages/cli/src/nonInteractive/control/controllers/systemController.test.ts b/packages/cli/src/nonInteractive/control/controllers/systemController.test.ts index d6dbaa0ac76..e96a018ffbc 100644 --- a/packages/cli/src/nonInteractive/control/controllers/systemController.test.ts +++ b/packages/cli/src/nonInteractive/control/controllers/systemController.test.ts @@ -345,6 +345,11 @@ describe('SystemController', () => { baseUrl: 'https://internal-proxy.corp/v1', envKey: 'SECRET_API_KEY', }, + { + id: 'qwen-image-2.0', + label: 'Qwen Image 2.0', + imageOnly: true, + }, ]); const controller = new SystemController( context, diff --git a/packages/cli/src/nonInteractive/control/controllers/systemController.ts b/packages/cli/src/nonInteractive/control/controllers/systemController.ts index 5f83dd99da0..7b3994d78a1 100644 --- a/packages/cli/src/nonInteractive/control/controllers/systemController.ts +++ b/packages/cli/src/nonInteractive/control/controllers/systemController.ts @@ -574,6 +574,7 @@ export class SystemController extends BaseController { try { const models = this.context.config .getAvailableModels() + .filter((model) => !model.imageOnly) .map(({ id, label, capabilities, contextWindowSize }) => ({ id, label, diff --git a/packages/cli/src/services/voice-model.ts b/packages/cli/src/services/voice-model.ts index 16a336d4b9f..087dfaabac0 100644 --- a/packages/cli/src/services/voice-model.ts +++ b/packages/cli/src/services/voice-model.ts @@ -37,6 +37,7 @@ export function isTranscribableVoiceModel(model: AvailableModel): boolean { return ( model.authType === AuthType.USE_OPENAI && model.isRuntimeModel !== true && + model.imageOnly !== true && typeof model.baseUrl === 'string' && model.baseUrl.trim().length > 0 ); diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 454c8e905bc..b50c8923282 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -1314,6 +1314,7 @@ export const AppContainer = (props: AppContainerProps) => { isFastModelMode, isVoiceModelMode, isVisionModelMode, + isImageModelMode, modelDialogPersistScope, openModelDialog, closeModelDialog, @@ -4066,6 +4067,7 @@ export const AppContainer = (props: AppContainerProps) => { isFastModelMode, isVoiceModelMode, isVisionModelMode, + isImageModelMode, modelDialogPersistScope, isTrustDialogOpen, activeArenaDialog, @@ -4209,6 +4211,7 @@ export const AppContainer = (props: AppContainerProps) => { isFastModelMode, isVoiceModelMode, isVisionModelMode, + isImageModelMode, modelDialogPersistScope, isTrustDialogOpen, activeArenaDialog, diff --git a/packages/cli/src/ui/commands/arenaCommand.test.ts b/packages/cli/src/ui/commands/arenaCommand.test.ts index 17eca36aef4..194032443df 100644 --- a/packages/cli/src/ui/commands/arenaCommand.test.ts +++ b/packages/cli/src/ui/commands/arenaCommand.test.ts @@ -86,6 +86,42 @@ describe('arenaCommand localization', () => { }); }); +describe('arenaCommand start subcommand', () => { + it('rejects image-only models passed explicitly', async () => { + const context = createMockCommandContext({ + executionMode: 'interactive', + services: { + config: { + getArenaManager: vi.fn(() => null), + getContentGeneratorConfig: vi.fn(() => undefined), + getModelsConfig: vi.fn(() => ({ + getAvailableModelsForAuthType: vi.fn(() => [ + { + id: 'qwen-image-2.0', + label: 'Qwen Image 2.0', + imageOnly: true, + }, + { id: 'qwen-plus', label: 'Qwen Plus' }, + ]), + })), + } as never, + }, + }); + + const result = await getArenaSubCommand('start').action!( + context, + '--models qwen-image-2.0,qwen-plus "build a feature"', + ); + + expect(result).toEqual({ + type: 'message', + messageType: 'error', + content: + "Image-only model 'qwen-image-2.0' cannot be used in an Arena session.", + }); + }); +}); + describe('arenaCommand stop subcommand', () => { let mockContext: CommandContext; let mockConfig: { diff --git a/packages/cli/src/ui/commands/arenaCommand.ts b/packages/cli/src/ui/commands/arenaCommand.ts index 73e577a028b..656e59ee5b1 100644 --- a/packages/cli/src/ui/commands/arenaCommand.ts +++ b/packages/cli/src/ui/commands/arenaCommand.ts @@ -144,17 +144,28 @@ function buildArenaExecutionInput( // Build ArenaModelConfig for each model, resolving display names from // the model registry when available. const modelsConfig = config.getModelsConfig(); - const models: ArenaModelConfig[] = parsed.models.map((parsedModel) => { + const models: ArenaModelConfig[] = []; + for (const parsedModel of parsed.models) { const authType = (parsedModel.authType as AuthType | undefined) ?? defaultAuthType; const registryModels = modelsConfig.getAvailableModelsForAuthType(authType); const resolved = registryModels.find((m) => m.id === parsedModel.modelId); - return { + if (resolved?.imageOnly) { + return { + type: 'message', + messageType: 'error', + content: t( + "Image-only model '{{model}}' cannot be used in an Arena session.", + { model: parsedModel.modelId }, + ), + }; + } + models.push({ modelId: parsedModel.modelId, authType, displayName: resolved?.label ?? parsedModel.modelId, - }; - }); + }); + } return { task: parsed.task, diff --git a/packages/cli/src/ui/commands/modelCommand.test.ts b/packages/cli/src/ui/commands/modelCommand.test.ts index 93ae32fabe2..b54c6c3585b 100644 --- a/packages/cli/src/ui/commands/modelCommand.test.ts +++ b/packages/cli/src/ui/commands/modelCommand.test.ts @@ -46,10 +46,32 @@ describe('modelCommand', () => { it('should have the correct name and description', () => { expect(modelCommand.name).toBe('model'); expect(modelCommand.description).toBe( - 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).', + 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --image for the image generation model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).', ); }); + it('should complete image models across providers', async () => { + mockContext.services.config = { + getAvailableModels: vi.fn().mockReturnValue([ + { + id: 'current-chat-model', + authType: AuthType.QWEN_OAUTH, + }, + ]), + getAllConfiguredModels: vi.fn().mockReturnValue([ + { + id: 'qwen-image-2.0', + authType: AuthType.USE_OPENAI, + imageOnly: true, + }, + ]), + } as unknown as Config; + + const result = await modelCommand.completion!(mockContext, '--image q'); + + expect(result).toEqual(['qwen-image-2.0']); + }); + it('should return error when config is not available', async () => { mockContext.services.config = null; @@ -1305,6 +1327,131 @@ describe('modelCommand', () => { }); }); + it('should open the image model dialog for /model --image', async () => { + const mockConfig = createMockConfig({ + model: 'qwen-plus', + authType: AuthType.USE_OPENAI, + }); + mockContext.services.config = mockConfig as Config; + + const result = await modelCommand.action!(mockContext, '--image'); + + expect(result).toEqual({ + type: 'dialog', + dialog: 'image-model', + }); + }); + + it('should return the current image model outside interactive mode', async () => { + mockContext = createMockCommandContext({ + executionMode: 'non_interactive', + invocation: { args: '--image' }, + services: { + config: createMockConfig({ + model: 'qwen-max', + authType: AuthType.USE_OPENAI, + }), + settings: { + merged: { + imageModel: + 'openai:qwen-image-2.0\0https://images.example.com/api/v1', + } as Record, + }, + }, + }); + + const result = await modelCommand.action!(mockContext, '--image'); + + expect(result).toEqual({ + type: 'message', + messageType: 'info', + content: + 'Current image model: openai:qwen-image-2.0 (https://images.example.com/api/v1)\nUse "/model --image " to set the image generation model.', + }); + }); + + it('should set an imageOnly model and hot-register its tool', async () => { + const setValue = vi.fn(); + const setImageModel = vi.fn().mockResolvedValue(undefined); + const baseUrl = 'https://images.example.com/api/v1'; + mockContext = createMockCommandContext({ + invocation: { + raw: '/model --image qwen-image-2.0', + name: 'model', + args: '--image qwen-image-2.0', + }, + services: { + config: { + getAllConfiguredModels: vi.fn().mockReturnValue([ + { + id: 'qwen-image-2.0', + label: 'Qwen Image 2.0', + authType: AuthType.USE_OPENAI, + baseUrl, + registryBaseUrl: baseUrl, + envKey: 'IMAGE_API_KEY', + imageOnly: true, + }, + ]), + resolveImageGenerationModel: vi.fn().mockReturnValue({ + model: 'qwen-image-2.0', + baseUrl, + apiKeyEnv: 'IMAGE_API_KEY', + }), + setImageModel, + }, + settings: createMockSettings(setValue), + }, + }); + + const result = await modelCommand.action!( + mockContext, + '--image qwen-image-2.0', + ); + + const persisted = `openai:qwen-image-2.0\0${baseUrl}`; + expect(setValue).toHaveBeenCalledWith( + expect.any(String), + 'imageModel', + persisted, + ); + expect(setImageModel).toHaveBeenCalledWith(persisted); + expect(result).toEqual({ + type: 'message', + messageType: 'info', + content: 'Image Model: qwen-image-2.0', + }); + }); + + it('should reject a chat model from /model --image', async () => { + const setValue = vi.fn(); + mockContext = createMockCommandContext({ + services: { + config: { + getAllConfiguredModels: vi.fn().mockReturnValue([ + { + id: 'qwen-plus', + label: 'Qwen Plus', + authType: AuthType.USE_OPENAI, + }, + ]), + }, + settings: createMockSettings(setValue), + }, + }); + + const result = await modelCommand.action!(mockContext, '--image qwen-plus'); + + expect(setValue).not.toHaveBeenCalled(); + expect(result).toEqual({ + type: 'message', + messageType: 'error', + content: expect.stringContaining( + "Image model 'qwen-plus' is not configured", + ), + }); + }); + it('should open the voice model dialog for /model --voice in interactive mode', async () => { const mockConfig = createMockConfig({ model: 'qwen-plus', @@ -1748,7 +1895,7 @@ describe('modelCommand', () => { }); }); - describe('fastOnly/voiceOnly filtering', () => { + describe('selector-only model filtering', () => { it('should reject fastOnly models from normal /model selection', async () => { mockContext = createMockCommandContext({ invocation: { @@ -1809,6 +1956,40 @@ describe('modelCommand', () => { }); }); + it('should reject image-generation-only models from normal /model selection', async () => { + mockContext = createMockCommandContext({ + invocation: { + raw: '/model qwen-image-2.0', + name: 'model', + args: 'qwen-image-2.0', + }, + services: { + config: { + getContentGeneratorConfig: vi.fn().mockReturnValue({ + model: 'main-model', + authType: AuthType.USE_OPENAI, + }), + getAvailableModelsForAuthType: vi.fn().mockReturnValue([ + { id: 'main-model', label: 'Main' }, + { + id: 'qwen-image-2.0', + label: 'Image', + imageOnly: true, + }, + ]), + }, + settings: createMockSettings(), + }, + }); + + const result = await modelCommand.action!(mockContext, 'qwen-image-2.0'); + expect(result).toMatchObject({ + type: 'message', + messageType: 'error', + content: expect.stringContaining('qwen-image-2.0'), + }); + }); + it('should allow fastOnly models in --fast selection', async () => { const setValue = vi.fn(); mockContext = createMockCommandContext({ @@ -2090,6 +2271,16 @@ describe('modelCommand', () => { }); }); + it('should include persistScope for /model --global --image dialog', async () => { + const ctx = setupContext(); + const result = await modelCommand.action!(ctx, '--global --image'); + expect(result).toEqual({ + type: 'dialog', + dialog: 'image-model', + persistScope: 'user', + }); + }); + it('should parse scope flags in any position', async () => { const ctx = setupContext(); const result = await modelCommand.action!(ctx, '--fast --project'); diff --git a/packages/cli/src/ui/commands/modelCommand.ts b/packages/cli/src/ui/commands/modelCommand.ts index 29d9f94da34..5efe5f55136 100644 --- a/packages/cli/src/ui/commands/modelCommand.ts +++ b/packages/cli/src/ui/commands/modelCommand.ts @@ -41,6 +41,9 @@ const FAST_MODEL_CONFIGURATION_HINT = const VISION_MODEL_CONFIGURATION_HINT = 'Configure an image-capable model in settings.modelProviders and ensure the required environment variables are set. Run /model --vision to set it, or leave it unset to auto-pick a same-provider vision model.'; +const IMAGE_MODEL_CONFIGURATION_HINT = + 'Configure a model with imageOnly: true, baseUrl, and envKey in settings.modelProviders. Run /model --image to select it.'; + /** * Parse --project / --global scope flags from the argument string. * Returns the resolved scope override and the remaining args with flags stripped. @@ -143,7 +146,7 @@ async function switchMainModel( } function formatUnavailableModelMessage( - kind: 'Model' | 'Fast model' | 'Vision model', + kind: 'Model' | 'Fast model' | 'Vision model' | 'Image model', modelName: string, authType: AuthType, availableModels: AvailableModel[], @@ -161,7 +164,9 @@ function formatUnavailableModelMessage( ? FAST_MODEL_CONFIGURATION_HINT : kind === 'Vision model' ? VISION_MODEL_CONFIGURATION_HINT - : MAIN_MODEL_CONFIGURATION_HINT; + : kind === 'Image model' + ? IMAGE_MODEL_CONFIGURATION_HINT + : MAIN_MODEL_CONFIGURATION_HINT; return ( `${kind} '${modelName}' is not available for auth type '${authType}'.\n` + @@ -170,10 +175,10 @@ function formatUnavailableModelMessage( ); } -// Fast and vision share the same "not configured for any auth type" message -// shape, differing only in the label and the configuration hint. +// Auxiliary model selectors share the same "not configured for any auth type" +// message shape, differing only in the label and configuration hint. function formatUnavailableAuxModelMessage( - label: 'Fast model' | 'Vision model', + label: 'Fast model' | 'Vision model' | 'Image model', modelName: string, availableModels: AvailableModel[], hint: string, @@ -217,6 +222,18 @@ function formatUnavailableVisionModelMessage( ); } +function formatUnavailableImageModelMessage( + modelName: string, + availableModels: AvailableModel[], +): string { + return formatUnavailableAuxModelMessage( + 'Image model', + modelName, + availableModels, + IMAGE_MODEL_CONFIGURATION_HINT, + ); +} + function formatAmbiguousVisionModelMessage( modelName: string, matchingModels: AvailableModel[], @@ -294,18 +311,24 @@ function formatUnavailableVoiceModelMessage( // Get an array of the available model IDs as strings, filtered by mode function getAvailableModelIds( context: CommandContext, - mode: 'main' | 'fast' | 'voice' | 'vision' = 'main', + mode: 'main' | 'fast' | 'voice' | 'vision' | 'image' = 'main', ) { const { services } = context; const { config } = services; if (!config) { return []; } - const availableModels = config.getAvailableModels().filter((m) => { - if (mode === 'fast') return !m.voiceOnly; - if (mode === 'voice') return !m.fastOnly; - // 'vision' and 'main' both exclude fast/voice-only models. - return !m.fastOnly && !m.voiceOnly; + const models = + mode === 'image' + ? config.getAllConfiguredModels() + : config.getAvailableModels(); + const availableModels = models.filter((m) => { + if (mode === 'image') + return m.imageOnly === true && !m.fastOnly && !m.voiceOnly; + if (mode === 'fast') return !m.voiceOnly && !m.imageOnly; + if (mode === 'voice') return !m.fastOnly && !m.imageOnly; + // 'vision' and 'main' both exclude selector-only models. + return !m.fastOnly && !m.voiceOnly && !m.imageOnly; }); return availableModels.map((model) => model.id); } @@ -315,11 +338,11 @@ export const modelCommand: SlashCommand = { completionPriority: 100, get description() { return t( - 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).', + 'Switch the model for this session (--fast for suggestion model, --voice for voice transcription model, --vision for the vision bridge model, --image for the image generation model, --project to persist to project settings, --global to persist to user settings, [model-id] to switch immediately, or [model-id] [prompt] to run a one-off prompt on another model; the inline prompt is sent verbatim without @file expansion).', ); }, argumentHint: - '[--fast|--voice|--vision] [--project|--global] [] | ', + '[--fast|--voice|--vision|--image] [--project|--global] [] | ', kind: CommandKind.BUILT_IN, supportedModes: ['interactive', 'non_interactive', 'acp'] as const, completion: async (context, partialArg) => { @@ -341,6 +364,10 @@ export const modelCommand: SlashCommand = { 'Set the image-capable model used to transcribe images for a text-only main model', ), }, + { + value: '--image', + description: t('Set the model used to generate images'), + }, { value: '--project', description: t( @@ -359,18 +386,20 @@ export const modelCommand: SlashCommand = { } const trimmed = partialArg.trim(); if (trimmed) { - let mode: 'main' | 'fast' | 'voice' | 'vision' = 'main'; + let mode: 'main' | 'fast' | 'voice' | 'vision' | 'image' = 'main'; // Strip all known flags to isolate the model prefix for completion const modelPrefix = trimmed .replace(/(?:^|\s)--fast(?:\s|$)/, ' ') .replace(/(?:^|\s)--voice(?:\s|$)/, ' ') .replace(/(?:^|\s)--vision(?:\s|$)/, ' ') + .replace(/(?:^|\s)--image(?:\s|$)/, ' ') .replace(/(?:^|\s)--project(?:\s|$)/, ' ') .replace(/(?:^|\s)--global(?:\s|$)/, ' ') .trim(); if (/(?:^|\s)--fast(?:\s|$)/.test(trimmed)) mode = 'fast'; else if (/(?:^|\s)--voice(?:\s|$)/.test(trimmed)) mode = 'voice'; else if (/(?:^|\s)--vision(?:\s|$)/.test(trimmed)) mode = 'vision'; + else if (/(?:^|\s)--image(?:\s|$)/.test(trimmed)) mode = 'image'; return getAvailableModelIds(context, mode).filter((id) => id.startsWith(modelPrefix), ); @@ -469,7 +498,7 @@ export const modelCommand: SlashCommand = { const availableModels = config .getAllConfiguredModels() - .filter((m) => !m.fastOnly); + .filter((m) => !m.fastOnly && !m.imageOnly); const matches = availableModels.filter((model) => model.id === modelName); if (matches.length === 0) { return { @@ -565,7 +594,7 @@ export const modelCommand: SlashCommand = { selector.authType ? config.getAvailableModelsForAuthType(selector.authType) : config.getAllConfiguredModels() - ).filter((m) => !m.voiceOnly); + ).filter((m) => !m.voiceOnly && !m.imageOnly); if (!availableModels.some((model) => model.id === selector.modelId)) { return { type: 'message', @@ -648,7 +677,7 @@ export const modelCommand: SlashCommand = { selector.authType ? config.getAvailableModelsForAuthType(selector.authType) : config.getAllConfiguredModels() - ).filter((m) => !m.fastOnly && !m.voiceOnly); + ).filter((m) => !m.fastOnly && !m.voiceOnly && !m.imageOnly); const matchingModels = availableModels.filter( (model) => model.id === selector.modelId, ); @@ -711,6 +740,119 @@ export const modelCommand: SlashCommand = { }; } + const isImageModelCommand = + args === '--image' || args.startsWith('--image '); + if (isImageModelCommand) { + const modelName = args.replace('--image', '').trim(); + if (!modelName) { + if (context.executionMode !== 'interactive') { + const imageModel = + context.services.settings?.merged?.imageModel?.trim(); + return { + type: 'message', + messageType: 'info', + content: t( + 'Current image model: {{imageModel}}\nUse "/model --image " to set the image generation model.', + { + imageModel: imageModel + ? formatVisionModelSettingForDisplay(imageModel) + : t('not set'), + }, + ), + }; + } + return { + type: 'dialog', + dialog: 'image-model', + ...persistScopeSpread(scopeOverride), + }; + } + if (!settings) { + return { + type: 'message', + messageType: 'error', + content: t('Settings service not available.'), + }; + } + + const selector = (() => { + try { + return resolveModelId(modelName); + } catch { + return undefined; + } + })(); + if (!selector) { + return { + type: 'message', + messageType: 'error', + content: formatUnavailableImageModelMessage(modelName, []), + }; + } + + const availableModels = ( + selector.authType + ? config.getAvailableModelsForAuthType(selector.authType) + : config.getAllConfiguredModels() + ).filter( + (model) => + model.imageOnly === true && !model.fastOnly && !model.voiceOnly, + ); + const matchingModels = availableModels.filter( + (model) => model.id === selector.modelId, + ); + if (matchingModels.length > 1) { + return { + type: 'message', + messageType: 'error', + content: t( + "Image model '{{modelName}}' matches multiple configured endpoints. Run /model --image without an argument and choose the exact endpoint.", + { modelName }, + ), + }; + } + const matched = matchingModels[0]; + if (!matched) { + return { + type: 'message', + messageType: 'error', + content: selector.authType + ? formatUnavailableModelMessage( + 'Image model', + selector.modelId, + selector.authType, + availableModels, + ) + : formatUnavailableImageModelMessage(modelName, availableModels), + }; + } + + const qualifiedModelName = `${ + selector.authType ?? matched.authType + }:${selector.modelId}`; + const imageModel = matched.baseUrl + ? `${qualifiedModelName}\0${matched.baseUrl}` + : qualifiedModelName; + if (!config.resolveImageGenerationModel(imageModel)) { + return { + type: 'message', + messageType: 'error', + content: t( + "Image model '{{modelName}}' must declare a valid HTTPS baseUrl and credential environment variable.", + { modelName }, + ), + }; + } + + persistSetting(settings, 'imageModel', imageModel, scopeOverride); + await config.setImageModel(imageModel); + return { + type: 'message', + messageType: 'info', + content: t('Image Model') + ': ' + modelName + scopeSuffix, + }; + } + const contentGeneratorConfig = config.getContentGeneratorConfig(); if (!contentGeneratorConfig) { return { @@ -743,7 +885,7 @@ export const modelCommand: SlashCommand = { const targetAuthType = parsed.authType ?? authType; const availableModels = config .getAvailableModelsForAuthType(targetAuthType) - .filter((m) => !m.fastOnly && !m.voiceOnly); + .filter((m) => !m.fastOnly && !m.voiceOnly && !m.imageOnly); if (!availableModels.some((model) => model.id === parsed.modelId)) { return { type: 'message', diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index 19714c33743..a41306770f7 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -189,6 +189,7 @@ export interface OpenDialogActionReturn { | 'fast-model' | 'voice-model' | 'vision-model' + | 'image-model' | 'subagent_create' | 'subagent_list' | 'skills_manage' diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index a611cae428b..86d2f127779 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -272,6 +272,7 @@ export const DialogManager = ({ isFastModelMode={uiState.isFastModelMode} isVoiceModelMode={uiState.isVoiceModelMode} isVisionModelMode={uiState.isVisionModelMode} + isImageModelMode={uiState.isImageModelMode} persistScope={uiState.modelDialogPersistScope} availableTerminalHeight={listDialogHeight} /> diff --git a/packages/cli/src/ui/components/ModelDialog.test.tsx b/packages/cli/src/ui/components/ModelDialog.test.tsx index db3d339355c..48ac5c34d3b 100644 --- a/packages/cli/src/ui/components/ModelDialog.test.tsx +++ b/packages/cli/src/ui/components/ModelDialog.test.tsx @@ -771,6 +771,65 @@ describe('', () => { expect(props.onClose).toHaveBeenCalledTimes(1); }); + it('shows only image-generation models and stores the exact provider route', async () => { + const setImageModel = vi.fn().mockResolvedValue(undefined); + const baseUrl = 'https://images.example.com/api/v1'; + const persisted = `openai:qwen-image-2.0\0${baseUrl}`; + const { props, mockSettings, getByText } = renderComponent( + { isImageModelMode: true }, + { + getAuthType: vi.fn(() => AuthType.USE_OPENAI), + getAllConfiguredModels: vi.fn(() => [ + { + id: 'qwen-plus', + label: 'Qwen Plus', + authType: AuthType.USE_OPENAI, + }, + { + id: 'qwen-image-2.0', + label: 'Qwen Image 2.0', + authType: AuthType.USE_OPENAI, + baseUrl, + envKey: 'IMAGE_API_KEY', + imageOnly: true, + }, + { + id: 'image-without-credentials', + label: 'Image without credentials', + authType: AuthType.USE_OPENAI, + baseUrl: 'https://invalid.example.com/api/v1', + imageOnly: true, + }, + ]), + resolveImageGenerationModel: vi.fn((selector: string) => + selector === persisted + ? { + model: 'qwen-image-2.0', + baseUrl, + apiKeyEnv: 'IMAGE_API_KEY', + } + : undefined, + ), + setImageModel, + } as unknown as Partial, + ); + + expect(getByText('Select Image Model')).toBeDefined(); + const selectProps = mockedSelect.mock.calls[0][0]; + expect(selectProps.items).toHaveLength(1); + await selectProps.onSelect( + `${AuthType.USE_OPENAI}::qwen-image-2.0\0${baseUrl}`, + ); + + expect(mockSettings.setValue).toHaveBeenCalledWith( + SettingScope.User, + 'imageModel', + persisted, + ); + expect(setImageModel).toHaveBeenCalledWith(persisted); + expect(props.onClose).toHaveBeenCalledTimes(1); + }); + it('keeps the selected baseUrl for same-provider duplicate vision model ids', async () => { const switchModel = vi.fn(); const setVisionModel = vi.fn(); diff --git a/packages/cli/src/ui/components/ModelDialog.tsx b/packages/cli/src/ui/components/ModelDialog.tsx index d8158e72cc8..a61e414755b 100644 --- a/packages/cli/src/ui/components/ModelDialog.tsx +++ b/packages/cli/src/ui/components/ModelDialog.tsx @@ -116,6 +116,7 @@ interface ModelDialogProps { isFastModelMode?: boolean; isVoiceModelMode?: boolean; isVisionModelMode?: boolean; + isImageModelMode?: boolean; /** Override which settings scope to persist the selection to. */ persistScope?: 'workspace' | 'user'; availableTerminalHeight?: number; @@ -285,6 +286,7 @@ export function ModelDialog({ isFastModelMode, isVoiceModelMode, isVisionModelMode, + isImageModelMode, persistScope, availableTerminalHeight, }: ModelDialogProps): React.JSX.Element { @@ -302,15 +304,26 @@ export function ModelDialog({ const allModels = config ? config.getAllConfiguredModels() : []; // Separate runtime models from registry models - const runtimeModels = allModels.filter((m) => m.isRuntimeModel); - const registryModels = allModels.filter( - (m) => + const runtimeModels = isImageModelMode + ? [] + : allModels.filter((m) => m.isRuntimeModel); + const registryModels = allModels.filter((m) => { + const imageModelSelector = encodeVisionModelSelector( + buildModelSelectionKey(m.authType, m.id, m.baseUrl), + ); + const isSelectableImageModel = isImageModelMode + ? m.imageOnly === true && + config?.resolveImageGenerationModel(imageModelSelector) !== undefined + : m.imageOnly !== true; + return ( !m.isRuntimeModel && (m.authType !== AuthType.QWEN_OAUTH || authType === AuthType.QWEN_OAUTH) && + isSelectableImageModel && (isFastModelMode || !m.fastOnly) && - (isVoiceModelMode || !m.voiceOnly), - ); + (isVoiceModelMode || !m.voiceOnly) + ); + }); // Group registry models by authType const modelsByAuthTypeMap = new Map(); @@ -363,7 +376,7 @@ export function ModelDialog({ } return result; - }, [authType, config, isFastModelMode, isVoiceModelMode]); + }, [authType, config, isFastModelMode, isImageModelMode, isVoiceModelMode]); const MODEL_OPTIONS = useMemo( () => @@ -462,7 +475,9 @@ export function ModelDialog({ const visionModelSetting = settings?.merged?.visionModel as | string | undefined; + const imageModelSetting = settings?.merged?.imageModel as string | undefined; const parsedVisionModelValue = parseVisionModelSetting(visionModelSetting); + const parsedImageModelValue = parseVisionModelSetting(imageModelSetting); const parsedFastModelSetting = useMemo(() => { if (!isFastModelMode) return undefined; try { @@ -479,17 +494,27 @@ export function ModelDialog({ return undefined; } }, [parsedVisionModelValue?.selector, isVisionModelMode]); + const parsedImageModelSetting = useMemo(() => { + if (!isImageModelMode) return undefined; + try { + return resolveModelId(parsedImageModelValue?.selector); + } catch { + return undefined; + } + }, [parsedImageModelValue?.selector, isImageModelMode]); const preferredModelId = isFastModelMode && parsedFastModelSetting ? parsedFastModelSetting.modelId : isVisionModelMode && parsedVisionModelSetting ? parsedVisionModelSetting.modelId - : config?.getModel() || MAINLINE_CODER_MODEL; + : isImageModelMode && parsedImageModelSetting + ? parsedImageModelSetting.modelId + : config?.getModel() || MAINLINE_CODER_MODEL; // Check if current model is a runtime model // Runtime snapshot ID is already in $runtime|${authType}|${modelId} format const activeRuntimeSnapshot = - isFastModelMode || isVoiceModelMode || isVisionModelMode - ? undefined // fast/voice/vision models are never runtime model selections + isFastModelMode || isVoiceModelMode || isVisionModelMode || isImageModelMode + ? undefined : config?.getActiveRuntimeModelSnapshot?.(); const currentBaseUrl = config ?.getModelsConfig() @@ -538,6 +563,23 @@ export function ModelDialog({ matchesVisionModelBaseUrl(model), ) : undefined; + const preferredImageModelEntry = + isImageModelMode && parsedImageModelSetting + ? parsedImageModelSetting.authType + ? availableModelEntries.find( + ({ authType: t2, model }) => + t2 === parsedImageModelSetting.authType && + model.id === parsedImageModelSetting.modelId && + (!parsedImageModelValue?.baseUrl || + model.baseUrl === parsedImageModelValue.baseUrl), + ) + : availableModelEntries.find( + ({ model }) => + model.id === parsedImageModelSetting.modelId && + (!parsedImageModelValue?.baseUrl || + model.baseUrl === parsedImageModelValue.baseUrl), + ) + : undefined; const preferredKey = activeRuntimeSnapshot ? activeRuntimeSnapshot.id : preferredVoiceModelEntry @@ -552,22 +594,35 @@ export function ModelDialog({ preferredVisionModelEntry.model.id, preferredVisionModelEntry.model.baseUrl, ) - : preferredFastModelEntry + : preferredImageModelEntry ? buildModelSelectionKey( - preferredFastModelEntry.authType, - preferredFastModelEntry.model.id, - preferredFastModelEntry.model.baseUrl, + preferredImageModelEntry.authType, + preferredImageModelEntry.model.id, + preferredImageModelEntry.model.baseUrl, ) - : authType - ? buildModelSelectionKey(authType, preferredModelId, currentBaseUrl) - : ''; + : preferredFastModelEntry + ? buildModelSelectionKey( + preferredFastModelEntry.authType, + preferredFastModelEntry.model.id, + preferredFastModelEntry.model.baseUrl, + ) + : authType + ? buildModelSelectionKey( + authType, + preferredModelId, + currentBaseUrl, + ) + : ''; useKeypress( (key) => { if ( key.name === 'escape' || (key.name === 'left' && - (isFastModelMode || isVoiceModelMode || isVisionModelMode)) + (isFastModelMode || + isVoiceModelMode || + isVisionModelMode || + isImageModelMode)) ) { onClose(); } @@ -730,6 +785,43 @@ export function ModelDialog({ return; } + if (isImageModelMode) { + if (!selectedEntry || !config) { + setErrorMessage(t('Selected image model is unavailable.')); + return; + } + const imageModel = encodeVisionModelSelector(selected); + const imageModelDisplay = + parseVisionModelSetting(imageModel)?.selector ?? imageModel; + if (!config.resolveImageGenerationModel(imageModel)) { + setErrorMessage( + t( + "'{{model}}' must declare a valid HTTPS baseUrl and credential environment variable.", + { model: imageModelDisplay }, + ), + ); + return; + } + const scope = resolvePersistScope(settings, persistScope); + settings.setValue(scope, 'imageModel', imageModel); + await config.setImageModel(imageModel); + const scopeSuffix = + persistScope === 'workspace' + ? t(' (this project)') + : persistScope === 'user' + ? t(' (global)') + : ''; + uiState?.historyManager.addItem( + { + type: 'success', + text: `${t('Image Model')}: ${imageModelDisplay}${scopeSuffix}`, + }, + Date.now(), + ); + onClose(); + return; + } + // Block selection of discontinued qwen-oauth models // (only block non-runtime OAuth; runtime OAuth models from existing // cached tokens are still allowed to work until the server rejects them) @@ -847,6 +939,7 @@ export function ModelDialog({ isFastModelMode, isVoiceModelMode, isVisionModelMode, + isImageModelMode, availableModelEntries, persistScope, ], @@ -867,9 +960,11 @@ export function ModelDialog({ ? t('Select Voice Model') : isVisionModelMode ? t('Select Vision Model') - : isFastModelMode - ? t('Select Fast Model') - : t('Select Model')) + + : isImageModelMode + ? t('Select Image Model') + : isFastModelMode + ? t('Select Fast Model') + : t('Select Model')) + (persistScope === 'workspace' ? t(' (this project)') : persistScope === 'user' diff --git a/packages/cli/src/ui/components/arena/ArenaStartDialog.tsx b/packages/cli/src/ui/components/arena/ArenaStartDialog.tsx index 7837ce043c8..71b768353fa 100644 --- a/packages/cli/src/ui/components/arena/ArenaStartDialog.tsx +++ b/packages/cli/src/ui/components/arena/ArenaStartDialog.tsx @@ -33,7 +33,9 @@ export function ArenaStartDialog({ const modelItems = useMemo(() => { const allModels = config.getAllConfiguredModels(); - const selectableModels = allModels.filter((model) => !model.isRuntimeModel); + const selectableModels = allModels.filter( + (model) => !model.isRuntimeModel && !model.imageOnly, + ); return selectableModels.map((model) => { const token = `${model.authType}:${model.id}`; diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index d4cfb202136..862b4f7ae39 100644 --- a/packages/cli/src/ui/contexts/UIActionsContext.tsx +++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx @@ -57,6 +57,7 @@ export interface UIActions { fastModelMode?: boolean; voiceModelMode?: boolean; visionModelMode?: boolean; + imageModelMode?: boolean; }) => void; openArenaDialog: (type: Exclude) => void; closeArenaDialog: () => void; diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 2cd2c4ac828..09d4aa55764 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -72,6 +72,7 @@ export interface UIState { isFastModelMode: boolean; isVoiceModelMode: boolean; isVisionModelMode: boolean; + isImageModelMode: boolean; modelDialogPersistScope: 'workspace' | 'user' | undefined; isTrustDialogOpen: boolean; activeArenaDialog: ArenaDialogType; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index de547b18876..11a6dab6588 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -618,6 +618,25 @@ describe('useSlashCommandProcessor', () => { }); }); + it('should handle "dialog: image-model" action', async () => { + const command = createTestCommand({ + name: 'imagemodelcmd', + action: vi + .fn() + .mockResolvedValue({ type: 'dialog', dialog: 'image-model' }), + }); + const result = setupProcessorHook([command]); + await waitFor(() => expect(result.current.slashCommands).toHaveLength(1)); + + await act(async () => { + await result.current.handleSlashCommand('/imagemodelcmd'); + }); + + expect(mockOpenModelDialog).toHaveBeenCalledWith({ + imageModelMode: true, + }); + }); + it('awaits direct resume session switching before returning handled', async () => { const actions = createMockActions(); let resolveResume: (() => void) | undefined; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index cfa59d4506c..862b7905500 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -134,6 +134,7 @@ export interface SlashCommandProcessorActions { fastModelMode?: boolean; voiceModelMode?: boolean; visionModelMode?: boolean; + imageModelMode?: boolean; persistScope?: 'workspace' | 'user'; }) => void; openTrustDialog: () => void; @@ -1119,6 +1120,12 @@ export const useSlashCommandProcessor = ( persistScope: result.persistScope, }); return { type: 'handled' }; + case 'image-model': + actions.openModelDialog({ + imageModelMode: true, + persistScope: result.persistScope, + }); + return { type: 'handled' }; case 'trust': actions.openTrustDialog(); return { type: 'handled' }; diff --git a/packages/cli/src/ui/hooks/useModelCommand.test.ts b/packages/cli/src/ui/hooks/useModelCommand.test.ts index f7c8d4ec47e..ced21d15ea7 100644 --- a/packages/cli/src/ui/hooks/useModelCommand.test.ts +++ b/packages/cli/src/ui/hooks/useModelCommand.test.ts @@ -53,6 +53,24 @@ describe('useModelCommand', () => { expect(result.current.isVoiceModelMode).toBe(false); }); + it('should open the model dialog in image model mode exclusively', () => { + const { result } = renderHook(() => useModelCommand()); + + act(() => { + result.current.openModelDialog({ + imageModelMode: true, + voiceModelMode: true, + visionModelMode: true, + fastModelMode: true, + }); + }); + + expect(result.current.isImageModelMode).toBe(true); + expect(result.current.isVisionModelMode).toBe(false); + expect(result.current.isVoiceModelMode).toBe(false); + expect(result.current.isFastModelMode).toBe(false); + }); + it('should close the model dialog when closeModelDialog is called', () => { const { result } = renderHook(() => useModelCommand()); @@ -85,4 +103,18 @@ describe('useModelCommand', () => { expect(result.current.isModelDialogOpen).toBe(false); expect(result.current.isVisionModelMode).toBe(false); }); + + it('should reset isImageModelMode on close', () => { + const { result } = renderHook(() => useModelCommand()); + + act(() => { + result.current.openModelDialog({ imageModelMode: true }); + }); + expect(result.current.isImageModelMode).toBe(true); + + act(() => { + result.current.closeModelDialog(); + }); + expect(result.current.isImageModelMode).toBe(false); + }); }); diff --git a/packages/cli/src/ui/hooks/useModelCommand.ts b/packages/cli/src/ui/hooks/useModelCommand.ts index 377fe5bd7c4..aae7dc6db70 100644 --- a/packages/cli/src/ui/hooks/useModelCommand.ts +++ b/packages/cli/src/ui/hooks/useModelCommand.ts @@ -13,11 +13,13 @@ interface UseModelCommandReturn { isFastModelMode: boolean; isVoiceModelMode: boolean; isVisionModelMode: boolean; + isImageModelMode: boolean; modelDialogPersistScope: ModelDialogPersistScope | undefined; openModelDialog: (options?: { fastModelMode?: boolean; voiceModelMode?: boolean; visionModelMode?: boolean; + imageModelMode?: boolean; persistScope?: ModelDialogPersistScope; }) => void; closeModelDialog: () => void; @@ -28,6 +30,7 @@ export const useModelCommand = (): UseModelCommandReturn => { const [isFastModelMode, setIsFastModelMode] = useState(false); const [isVoiceModelMode, setIsVoiceModelMode] = useState(false); const [isVisionModelMode, setIsVisionModelMode] = useState(false); + const [isImageModelMode, setIsImageModelMode] = useState(false); const [modelDialogPersistScope, setModelDialogPersistScope] = useState< ModelDialogPersistScope | undefined >(undefined); @@ -37,20 +40,23 @@ export const useModelCommand = (): UseModelCommandReturn => { fastModelMode?: boolean; voiceModelMode?: boolean; visionModelMode?: boolean; + imageModelMode?: boolean; persistScope?: ModelDialogPersistScope; }) => { const voiceModelMode = options?.voiceModelMode ?? false; const visionModelMode = options?.visionModelMode ?? false; + const imageModelMode = options?.imageModelMode ?? false; // Modes are mutually exclusive; a specialized mode suppresses fast mode. setIsFastModelMode( - voiceModelMode || visionModelMode + voiceModelMode || visionModelMode || imageModelMode ? false : (options?.fastModelMode ?? false), ); - // Vision wins over voice when both are passed, so the dialog can't end up - // in two specialized modes at once (mismatched title vs. highlighted row). - setIsVoiceModelMode(visionModelMode ? false : voiceModelMode); - setIsVisionModelMode(visionModelMode); + setIsVoiceModelMode( + visionModelMode || imageModelMode ? false : voiceModelMode, + ); + setIsVisionModelMode(imageModelMode ? false : visionModelMode); + setIsImageModelMode(imageModelMode); setModelDialogPersistScope(options?.persistScope); setIsModelDialogOpen(true); }, @@ -62,6 +68,7 @@ export const useModelCommand = (): UseModelCommandReturn => { setIsFastModelMode(false); setIsVoiceModelMode(false); setIsVisionModelMode(false); + setIsImageModelMode(false); setModelDialogPersistScope(undefined); }, []); @@ -70,6 +77,7 @@ export const useModelCommand = (): UseModelCommandReturn => { isFastModelMode, isVoiceModelMode, isVisionModelMode, + isImageModelMode, modelDialogPersistScope, openModelDialog, closeModelDialog, diff --git a/packages/cli/src/ui/models/availableModels.test.ts b/packages/cli/src/ui/models/availableModels.test.ts index 767fb6f0604..bdbcf24cb74 100644 --- a/packages/cli/src/ui/models/availableModels.test.ts +++ b/packages/cli/src/ui/models/availableModels.test.ts @@ -94,6 +94,12 @@ describe('availableModels', () => { authType: AuthType.QWEN_OAUTH, isVision: false, }, + { + id: 'image', + label: 'Image', + authType: AuthType.QWEN_OAUTH, + imageOnly: true, + }, ]), } as unknown as Config; diff --git a/packages/cli/src/ui/models/availableModels.ts b/packages/cli/src/ui/models/availableModels.ts index def4f12a7a5..da5f9e5eaec 100644 --- a/packages/cli/src/ui/models/availableModels.ts +++ b/packages/cli/src/ui/models/availableModels.ts @@ -97,7 +97,9 @@ export function getAvailableModelsForAuthType( // Use config's model registry when available if (config) { try { - const models = config.getAvailableModelsForAuthType(authType); + const models = config + .getAvailableModelsForAuthType(authType) + .filter((model) => !model.imageOnly); if (models.length > 0) { return models.map(convertCoreModelToCliModel); } diff --git a/packages/cli/src/ui/voice/voice-model.test.ts b/packages/cli/src/ui/voice/voice-model.test.ts index d54536d4f1c..5d073e74353 100644 --- a/packages/cli/src/ui/voice/voice-model.test.ts +++ b/packages/cli/src/ui/voice/voice-model.test.ts @@ -33,6 +33,7 @@ describe('voice model guards', () => { false, ); expect(isTranscribableVoiceModel(model({ baseUrl: '' }))).toBe(false); + expect(isTranscribableVoiceModel(model({ imageOnly: true }))).toBe(false); }); it('isSelectableVoiceModel accepts ids with a real ASR transport', () => { diff --git a/packages/cli/src/utils/acpModelUtils.test.ts b/packages/cli/src/utils/acpModelUtils.test.ts index 99483a1bafa..9291ec67238 100644 --- a/packages/cli/src/utils/acpModelUtils.test.ts +++ b/packages/cli/src/utils/acpModelUtils.test.ts @@ -334,13 +334,15 @@ describe('acpModelUtils', () => { expect(isInlineModelOverrideAllowed(config, 'missing')).toBe(false); }); - it('does not match fast-only or voice-only models', () => { + it('does not match selector-only models', () => { const config = makeConfig({ authType: AuthType.QWEN_OAUTH }, [ { id: 'qwen-fast', authType: AuthType.QWEN_OAUTH, fastOnly: true }, { id: 'qwen-voice', authType: AuthType.QWEN_OAUTH, voiceOnly: true }, + { id: 'qwen-image', authType: AuthType.QWEN_OAUTH, imageOnly: true }, ]); expect(isInlineModelOverrideAllowed(config, 'qwen-fast')).toBe(false); expect(isInlineModelOverrideAllowed(config, 'qwen-voice')).toBe(false); + expect(isInlineModelOverrideAllowed(config, 'qwen-image')).toBe(false); }); it('rejects when no active auth type is available', () => { diff --git a/packages/cli/src/utils/acpModelUtils.ts b/packages/cli/src/utils/acpModelUtils.ts index b9b3e09829a..de14c6f0aad 100644 --- a/packages/cli/src/utils/acpModelUtils.ts +++ b/packages/cli/src/utils/acpModelUtils.ts @@ -52,7 +52,12 @@ export function buildAcpModelOptions( models: readonly AvailableModel[], ): AcpModelOption[] { const candidates = models - .filter((model) => model.fastOnly !== true && model.voiceOnly !== true) + .filter( + (model) => + model.fastOnly !== true && + model.voiceOnly !== true && + model.imageOnly !== true, + ) .map((model) => { const effectiveModelId = model.isRuntimeModel && model.runtimeSnapshotId @@ -296,7 +301,7 @@ export function isInlineModelOverrideAllowed( const activeEnvKey = contentGeneratorConfig.apiKeyEnvKey; return config .getAvailableModelsForAuthType(authType) - .filter((m) => !m.fastOnly && !m.voiceOnly) + .filter((m) => !m.fastOnly && !m.voiceOnly && !m.imageOnly) .some( (m) => m.id === modelId && diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 2cd401ce454..358575e1f20 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2889,6 +2889,137 @@ describe('Server Config (config.ts)', () => { expect(registeredNames).not.toContain(ToolNames.RECORD_ARTIFACT); }); + it('registers image_gen when an image-only model route is selected', async () => { + const baseUrl = 'https://images.example.com/api/v1'; + const config = new Config({ + ...baseParams, + modelProvidersConfig: { + openai: [ + { + id: 'qwen-image-2.0', + baseUrl, + envKey: 'TEST_IMAGE_GENERATION_KEY', + imageOnly: true, + }, + ], + }, + imageModel: `openai:qwen-image-2.0\0${baseUrl}`, + }); + await config.initialize(); + + const registeredNames = ( + ToolRegistry.prototype.registerFactory as Mock + ).mock.calls.map((call) => call[0]); + expect(registeredNames).toContain(ToolNames.IMAGE_GEN); + expect(config.getImageGenerationConfig()).toEqual({ + model: 'qwen-image-2.0', + baseUrl, + apiKeyEnv: 'TEST_IMAGE_GENERATION_KEY', + }); + }); + + it('does not register image_gen without an image model selection', async () => { + const config = new Config({ + ...baseParams, + modelProvidersConfig: { + openai: [ + { + id: 'qwen-image-2.0', + baseUrl: 'https://images.example.com/api/v1', + envKey: 'TEST_IMAGE_GENERATION_KEY', + imageOnly: true, + }, + ], + }, + }); + await config.initialize(); + + const registeredNames = ( + ToolRegistry.prototype.registerFactory as Mock + ).mock.calls.map((call) => call[0]); + expect(registeredNames).not.toContain(ToolNames.IMAGE_GEN); + }); + + it('does not use a protocol default as the image generation endpoint', () => { + const config = new Config({ + ...baseParams, + modelProvidersConfig: { + openai: [ + { + id: 'qwen-image-2.0', + envKey: 'TEST_IMAGE_GENERATION_KEY', + imageOnly: true, + }, + ], + }, + imageModel: 'openai:qwen-image-2.0', + }); + + expect(config.getImageGenerationConfig()).toBeUndefined(); + }); + + it('registers image_gen immediately when the image model changes at runtime', async () => { + const baseUrl = 'https://images.example.com/api/v1'; + const config = new Config({ + ...baseParams, + modelProvidersConfig: { + openai: [ + { + id: 'qwen-image-2.0', + baseUrl, + envKey: 'TEST_IMAGE_GENERATION_KEY', + imageOnly: true, + }, + ], + }, + }); + await config.initialize(); + vi.mocked(ToolRegistry.prototype.registerFactory).mockClear(); + + await config.setImageModel(`openai:qwen-image-2.0\0${baseUrl}`); + + expect(ToolRegistry.prototype.registerFactory).toHaveBeenCalledWith( + ToolNames.IMAGE_GEN, + expect.any(Function), + ); + expect(ToolRegistry.prototype.ensureTool).toHaveBeenCalledWith( + ToolNames.IMAGE_GEN, + ); + }); + + it('does not register image_gen when the permission manager disables it', async () => { + const baseUrl = 'https://images.example.com/api/v1'; + const config = new Config({ + ...baseParams, + modelProvidersConfig: { + openai: [ + { + id: 'qwen-image-2.0', + baseUrl, + envKey: 'TEST_IMAGE_GENERATION_KEY', + imageOnly: true, + }, + ], + }, + }); + await config.initialize(); + vi.mocked(ToolRegistry.prototype.registerFactory).mockClear(); + ( + config as unknown as { + permissionManager: { isToolEnabled: () => Promise }; + } + ).permissionManager = { + isToolEnabled: vi.fn().mockResolvedValue(false), + }; + + await config.setImageModel(`openai:qwen-image-2.0\0${baseUrl}`); + + expect(ToolRegistry.prototype.registerFactory).not.toHaveBeenCalledWith( + ToolNames.IMAGE_GEN, + expect.any(Function), + ); + }); + it('registers both artifact tools by default for interactive sessions', async () => { const config = new Config({ ...baseParams, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 8240501be8a..53201c3bd14 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -107,6 +107,7 @@ import { SubagentManager } from '../subagents/subagent-manager.js'; import type { SubagentConfig } from '../subagents/types.js'; import { BackgroundTaskRegistry } from '../agents/background-tasks.js'; import { MonitorRegistry } from '../services/monitorRegistry.js'; +import { normalizeImageGenerationBaseUrl } from '../services/image-generation-service.js'; import { BackgroundAgentResumeService } from '../agents/background-agent-resume.js'; import { BackgroundShellRegistry } from '../services/backgroundShellRegistry.js'; import { WorkflowRunRegistry } from '../agents/workflow-run-registry.js'; @@ -1044,6 +1045,8 @@ export interface ConfigParameters { artifactPublisher?: 'local' | 'host' | 'oss'; artifactHost?: ArtifactHostConfig; artifactOss?: ArtifactOssConfig; + /** Image generation model selected through `/model --image`. */ + imageModel?: string; /** * P5 T7: suppress the one-time `Workflow` tool usage-warning banner. * When `true`, the registry-side warning latch is bypassed and the @@ -1289,6 +1292,12 @@ export interface ConfigParameters { settingsWatcher?: { stopWatching(): void }; } +export interface ImageGenerationConfig { + model: string; + baseUrl: string; + apiKeyEnv: string; +} + function normalizeConfigOutputFormat( format: OutputFormat | undefined, ): OutputFormat | undefined { @@ -1892,6 +1901,7 @@ export class Config { private readonly webSearchSettings?: WebSearchSettings; private webSearchNoticeEmitted = false; private visionModel?: string; + private imageModel?: string; private readonly visionBridgeTimeoutMs: number | undefined; private readonly modelFallbacks: string[]; private readonly disableAllHooks: boolean; @@ -2309,6 +2319,7 @@ export class Config { this.fastModel = params.fastModel || undefined; this.webSearchSettings = params.webSearch; this.visionModel = params.visionModel || undefined; + this.imageModel = params.imageModel || undefined; // Guard: nothing validates settings.json on the load path, so this is the // only real gate. `AbortSignal.timeout()` requires an integer in // [0, 2^31-1] — a fractional or out-of-range value (which the number-typed @@ -3678,7 +3689,12 @@ export class Config { const available = selector.authType ? this.getAllConfiguredModels([selector.authType]) : this.getAllConfiguredModels(); - if (!available.some((m) => m.id === selector.modelId)) { + if ( + !available.some( + (model) => + model.id === selector.modelId && !model.voiceOnly && !model.imageOnly, + ) + ) { return undefined; } @@ -3738,6 +3754,19 @@ export class Config { this.visionModel = model || undefined; } + /** + * Update the image generation model and make the tool available immediately + * when the selected provider route is valid. + */ + async setImageModel(model: string | undefined): Promise { + this.imageModel = model || undefined; + if (!this.initialized || !this.isImageGenerationEnabled()) { + return; + } + await this.registerImageGenerationTool(this.toolRegistry); + await this.toolRegistry.ensureTool(ToolNames.IMAGE_GEN); + } + /** * Return the ordered list of fallback model IDs configured for this session. * The list is already normalized (deduplicated, capped at 3, blanks removed). @@ -3865,7 +3894,7 @@ export class Config { return undefined; } // Each guard below silently drops the pin (the hardest failure mode to - // debug, hence the warn): skip fast/voice-only models (a `settings.json` + // debug, hence the warn): skip selector-only models (a `settings.json` // pin can bypass the slash command's filter), and never route the bridge at // the primary entry itself (the text-only model the bridge works around) — // via the provider-aware identity check so a cross-provider namesake stays @@ -3877,6 +3906,7 @@ export class Config { (!parsedSetting.baseUrl || m.baseUrl === parsedSetting.baseUrl) && !m.fastOnly && !m.voiceOnly && + !m.imageOnly && !this.isCurrentPrimaryModel(m), ); if (routeMatches.length > 1) { @@ -3889,7 +3919,7 @@ export class Config { if (!match) { this.debugLogger.warn( `vision model pin '${visionModelForLog}' did not match a usable configured model ` + - `(removed, mistyped, fast/voice-only, or the primary itself); falling back to auto-select`, + `(removed, mistyped, selector-only, or the primary itself); falling back to auto-select`, ); return undefined; } @@ -5752,6 +5782,57 @@ export class Config { return this.artifactOss; } + resolveImageGenerationModel( + setting: string | undefined, + ): ImageGenerationConfig | undefined { + const parsedSetting = parseVisionModelSetting(setting); + if (!parsedSetting) return undefined; + + let selector; + try { + selector = resolveModelId(parsedSetting.selector); + } catch { + return undefined; + } + if (!selector) return undefined; + + const routeMatches = this.getAllConfiguredModels().filter( + (model) => + model.imageOnly === true && + !model.fastOnly && + !model.voiceOnly && + model.id === selector.modelId && + (!selector.authType || model.authType === selector.authType) && + (!parsedSetting.baseUrl || model.baseUrl === parsedSetting.baseUrl), + ); + if (routeMatches.length !== 1) return undefined; + + const match = routeMatches[0]!; + const apiKeyEnv = match.envKey?.trim(); + const configuredBaseUrl = match.registryBaseUrl?.trim(); + if (!apiKeyEnv || !configuredBaseUrl) return undefined; + + const baseUrl = normalizeImageGenerationBaseUrl( + parsedSetting.baseUrl ?? configuredBaseUrl, + ); + if (!baseUrl) return undefined; + + return { + model: match.id, + baseUrl, + apiKeyEnv, + }; + } + + getImageGenerationConfig(): ImageGenerationConfig | undefined { + if (this.bareMode || this.safeMode) return undefined; + return this.resolveImageGenerationModel(this.imageModel); + } + + isImageGenerationEnabled(): boolean { + return this.getImageGenerationConfig() !== undefined; + } + shouldAutoOpenArtifact(): boolean { if (process.env['QWEN_ARTIFACT_NO_AUTO_OPEN'] === '1') return false; return this.artifactAutoOpen && !this.isBrowserLaunchSuppressed(); @@ -6774,6 +6855,35 @@ export class Config { return this.onPersistPermissionRuleCallback; } + private async registerImageGenerationTool( + registry: ToolRegistry, + ): Promise { + if ( + !this.isImageGenerationEnabled() || + registry.getAllToolNames().includes(ToolNames.IMAGE_GEN) + ) { + return; + } + let enabled = true; + try { + enabled = this.permissionManager + ? await this.permissionManager.isToolEnabled(ToolNames.IMAGE_GEN) + : true; + } catch (error) { + this.debugLogger.warn( + `Failed to check permissions for tool "${ToolNames.IMAGE_GEN}", skipping registration:`, + error, + ); + return; + } + if (!enabled) return; + + registry.registerFactory(ToolNames.IMAGE_GEN, async () => { + const { ImageGenTool } = await import('../tools/image-gen.js'); + return new ImageGenTool(this); + }); + } + async createToolRegistry( sendSdkMcpMessage?: SendSdkMcpMessage, options?: { skipDiscovery?: boolean; forSubAgent?: boolean }, @@ -7018,6 +7128,7 @@ export class Config { this.warnings.push(gate.notice); } } + await this.registerImageGenerationTool(registry); if (this.isArtifactEnabled()) { await registerLazy(ToolNames.ARTIFACT, async () => { const { ArtifactTool } = await import( diff --git a/packages/core/src/models/content-generator-config.test.ts b/packages/core/src/models/content-generator-config.test.ts index a028a007991..fdc13da0a25 100644 --- a/packages/core/src/models/content-generator-config.test.ts +++ b/packages/core/src/models/content-generator-config.test.ts @@ -241,6 +241,21 @@ describe('buildAgentContentGeneratorConfig', () => { expect(result.thinkingMandatory).toBeUndefined(); }); + + it('rejects image-only models for agent content generation', () => { + const config = createMockConfig(parentConfig, { + ...resolvedModel, + imageOnly: true, + }); + + expect(() => + buildAgentContentGeneratorConfig(config, 'registry-model-id', { + authType: 'anthropic', + }), + ).toThrow( + "Image-only model 'registry-model-id' cannot be used for content generation", + ); + }); }); describe('edge cases', () => { diff --git a/packages/core/src/models/content-generator-config.ts b/packages/core/src/models/content-generator-config.ts index 2fc3200ae71..751d7aa7e13 100644 --- a/packages/core/src/models/content-generator-config.ts +++ b/packages/core/src/models/content-generator-config.ts @@ -56,6 +56,11 @@ export function buildAgentContentGeneratorConfig( authOverrides.baseUrl, ) : undefined; + if (resolvedModel?.imageOnly) { + throw new Error( + `Image-only model '${resolvedModel.id}' cannot be used for content generation`, + ); + } const nextConfig: ContentGeneratorConfig = { ...parentConfig, diff --git a/packages/core/src/models/modelRegistry.test.ts b/packages/core/src/models/modelRegistry.test.ts index 28e0e78d232..aa3a9b07740 100644 --- a/packages/core/src/models/modelRegistry.test.ts +++ b/packages/core/src/models/modelRegistry.test.ts @@ -1016,6 +1016,21 @@ describe('fastOnly and voiceOnly flags', () => { expect(models.find((m) => m.id === 'whisper-1')?.voiceOnly).toBe(true); }); + it('should propagate imageOnly flag to AvailableModel', () => { + const config: ModelProvidersConfig = { + openai: [ + { + id: 'qwen-image-2.0', + imageOnly: true, + }, + ], + }; + const registry = new ModelRegistry(config); + expect( + registry.getModelsForAuthType(AuthType.USE_OPENAI)[0]?.imageOnly, + ).toBe(true); + }); + it('should warn when both fastOnly and voiceOnly are set', () => { const config: ModelProvidersConfig = { openai: [ diff --git a/packages/core/src/models/modelRegistry.ts b/packages/core/src/models/modelRegistry.ts index 225eb343e21..ec2b4c3a0ea 100644 --- a/packages/core/src/models/modelRegistry.ts +++ b/packages/core/src/models/modelRegistry.ts @@ -235,6 +235,7 @@ export class ModelRegistry { envKey: model.envKey, fastOnly: model.fastOnly, voiceOnly: model.voiceOnly, + imageOnly: model.imageOnly, })); } @@ -284,7 +285,7 @@ export class ModelRegistry { /** * Get default model for an authType. * For qwen-oauth, returns the coder model. - * For others, returns the first configured model. + * For others, returns the first configured primary-capable model. */ getDefaultModelForAuthType( authType: AuthType, @@ -294,7 +295,7 @@ export class ModelRegistry { } const models = this.modelsByAuthType.get(authType); if (!models || models.size === 0) return undefined; - return Array.from(models.values())[0]; + return Array.from(models.values()).find((model) => !model.imageOnly); } /** @@ -338,9 +339,14 @@ export class ModelRegistry { `Model config in authType '${authType}' missing required field: id`, ); } - if (config.fastOnly && config.voiceOnly) { + const selectorOnlyCount = [ + config.fastOnly, + config.voiceOnly, + config.imageOnly, + ].filter(Boolean).length; + if (selectorOnlyCount > 1) { debugLogger.warn( - `Model "${config.id}" in authType "${authType}" has both fastOnly and voiceOnly set. It will be unreachable in all model selectors.`, + `Model "${config.id}" in authType "${authType}" has multiple selector-only flags. It will be unreachable in at least one model selector.`, ); } } diff --git a/packages/core/src/models/modelsConfig.test.ts b/packages/core/src/models/modelsConfig.test.ts index 9357ba7192f..f7f43226632 100644 --- a/packages/core/src/models/modelsConfig.test.ts +++ b/packages/core/src/models/modelsConfig.test.ts @@ -35,6 +35,53 @@ describe('ModelsConfig', () => { return modelsConfig.getGenerationConfig() as ContentGeneratorConfig; } + it('rejects image-only models as the primary model', async () => { + const modelsConfig = new ModelsConfig({ + initialAuthType: AuthType.USE_OPENAI, + modelProvidersConfig: { + openai: [{ id: 'chat-model' }, { id: 'image-model', imageOnly: true }], + }, + }); + await modelsConfig.switchModel(AuthType.USE_OPENAI, 'chat-model'); + + await expect( + modelsConfig.switchModel(AuthType.USE_OPENAI, 'image-model'), + ).rejects.toThrow( + "Image-only model 'image-model' cannot be used as the primary model", + ); + expect(modelsConfig.getModel()).toBe('chat-model'); + }); + + it('rejects an image-only model during auth refresh without changing state', () => { + const modelsConfig = new ModelsConfig({ + initialAuthType: AuthType.USE_ANTHROPIC, + modelProvidersConfig: { + openai: [{ id: 'image-model', imageOnly: true }], + }, + generationConfig: { model: 'previous-model' }, + }); + + expect(() => + modelsConfig.syncAfterAuthRefresh(AuthType.USE_OPENAI, 'image-model'), + ).toThrow( + "Image-only model 'image-model' cannot be used as the primary model", + ); + expect(modelsConfig.getCurrentAuthType()).toBe(AuthType.USE_ANTHROPIC); + expect(modelsConfig.getModel()).toBe('previous-model'); + }); + + it('does not choose an image-only model as the auth default', () => { + const modelsConfig = new ModelsConfig({ + modelProvidersConfig: { + openai: [{ id: 'image-model', imageOnly: true }, { id: 'chat-model' }], + }, + }); + + modelsConfig.syncAfterAuthRefresh(AuthType.USE_OPENAI, 'missing-model'); + + expect(modelsConfig.getModel()).toBe('chat-model'); + }); + it('should fully rollback state when switchModel fails after applying defaults (authType change)', async () => { const modelProvidersConfig: ModelProvidersConfig = { openai: [ diff --git a/packages/core/src/models/modelsConfig.ts b/packages/core/src/models/modelsConfig.ts index eeb488b2ae9..28f0a219588 100644 --- a/packages/core/src/models/modelsConfig.ts +++ b/packages/core/src/models/modelsConfig.ts @@ -507,6 +507,11 @@ export class ModelsConfig { `Model '${modelId}' not found for authType '${authType}'`, ); } + if (model.imageOnly) { + throw new Error( + `Image-only model '${modelId}' cannot be used as the primary model`, + ); + } const previousModelId = rollbackSnapshot.generationConfig.model || ''; const previousModel = @@ -1013,9 +1018,7 @@ export class ModelsConfig { modelId?: string, providerBaseUrlOverride?: string, ): void { - this.strictModelProviderSelection = false; const previousAuthType = this.currentAuthType; - this.currentAuthType = authType; // Step 1: If modelId exists in registry, always use config from modelRegistry // Manual credentials won't have a modelId that matches a provider model (the /auth provider-setup flow prevents it), @@ -1035,6 +1038,14 @@ export class ModelsConfig { ? (this.modelRegistry.getModel(authType, modelId, providerBaseUrl) ?? this.modelRegistry.getModel(authType, modelId)) : undefined; + if (resolved?.imageOnly) { + throw new Error( + `Image-only model '${modelId}' cannot be used as the primary model`, + ); + } + + this.strictModelProviderSelection = false; + this.currentAuthType = authType; if (resolved) { // When authType and modelId haven't changed (startup/restart scenario), // the current apiKey was already correctly resolved by diff --git a/packages/core/src/models/types.ts b/packages/core/src/models/types.ts index 668eebe6ccf..cd77e58fcf0 100644 --- a/packages/core/src/models/types.ts +++ b/packages/core/src/models/types.ts @@ -68,6 +68,8 @@ export interface ModelConfig { fastOnly?: boolean; /** When true, this model only appears in the voice model selector, not the main model list */ voiceOnly?: boolean; + /** When true, this model only appears in the image generation model selector */ + imageOnly?: boolean; } /** @@ -134,6 +136,8 @@ export interface AvailableModel { fastOnly?: boolean; /** When true, this model only appears in the voice model selector */ voiceOnly?: boolean; + /** When true, this model only appears in the image generation model selector */ + imageOnly?: boolean; /** Whether this is a runtime model (not from modelProviders) */ isRuntimeModel?: boolean; diff --git a/packages/core/src/providers/__tests__/provider-config.test.ts b/packages/core/src/providers/__tests__/provider-config.test.ts index 9dafe4f9aa0..9c2ff332ec0 100644 --- a/packages/core/src/providers/__tests__/provider-config.test.ts +++ b/packages/core/src/providers/__tests__/provider-config.test.ts @@ -253,6 +253,14 @@ describe('specToModelConfig (via buildProviderTemplate)', () => { const template = buildProviderTemplate(config); expect(template[0]?.description).toBe('A model'); }); + + it('preserves image-only model metadata in the provider template', () => { + const config = makeConfig({ + models: [{ id: 'image-model', imageOnly: true }], + }); + + expect(buildProviderTemplate(config)[0]?.imageOnly).toBe(true); + }); }); describe('resolveOwnsModel (via buildInstallPlan)', () => { diff --git a/packages/core/src/providers/provider-config.ts b/packages/core/src/providers/provider-config.ts index 5d5603155c8..88577e24b77 100644 --- a/packages/core/src/providers/provider-config.ts +++ b/packages/core/src/providers/provider-config.ts @@ -122,6 +122,7 @@ function specToModelConfig( ...(spec.description ? { description: spec.description } : {}), baseUrl, envKey, + ...(spec.imageOnly ? { imageOnly: true } : {}), ...(genConfig ? { generationConfig: genConfig } : {}), }; } diff --git a/packages/core/src/providers/types.ts b/packages/core/src/providers/types.ts index b871283a524..aac55346515 100644 --- a/packages/core/src/providers/types.ts +++ b/packages/core/src/providers/types.ts @@ -23,6 +23,7 @@ export interface ModelSpec { thinkingMandatory?: boolean; modalities?: InputModalities; description?: string; + imageOnly?: boolean; } export interface BaseUrlOption { diff --git a/packages/core/src/services/image-generation-service.test.ts b/packages/core/src/services/image-generation-service.test.ts new file mode 100644 index 00000000000..27bc9347239 --- /dev/null +++ b/packages/core/src/services/image-generation-service.test.ts @@ -0,0 +1,619 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { + generateImage, + normalizeImageGenerationBaseUrl, +} from './image-generation-service.js'; + +const networkPolicyMocks = vi.hoisted(() => ({ + resolveNetworkTarget: vi.fn(), +})); + +vi.mock('../extension/network-policy.js', () => networkPolicyMocks); + +const PNG_BYTES = new Uint8Array([ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, +]); + +beforeEach(() => { + networkPolicyMocks.resolveNetworkTarget.mockImplementation( + async (value: string | URL) => ({ + url: value instanceof URL ? value : new URL(value), + }), + ); +}); + +describe('normalizeImageGenerationBaseUrl', () => { + it('accepts a user-configured HTTPS endpoint', () => { + expect( + normalizeImageGenerationBaseUrl('https://images.example.com/api/v1/'), + ).toBe('https://images.example.com/api/v1'); + }); + + it('accepts a full multimodal generation endpoint', () => { + expect( + normalizeImageGenerationBaseUrl( + 'https://gateway.example.com/api/v1/services/aigc/multimodal-generation/generation', + ), + ).toBe( + 'https://gateway.example.com/api/v1/services/aigc/multimodal-generation/generation', + ); + }); + + it('removes repeated trailing slashes from the configured endpoint', () => { + expect( + normalizeImageGenerationBaseUrl('https://images.example.com/api/v1///'), + ).toBe('https://images.example.com/api/v1'); + }); + + it('rejects unsafe or malformed endpoints', () => { + expect( + normalizeImageGenerationBaseUrl('http://images.example.com/api/v1'), + ).toBeUndefined(); + expect( + normalizeImageGenerationBaseUrl( + 'https://user:secret@images.example.com/api/v1', + ), + ).toBeUndefined(); + }); +}); + +describe('generateImage', () => { + it('returns verified image bytes from a synchronous image endpoint', async () => { + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + request_id: 'request-1', + output: { + choices: [ + { + message: { + content: [ + { + image: 'https://cdn.example.com/generated/image.png', + }, + ], + }, + }, + ], + }, + }), + { + status: 200, + headers: { 'content-type': 'application/json' }, + }, + ), + ) + .mockResolvedValueOnce( + new Response(PNG_BYTES, { + status: 200, + headers: { 'content-type': 'image/png' }, + }), + ); + + const result = await generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'A Qwen Code poster', + size: '1536*864', + signal: new AbortController().signal, + fetchFn, + }); + + expect(result).toEqual({ + bytes: Buffer.from(PNG_BYTES), + mimeType: 'image/png', + requestId: 'request-1', + }); + expect(fetchFn).toHaveBeenCalledTimes(2); + expect(fetchFn.mock.calls[0]?.[0]).toBe( + 'https://images.example.com/api/v1/services/aigc/multimodal-generation/generation', + ); + const requestInit = fetchFn.mock.calls[0]?.[1]; + expect(requestInit?.method).toBe('POST'); + expect(requestInit?.headers).toEqual({ + Authorization: 'Bearer secret', + 'Content-Type': 'application/json', + }); + expect(fetchFn.mock.calls[1]?.[1]?.headers).toEqual({ + Accept: 'image/png', + }); + expect(JSON.parse(String(requestInit?.body))).toEqual({ + model: 'qwen-image-2.0', + input: { + messages: [ + { + role: 'user', + content: [{ text: 'A Qwen Code poster' }], + }, + ], + }, + parameters: { + n: 1, + prompt_extend: true, + size: '1536*864', + watermark: false, + }, + }); + }); + + it('pins the validated result hostname for the download connection', async () => { + const lookup = vi.fn(); + networkPolicyMocks.resolveNetworkTarget.mockResolvedValueOnce({ + url: new URL('https://cdn.example.com/generated/image.png'), + lookup, + }); + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [ + { + image: 'https://cdn.example.com/generated/image.png', + }, + ], + }, + }, + ], + }, + }), + { status: 200 }, + ), + ) + .mockResolvedValueOnce(new Response(PNG_BYTES, { status: 200 })); + + await generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }); + + expect(fetchFn.mock.calls[1]?.[1]).toEqual( + expect.objectContaining({ dispatcher: expect.anything() }), + ); + }); + + it('reports throttling without attempting a download', async () => { + const fetchFn = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + code: 'Throttling', + message: 'Requests rate limit exceeded', + request_id: 'request-2', + }), + { + status: 429, + headers: { 'content-type': 'application/json' }, + }, + ), + ); + + await expect( + generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }), + ).rejects.toThrow(/rate limit/i); + expect(fetchFn).toHaveBeenCalledTimes(1); + }); + + it('rejects an unsafe result URL before downloading it', async () => { + const fetchFn = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [{ image: 'http://127.0.0.1/private.png' }], + }, + }, + ], + }, + }), + { + status: 200, + headers: { 'content-type': 'application/json' }, + }, + ), + ); + + await expect( + generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }), + ).rejects.toThrow(/safe public HTTPS/i); + expect(fetchFn).toHaveBeenCalledTimes(1); + }); + + it('rejects a result hostname that resolves to a blocked address', async () => { + networkPolicyMocks.resolveNetworkTarget.mockRejectedValueOnce( + new Error('host resolved to 169.254.169.254'), + ); + const fetchFn = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [ + { image: 'https://images.example.com/private.png' }, + ], + }, + }, + ], + }, + }), + { + status: 200, + headers: { 'content-type': 'application/json' }, + }, + ), + ); + + await expect( + generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }), + ).rejects.toThrow(/safe public HTTPS/i); + expect(fetchFn).toHaveBeenCalledTimes(1); + }); + + it('rejects a download that is not a PNG image', async () => { + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [{ image: 'https://cdn.example.com/image.png' }], + }, + }, + ], + }, + }), + { + status: 200, + headers: { 'content-type': 'application/json' }, + }, + ), + ) + .mockResolvedValueOnce( + new Response('not an image', { + status: 200, + headers: { 'content-type': 'text/plain' }, + }), + ); + + await expect( + generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }), + ).rejects.toThrow(/valid PNG/i); + }); + + it('rejects a download with only a partial PNG signature', async () => { + const partialPngSignature = new Uint8Array([ + 0x89, 0x50, 0x4e, 0x47, 0, 0, 0, 0, + ]); + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [{ image: 'https://cdn.example.com/image.png' }], + }, + }, + ], + }, + }), + { + status: 200, + headers: { 'content-type': 'application/json' }, + }, + ), + ) + .mockResolvedValueOnce( + new Response(partialPngSignature, { + status: 200, + headers: { 'content-type': 'image/png' }, + }), + ); + + await expect( + generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }), + ).rejects.toThrow(/valid PNG/i); + }); + + it('rejects an image response above the download byte limit', async () => { + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [{ image: 'https://cdn.example.com/image.png' }], + }, + }, + ], + }, + }), + { + status: 200, + headers: { 'content-type': 'application/json' }, + }, + ), + ) + .mockResolvedValueOnce( + new Response(null, { + status: 200, + headers: { + 'content-length': String(10 * 1024 * 1024 + 1), + 'content-type': 'image/png', + }, + }), + ); + + await expect( + generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }), + ).rejects.toThrow(/byte limit/i); + }); + + it('does not expose a signed result URL when its download fails', async () => { + const signedUrl = + 'https://cdn.example.com/image.png?signature=temporary-secret'; + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [{ image: signedUrl }], + }, + }, + ], + }, + }), + { + status: 200, + headers: { 'content-type': 'application/json' }, + }, + ), + ) + .mockRejectedValueOnce(new Error(`Failed to fetch ${signedUrl}`)); + + const request = generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }); + + await expect(request).rejects.toThrow( + 'Generated image download failed before completion.', + ); + await expect(request).rejects.not.toThrow(signedUrl); + }); + + it('does not expose a signed result URL when its response stream fails', async () => { + const signedUrl = + 'https://cdn.example.com/image.png?signature=temporary-secret'; + const body = new ReadableStream({ + start(controller) { + controller.error(new Error(`Failed to read ${signedUrl}`)); + }, + }); + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [{ image: signedUrl }], + }, + }, + ], + }, + }), + { + status: 200, + headers: { 'content-type': 'application/json' }, + }, + ), + ) + .mockResolvedValueOnce(new Response(body, { status: 200 })); + + const request = generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }); + + await expect(request).rejects.toThrow( + 'Generated image download failed before completion.', + ); + await expect(request).rejects.not.toThrow(signedUrl); + }); +}); + +describe('generateImage redirect handling', () => { + it('follows a valid 302 → 200 redirect chain and returns the PNG', async () => { + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [{ image: 'https://api.example.com/image.png' }], + }, + }, + ], + }, + }), + { status: 200 }, + ), + ) + .mockResolvedValueOnce( + new Response(null, { + status: 302, + headers: { location: 'https://cdn.example.com/final.png' }, + }), + ) + .mockResolvedValueOnce( + new Response(PNG_BYTES, { + status: 200, + headers: { 'content-type': 'image/png' }, + }), + ); + + const result = await generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }); + + expect(result.bytes).toEqual(Buffer.from(PNG_BYTES)); + expect(fetchFn).toHaveBeenCalledTimes(3); + }); + + it('rejects when redirects exceed the maximum allowed', async () => { + const redirectResponse = () => + new Response(null, { + status: 302, + headers: { location: 'https://cdn.example.com/next.png' }, + }); + const fetchFn = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + output: { + choices: [ + { + message: { + content: [{ image: 'https://api.example.com/image.png' }], + }, + }, + ], + }, + }), + { status: 200 }, + ), + ) + // MAX_DOWNLOAD_REDIRECTS + 1 consecutive redirects + .mockResolvedValueOnce(redirectResponse()) + .mockResolvedValueOnce(redirectResponse()) + .mockResolvedValueOnce(redirectResponse()) + .mockResolvedValueOnce(redirectResponse()); + + await expect( + generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }), + ).rejects.toThrow(/exceeded.*redirects/i); + }); +}); + +describe('generateImage error body handling', () => { + it('reports HTTP status when the error body is not JSON', async () => { + const fetchFn = vi.fn().mockResolvedValue( + new Response('502 Bad Gateway', { + status: 502, + headers: { 'content-type': 'text/html' }, + }), + ); + + await expect( + generateImage({ + baseUrl: 'https://images.example.com/api/v1', + apiKey: 'secret', + model: 'qwen-image-2.0', + prompt: 'poster', + signal: new AbortController().signal, + fetchFn, + }), + ).rejects.toThrow(/HTTP 502/); + }); +}); diff --git a/packages/core/src/services/image-generation-service.ts b/packages/core/src/services/image-generation-service.ts new file mode 100644 index 00000000000..83d807ae17e --- /dev/null +++ b/packages/core/src/services/image-generation-service.ts @@ -0,0 +1,386 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { isPrivateHost } from '../utils/fetch.js'; +import { getErrorMessage } from '../utils/errors.js'; +import { + resolveNetworkTarget, + type ResolvedNetworkTarget, +} from '../extension/network-policy.js'; +import { loadUndici } from '../utils/runtimeFetchOptions.js'; + +const GENERATION_TIMEOUT_MS = 240_000; +const DOWNLOAD_TIMEOUT_MS = 120_000; +const MAX_API_RESPONSE_BYTES = 1024 * 1024; +const MAX_IMAGE_BYTES = 10 * 1024 * 1024; +const MAX_DOWNLOAD_REDIRECTS = 3; +const REDIRECT_STATUSES = new Set([301, 302, 303, 307, 308]); +const PNG_SIGNATURE = Buffer.from([ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, +]); + +class ResponseSizeLimitError extends Error {} + +export interface ImageGenerationRequest { + baseUrl: string; + apiKey: string; + model: string; + prompt: string; + size?: string; + signal: AbortSignal; + fetchFn?: typeof fetch; +} + +export interface GeneratedImage { + bytes: Buffer; + mimeType: 'image/png'; + requestId?: string; +} + +export type GenerateImage = ( + request: ImageGenerationRequest, +) => Promise; + +export function normalizeImageGenerationBaseUrl( + value: string | undefined, +): string | undefined { + const baseUrl = value?.trim(); + if (!baseUrl) return undefined; + let parsed: URL; + try { + parsed = new URL(baseUrl); + } catch { + return undefined; + } + if ( + parsed.protocol !== 'https:' || + parsed.username || + parsed.password || + parsed.search || + parsed.hash + ) { + return undefined; + } + return parsed.toString().replace(/\/+$/, ''); +} + +export async function generateImage( + request: ImageGenerationRequest, +): Promise { + const fetchFn = request.fetchFn ?? fetch; + const baseUrl = normalizeImageGenerationBaseUrl(request.baseUrl); + if (!baseUrl) { + throw new Error( + 'Image generation baseUrl must be a valid HTTPS URL without credentials, query, or fragment.', + ); + } + const generationUrl = baseUrl.endsWith( + '/services/aigc/multimodal-generation/generation', + ) + ? baseUrl + : `${baseUrl}/services/aigc/multimodal-generation/generation`; + const parameters: Record = { + n: 1, + prompt_extend: true, + watermark: false, + }; + if (request.size) { + parameters['size'] = request.size; + } + + let response: Response; + try { + response = await fetchFn(generationUrl, { + method: 'POST', + headers: { + Authorization: `Bearer ${request.apiKey}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + model: request.model, + input: { + messages: [ + { + role: 'user', + content: [{ text: request.prompt }], + }, + ], + }, + parameters, + }), + redirect: 'error', + signal: combineWithTimeout(request.signal, GENERATION_TIMEOUT_MS), + }); + } catch (error) { + throw new Error( + `Image generation request failed: ${getErrorMessage(error)}`, + { cause: error }, + ); + } + + if (!response.ok) { + let payload: unknown = {}; + try { + payload = await readJsonResponse(response, MAX_API_RESPONSE_BYTES); + } catch { + // non-JSON error body — formatImageGenerationError handles missing fields + } + throw new Error(formatImageGenerationError(response.status, payload)); + } + const payload = await readJsonResponse(response, MAX_API_RESPONSE_BYTES); + + const imageUrl = findGeneratedImageUrl(payload); + if (!imageUrl) { + throw new Error('Image generation response did not contain an image URL.'); + } + + const bytes = await downloadPng(imageUrl, fetchFn, request.signal); + const requestId = readString(payload, 'request_id', 'requestId'); + return { + bytes, + mimeType: 'image/png', + ...(requestId ? { requestId } : {}), + }; +} + +async function readJsonResponse( + response: Response, + maxBytes: number, +): Promise { + const bytes = await readBoundedBody(response, maxBytes); + if (bytes.length === 0) { + return {}; + } + try { + return JSON.parse(bytes.toString('utf8')) as unknown; + } catch { + throw new Error( + `Image generation endpoint returned malformed JSON (HTTP ${response.status}).`, + ); + } +} + +async function readBoundedBody( + response: Response, + maxBytes: number, +): Promise { + const contentLength = Number(response.headers.get('content-length')); + if (Number.isFinite(contentLength) && contentLength > maxBytes) { + await response.body?.cancel().catch(() => {}); + throw new ResponseSizeLimitError( + `Response exceeds the ${maxBytes}-byte limit.`, + ); + } + if (!response.body) { + return Buffer.alloc(0); + } + + const chunks: Uint8Array[] = []; + let total = 0; + const reader = response.body.getReader(); + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + total += value.byteLength; + if (total > maxBytes) { + throw new ResponseSizeLimitError( + `Response exceeds the ${maxBytes}-byte limit.`, + ); + } + chunks.push(value); + } + } finally { + await reader.cancel().catch(() => {}); + reader.releaseLock(); + } + return Buffer.concat(chunks, total); +} + +function formatImageGenerationError(status: number, payload: unknown): string { + const code = readString(payload, 'code'); + const message = readString(payload, 'message'); + const suffix = [code, message].filter(Boolean).join(': '); + + if (status === 429 || /throttl|rate.?limit/i.test(`${code} ${message}`)) { + return `Image generation rate limit reached${suffix ? ` (${suffix})` : ''}.`; + } + if ( + status === 401 || + status === 403 || + /access|permission/i.test(code ?? '') + ) { + return `Image generation access denied${suffix ? ` (${suffix})` : ''}. Check the API key, endpoint, and model access.`; + } + if (/DataInspectionFailed/i.test(code ?? '')) { + return `The image generation endpoint blocked the prompt during content moderation${message ? `: ${message}` : '.'}`; + } + return `Image generation failed with HTTP ${status}${suffix ? ` (${suffix})` : ''}.`; +} + +function findGeneratedImageUrl(payload: unknown): string | undefined { + if (!isRecord(payload)) return undefined; + const output = payload['output']; + if (!isRecord(output) || !Array.isArray(output['choices'])) return undefined; + + for (const choice of output['choices']) { + if (!isRecord(choice)) continue; + const message = choice['message']; + if (!isRecord(message) || !Array.isArray(message['content'])) continue; + for (const part of message['content']) { + if (!isRecord(part)) continue; + const image = part['image']; + if (typeof image === 'string' && image.trim()) { + return image.trim(); + } + } + } + return undefined; +} + +async function downloadPng( + imageUrl: string, + fetchFn: typeof fetch, + signal: AbortSignal, +): Promise { + const combinedSignal = combineWithTimeout(signal, DOWNLOAD_TIMEOUT_MS); + let currentTarget = await validateResultUrl(imageUrl, combinedSignal); + + for ( + let redirectCount = 0; + redirectCount <= MAX_DOWNLOAD_REDIRECTS; + redirectCount++ + ) { + let dispatcher: import('undici').Dispatcher | undefined; + if (currentTarget.lookup) { + const undici = await loadUndici(); + const { Agent } = undici as unknown as typeof import('undici'); + dispatcher = new Agent({ connect: { lookup: currentTarget.lookup } }); + } + let response: Response; + try { + response = await fetchFn(currentTarget.url.toString(), { + method: 'GET', + headers: { Accept: 'image/png' }, + redirect: 'manual', + signal: combinedSignal, + ...(dispatcher ? { dispatcher } : {}), + } as RequestInit); + } catch (error) { + await dispatcher?.close(); + throw new Error('Generated image download failed before completion.', { + cause: error, + }); + } + + try { + if (REDIRECT_STATUSES.has(response.status)) { + const location = response.headers.get('location'); + await response.body?.cancel().catch(() => {}); + if (!location) { + throw new Error( + 'Generated image redirect is missing a Location header.', + ); + } + let redirectUrl: string; + try { + redirectUrl = new URL( + location, + currentTarget.url.toString(), + ).toString(); + } catch { + throw new Error('Generated image redirect URL is invalid.'); + } + currentTarget = await validateResultUrl(redirectUrl, combinedSignal); + continue; + } + if (!response.ok) { + await response.body?.cancel().catch(() => {}); + throw new Error( + `Generated image download failed with HTTP ${response.status}.`, + ); + } + + let bytes: Buffer; + try { + bytes = await readBoundedBody(response, MAX_IMAGE_BYTES); + } catch (error) { + if (error instanceof ResponseSizeLimitError) { + throw error; + } + throw new Error('Generated image download failed before completion.', { + cause: error, + }); + } + if ( + bytes.length < PNG_SIGNATURE.length || + !bytes.subarray(0, PNG_SIGNATURE.length).equals(PNG_SIGNATURE) + ) { + throw new Error('Downloaded result is not a valid PNG image.'); + } + return bytes; + } finally { + await dispatcher?.close(); + } + } + + throw new Error( + `Generated image download exceeded ${MAX_DOWNLOAD_REDIRECTS} redirects.`, + ); +} + +async function validateResultUrl( + value: string, + signal: AbortSignal, +): Promise { + let parsed: URL; + try { + parsed = new URL(value); + } catch { + throw new Error('Image generation returned an invalid image URL.'); + } + if ( + parsed.protocol !== 'https:' || + parsed.username || + parsed.password || + isPrivateHost(parsed.toString()) + ) { + throw new Error( + 'Image generation returned an image URL that is not a safe public HTTPS URL.', + ); + } + try { + return await resolveNetworkTarget(parsed, 'public', signal); + } catch (error) { + signal.throwIfAborted(); + throw new Error( + 'Image generation returned an image URL that is not a safe public HTTPS URL.', + { cause: error }, + ); + } +} + +function combineWithTimeout( + signal: AbortSignal, + timeoutMs: number, +): AbortSignal { + return AbortSignal.any([signal, AbortSignal.timeout(timeoutMs)]); +} + +function readString(value: unknown, ...keys: string[]): string | undefined { + if (!isRecord(value)) return undefined; + for (const key of keys) { + const candidate = value[key]; + if (typeof candidate === 'string' && candidate.trim()) { + return candidate.trim(); + } + } + return undefined; +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} diff --git a/packages/core/src/services/visionBridge/vision-bridge-service.test.ts b/packages/core/src/services/visionBridge/vision-bridge-service.test.ts index b6dd3546081..8e8b53e9c43 100644 --- a/packages/core/src/services/visionBridge/vision-bridge-service.test.ts +++ b/packages/core/src/services/visionBridge/vision-bridge-service.test.ts @@ -15,6 +15,7 @@ import { runVisionBridge, selectVisionBridgeModel, isImageCapable, + isFullTurnVisionCapable, type VisionModelCandidate, } from './vision-bridge-service.js'; import type { Config } from '../../config/config.js'; @@ -1104,3 +1105,26 @@ describe('isImageCapable', () => { expect(isImageCapable({ id: 'qwen-text-max' })).toBe(false); }); }); + +describe('isFullTurnVisionCapable', () => { + it('excludes an image-only model even when agent-capable', () => { + expect( + isFullTurnVisionCapable({ + id: 'qwen-image-2.0', + imageOnly: true, + isVision: true, + capabilities: { agent: true }, + }), + ).toBe(false); + }); + + it('includes a non-image-only agent-capable vision model', () => { + expect( + isFullTurnVisionCapable({ + id: 'qwen3-vl-plus', + isVision: true, + capabilities: { agent: true }, + }), + ).toBe(true); + }); +}); diff --git a/packages/core/src/services/visionBridge/vision-bridge-service.ts b/packages/core/src/services/visionBridge/vision-bridge-service.ts index aefb7aef5bb..8443856a534 100644 --- a/packages/core/src/services/visionBridge/vision-bridge-service.ts +++ b/packages/core/src/services/visionBridge/vision-bridge-service.ts @@ -37,6 +37,7 @@ export interface VisionModelCandidate { capabilities?: { agent?: boolean }; fastOnly?: boolean; voiceOnly?: boolean; + imageOnly?: boolean; } /** The model/endpoint selected for a vision bridge call. */ @@ -63,6 +64,7 @@ export function isFullTurnVisionCapable(model: VisionModelCandidate): boolean { return ( !model.fastOnly && !model.voiceOnly && + !model.imageOnly && model.capabilities?.agent === true && isImageCapable(model) ); diff --git a/packages/core/src/tools/image-gen.test.ts b/packages/core/src/tools/image-gen.test.ts new file mode 100644 index 00000000000..0ccdb3b6f38 --- /dev/null +++ b/packages/core/src/tools/image-gen.test.ts @@ -0,0 +1,270 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { + access, + mkdir, + mkdtemp, + readFile, + rm, + symlink, +} from 'node:fs/promises'; +import os from 'node:os'; +import path from 'node:path'; +import type { Part } from '@google/genai'; +import type { Config } from '../config/config.js'; +import { ImageGenTool } from './image-gen.js'; +import { ToolErrorType } from './tool-error.js'; + +const PNG_BYTES = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); +const workspaces: string[] = []; + +function createConfig( + workspace: string, + imageCapable: boolean, +): Pick< + Config, + | 'getEffectiveInputModalities' + | 'getImageGenerationConfig' + | 'getSessionId' + | 'getTargetDir' +> { + return { + getEffectiveInputModalities: () => (imageCapable ? { image: true } : {}), + getImageGenerationConfig: () => ({ + model: 'qwen-image-2.0', + baseUrl: 'https://images.example.com/api/v1', + apiKeyEnv: 'TEST_IMAGE_API_KEY', + }), + getSessionId: () => 'session-1', + getTargetDir: () => workspace, + }; +} + +afterEach(async () => { + delete process.env['TEST_IMAGE_API_KEY']; + await Promise.all( + workspaces + .splice(0) + .map((workspace) => rm(workspace, { recursive: true, force: true })), + ); +}); + +describe('ImageGenTool', () => { + it('persists a workspace image artifact and returns it to an image-capable model', async () => { + const workspace = await mkdtemp(path.join(os.tmpdir(), 'image-gen-')); + workspaces.push(workspace); + process.env['TEST_IMAGE_API_KEY'] = 'secret'; + const generateImage = vi.fn().mockResolvedValue({ + bytes: PNG_BYTES, + mimeType: 'image/png', + requestId: 'request-1', + }); + const tool = new ImageGenTool( + createConfig(workspace, true) as Config, + generateImage, + ); + + const result = await tool.buildAndExecute( + { prompt: 'A Qwen Code poster', size: '1536*864' }, + new AbortController().signal, + ); + + expect(generateImage).toHaveBeenCalledWith( + expect.objectContaining({ + model: 'qwen-image-2.0', + prompt: 'A Qwen Code poster', + size: '1536*864', + apiKey: 'secret', + }), + ); + expect(result.resultFilePaths).toHaveLength(1); + const outputPath = result.resultFilePaths?.[0]; + expect(outputPath).toBeDefined(); + await expect(readFile(outputPath!)).resolves.toEqual(PNG_BYTES); + expect(result.artifacts).toEqual([ + expect.objectContaining({ + kind: 'image', + storage: 'workspace', + workspacePath: expect.stringMatching( + /^\.qwen\/generated-images\/session-1\/.+\.png$/, + ), + mimeType: 'image/png', + sizeBytes: PNG_BYTES.length, + metadata: { + model: 'qwen-image-2.0', + requestId: 'request-1', + size: '1536*864', + }, + }), + ]); + const parts = result.llmContent as Part[]; + expect( + parts.some((part) => part.inlineData?.mimeType === 'image/png'), + ).toBe(true); + expect( + parts.some( + (part) => part.inlineData?.data === PNG_BYTES.toString('base64'), + ), + ).toBe(true); + }); + + it('returns only the saved path to a text-only primary model', async () => { + const workspace = await mkdtemp(path.join(os.tmpdir(), 'image-gen-')); + workspaces.push(workspace); + process.env['TEST_IMAGE_API_KEY'] = 'secret'; + const tool = new ImageGenTool( + createConfig(workspace, false) as Config, + vi.fn().mockResolvedValue({ + bytes: PNG_BYTES, + mimeType: 'image/png', + requestId: 'request-2', + }), + ); + + const result = await tool.buildAndExecute( + { prompt: 'A poster' }, + new AbortController().signal, + ); + + expect(JSON.stringify(result.llmContent)).not.toContain('inlineData'); + expect(result.resultFilePaths).toHaveLength(1); + expect(result.artifacts?.[0]?.kind).toBe('image'); + }); + + it('rejects image sizes outside the documented total-pixel range', () => { + const tool = new ImageGenTool( + createConfig('/workspace', true) as Config, + vi.fn(), + ); + + expect(() => tool.build({ prompt: 'poster', size: '100*100' })).toThrow( + /total pixels/i, + ); + expect(() => + tool.build({ prompt: 'poster', size: '2688*1536' }), + ).not.toThrow(); + }); + + it('requires approval because generation is a billable external request', async () => { + const tool = new ImageGenTool( + createConfig('/workspace', true) as Config, + vi.fn(), + ); + + await expect( + tool.build({ prompt: 'poster' }).getDefaultPermission(), + ).resolves.toBe('ask'); + }); + + it('does not start a billable request when already cancelled', async () => { + const workspace = await mkdtemp(path.join(os.tmpdir(), 'image-gen-')); + workspaces.push(workspace); + process.env['TEST_IMAGE_API_KEY'] = 'secret'; + const generateImage = vi.fn(); + const controller = new AbortController(); + controller.abort(); + const tool = new ImageGenTool( + createConfig(workspace, true) as Config, + generateImage, + ); + + const result = await tool + .build({ prompt: 'A poster' }) + .execute(controller.signal); + + expect(result.error?.type).toBe(ToolErrorType.EXECUTION_FAILED); + expect(generateImage).not.toHaveBeenCalled(); + }); + + it('rejects an unsafe output directory before starting generation', async () => { + const workspace = await mkdtemp(path.join(os.tmpdir(), 'image-gen-')); + const outside = await mkdtemp(path.join(os.tmpdir(), 'image-gen-outside-')); + workspaces.push(workspace, outside); + process.env['TEST_IMAGE_API_KEY'] = 'secret'; + const generatedImagesDir = path.join( + workspace, + '.qwen', + 'generated-images', + ); + await mkdir(generatedImagesDir, { recursive: true }); + await symlink( + outside, + path.join(generatedImagesDir, 'session-1'), + process.platform === 'win32' ? 'junction' : 'dir', + ); + const generateImage = vi.fn(); + const tool = new ImageGenTool( + createConfig(workspace, true) as Config, + generateImage, + ); + + const result = await tool.buildAndExecute( + { prompt: 'A poster' }, + new AbortController().signal, + ); + + expect(result.error?.type).toBe(ToolErrorType.EXECUTION_FAILED); + expect(result.error?.message).toMatch(/inside the workspace/i); + expect(generateImage).not.toHaveBeenCalled(); + }); + + it('does not write an image when cancellation wins before persistence', async () => { + const workspace = await mkdtemp(path.join(os.tmpdir(), 'image-gen-')); + workspaces.push(workspace); + process.env['TEST_IMAGE_API_KEY'] = 'secret'; + const controller = new AbortController(); + const tool = new ImageGenTool( + createConfig(workspace, true) as Config, + vi.fn().mockImplementation(async () => { + controller.abort(); + return { + bytes: PNG_BYTES, + mimeType: 'image/png', + requestId: 'request-3', + }; + }), + ); + const invocation = tool.build({ prompt: 'A poster' }); + const outputPath = invocation.toolLocations()[0]?.path; + + const result = await invocation.execute(controller.signal); + + expect(result.error?.type).toBe(ToolErrorType.EXECUTION_FAILED); + expect(outputPath).toBeDefined(); + await expect(access(outputPath!)).rejects.toThrow(); + }); + + it('does not leak signed URLs from the error cause chain', async () => { + const workspace = await mkdtemp(path.join(os.tmpdir(), 'image-gen-')); + workspaces.push(workspace); + process.env['TEST_IMAGE_API_KEY'] = 'secret'; + const signedUrl = + 'https://cdn.example.com/image.png?signature=temporary-secret&token=abc123'; + const tool = new ImageGenTool( + createConfig(workspace, true) as Config, + vi.fn().mockRejectedValue( + new Error('Generated image download failed before completion.', { + cause: new Error(`GET ${signedUrl} failed: ECONNREFUSED`), + }), + ), + ); + + const result = await tool.buildAndExecute( + { prompt: 'A poster' }, + new AbortController().signal, + ); + + expect(result.error?.type).toBe(ToolErrorType.EXECUTION_FAILED); + expect(result.error?.message).toBe( + 'Generated image download failed before completion.', + ); + expect(JSON.stringify(result.llmContent)).not.toContain('signature'); + expect(JSON.stringify(result.llmContent)).not.toContain('temporary-secret'); + expect(JSON.stringify(result.returnDisplay)).not.toContain('signature'); + }); +}); diff --git a/packages/core/src/tools/image-gen.ts b/packages/core/src/tools/image-gen.ts new file mode 100644 index 00000000000..d12837c9fc4 --- /dev/null +++ b/packages/core/src/tools/image-gen.ts @@ -0,0 +1,252 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { randomUUID } from 'node:crypto'; +import { mkdir } from 'node:fs/promises'; +import path from 'node:path'; +import type { Part } from '@google/genai'; +import type { Config } from '../config/config.js'; +import { Storage } from '../config/storage.js'; +import { + generateImage as generateConfiguredImage, + type GenerateImage, +} from '../services/image-generation-service.js'; +import { atomicWriteFile } from '../utils/atomicFileWrite.js'; +import { getErrorMessage } from '../utils/errors.js'; +import { ToolErrorType } from './tool-error.js'; +import { ToolDisplayNames, ToolNames } from './tool-names.js'; +import type { ToolInvocation, ToolLocation, ToolResult } from './tools.js'; +import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js'; + +const MIN_TOTAL_PIXELS = 512 * 512; +const MAX_TOTAL_PIXELS = 2048 * 2048; +const MAX_PROMPT_CHARS = 10_000; + +export interface ImageGenParams { + prompt: string; + size?: string; +} + +class ImageGenInvocation extends BaseToolInvocation< + ImageGenParams, + ToolResult +> { + private readonly outputPath: string; + + constructor( + private readonly config: Config, + private readonly generateImage: GenerateImage, + params: ImageGenParams, + ) { + super(params); + const sessionDir = Storage.sanitizePlanSessionId(config.getSessionId()); + this.outputPath = path.join( + config.getTargetDir(), + '.qwen', + 'generated-images', + sessionDir, + `${randomUUID()}.png`, + ); + } + + override toolLocations(): ToolLocation[] { + return [{ path: this.outputPath }]; + } + + override getDescription(): string { + const imageConfig = this.config.getImageGenerationConfig(); + const size = this.params.size ? ` at ${this.params.size}` : ''; + return `Generate an image with ${imageConfig?.model ?? 'the configured model'}${size}: ${this.params.prompt}`; + } + + override getDefaultPermission(): Promise<'ask'> { + return Promise.resolve('ask'); + } + + async execute(signal: AbortSignal): Promise { + const imageConfig = this.config.getImageGenerationConfig(); + if (!imageConfig) { + return failureResult( + 'Image generation is not configured with a valid endpoint.', + ); + } + const apiKey = process.env[imageConfig.apiKeyEnv]?.trim(); + if (!apiKey) { + return failureResult( + `Image generation requires the ${imageConfig.apiKeyEnv} environment variable.`, + ); + } + + try { + signal.throwIfAborted(); + const outputDir = path.dirname(this.outputPath); + Storage.assertPathWithinDirectory( + outputDir, + this.config.getTargetDir(), + 'Generated image path must stay inside the workspace.', + ); + await mkdir(outputDir, { recursive: true }); + Storage.assertPathWithinDirectory( + outputDir, + this.config.getTargetDir(), + 'Generated image path must stay inside the workspace.', + ); + signal.throwIfAborted(); + const generated = await this.generateImage({ + baseUrl: imageConfig.baseUrl, + apiKey, + model: imageConfig.model, + prompt: this.params.prompt, + size: this.params.size, + signal, + }); + signal.throwIfAborted(); + Storage.assertPathWithinDirectory( + outputDir, + this.config.getTargetDir(), + 'Generated image path must stay inside the workspace.', + ); + await atomicWriteFile(this.outputPath, generated.bytes, { + mode: 0o600, + noFollow: true, + }); + + const workspacePath = path + .relative(this.config.getTargetDir(), this.outputPath) + .split(path.sep) + .join('/'); + const metadata: Record = { + model: imageConfig.model, + ...(generated.requestId ? { requestId: generated.requestId } : {}), + ...(this.params.size ? { size: this.params.size } : {}), + }; + const llmContent: Part[] = [ + { + text: `Generated image saved to ${this.outputPath}.`, + }, + ]; + if (this.config.getEffectiveInputModalities().image === true) { + llmContent.push({ + inlineData: { + mimeType: generated.mimeType, + data: generated.bytes.toString('base64'), + }, + }); + } + + return { + llmContent, + returnDisplay: `Generated image saved to **${this.outputPath}**.`, + resultFilePaths: [this.outputPath], + artifacts: [ + { + title: 'Generated image', + kind: 'image', + storage: 'workspace', + workspacePath, + mimeType: generated.mimeType, + sizeBytes: generated.bytes.length, + metadata, + }, + ], + }; + } catch (error) { + return failureResult( + error instanceof Error ? error.message : getErrorMessage(error), + ); + } + } +} + +export class ImageGenTool extends BaseDeclarativeTool< + ImageGenParams, + ToolResult +> { + static readonly Name = ToolNames.IMAGE_GEN; + + constructor( + private readonly config: Config, + private readonly generateImage: GenerateImage = generateConfiguredImage, + ) { + super( + ImageGenTool.Name, + ToolDisplayNames.IMAGE_GEN, + 'Generates a PNG image with the configured image model and saves it as a workspace artifact. Use size in width*height form when the user requests a specific aspect ratio.', + Kind.Execute, + { + type: 'object', + properties: { + prompt: { + type: 'string', + minLength: 1, + maxLength: MAX_PROMPT_CHARS, + description: 'Detailed text description of the image to generate.', + }, + size: { + type: 'string', + pattern: '^\\d+\\*\\d+$', + description: + 'Optional output size in width*height form, for example 1536*864.', + }, + }, + required: ['prompt'], + }, + true, + false, + false, + false, + 'image generation picture poster illustration', + ); + } + + protected override validateToolParamValues( + params: ImageGenParams, + ): string | null { + params.prompt = params.prompt.trim(); + if (!params.prompt) { + return 'The image prompt must be non-empty.'; + } + if (params.prompt.length > MAX_PROMPT_CHARS) { + return `The image prompt must not exceed ${MAX_PROMPT_CHARS} characters.`; + } + if (!params.size) { + return null; + } + + const match = /^(\d+)\*(\d+)$/.exec(params.size); + if (!match) { + return 'Image size must use width*height form, for example 1536*864.'; + } + const width = Number(match[1]); + const height = Number(match[2]); + const totalPixels = width * height; + if ( + !Number.isSafeInteger(totalPixels) || + totalPixels < MIN_TOTAL_PIXELS || + totalPixels > MAX_TOTAL_PIXELS + ) { + return `Image size total pixels must be between 512*512 and 2048*2048.`; + } + return null; + } + + protected createInvocation( + params: ImageGenParams, + ): ToolInvocation { + return new ImageGenInvocation(this.config, this.generateImage, params); + } +} + +function failureResult(message: string): ToolResult { + return { + llmContent: `Image generation failed: ${message}`, + returnDisplay: `Image generation failed: ${message}`, + error: { + message, + type: ToolErrorType.EXECUTION_FAILED, + }, + }; +} diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index 21668ed38aa..f685ddf543f 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -32,6 +32,7 @@ export const ToolNames = { ENTER_PLAN_MODE: 'enter_plan_mode', WEB_FETCH: 'web_fetch', WEB_SEARCH: 'web_search', + IMAGE_GEN: 'image_gen', LS: 'list_directory', LSP: 'lsp', ASK_USER_QUESTION: 'ask_user_question', @@ -87,6 +88,7 @@ export const ToolDisplayNames = { ENTER_PLAN_MODE: 'EnterPlanMode', WEB_FETCH: 'WebFetch', WEB_SEARCH: 'WebSearch', + IMAGE_GEN: 'ImageGen', LS: 'ListFiles', LSP: 'Lsp', ASK_USER_QUESTION: 'AskUserQuestion', diff --git a/packages/vscode-ide-companion/schemas/settings.schema.json b/packages/vscode-ide-companion/schemas/settings.schema.json index e289910eb8b..383305b2d66 100644 --- a/packages/vscode-ide-companion/schemas/settings.schema.json +++ b/packages/vscode-ide-companion/schemas/settings.schema.json @@ -563,6 +563,11 @@ "type": "string", "default": "" }, + "imageModel": { + "description": "Model used by the built-in image_gen tool. Set with /model --image. The selected model must be marked imageOnly in modelProviders.", + "type": "string", + "default": "" + }, "visionBridgeTimeoutMs": { "description": "Per-attempt timeout in milliseconds for the vision bridge image transcription call (a positive integer up to 2147483647). Unset uses the built-in 30s. Raise for slow or proxied vision endpoints.", "type": "integer", diff --git a/packages/web-shell/client/components/messages/toolFormatting.ts b/packages/web-shell/client/components/messages/toolFormatting.ts index c38cc6408f1..23ff41fe263 100644 --- a/packages/web-shell/client/components/messages/toolFormatting.ts +++ b/packages/web-shell/client/components/messages/toolFormatting.ts @@ -54,6 +54,7 @@ export const TOOL_DISPLAY_NAMES: Record = { artifact: 'Artifact', record_artifact: 'RecordArtifact', web_search: 'WebSearch', + image_gen: 'ImageGen', bash: 'Shell', shell: 'Shell Command', read: 'ReadFile', diff --git a/packages/web-shell/client/i18n.tsx b/packages/web-shell/client/i18n.tsx index 8476d531d19..797ab29a093 100644 --- a/packages/web-shell/client/i18n.tsx +++ b/packages/web-shell/client/i18n.tsx @@ -2384,6 +2384,7 @@ const ZH: Messages = { 'toolName.read_mcp_resource': '读取 MCP 资源', 'toolName.artifact': '制品', 'toolName.record_artifact': '记录制品', + 'toolName.image_gen': '生成图片', // web-shell-only wire aliases (see TOOL_DISPLAY_NAMES in toolFormatting.ts) 'toolName.bash': '运行命令', 'toolName.shell': 'Shell 命令',