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..39d9237c55fe 100644 --- a/webview-ui/src/components/modes/ModesView.tsx +++ b/webview-ui/src/components/modes/ModesView.tsx @@ -55,6 +55,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 +188,28 @@ const ModesView = ({ onDone }: ModesViewProps) => { [visualMode, switchMode], ) + // Keep latest handleModeSwitch and customModes available inside window message handler + const handleModeSwitchRef = useRef(handleModeSwitch) + useEffect(() => { + handleModeSwitchRef.current = handleModeSwitch + }, [handleModeSwitch]) + + const customModesRef = useRef(customModes) + useEffect(() => { + customModesRef.current = customModes + }, [customModes]) + + // Keep latest switchMode available inside window message handler + const switchModeRef = useRef(switchMode) + 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 +484,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: switch by slug to keep backend in sync and update visual selection + setVisualMode(slug) + switchModeRef.current?.(slug) + } + } + } 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..8be38e4cce69 --- /dev/null +++ b/webview-ui/src/components/modes/__tests__/ModesView.import-switch.spec.tsx @@ -0,0 +1,118 @@ +import { render, screen, waitFor } from "@/utils/test-utils" +import ModesView from "../ModesView" +import { ExtensionStateContext } from "@src/context/ExtensionStateContext" +import { vscode } from "@src/utils/vscode" + +vitest.mock("@src/utils/vscode", () => ({ + vscode: { + postMessage: vitest.fn(), + }, +})) + +const baseState = { + customModePrompts: {}, + listApiConfigMeta: [], + enhancementApiConfigId: "", + setEnhancementApiConfigId: vitest.fn(), + mode: "code", + customModes: [], + customSupportPrompts: [], + currentApiConfigName: "", + customInstructions: "", + setCustomInstructions: vitest.fn(), +} + +describe("ModesView - auto switch after import", () => { + beforeEach(() => { + vitest.clearAllMocks() + }) + + it("switches to imported mode when import succeeds and slug is provided", async () => { + const importedMode = { + slug: "imported-mode", + name: "Imported Mode", + roleDefinition: "Role", + groups: ["read"] as const, + source: "global" as const, + } + + render( + + + , + ) + + const trigger = screen.getByTestId("mode-select-trigger") + expect(trigger).toHaveTextContent("Code") + + // Simulate extension sending successful import result with slug + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "importModeResult", success: true, slug: "imported-mode" }, + }), + ) + + // Backend switch message sent + await waitFor(() => { + expect(vscode.postMessage).toHaveBeenCalledWith({ type: "mode", text: "imported-mode" }) + }) + + // UI reflects new mode selection + await waitFor(() => { + expect(trigger).toHaveTextContent("Imported Mode") + }) + }) + + it("does not switch when import fails or slug missing", async () => { + render( + + + , + ) + + const trigger = screen.getByTestId("mode-select-trigger") + expect(trigger).toHaveTextContent("Code") + + // Import failure + window.dispatchEvent( + new MessageEvent("message", { data: { type: "importModeResult", success: false, error: "x" } }), + ) + + await waitFor(() => { + expect(vscode.postMessage).not.toHaveBeenCalledWith({ type: "mode", text: expect.any(String) }) + }) + expect(trigger).toHaveTextContent("Code") + + // Success but no slug provided + window.dispatchEvent(new MessageEvent("message", { data: { type: "importModeResult", success: true } })) + + await waitFor(() => { + expect(vscode.postMessage).not.toHaveBeenCalledWith({ type: "mode", text: expect.any(String) }) + }) + expect(trigger).toHaveTextContent("Code") + }) + + it("uses fallback branch when imported slug not yet present in customModes", async () => { + // Render with empty customModes - imported mode hasn't been added to state yet + render( + + + , + ) + + const trigger = screen.getByTestId("mode-select-trigger") + expect(trigger).toHaveTextContent("Code") + + // Simulate successful import for a slug not yet in customModes (timing race condition) + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "importModeResult", success: true, slug: "not-yet-loaded-mode" }, + }), + ) + + // Fallback branch should send backend switch message + await waitFor(() => { + expect(vscode.postMessage).toHaveBeenCalledWith({ type: "mode", text: "not-yet-loaded-mode" }) + }) + }) +})