Cloud account switcher (#8223)

* Cloud account switcher

* Bare metal evals fixes (#8224)

Co-authored-by: Roo Code <[email protected]>

* Rounded icons in chat account switcher

* Visual tweaks to CloudView

* Remove hardcoded timeout

* Safer check

* PR feedback

* Fix test

* Check for org mismatches in handleCredentialsChange

* Cloud: use the existing auth event flow to handle org switching

* Cleanup: broadcast() might be confusingly named.

---------

Co-authored-by: Chris Estreich <[email protected]>
Co-authored-by: Roo Code <[email protected]>
Co-authored-by: Bruno Bergher <[email protected]>
Co-authored-by: John Richmond <[email protected]>

Author: 5 people

packages/cloud/src/CloudService.ts

@@ -8,6 +8,7 @@ import type {
 	AuthService,
 	SettingsService,
 	CloudUserInfo,
+	CloudOrganizationMembership,
 	OrganizationAllowList,
 	OrganizationSettings,
 	ShareVisibility,
@@ -242,6 +243,21 @@ 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()
+
+		// Perform the organization switch
+		// StaticTokenAuthService will throw an error if organization switching is not supported
+		await this.authService!.switchOrganization(organizationId)
+	}
+
+	public async getOrganizationMemberships(): Promise<CloudOrganizationMembership[]> {
+		this.ensureInitialized()
+
+		// StaticTokenAuthService will throw an error if organization memberships are not supported
+		return await this.authService!.getOrganizationMemberships()
+	}
+
 	// SettingsService
 
 	public getAllowList(): OrganizationAllowList {

packages/cloud/src/StaticTokenAuthService.ts

@@ -63,6 +63,14 @@ export class StaticTokenAuthService extends EventEmitter<AuthServiceEvents> impl
 		throw new Error("Authentication methods are disabled in StaticTokenAuthService")
 	}
 
+	public async switchOrganization(_organizationId: string | null): Promise<void> {
+		throw new Error("Authentication methods are disabled in StaticTokenAuthService")
+	}
+
+	public async getOrganizationMemberships(): Promise<import("@roo-code/types").CloudOrganizationMembership[]> {
+		throw new Error("Authentication methods are disabled in StaticTokenAuthService")
+	}
+
 	public getState(): AuthState {
 		return this.state
 	}

packages/cloud/src/WebAuthService.ts

@@ -141,7 +141,8 @@ export class WebAuthService extends EventEmitter<AuthServiceEvents> implements A
 				if (
 					this.credentials === null ||
 					this.credentials.clientToken !== credentials.clientToken ||
-					this.credentials.sessionId !== credentials.sessionId
+					this.credentials.sessionId !== credentials.sessionId ||
+					this.credentials.organizationId !== credentials.organizationId
 				) {
 					this.transitionToAttemptingSession(credentials)
 				}
@@ -174,6 +175,7 @@ export class WebAuthService extends EventEmitter<AuthServiceEvents> implements A
 
 		this.changeState("attempting-session")
 
+		this.timer.stop()
 		this.timer.start()
 	}
 
@@ -469,6 +471,42 @@ 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, handleCredentialsChange will handle the update
+		await this.storeCredentials(updatedCredentials)
+	}
+
+	/**
+	 * 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")
@@ -653,9 +691,14 @@ export class WebAuthService extends EventEmitter<AuthServiceEvents> implements A
 	}
 
 	private async clerkGetOrganizationMemberships(): Promise<CloudOrganizationMembership[]> {
+		if (!this.credentials) {
+			this.log("[auth] Cannot get organization memberships: missing credentials")
+			return []
+		}
+
 		const response = await fetch(`${getClerkBaseUrl()}/v1/me/organization_memberships`, {
 			headers: {
-				Authorization: `Bearer ${this.credentials!.clientToken}`,
+				Authorization: `Bearer ${this.credentials.clientToken}`,
 				"User-Agent": this.userAgent(),
 			},
 			signal: AbortSignal.timeout(10000),

packages/types/src/cloud.ts

@@ -242,6 +242,7 @@ export interface AuthService extends EventEmitter<AuthServiceEvents> {
 	login(landingPageSlug?: string): Promise<void>
 	logout(): Promise<void>
 	handleCallback(code: string | null, state: string | null, organizationId?: string | null): Promise<void>
+	switchOrganization(organizationId: string | null): Promise<void>
 
 	// State methods
 	getState(): AuthState
@@ -253,6 +254,9 @@ export interface AuthService extends EventEmitter<AuthServiceEvents> {
 	getSessionToken(): string | undefined
 	getUserInfo(): CloudUserInfo | null
 	getStoredOrganizationId(): string | null
+
+	// Organization management
+	getOrganizationMemberships(): Promise<CloudOrganizationMembership[]>
 }
 
 /**

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,38 @@ 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()
+
+				// Send success response back to webview
+				await provider.postMessageToWebview({
+					type: "organizationSwitchResult",
+					success: true,
+					organizationId: organizationId,
+				})
+			} catch (error) {
+				provider.log(`Organization switch failed: ${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
+		}
 
 		case "saveCodeIndexSettingsAtomic": {
 			if (!message.codeIndexSettings) {

src/shared/ExtensionMessage.ts

@@ -10,6 +10,7 @@ import type {
 	MarketplaceItem,
 	TodoItem,
 	CloudUserInfo,
+	CloudOrganizationMembership,
 	OrganizationAllowList,
 	ShareVisibility,
 	QueuedMessage,
@@ -124,6 +125,7 @@ export interface ExtensionMessage {
 		| "commands"
 		| "insertTextIntoTextarea"
 		| "dismissedUpsells"
+		| "organizationSwitchResult"
 	text?: string
 	payload?: any // Add a generic payload for now, can refine later
 	action?:
@@ -202,6 +204,7 @@ export interface ExtensionMessage {
 	commands?: Command[]
 	queuedMessages?: QueuedMessage[]
 	list?: string[] // For dismissedUpsells
+	organizationId?: string | null // For organizationSwitchResult
 }
 
 export type ExtensionState = Pick<
@@ -327,6 +330,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>