fix(auth): handle mixed OAuth project scopes

Generated-By: PostHog Code
Task-Id: 07b377f9-4583-463c-b256-cc29621375e6

Author: dmarticus

packages/core/src/auth/auth.test.ts

@@ -1524,6 +1524,139 @@ describe("AuthService", () => {
         },
       });
     });
+
+    it("merges organization and project scopes", async () => {
+      vi.stubGlobal(
+        "fetch",
+        vi.fn(async (input: string | Request) => {
+          const url = typeof input === "string" ? input : input.url;
+          if (url.includes("/api/users/@me/")) {
+            return {
+              ok: true,
+              json: vi.fn().mockResolvedValue({
+                uuid: "user-1",
+                organization: { id: "org-1", name: "Org 1" },
+              }),
+            } as unknown as Response;
+          }
+          if (url.endsWith("/api/organizations/org-1/")) {
+            return {
+              ok: true,
+              json: vi.fn().mockResolvedValue({
+                name: "Org 1",
+                teams: [{ id: 11, name: "Organization Project" }],
+              }),
+            } as unknown as Response;
+          }
+          if (url.endsWith("/api/projects/42/")) {
+            return {
+              ok: true,
+              json: vi.fn().mockResolvedValue({
+                id: 42,
+                name: "Scoped Project",
+                organization: "org-1",
+              }),
+            } as unknown as Response;
+          }
+          return {
+            ok: true,
+            json: vi.fn().mockResolvedValue({ has_access: true }),
+          } as unknown as Response;
+        }) as unknown as typeof fetch,
+      );
+      oauthFlow.startFlow.mockResolvedValue(
+        mockTokenResponse({ scopedOrgs: ["org-1"], scopedTeams: [42] }),
+      );
+
+      await service.login("us");
+
+      expect(service.getState().orgProjectsMap["org-1"].projects).toEqual([
+        { id: 11, name: "Organization Project" },
+        { id: 42, name: "Scoped Project" },
+      ]);
+    });
+
+    it("aligns the current organization with the restored scoped project", async () => {
+      seedStoredSession({
+        selectedProjectId: 84,
+        scopeVersion: OAUTH_SCOPE_VERSION - 1,
+      });
+      await service.initialize();
+      vi.stubGlobal(
+        "fetch",
+        vi.fn(async (input: string | Request) => {
+          const url = typeof input === "string" ? input : input.url;
+          if (url.includes("/api/users/@me/")) {
+            return {
+              ok: true,
+              json: vi.fn().mockResolvedValue({
+                uuid: "user-1",
+                organization: { id: "org-1", name: "Org 1" },
+              }),
+            } as unknown as Response;
+          }
+          const projectMatch = url.match(/\/api\/projects\/(42|84)\/$/);
+          if (projectMatch) {
+            const projectId = Number(projectMatch[1]);
+            return {
+              ok: true,
+              json: vi.fn().mockResolvedValue({
+                id: projectId,
+                name: `Project ${projectId}`,
+                organization: projectId === 42 ? "org-1" : "org-2",
+              }),
+            } as unknown as Response;
+          }
+          return {
+            ok: true,
+            json: vi.fn().mockResolvedValue({ has_access: true }),
+          } as unknown as Response;
+        }) as unknown as typeof fetch,
+      );
+      oauthFlow.startFlow.mockResolvedValue(
+        mockTokenResponse({ scopedOrgs: [], scopedTeams: [42, 84] }),
+      );
+
+      await service.login("us");
+
+      expect(service.getState()).toMatchObject({
+        currentOrgId: "org-2",
+        currentProjectId: 84,
+      });
+    });
+
+    it("rejects authentication when a scoped project cannot be loaded permanently", async () => {
+      vi.stubGlobal(
+        "fetch",
+        vi.fn(async (input: string | Request) => {
+          const url = typeof input === "string" ? input : input.url;
+          if (url.includes("/api/users/@me/")) {
+            return {
+              ok: true,
+              json: vi.fn().mockResolvedValue({
+                uuid: "user-1",
+                organization: { id: "org-1", name: "Org 1" },
+              }),
+            } as unknown as Response;
+          }
+          if (url.endsWith("/api/projects/42/")) {
+            return { ok: false, status: 403 } as Response;
+          }
+          return {
+            ok: true,
+            json: vi.fn().mockResolvedValue({ has_access: true }),
+          } as unknown as Response;
+        }) as unknown as typeof fetch,
+      );
+      oauthFlow.startFlow.mockResolvedValue(
+        mockTokenResponse({ scopedOrgs: [], scopedTeams: [42] }),
+      );
+
+      await expect(service.login("us")).rejects.toThrow(
+        "Unable to load OAuth-scoped project 42",
+      );
+      expect(service.getState().status).toBe("anonymous");
+    });
   });
 
   describe("switchOrg", () => {

packages/core/src/auth/auth.ts

@@ -57,6 +57,7 @@ interface InMemorySession {
   orgProjectsMap: OrgProjectsMap;
   currentOrgId: string | null;
   currentProjectId: number | null;
+  scopedOrgIds: string[];
   scopedProjectIds: number[];
   orgProjectsIncomplete: boolean;
 }
@@ -647,28 +648,15 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
       );
     const previousMap = this.session?.orgProjectsMap ?? {};
     const { map: orgProjectsMap, incomplete: orgProjectsIncomplete } =
-      scopedOrgIds.length > 0
-        ? await this.buildOrgProjectsMap(
-            tokenResponse.access_token,
-            options.cloudRegion,
-            scopedOrgIds,
-            previousMap,
-          )
-        : scopedProjectIds.length > 0
-          ? await this.buildScopedProjectMap(
-              tokenResponse.access_token,
-              options.cloudRegion,
-              scopedProjectIds,
-              currentOrgId,
-              currentOrgName,
-              previousMap,
-            )
-          : await this.buildOrgProjectsMap(
-              tokenResponse.access_token,
-              options.cloudRegion,
-              currentOrgId ? [currentOrgId] : [],
-              previousMap,
-            );
+      await this.buildAuthorizedProjectsMap({
+        accessToken: tokenResponse.access_token,
+        cloudRegion: options.cloudRegion,
+        scopedOrgIds,
+        scopedProjectIds,
+        currentOrgId,
+        currentOrgName,
+        previousMap,
+      });
     const lastPrefs = accountKey
       ? this.authPreference.get(accountKey, options.cloudRegion)
       : null;
@@ -680,8 +668,10 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
       lastSelectedOrgId: lastPrefs?.lastSelectedOrgId ?? null,
     });
 
-    const resolvedCurrentOrgId =
-      currentOrgId && orgProjectsMap[currentOrgId]
+    const resolvedCurrentOrgId = currentProjectId
+      ? (findOrgForProject(orgProjectsMap, currentProjectId, currentOrgId) ??
+        currentOrgId)
+      : currentOrgId && orgProjectsMap[currentOrgId]
         ? currentOrgId
         : (Object.keys(orgProjectsMap)[0] ?? currentOrgId);
     const session: InMemorySession = {
@@ -693,12 +683,74 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
       orgProjectsMap,
       currentOrgId: resolvedCurrentOrgId,
       currentProjectId,
+      scopedOrgIds,
       scopedProjectIds,
       orgProjectsIncomplete,
     };
 
     return session;
   }
+  private async buildAuthorizedProjectsMap(input: {
+    accessToken: string;
+    cloudRegion: CloudRegion;
+    scopedOrgIds: string[];
+    scopedProjectIds: number[];
+    currentOrgId: string | null;
+    currentOrgName: string | null;
+    previousMap: OrgProjectsMap;
+  }): Promise<{ map: OrgProjectsMap; incomplete: boolean }> {
+    const orgIdsToFetch =
+      input.scopedOrgIds.length > 0
+        ? input.scopedOrgIds
+        : input.scopedProjectIds.length === 0 && input.currentOrgId
+          ? [input.currentOrgId]
+          : [];
+    const [orgResult, projectResult] = await Promise.all([
+      this.buildOrgProjectsMap(
+        input.accessToken,
+        input.cloudRegion,
+        orgIdsToFetch,
+        input.previousMap,
+      ),
+      input.scopedProjectIds.length > 0
+        ? this.buildScopedProjectMap(
+            input.accessToken,
+            input.cloudRegion,
+            input.scopedProjectIds,
+            input.currentOrgId,
+            input.currentOrgName,
+            input.previousMap,
+          )
+        : Promise.resolve({ map: {}, incomplete: false }),
+    ]);
+
+    return {
+      map: this.mergeOrgProjectsMaps(orgResult.map, projectResult.map),
+      incomplete: orgResult.incomplete || projectResult.incomplete,
+    };
+  }
+  private mergeOrgProjectsMaps(...maps: OrgProjectsMap[]): OrgProjectsMap {
+    const merged: OrgProjectsMap = {};
+    for (const map of maps) {
+      for (const [orgId, org] of Object.entries(map)) {
+        const existing = merged[orgId];
+        const projects = new Map(
+          existing?.projects.map((project) => [project.id, project]) ?? [],
+        );
+        for (const project of org.projects) {
+          projects.set(project.id, project);
+        }
+        merged[orgId] = {
+          orgName:
+            existing?.orgName && existing.orgName !== "(unknown)"
+              ? existing.orgName
+              : org.orgName,
+          projects: [...projects.values()],
+        };
+      }
+    }
+    return merged;
+  }
   private async buildScopedProjectMap(
     accessToken: string,
     cloudRegion: CloudRegion,
@@ -718,7 +770,10 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
           projectId,
         );
         if (!result.ok) {
-          incomplete ||= result.retryable;
+          if (!result.retryable) {
+            throw new Error(`Unable to load OAuth-scoped project ${projectId}`);
+          }
+          incomplete = true;
           return;
         }
 
@@ -1277,24 +1332,16 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
 
       if (!session.orgProjectsIncomplete) return;
 
-      const orgIds = Object.keys(session.orgProjectsMap);
-      const { map, incomplete } =
-        session.scopedProjectIds.length > 0
-          ? await this.buildScopedProjectMap(
-              session.accessToken,
-              session.cloudRegion,
-              session.scopedProjectIds,
-              session.currentOrgId,
-              session.orgProjectsMap[session.currentOrgId ?? ""]?.orgName ??
-                null,
-              session.orgProjectsMap,
-            )
-          : await this.buildOrgProjectsMap(
-              session.accessToken,
-              session.cloudRegion,
-              orgIds,
-              session.orgProjectsMap,
-            );
+      const { map, incomplete } = await this.buildAuthorizedProjectsMap({
+        accessToken: session.accessToken,
+        cloudRegion: session.cloudRegion,
+        scopedOrgIds: session.scopedOrgIds,
+        scopedProjectIds: session.scopedProjectIds,
+        currentOrgId: session.currentOrgId,
+        currentOrgName:
+          session.orgProjectsMap[session.currentOrgId ?? ""]?.orgName ?? null,
+        previousMap: session.orgProjectsMap,
+      });
 
       // The session may have been replaced (logout, re-login) while the fetch
       // was in flight; committing the stale one would resurrect it.