Store the organization id in credentials

Author: mrubens

packages/cloud/src/AuthService.ts

@@ -21,6 +21,7 @@ export interface AuthServiceEvents {
 const authCredentialsSchema = z.object({
 	clientToken: z.string().min(1, "Client token cannot be empty"),
 	sessionId: z.string().min(1, "Session ID cannot be empty"),
+	organizationId: z.string().nullable().optional(),
 })
 
 type AuthCredentials = z.infer<typeof authCredentialsSchema>
@@ -220,7 +221,16 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
 
 		try {
 			const parsedJson = JSON.parse(credentialsJson)
-			return authCredentialsSchema.parse(parsedJson)
+			const credentials = authCredentialsSchema.parse(parsedJson)
+
+			// Migration: If no organizationId but we have userInfo, add it
+			if (credentials.organizationId === undefined && this.userInfo?.organizationId) {
+				credentials.organizationId = this.userInfo.organizationId
+				await this.storeCredentials(credentials)
+				this.log("[auth] Migrated credentials with organizationId")
+			}
+
+			return credentials
 		} catch (error) {
 			if (error instanceof z.ZodError) {
 				this.log("[auth] Invalid credentials format:", error.errors)
@@ -269,8 +279,13 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
 	 *
 	 * @param code The authorization code from the callback
 	 * @param state The state parameter from the callback
+	 * @param organizationId The organization ID from the callback (null for personal accounts)
 	 */
-	public async handleCallback(code: string | null, state: string | null): Promise<void> {
+	public async handleCallback(
+		code: string | null,
+		state: string | null,
+		organizationId?: string | null,
+	): Promise<void> {
 		if (!code || !state) {
 			vscode.window.showInformationMessage("Invalid Roo Code Cloud sign in url")
 			return
@@ -287,6 +302,9 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
 
 			const credentials = await this.clerkSignIn(code)
 
+			// Set organizationId (null for personal accounts)
+			credentials.organizationId = organizationId || null
+
 			await this.storeCredentials(credentials)
 
 			vscode.window.showInformationMessage("Successfully authenticated with Roo Code Cloud")
@@ -417,6 +435,15 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
 		return this.userInfo
 	}
 
+	/**
+	 * Get the stored organization ID from credentials
+	 *
+	 * @returns The stored organization ID, null for personal accounts or if no credentials exist
+	 */
+	public getStoredOrganizationId(): string | null {
+		return this.credentials?.organizationId || null
+	}
+
 	private async clerkSignIn(ticket: string): Promise<AuthCredentials> {
 		const formData = new URLSearchParams()
 		formData.append("strategy", "ticket")
@@ -454,6 +481,12 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
 		const formData = new URLSearchParams()
 		formData.append("_is_native", "1")
 
+		// Only add organization_id if not null (personal accounts)
+		const organizationId = this.getStoredOrganizationId()
+		if (organizationId !== null) {
+			formData.append("organization_id", organizationId)
+		}
+
 		const response = await fetch(`${getClerkBaseUrl()}/v1/client/sessions/${this.credentials!.sessionId}/tokens`, {
 			method: "POST",
 			headers: {

packages/cloud/src/CloudService.ts

@@ -118,14 +118,28 @@ export class CloudService {
 		return userInfo?.organizationRole || null
 	}
 
+	public hasStoredOrganizationId(): boolean {
+		this.ensureInitialized()
+		return this.authService!.getStoredOrganizationId() !== null
+	}
+
+	public getStoredOrganizationId(): string | null {
+		this.ensureInitialized()
+		return this.authService!.getStoredOrganizationId()
+	}
+
 	public getAuthState(): string {
 		this.ensureInitialized()
 		return this.authService!.getState()
 	}
 
-	public async handleAuthCallback(code: string | null, state: string | null): Promise<void> {
+	public async handleAuthCallback(
+		code: string | null,
+		state: string | null,
+		organizationId?: string | null,
+	): Promise<void> {
 		this.ensureInitialized()
-		return this.authService!.handleCallback(code, state)
+		return this.authService!.handleCallback(code, state, organizationId)
 	}
 
 	// SettingsService

packages/cloud/src/__tests__/AuthService.spec.ts

@@ -328,7 +328,7 @@ describe("AuthService", () => {
 
 			expect(mockContext.secrets.store).toHaveBeenCalledWith(
 				"clerk-auth-credentials",
-				JSON.stringify({ clientToken: "Bearer token-123", sessionId: "session-123" }),
+				JSON.stringify({ clientToken: "Bearer token-123", sessionId: "session-123", organizationId: null }),
 			)
 			expect(mockShowInfo).toHaveBeenCalledWith("Successfully authenticated with Roo Code Cloud")
 		})

packages/cloud/src/__tests__/CloudService.test.ts

@@ -40,6 +40,7 @@ describe("CloudService", () => {
 		getState: ReturnType<typeof vi.fn>
 		getSessionToken: ReturnType<typeof vi.fn>
 		handleCallback: ReturnType<typeof vi.fn>
+		getStoredOrganizationId: ReturnType<typeof vi.fn>
 		on: ReturnType<typeof vi.fn>
 		off: ReturnType<typeof vi.fn>
 		once: ReturnType<typeof vi.fn>
@@ -88,6 +89,7 @@ describe("CloudService", () => {
 			getState: vi.fn().mockReturnValue("logged-out"),
 			getSessionToken: vi.fn(),
 			handleCallback: vi.fn(),
+			getStoredOrganizationId: vi.fn().mockReturnValue(null),
 			on: vi.fn(),
 			off: vi.fn(),
 			once: vi.fn(),
@@ -255,7 +257,41 @@ describe("CloudService", () => {
 
 		it("should delegate handleAuthCallback to AuthService", async () => {
 			await cloudService.handleAuthCallback("code", "state")
-			expect(mockAuthService.handleCallback).toHaveBeenCalledWith("code", "state")
+			expect(mockAuthService.handleCallback).toHaveBeenCalledWith("code", "state", undefined)
+		})
+
+		it("should delegate handleAuthCallback with organizationId to AuthService", async () => {
+			await cloudService.handleAuthCallback("code", "state", "org_123")
+			expect(mockAuthService.handleCallback).toHaveBeenCalledWith("code", "state", "org_123")
+		})
+
+		it("should return stored organization ID from AuthService", () => {
+			mockAuthService.getStoredOrganizationId.mockReturnValue("org_456")
+
+			const result = cloudService.getStoredOrganizationId()
+			expect(mockAuthService.getStoredOrganizationId).toHaveBeenCalled()
+			expect(result).toBe("org_456")
+		})
+
+		it("should return null when no stored organization ID available", () => {
+			mockAuthService.getStoredOrganizationId.mockReturnValue(null)
+
+			const result = cloudService.getStoredOrganizationId()
+			expect(result).toBe(null)
+		})
+
+		it("should return true when stored organization ID exists", () => {
+			mockAuthService.getStoredOrganizationId.mockReturnValue("org_789")
+
+			const result = cloudService.hasStoredOrganizationId()
+			expect(result).toBe(true)
+		})
+
+		it("should return false when no stored organization ID exists", () => {
+			mockAuthService.getStoredOrganizationId.mockReturnValue(null)
+
+			const result = cloudService.hasStoredOrganizationId()
+			expect(result).toBe(false)
 		})
 	})
 

src/activate/handleUri.ts

@@ -38,7 +38,13 @@ export const handleUri = async (uri: vscode.Uri) => {
 		case "/auth/clerk/callback": {
 			const code = query.get("code")
 			const state = query.get("state")
-			await CloudService.instance.handleAuthCallback(code, state)
+			const organizationId = query.get("organizationId")
+
+			await CloudService.instance.handleAuthCallback(
+				code,
+				state,
+				organizationId === "null" ? null : organizationId,
+			)
 			break
 		}
 		default:

src/services/mdm/MdmService.ts

@@ -6,6 +6,7 @@ import { z } from "zod"
 
 import { CloudService, getClerkBaseUrl, PRODUCTION_CLERK_BASE_URL } from "@roo-code/cloud"
 import { Package } from "../../shared/package"
+import { t } from "../../i18n"
 
 // MDM Configuration Schema
 const mdmConfigSchema = z.object({
@@ -89,26 +90,43 @@ export class MdmService {
 		if (!CloudService.hasInstance() || !CloudService.instance.hasOrIsAcquiringActiveSession()) {
 			return {
 				compliant: false,
-				reason: "Your organization requires Roo Code Cloud authentication. Please sign in to continue.",
+				reason: t("mdm.errors.cloud_auth_required"),
 			}
 		}
 
 		// Check organization match if specified
 		const requiredOrgId = this.getRequiredOrganizationId()
 		if (requiredOrgId) {
 			try {
-				const currentOrgId = CloudService.instance.getOrganizationId()
+				// First try to get from active session
+				let currentOrgId = CloudService.instance.getOrganizationId()
+
+				// If no active session, check stored credentials
+				if (!currentOrgId) {
+					const storedOrgId = CloudService.instance.getStoredOrganizationId()
+
+					// null means personal account, which is not compliant for org requirements
+					if (storedOrgId === null || storedOrgId !== requiredOrgId) {
+						return {
+							compliant: false,
+							reason: t("mdm.errors.organization_mismatch"),
+						}
+					}
+
+					currentOrgId = storedOrgId
+				}
+
 				if (currentOrgId !== requiredOrgId) {
 					return {
 						compliant: false,
-						reason: "You must be authenticated with your organization's Roo Code Cloud account.",
+						reason: t("mdm.errors.organization_mismatch"),
 					}
 				}
 			} catch (error) {
 				this.log("[MDM] Error checking organization ID:", error)
 				return {
 					compliant: false,
-					reason: "Unable to verify organization authentication.",
+					reason: t("mdm.errors.verification_failed"),
 				}
 			}
 		}