Cloud account switcher

Author: mrubens

packages/cloud/src/CloudService.ts

@@ -8,6 +8,7 @@ import type {
 	AuthService,
 	SettingsService,
 	CloudUserInfo,
+	CloudOrganizationMembership,
 	OrganizationAllowList,
 	OrganizationSettings,
 	ShareVisibility,
@@ -242,6 +243,35 @@ export class CloudService extends EventEmitter<CloudServiceEvents> implements Di
 		return this.authService!.handleCallback(code, state, organizationId)
 	}
 
+	public async switchOrganization(organizationId: string | null): Promise<void> {
+		this.ensureInitialized()
+
+		// Check if we have a WebAuthService (not StaticTokenAuthService)
+		if (!(this.authService instanceof WebAuthService)) {
+			throw new Error("Organization switching is only available with web authentication")
+		}
+
+		const webAuthService = this.authService as WebAuthService
+
+		// Perform the organization switch
+		await webAuthService.switchOrganization(organizationId)
+
+		// Broadcast the change
+		await this.authService.broadcast()
+	}
+
+	public async getOrganizationMemberships(): Promise<CloudOrganizationMembership[]> {
+		this.ensureInitialized()
+
+		// Check if we have a WebAuthService (not StaticTokenAuthService)
+		if (!(this.authService instanceof WebAuthService)) {
+			return []
+		}
+
+		const webAuthService = this.authService as WebAuthService
+		return await webAuthService.getOrganizationMemberships()
+	}
+
 	// SettingsService
 
 	public getAllowList(): OrganizationAllowList {

packages/cloud/src/WebAuthService.ts

@@ -469,6 +469,60 @@ export class WebAuthService extends EventEmitter<AuthServiceEvents> implements A
 		return this.credentials?.organizationId || null
 	}
 
+	/**
+	 * Switch to a different organization context
+	 * @param organizationId The organization ID to switch to, or null for personal account
+	 */
+	public async switchOrganization(organizationId: string | null): Promise<void> {
+		if (!this.credentials) {
+			throw new Error("Cannot switch organization: not authenticated")
+		}
+
+		// Update the stored credentials with the new organization ID
+		const updatedCredentials: AuthCredentials = {
+			...this.credentials,
+			organizationId: organizationId,
+		}
+
+		// Store the updated credentials
+		await this.storeCredentials(updatedCredentials)
+
+		// Update the local credentials
+		this.credentials = updatedCredentials
+
+		// Clear the current session token to force a refresh with new org context
+		this.sessionToken = null
+
+		// Trigger a session refresh to get a new token with the correct org context
+		try {
+			await this.refreshSession()
+			// Fetch updated user info after organization switch to reflect new context
+			await this.fetchUserInfo()
+		} catch (error) {
+			this.log(`[auth] Failed to refresh session after organization switch: ${error}`)
+			// Even if refresh fails, the credentials are updated for next attempt
+		}
+
+		this.log(`[auth] Switched organization context to: ${organizationId || "personal account"}`)
+	}
+
+	/**
+	 * Get all organization memberships for the current user
+	 * @returns Array of organization memberships
+	 */
+	public async getOrganizationMemberships(): Promise<CloudOrganizationMembership[]> {
+		if (!this.credentials) {
+			return []
+		}
+
+		try {
+			return await this.clerkGetOrganizationMemberships()
+		} catch (error) {
+			this.log(`[auth] Failed to get organization memberships: ${error}`)
+			return []
+		}
+	}
+
 	private async clerkSignIn(ticket: string): Promise<AuthCredentials> {
 		const formData = new URLSearchParams()
 		formData.append("strategy", "ticket")

src/core/webview/ClineProvider.ts

@@ -30,6 +30,7 @@ import {
 	type TerminalActionPromptType,
 	type HistoryItem,
 	type CloudUserInfo,
+	type CloudOrganizationMembership,
 	type CreateTaskOptions,
 	type TokenUsage,
 	RooCodeEventName,
@@ -1815,6 +1816,16 @@ export class ClineProvider
 			featureRoomoteControlEnabled,
 		} = await this.getState()
 
+		let cloudOrganizations: CloudOrganizationMembership[] = []
+
+		try {
+			cloudOrganizations = await CloudService.instance.getOrganizationMemberships()
+		} catch (error) {
+			console.error(
+				`[getStateToPostToWebview] failed to get cloud organizations: ${error instanceof Error ? error.message : String(error)}`,
+			)
+		}
+
 		const telemetryKey = process.env.POSTHOG_API_KEY
 		const machineId = vscode.env.machineId
 		const mergedAllowedCommands = this.mergeAllowedCommands(allowedCommands)
@@ -1916,6 +1927,7 @@ export class ClineProvider
 			historyPreviewCollapsed: historyPreviewCollapsed ?? false,
 			cloudUserInfo,
 			cloudIsAuthenticated: cloudIsAuthenticated ?? false,
+			cloudOrganizations,
 			sharingEnabled: sharingEnabled ?? false,
 			organizationAllowList,
 			organizationSettingsVersion,

src/core/webview/webviewMessageHandler.ts

@@ -2379,6 +2379,23 @@ export const webviewMessageHandler = async (
 
 			break
 		}
+		case "switchOrganization": {
+			try {
+				const organizationId = message.organizationId ?? null
+
+				// Switch to the new organization context
+				await CloudService.instance.switchOrganization(organizationId)
+
+				// Refresh the state to update UI
+				await provider.postStateToWebview()
+			} catch (error) {
+				provider.log(`Organization switch failed: ${error}`)
+				vscode.window.showErrorMessage(
+					`Failed to switch organization: ${error instanceof Error ? error.message : String(error)}`,
+				)
+			}
+			break
+		}
 
 		case "saveCodeIndexSettingsAtomic": {
 			if (!message.codeIndexSettings) {

src/shared/ExtensionMessage.ts

@@ -10,6 +10,7 @@ import type {
 	MarketplaceItem,
 	TodoItem,
 	CloudUserInfo,
+	CloudOrganizationMembership,
 	OrganizationAllowList,
 	ShareVisibility,
 	QueuedMessage,
@@ -327,6 +328,7 @@ export type ExtensionState = Pick<
 	cloudUserInfo: CloudUserInfo | null
 	cloudIsAuthenticated: boolean
 	cloudApiUrl?: string
+	cloudOrganizations?: CloudOrganizationMembership[]
 	sharingEnabled: boolean
 	organizationAllowList: OrganizationAllowList
 	organizationSettingsVersion?: number

src/shared/WebviewMessage.ts

@@ -184,6 +184,7 @@ export interface WebviewMessage {
 		| "cloudLandingPageSignIn"
 		| "rooCloudSignOut"
 		| "rooCloudManualUrl"
+		| "switchOrganization"
 		| "condenseTaskContextRequest"
 		| "requestIndexingStatus"
 		| "startIndexing"
@@ -274,6 +275,7 @@ export interface WebviewMessage {
 	checkOnly?: boolean // For deleteCustomMode check
 	upsellId?: string // For dismissUpsell
 	list?: string[] // For dismissedUpsells response
+	organizationId?: string | null // For organization switching
 	codeIndexSettings?: {
 		// Global state settings
 		codebaseIndexEnabled: boolean

webview-ui/src/App.tsx

@@ -76,6 +76,7 @@ const App = () => {
 		cloudUserInfo,
 		cloudIsAuthenticated,
 		cloudApiUrl,
+		cloudOrganizations,
 		renderContext,
 		mdmCompliant,
 	} = useExtensionState()
@@ -267,6 +268,7 @@ const App = () => {
 					userInfo={cloudUserInfo}
 					isAuthenticated={cloudIsAuthenticated}
 					cloudApiUrl={cloudApiUrl}
+					organizations={cloudOrganizations}
 					onDone={() => switchTab("chat")}
 				/>
 			)}

webview-ui/src/components/chat/ChatTextArea.tsx

@@ -31,6 +31,7 @@ import { MAX_IMAGES_PER_MESSAGE } from "./ChatView"
 import ContextMenu from "./ContextMenu"
 import { IndexingStatusBadge } from "./IndexingStatusBadge"
 import { usePromptHistory } from "./hooks/usePromptHistory"
+import { CloudAccountSwitcher } from "../cloud/CloudAccountSwitcher"
 
 interface ChatTextAreaProps {
 	inputValue: string
@@ -87,6 +88,8 @@ export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 			taskHistory,
 			clineMessages,
 			commands,
+			cloudUserInfo,
+			cloudOrganizations,
 		} = useExtensionState()
 
 		// Find the ID and display text for the currently selected API configuration.
@@ -1232,7 +1235,13 @@ export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 						/>
 						<AutoApproveDropdown triggerClassName="min-w-[28px] text-ellipsis overflow-hidden flex-shrink" />
 					</div>
-					<div className="flex flex-shrink-0 items-center gap-0.5 pr-2">
+					<div
+						className={cn(
+							"flex flex-shrink-0 items-center gap-0.5",
+							!isEditMode && cloudOrganizations && cloudOrganizations.length > 0 && cloudUserInfo
+								? ""
+								: "pr-2",
+						)}>
 						{isTtsPlaying && (
 							<StandardTooltip content={t("chat:stopTts")}>
 								<button
@@ -1254,6 +1263,9 @@ export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 							</StandardTooltip>
 						)}
 						{!isEditMode ? <IndexingStatusBadge /> : null}
+						{!isEditMode && cloudOrganizations && cloudOrganizations.length > 0 && cloudUserInfo && (
+							<CloudAccountSwitcher />
+						)}
 					</div>
 				</div>
 			</div>