From c327badc0da08977a7115b2116ab972ba3a5c76c Mon Sep 17 00:00:00 2001 From: Matt Rubens Date: Tue, 17 Jun 2025 11:34:47 -0400 Subject: [PATCH] Handle initial session refresh when checking compliance --- packages/cloud/src/AuthService.ts | 23 +++- packages/cloud/src/CloudService.ts | 8 ++ .../cloud/src/__tests__/AuthService.test.ts | 125 ++++++++++++++++++ .../cloud/src/__tests__/CloudService.test.ts | 8 ++ src/services/mdm/MdmService.ts | 8 +- src/services/mdm/__tests__/MdmService.spec.ts | 62 +++++++++ 6 files changed, 226 insertions(+), 8 deletions(-) create mode 100644 packages/cloud/src/__tests__/AuthService.test.ts diff --git a/packages/cloud/src/AuthService.ts b/packages/cloud/src/AuthService.ts index 68036ce3c9..6830651add 100644 --- a/packages/cloud/src/AuthService.ts +++ b/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 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 { private context: vscode.ExtensionContext @@ -277,6 +278,10 @@ export class AuthService extends EventEmitter { 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 { 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 diff --git a/packages/cloud/src/CloudService.ts b/packages/cloud/src/CloudService.ts index fe3bad970c..427822ec88 100644 --- a/packages/cloud/src/CloudService.ts +++ b/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) } diff --git a/packages/cloud/src/__tests__/AuthService.test.ts b/packages/cloud/src/__tests__/AuthService.test.ts new file mode 100644 index 0000000000..bcfae055e1 --- /dev/null +++ b/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 + 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 as vscode.SecretStorage, + globalState: { + update: vi.fn(), + get: vi.fn(), + keys: vi.fn(() => []), + setKeysForSync: vi.fn(), + } as Partial as vscode.Memento & { + setKeysForSync(keys: readonly string[]): void + }, + subscriptions: [], + extension: { + packageJSON: { + publisher: "test", + name: "test-extension", + }, + } as Partial> as vscode.Extension, + } + + 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 + }) + }) +}) diff --git a/packages/cloud/src/__tests__/CloudService.test.ts b/packages/cloud/src/__tests__/CloudService.test.ts index 03b28568d5..929e852c72 100644 --- a/packages/cloud/src/__tests__/CloudService.test.ts +++ b/packages/cloud/src/__tests__/CloudService.test.ts @@ -36,6 +36,7 @@ describe("CloudService", () => { logout: ReturnType isAuthenticated: ReturnType hasActiveSession: ReturnType + isRefreshingSession: ReturnType getUserInfo: ReturnType getState: ReturnType getSessionToken: ReturnType @@ -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() diff --git a/src/services/mdm/MdmService.ts b/src/services/mdm/MdmService.ts index b649f4d4a2..104ca60fb9 100644 --- a/src/services/mdm/MdmService.ts +++ b/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.", diff --git a/src/services/mdm/__tests__/MdmService.spec.ts b/src/services/mdm/__tests__/MdmService.spec.ts index 79ce83c3b5..a5b3f2ca4f 100644 --- a/src/services/mdm/__tests__/MdmService.spec.ts +++ b/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,6 +306,8 @@ 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() @@ -307,6 +315,60 @@ describe("MdmService", () => { 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", () => {