Handle initial session refresh when checking compliance

Author: mrubens

packages/cloud/src/AuthService.ts

@@ -14,6 +14,7 @@ import { getUserAgent } from "./utils"
 export interface AuthServiceEvents {
 	"inactive-session": [data: { previousState: AuthState }]
 	"active-session": [data: { previousState: AuthState }]
+	"refreshing-session": [data: { previousState: AuthState }]
 	"logged-out": [data: { previousState: AuthState }]
 	"user-info": [data: { userInfo: CloudUserInfo }]
 }
@@ -28,7 +29,7 @@ type AuthCredentials = z.infer<typeof authCredentialsSchema>
 const AUTH_CREDENTIALS_KEY = "clerk-auth-credentials"
 const AUTH_STATE_KEY = "clerk-auth-state"
 
-type AuthState = "initializing" | "logged-out" | "active-session" | "inactive-session"
+type AuthState = "initializing" | "logged-out" | "active-session" | "inactive-session" | "refreshing-session"
 
 export class AuthService extends EventEmitter<AuthServiceEvents> {
 	private context: vscode.ExtensionContext
@@ -277,6 +278,10 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
 		return this.state === "active-session"
 	}
 
+	public isRefreshingSession(): boolean {
+		return this.state === "refreshing-session"
+	}
+
 	/**
 	 * Refresh the session
 	 *
@@ -291,14 +296,20 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
 
 		try {
 			const previousState = this.state
+
+			// Transition to refreshing state
+			if (this.state !== "refreshing-session") {
+				this.state = "refreshing-session"
+				this.emit("refreshing-session", { previousState })
+				this.log("[auth] Transitioned to refreshing-session state")
+			}
+
 			this.sessionToken = await this.clerkCreateSessionToken()
 			this.state = "active-session"
 
-			if (previousState !== "active-session") {
-				this.log("[auth] Transitioned to active-session state")
-				this.emit("active-session", { previousState })
-				this.fetchUserInfo()
-			}
+			this.log("[auth] Transitioned to active-session state")
+			this.emit("active-session", { previousState: "refreshing-session" })
+			this.fetchUserInfo()
 		} catch (error) {
 			this.log("[auth] Failed to refresh session", error)
 			throw error

packages/cloud/src/CloudService.ts

@@ -41,6 +41,7 @@ export class CloudService {
 
 			this.authService.on("inactive-session", this.authListener)
 			this.authService.on("active-session", this.authListener)
+			this.authService.on("refreshing-session", this.authListener)
 			this.authService.on("logged-out", this.authListener)
 			this.authService.on("user-info", this.authListener)
 
@@ -87,6 +88,11 @@ export class CloudService {
 		return this.authService!.hasActiveSession()
 	}
 
+	public isRefreshingSession(): boolean {
+		this.ensureInitialized()
+		return this.authService!.isRefreshingSession()
+	}
+
 	public getUserInfo(): CloudUserInfo | null {
 		this.ensureInitialized()
 		return this.authService!.getUserInfo()
@@ -150,7 +156,9 @@ export class CloudService {
 
 	public dispose(): void {
 		if (this.authService) {
+			this.authService.off("inactive-session", this.authListener)
 			this.authService.off("active-session", this.authListener)
+			this.authService.off("refreshing-session", this.authListener)
 			this.authService.off("logged-out", this.authListener)
 			this.authService.off("user-info", this.authListener)
 		}

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

@@ -0,0 +1,125 @@
+// npx vitest run src/__tests__/AuthService.test.ts
+
+import { describe, it, expect, vi, beforeEach } from "vitest"
+import * as vscode from "vscode"
+import { AuthService } from "../AuthService"
+
+// Mock vscode
+vi.mock("vscode", () => ({
+	ExtensionContext: vi.fn(),
+	window: {
+		showInformationMessage: vi.fn(),
+	},
+	env: {
+		openExternal: vi.fn(),
+		uriScheme: "vscode",
+	},
+	Uri: {
+		parse: vi.fn(),
+	},
+}))
+
+// Mock axios
+vi.mock("axios", () => ({
+	default: {
+		post: vi.fn(),
+		get: vi.fn(),
+	},
+}))
+
+// Mock other dependencies
+vi.mock("../Config", () => ({
+	getClerkBaseUrl: vi.fn(() => "https://clerk.test"),
+	getRooCodeApiUrl: vi.fn(() => "https://api.test"),
+}))
+
+vi.mock("../RefreshTimer", () => ({
+	RefreshTimer: vi.fn().mockImplementation(() => ({
+		start: vi.fn(),
+		stop: vi.fn(),
+	})),
+}))
+
+vi.mock("../utils", () => ({
+	getUserAgent: vi.fn(() => "test-agent"),
+}))
+
+describe("AuthService", () => {
+	let mockContext: Partial<vscode.ExtensionContext>
+	let authService: AuthService
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+
+		mockContext = {
+			secrets: {
+				store: vi.fn(),
+				get: vi.fn(),
+				delete: vi.fn(),
+				onDidChange: vi.fn(() => ({ dispose: vi.fn() })),
+			} as Partial<vscode.SecretStorage> as vscode.SecretStorage,
+			globalState: {
+				update: vi.fn(),
+				get: vi.fn(),
+				keys: vi.fn(() => []),
+				setKeysForSync: vi.fn(),
+			} as Partial<vscode.Memento & { setKeysForSync(keys: readonly string[]): void }> as vscode.Memento & {
+				setKeysForSync(keys: readonly string[]): void
+			},
+			subscriptions: [],
+			extension: {
+				packageJSON: {
+					publisher: "test",
+					name: "test-extension",
+				},
+			} as Partial<vscode.Extension<unknown>> as vscode.Extension<unknown>,
+		}
+
+		authService = new AuthService(mockContext as vscode.ExtensionContext)
+	})
+
+	describe("State Management", () => {
+		it("should initialize with 'initializing' state", () => {
+			expect(authService.getState()).toBe("initializing")
+		})
+
+		it("should have isRefreshingSession method that returns false initially", () => {
+			expect(authService.isRefreshingSession()).toBe(false)
+		})
+
+		it("should include refreshing-session in AuthState type", () => {
+			// This test verifies that the new state is properly typed
+			// by checking that the method exists and returns a boolean
+			expect(typeof authService.isRefreshingSession).toBe("function")
+			expect(typeof authService.isRefreshingSession()).toBe("boolean")
+		})
+	})
+
+	describe("Event Emission", () => {
+		it("should emit refreshing-session event when transitioning to refreshing state", async () => {
+			// Set up the auth service to have credentials
+			const mockCredentials = {
+				clientToken: "test-token",
+				sessionId: "test-session",
+			}
+
+			// Mock the secrets.get to return credentials
+			vi.mocked(mockContext.secrets!.get).mockResolvedValue(JSON.stringify(mockCredentials))
+
+			// Create a promise to wait for the event
+			const eventPromise = new Promise((resolve) => {
+				authService.on("refreshing-session", (data) => {
+					expect(data).toHaveProperty("previousState")
+					resolve(data)
+				})
+			})
+
+			// This would trigger the refresh process in a real scenario
+			// For this test, we're just verifying the event structure exists
+			// We can manually emit the event to test the interface
+			authService.emit("refreshing-session", { previousState: "inactive-session" })
+
+			await eventPromise
+		})
+	})
+})

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

@@ -36,6 +36,7 @@ describe("CloudService", () => {
 		logout: ReturnType<typeof vi.fn>
 		isAuthenticated: ReturnType<typeof vi.fn>
 		hasActiveSession: ReturnType<typeof vi.fn>
+		isRefreshingSession: ReturnType<typeof vi.fn>
 		getUserInfo: ReturnType<typeof vi.fn>
 		getState: ReturnType<typeof vi.fn>
 		getSessionToken: ReturnType<typeof vi.fn>
@@ -84,6 +85,7 @@ describe("CloudService", () => {
 			logout: vi.fn(),
 			isAuthenticated: vi.fn().mockReturnValue(false),
 			hasActiveSession: vi.fn().mockReturnValue(false),
+			isRefreshingSession: vi.fn().mockReturnValue(false),
 			getUserInfo: vi.fn(),
 			getState: vi.fn().mockReturnValue("logged-out"),
 			getSessionToken: vi.fn(),
@@ -179,6 +181,12 @@ describe("CloudService", () => {
 			expect(result).toBe(false)
 		})
 
+		it("should delegate isRefreshingSession to AuthService", () => {
+			const result = cloudService.isRefreshingSession()
+			expect(mockAuthService.isRefreshingSession).toHaveBeenCalled()
+			expect(result).toBe(false)
+		})
+
 		it("should delegate getUserInfo to AuthService", async () => {
 			await cloudService.getUserInfo()
 			expect(mockAuthService.getUserInfo).toHaveBeenCalled()

src/services/mdm/MdmService.ts

@@ -85,8 +85,12 @@ export class MdmService {
 			return { compliant: true }
 		}
 
-		// Check if cloud service is available and authenticated
-		if (!CloudService.hasInstance() || !CloudService.instance.hasActiveSession()) {
+		const cloudService = CloudService.instance
+		const hasActiveSession = cloudService?.hasActiveSession()
+		const isRefreshingSession = cloudService?.isRefreshingSession()
+
+		// Allow only if user has active session or is refreshing session
+		if (!hasActiveSession && !isRefreshingSession) {
 			return {
 				compliant: false,
 				reason: "Your organization requires Roo Code Cloud authentication. Please sign in to continue.",

src/services/mdm/__tests__/MdmService.spec.ts

@@ -16,6 +16,8 @@ vi.mock("@roo-code/cloud", () => ({
 		hasInstance: vi.fn(),
 		instance: {
 			hasActiveSession: vi.fn(),
+			isRefreshingSession: vi.fn(),
+			isAuthenticated: vi.fn(),
 			getOrganizationId: vi.fn(),
 		},
 	},
@@ -244,6 +246,8 @@ describe("MdmService", () => {
 
 			mockCloudService.hasInstance.mockReturnValue(true)
 			mockCloudService.instance.hasActiveSession.mockReturnValue(true)
+			mockCloudService.instance.isRefreshingSession.mockReturnValue(false)
+			mockCloudService.instance.isAuthenticated.mockReturnValue(true)
 
 			const service = await MdmService.createInstance()
 			const compliance = service.isCompliant()
@@ -279,6 +283,8 @@ describe("MdmService", () => {
 			// Mock CloudService to have instance and active session but wrong org
 			mockCloudService.hasInstance.mockReturnValue(true)
 			mockCloudService.instance.hasActiveSession.mockReturnValue(true)
+			mockCloudService.instance.isRefreshingSession.mockReturnValue(false)
+			mockCloudService.instance.isAuthenticated.mockReturnValue(true)
 			mockCloudService.instance.getOrganizationId.mockReturnValue("different-org-456")
 
 			const service = await MdmService.createInstance()
@@ -300,13 +306,69 @@ describe("MdmService", () => {
 
 			mockCloudService.hasInstance.mockReturnValue(true)
 			mockCloudService.instance.hasActiveSession.mockReturnValue(true)
+			mockCloudService.instance.isRefreshingSession.mockReturnValue(false)
+			mockCloudService.instance.isAuthenticated.mockReturnValue(true)
 			mockCloudService.instance.getOrganizationId.mockReturnValue("correct-org-123")
 
 			const service = await MdmService.createInstance()
 			const compliance = service.isCompliant()
 
 			expect(compliance.compliant).toBe(true)
 		})
+
+		it("should be compliant when refreshing session", async () => {
+			const mockConfig = { requireCloudAuth: true }
+			mockFs.existsSync.mockReturnValue(true)
+			mockFs.readFileSync.mockReturnValue(JSON.stringify(mockConfig))
+
+			mockCloudService.hasInstance.mockReturnValue(true)
+			mockCloudService.instance.hasActiveSession.mockReturnValue(false)
+			mockCloudService.instance.isRefreshingSession.mockReturnValue(true)
+			mockCloudService.instance.isAuthenticated.mockReturnValue(true)
+
+			const service = await MdmService.createInstance()
+			const compliance = service.isCompliant()
+
+			expect(compliance.compliant).toBe(true)
+		})
+
+		it("should be non-compliant when authenticated but not active or refreshing", async () => {
+			const mockConfig = { requireCloudAuth: true }
+			mockFs.existsSync.mockReturnValue(true)
+			mockFs.readFileSync.mockReturnValue(JSON.stringify(mockConfig))
+
+			mockCloudService.hasInstance.mockReturnValue(true)
+			mockCloudService.instance.hasActiveSession.mockReturnValue(false)
+			mockCloudService.instance.isRefreshingSession.mockReturnValue(false)
+			mockCloudService.instance.isAuthenticated.mockReturnValue(true)
+
+			const service = await MdmService.createInstance()
+			const compliance = service.isCompliant()
+
+			expect(compliance.compliant).toBe(false)
+			if (!compliance.compliant) {
+				expect(compliance.reason).toContain("requires Roo Code Cloud authentication")
+			}
+		})
+
+		it("should be non-compliant when not authenticated, not active, and not refreshing", async () => {
+			const mockConfig = { requireCloudAuth: true }
+			mockFs.existsSync.mockReturnValue(true)
+			mockFs.readFileSync.mockReturnValue(JSON.stringify(mockConfig))
+
+			mockCloudService.hasInstance.mockReturnValue(true)
+			mockCloudService.instance.hasActiveSession.mockReturnValue(false)
+			mockCloudService.instance.isRefreshingSession.mockReturnValue(false)
+			mockCloudService.instance.isAuthenticated.mockReturnValue(false)
+
+			const service = await MdmService.createInstance()
+			const compliance = service.isCompliant()
+
+			expect(compliance.compliant).toBe(false)
+			if (!compliance.compliant) {
+				expect(compliance.reason).toContain("requires Roo Code Cloud authentication")
+			}
+		})
 	})
 
 	describe("cloud enablement", () => {