diff --git a/src/core/config/CustomModesManager.ts b/src/core/config/CustomModesManager.ts index a9a2e6a6b55a..a243a9236bec 100644 --- a/src/core/config/CustomModesManager.ts +++ b/src/core/config/CustomModesManager.ts @@ -40,6 +40,7 @@ interface ExportResult { interface ImportResult { success: boolean + slug?: string error?: string } @@ -411,7 +412,7 @@ export class CustomModesManager { const errorMessage = `Invalid mode configuration: ${errorMessages}` logger.error("Mode validation failed", { slug, errors: validationResult.error.errors }) vscode.window.showErrorMessage(t("common:customModes.errors.updateFailed", { error: errorMessage })) - return + throw new Error(errorMessage) } const isProjectMode = config.source === "project" @@ -457,6 +458,7 @@ export class CustomModesManager { const errorMessage = error instanceof Error ? error.message : String(error) logger.error("Failed to update custom mode", { slug, error: errorMessage }) vscode.window.showErrorMessage(t("common:customModes.errors.updateFailed", { error: errorMessage })) + throw error } } @@ -989,7 +991,8 @@ export class CustomModesManager { // Refresh the modes after import await this.refreshMergedState() - return { success: true } + // Return the imported mode's slug so the UI can activate it + return { success: true, slug: importData.customModes[0]?.slug } } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) logger.error("Failed to import mode with rules", { error: errorMessage }) diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index af5f9925c353..142520dbabfe 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -1981,6 +1981,7 @@ export const webviewMessageHandler = async ( break case "updateCustomMode": if (message.modeConfig) { + try { // Check if this is a new mode or an update to an existing mode const existingModes = await provider.customModesManager.getCustomModes() const isNewMode = !existingModes.some((mode) => mode.slug === message.modeConfig?.slug) @@ -2016,6 +2017,10 @@ export const webviewMessageHandler = async ( } } } + } catch (error) { + // Error already shown to user by updateCustomMode + // Just prevent unhandled rejection and skip state updates + } } break case "deleteCustomMode": @@ -2222,10 +2227,11 @@ export const webviewMessageHandler = async ( await updateGlobalState("customModes", customModes) await provider.postStateToWebview() - // Send success message to webview + // Send success message to webview, include the imported slug so UI can switch provider.postMessageToWebview({ type: "importModeResult", success: true, + slug: result.slug, }) // Show success message diff --git a/webview-ui/src/components/modes/ModesView.tsx b/webview-ui/src/components/modes/ModesView.tsx index c50996585fe7..1107eb44e10a 100644 --- a/webview-ui/src/components/modes/ModesView.tsx +++ b/webview-ui/src/components/modes/ModesView.tsx @@ -20,6 +20,7 @@ import { getCustomInstructions, getAllModes, findModeBySlug as findCustomModeBySlug, + defaultModeSlug, } from "@roo/modes" import { TOOL_GROUPS } from "@roo/tools" @@ -55,6 +56,8 @@ const availableGroups = (Object.keys(TOOL_GROUPS) as ToolGroup[]).filter((group) type ModeSource = "global" | "project" +type ImportModeResult = { type: "importModeResult"; success: boolean; slug?: string; error?: string } + type ModesViewProps = { onDone: () => void } @@ -186,6 +189,29 @@ const ModesView = ({ onDone }: ModesViewProps) => { [visualMode, switchMode], ) + // Refs to track latest state/functions for message handler (which has no dependencies) + const handleModeSwitchRef = useRef(handleModeSwitch) + const customModesRef = useRef(customModes) + const switchModeRef = useRef(switchMode) + + // Update refs when dependencies change + useEffect(() => { + handleModeSwitchRef.current = handleModeSwitch + }, [handleModeSwitch]) + + useEffect(() => { + customModesRef.current = customModes + }, [customModes]) + + useEffect(() => { + switchModeRef.current = switchMode + }, [switchMode]) + + // Sync visualMode with backend mode changes to prevent desync + useEffect(() => { + setVisualMode(mode) + }, [mode]) + // Handler for popover open state change const onOpenChange = useCallback((open: boolean) => { setOpen(open) @@ -460,7 +486,21 @@ const ModesView = ({ onDone }: ModesViewProps) => { setIsImporting(false) setShowImportDialog(false) - if (!message.success) { + if (message.success) { + const { slug } = message as ImportModeResult + if (slug) { + // Try switching using the freshest mode list available + const all = getAllModes(customModesRef.current) + const importedMode = all.find((m) => m.slug === slug) + if (importedMode) { + handleModeSwitchRef.current(importedMode) + } else { + // Fallback: slug not yet in state (race condition) - select default mode + setVisualMode(defaultModeSlug) + switchModeRef.current?.(defaultModeSlug) + } + } + } else { // Only log error if it's not a cancellation if (message.error !== "cancelled") { console.error("Failed to import mode:", message.error) diff --git a/webview-ui/src/components/modes/__tests__/ModesView.import-switch.spec.tsx b/webview-ui/src/components/modes/__tests__/ModesView.import-switch.spec.tsx new file mode 100644 index 000000000000..b62fc8856368 --- /dev/null +++ b/webview-ui/src/components/modes/__tests__/ModesView.import-switch.spec.tsx @@ -0,0 +1,156 @@ +// npx vitest src/components/modes/__tests__/ModesView.import-switch.spec.tsx + +import { render, waitFor } from "@/utils/test-utils" +import ModesView from "../ModesView" +import { ExtensionStateContext } from "@src/context/ExtensionStateContext" +import { vscode } from "@src/utils/vscode" +import { defaultModeSlug } from "@roo/modes" + +// Mock vscode API +vitest.mock("@src/utils/vscode", () => ({ + vscode: { + postMessage: vitest.fn(), + }, +})) + +const mockExtensionState = { + customModePrompts: {}, + listApiConfigMeta: [ + { id: "config1", name: "Config 1" }, + { id: "config2", name: "Config 2" }, + ], + enhancementApiConfigId: "", + setEnhancementApiConfigId: vitest.fn(), + mode: "code", + customModes: [], + customSupportPrompts: [], + currentApiConfigName: "", + customInstructions: "", + setCustomInstructions: vitest.fn(), +} + +const renderModesView = (props = {}) => { + const mockOnDone = vitest.fn() + return render( + + + , + ) +} + +Element.prototype.scrollIntoView = vitest.fn() + +describe("ModesView Import Auto-Switch", () => { + beforeEach(() => { + vitest.clearAllMocks() + }) + + it("should auto-switch to imported mode when found in current state", async () => { + const importedModeSlug = "custom-test-mode" + const customModes = [ + { + slug: importedModeSlug, + name: "Custom Test Mode", + roleDefinition: "Test role", + groups: [], + }, + ] + + renderModesView({ customModes }) + + // Simulate successful import message with the mode already in state + const importMessage = { + data: { + type: "importModeResult", + success: true, + slug: importedModeSlug, + }, + } + + window.dispatchEvent(new MessageEvent("message", importMessage)) + + // Wait for the mode switch message to be sent + await waitFor(() => { + expect(vscode.postMessage).toHaveBeenCalledWith({ + type: "mode", + text: importedModeSlug, + }) + }) + }) + + it("should fallback to architect mode when imported slug not yet in state (race condition)", async () => { + const importedModeSlug = "custom-new-mode" + + // Render without the imported mode in customModes (simulating race condition) + renderModesView({ customModes: [] }) + + // Simulate successful import message but mode not yet in state + const importMessage = { + data: { + type: "importModeResult", + success: true, + slug: importedModeSlug, + }, + } + + window.dispatchEvent(new MessageEvent("message", importMessage)) + + // Wait for the fallback to default mode (architect) + await waitFor(() => { + expect(vscode.postMessage).toHaveBeenCalledWith({ + type: "mode", + text: defaultModeSlug, + }) + }) + }) + + it("should not switch modes on import failure", async () => { + renderModesView() + + // Simulate failed import message + const importMessage = { + data: { + type: "importModeResult", + success: false, + error: "Import failed", + }, + } + + window.dispatchEvent(new MessageEvent("message", importMessage)) + + // Wait a bit to ensure no mode switch happens + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Verify no mode switch message was sent + expect(vscode.postMessage).not.toHaveBeenCalledWith( + expect.objectContaining({ + type: "mode", + }), + ) + }) + + it("should not switch modes on cancelled import", async () => { + renderModesView() + + // Simulate cancelled import message + const importMessage = { + data: { + type: "importModeResult", + success: false, + error: "cancelled", + }, + } + + window.dispatchEvent(new MessageEvent("message", importMessage)) + + // Wait a bit to ensure no mode switch happens + await new Promise((resolve) => setTimeout(resolve, 100)) + + // Verify no mode switch message was sent + expect(vscode.postMessage).not.toHaveBeenCalledWith( + expect.objectContaining({ + type: "mode", + }), + ) + }) +})