Remove hardcoded timeout

Author: mrubens

src/core/webview/webviewMessageHandler.ts

@@ -2388,11 +2388,26 @@ export const webviewMessageHandler = async (
 
 				// Refresh the state to update UI
 				await provider.postStateToWebview()
+
+				// Send success response back to webview
+				await provider.postMessageToWebview({
+					type: "organizationSwitchResult",
+					success: true,
+					organizationId: organizationId,
+				})
 			} catch (error) {
 				provider.log(`Organization switch failed: ${error}`)
-				vscode.window.showErrorMessage(
-					`Failed to switch organization: ${error instanceof Error ? error.message : String(error)}`,
-				)
+				const errorMessage = error instanceof Error ? error.message : String(error)
+
+				// Send error response back to webview
+				await provider.postMessageToWebview({
+					type: "organizationSwitchResult",
+					success: false,
+					error: errorMessage,
+					organizationId: message.organizationId ?? null,
+				})
+
+				vscode.window.showErrorMessage(`Failed to switch organization: ${errorMessage}`)
 			}
 			break
 		}

src/shared/ExtensionMessage.ts

@@ -125,6 +125,7 @@ export interface ExtensionMessage {
 		| "commands"
 		| "insertTextIntoTextarea"
 		| "dismissedUpsells"
+		| "organizationSwitchResult"
 	text?: string
 	payload?: any // Add a generic payload for now, can refine later
 	action?:
@@ -203,6 +204,7 @@ export interface ExtensionMessage {
 	commands?: Command[]
 	queuedMessages?: QueuedMessage[]
 	list?: string[] // For dismissedUpsells
+	organizationId?: string | null // For organizationSwitchResult
 }
 
 export type ExtensionState = Pick<

webview-ui/src/components/cloud/OrganizationSwitcher.tsx

@@ -4,6 +4,7 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue, SelectSe
 import { type CloudUserInfo, type CloudOrganizationMembership } from "@roo-code/types"
 import { useAppTranslation } from "@src/i18n/TranslationContext"
 import { vscode } from "@src/utils/vscode"
+import { type ExtensionMessage } from "@roo/ExtensionMessage"
 
 type OrganizationSwitcherProps = {
 	userInfo: CloudUserInfo
@@ -21,6 +22,28 @@ export const OrganizationSwitcher = ({ userInfo, organizations, onOrganizationCh
 		setSelectedOrgId(userInfo.organizationId || null)
 	}, [userInfo.organizationId])
 
+	// Listen for organization switch results
+	useEffect(() => {
+		const handleMessage = (event: MessageEvent) => {
+			const message = event.data as ExtensionMessage
+			if (message.type === "organizationSwitchResult") {
+				// Reset loading state when we receive the result
+				setIsLoading(false)
+
+				if (message.success) {
+					// Update selected org based on the result
+					setSelectedOrgId(message.organizationId ?? null)
+				} else {
+					// Revert to the previous organization on error
+					setSelectedOrgId(userInfo.organizationId || null)
+				}
+			}
+		}
+
+		window.addEventListener("message", handleMessage)
+		return () => window.removeEventListener("message", handleMessage)
+	}, [userInfo.organizationId])
+
 	const handleOrganizationChange = async (value: string) => {
 		const newOrgId = value === "personal" ? null : value
 
@@ -44,11 +67,6 @@ export const OrganizationSwitcher = ({ userInfo, organizations, onOrganizationCh
 		if (onOrganizationChange) {
 			onOrganizationChange(newOrgId)
 		}
-
-		// Reset loading state after a delay
-		setTimeout(() => {
-			setIsLoading(false)
-		}, 1000)
 	}
 
 	// If user has no organizations, don't show the switcher