Store the organization id in credentials

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: {
@@ -505,23 +538,37 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
 
 		userInfo.picture = userData.image_url
 
-		// Fetch organization memberships separately
+		// Fetch organization info separately - but only populate if user is in organization context
 		try {
-			const orgMemberships = await this.clerkGetOrganizationMemberships()
-			if (orgMemberships && orgMemberships.length > 0) {
-				// Get the first (or active) organization membership
-				const primaryOrgMembership = orgMemberships[0]
-				const organization = primaryOrgMembership?.organization
-
-				if (organization) {
-					userInfo.organizationId = organization.id
-					userInfo.organizationName = organization.name
-					userInfo.organizationRole = primaryOrgMembership.role
-					userInfo.organizationImageUrl = organization.image_url
+			const storedOrgId = this.getStoredOrganizationId()
+
+			if (storedOrgId !== null) {
+				// User is in organization context - fetch user's memberships and filter
+				const orgMemberships = await this.clerkGetOrganizationMemberships()
+
+				// Find the user's membership in this organization
+				const userMembership = orgMemberships?.find(
+					(membership: CloudOrganizationMembership) => membership.organization.id === storedOrgId,
+				)
+
+				if (userMembership) {
+					userInfo.organizationId = userMembership.organization.id
+					userInfo.organizationName = userMembership.organization.name
+					userInfo.organizationRole = userMembership.role
+					userInfo.organizationImageUrl = userMembership.organization.image_url
+					this.log("[auth] User in organization context:", {
+						id: userMembership.organization.id,
+						name: userMembership.organization.name,
+						role: userMembership.role,
+					})
+				} else {
+					this.log("[auth] Warning: User not found in stored organization:", storedOrgId)
 				}
+			} else {
+				this.log("[auth] User in personal account context - not setting organization info")
 			}
 		} catch (error) {
-			this.log("[auth] Failed to fetch organization memberships:", error)
+			this.log("[auth] Failed to fetch organization info:", error)
 			// Don't throw - organization info is optional
 		}
 

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")
 		})
@@ -633,9 +633,55 @@ describe("AuthService", () => {
 			expect(authService.getUserInfo()).toBeNull()
 		})
 
-		it("should parse user info correctly", async () => {
-			// Set up with credentials
-			const credentials = { clientToken: "test-token", sessionId: "test-session" }
+		it("should parse user info correctly for personal accounts", async () => {
+			// Set up with credentials for personal account (no organizationId)
+			const credentials = { clientToken: "test-token", sessionId: "test-session", organizationId: null }
+			mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials))
+			await authService.initialize()
+
+			// Clear previous mock calls
+			mockFetch.mockClear()
+
+			// Mock successful responses
+			mockFetch
+				.mockResolvedValueOnce({
+					ok: true,
+					json: () => Promise.resolve({ jwt: "jwt-token" }),
+				})
+				.mockResolvedValueOnce({
+					ok: true,
+					json: () =>
+						Promise.resolve({
+							response: {
+								first_name: "Jane",
+								last_name: "Smith",
+								image_url: "https://example.com/jane.jpg",
+								primary_email_address_id: "email-2",
+								email_addresses: [
+									{ id: "email-1", email_address: "jane.old@example.com" },
+									{ id: "email-2", email_address: "jane@example.com" },
+								],
+							},
+						}),
+				})
+
+			const timerCallback = vi.mocked(RefreshTimer).mock.calls[0][0].callback
+			await timerCallback()
+
+			// Wait for async operations to complete
+			await new Promise((resolve) => setTimeout(resolve, 0))
+
+			const userInfo = authService.getUserInfo()
+			expect(userInfo).toEqual({
+				name: "Jane Smith",
+				email: "jane@example.com",
+				picture: "https://example.com/jane.jpg",
+			})
+		})
+
+		it("should parse user info correctly for organization accounts", async () => {
+			// Set up with credentials for organization account
+			const credentials = { clientToken: "test-token", sessionId: "test-session", organizationId: "org_1" }
 			mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials))
 			await authService.initialize()
 
@@ -699,8 +745,8 @@ describe("AuthService", () => {
 		})
 
 		it("should handle missing user info fields", async () => {
-			// Set up with credentials
-			const credentials = { clientToken: "test-token", sessionId: "test-session" }
+			// Set up with credentials for personal account (no organizationId)
+			const credentials = { clientToken: "test-token", sessionId: "test-session", organizationId: null }
 			mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials))
 			await authService.initialize()
 

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)
 		})
 	})
 

packages/types/src/cloud.ts

@@ -125,6 +125,7 @@ export const ORGANIZATION_ALLOW_ALL: OrganizationAllowList = {
 export const ORGANIZATION_DEFAULT: OrganizationSettings = {
 	version: 0,
 	cloudSettings: {
+		recordTaskMessages: true,
 		enableTaskSharing: true,
 		taskShareExpirationDays: 30,
 	},

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"),
 				}
 			}
 		}