fix(api, web): More fixes on env and org switching (#6091)

apps/api/src/app/auth/services/passport/jwt.strategy.ts

@@ -4,22 +4,34 @@ import { PassportStrategy } from '@nestjs/passport';
 import { Injectable, UnauthorizedException } from '@nestjs/common';
 import { ApiAuthSchemeEnum, HttpRequestHeaderKeysEnum, UserSessionData } from '@novu/shared';
 import { AuthService, Instrument } from '@novu/application-generic';
+import { EnvironmentRepository } from '@novu/dal';
 
 @Injectable()
 export class JwtStrategy extends PassportStrategy(Strategy) {
-  constructor(private readonly authService: AuthService) {
+  constructor(private readonly authService: AuthService, private environmentRepository: EnvironmentRepository) {
     super({
       jwtFromRequest: ExtractJwt.fromAuthHeaderAsBearerToken(),
       secretOrKey: process.env.JWT_SECRET,
       passReqToCallback: true,
     });
   }
-
   @Instrument()
   async validate(req: http.IncomingMessage, session: UserSessionData) {
     // Set the scheme to Bearer, meaning the user is authenticated via a JWT coming from Dashboard
     session.scheme = ApiAuthSchemeEnum.BEARER;
 
+    const user = await this.authService.validateUser(session);
+    if (!user) {
+      throw new UnauthorizedException();
+    }
+
+    await this.resolveEnvironmentId(req, session);
+
+    return session;
+  }
+
+  @Instrument()
+  async resolveEnvironmentId(req: http.IncomingMessage, session: UserSessionData) {
     // Fetch the environmentId from the request header
     const environmentIdFromHeader =
       (req.headers[HttpRequestHeaderKeysEnum.NOVU_ENVIRONMENT_ID.toLowerCase()] as string) || '';
@@ -29,13 +41,24 @@ export class JwtStrategy extends PassportStrategy(Strategy) {
      * or cached SPA versions of Dashboard as there is no guarantee all current users
      * will have environmentId in localStorage instantly after the deployment.
      */
-    session.environmentId = session.environmentId || environmentIdFromHeader;
+    const environmentIdFromLegacyAuthToken = session.environmentId;
 
-    const user = await this.authService.validateUser(session);
-    if (!user) {
-      throw new UnauthorizedException();
+    let currentEnvironmentId = '';
+
+    if (environmentIdFromLegacyAuthToken) {
+      currentEnvironmentId = environmentIdFromLegacyAuthToken;
+    } else {
+      const environments = await this.environmentRepository.findOrganizationEnvironments(session.organizationId);
+      const environmentIds = environments.map((env) => env._id);
+      const developmentEnvironmentId = environments.find((env) => env.name === 'Development')?._id || '';
+
+      currentEnvironmentId = developmentEnvironmentId;
+
+      if (environmentIds.includes(environmentIdFromHeader)) {
+        currentEnvironmentId = environmentIdFromHeader;
+      }
     }
 
-    return session;
+    session.environmentId = currentEnvironmentId;
   }
 }

apps/web/src/components/providers/CommunityAuthProvider.tsx

@@ -134,6 +134,9 @@ export const CommunityAuthProvider = ({ children }: { children: React.ReactNode
         return;
       }
 
+      // TODO: Revise storing environment id in local storage to avoid having to clear it during org or env switching
+      clearEnvironmentId();
+
       saveToken(newToken);
       await refetchOrganizations();
       /*
@@ -205,17 +208,15 @@ export const CommunityAuthProvider = ({ children }: { children: React.ReactNode
   const switchOrganization = useCallback(
     async (orgId: string) => {
       if (!orgId) {
-        return;
+        throw new Error('Organization ID is required');
       }
 
       if (orgId === currentOrganization?._id) {
         return;
       }
 
       // TODO: Revise storing environment id in local storage to avoid having to clear it during org or env switching
-      if (currentOrganization) {
-        clearEnvironmentId();
-      }
+      clearEnvironmentId();
 
       const token = await apiSwitchOrganization(orgId);
       await login(token);
@@ -225,11 +226,9 @@ export const CommunityAuthProvider = ({ children }: { children: React.ReactNode
   );
 
   useEffect(() => {
-    (async () => {
-      if (organizations && !currentOrganization) {
-        await switchOrganization(getTokenClaims()?.organizationId || '');
-      }
-    })();
+    if (organizations) {
+      setCurrentOrganization(selectOrganization(organizations, getTokenClaims()?.organizationId));
+    }
   }, [organizations, currentOrganization, switchOrganization]);
 
   useEffect(() => {

apps/web/src/components/providers/EnvironmentProvider.tsx

@@ -26,9 +26,9 @@ export function clearEnvironmentId() {
 }
 
 type EnvironmentContextValue = {
-  currentEnvironment?: IEnvironment;
+  currentEnvironment?: IEnvironment | null;
   // @deprecated use currentEnvironment instead;
-  environment?: IEnvironment;
+  environment?: IEnvironment | null;
   environments?: IEnvironment[];
   refetchEnvironments: () => Promise<void>;
   switchEnvironment: (params: Partial<{ environmentId: string; redirectUrl: string }>) => Promise<void>;
@@ -40,11 +40,11 @@ type EnvironmentContextValue = {
 
 const [EnvironmentCtx, useEnvironmentCtx] = createContextAndHook<EnvironmentContextValue>('Environment Context');
 
-function selectEnvironment(environments: IEnvironment[] | undefined, selectedEnvironmentId?: string) {
-  let e: IEnvironment | undefined;
+function selectEnvironment(environments: IEnvironment[] | undefined | null, selectedEnvironmentId?: string) {
+  let e: IEnvironment | undefined | null = null;
 
   if (!environments) {
-    return;
+    return null;
   }
 
   // Find the environment based on the current user's last environment
@@ -80,12 +80,16 @@ export function EnvironmentProvider({ children }: { children: React.ReactNode })
     staleTime: Infinity,
   });
 
-  const [currentEnvironment, setCurrentEnvironment] = useState<IEnvironment | undefined>(
+  const [currentEnvironment, setCurrentEnvironment] = useState<IEnvironment | null>(
     selectEnvironment(environments, getEnvironmentId())
   );
 
   const switchEnvironment = useCallback(
     async ({ environmentId, redirectUrl }: Partial<{ environmentId: string; redirectUrl: string }> = {}) => {
+      if (currentEnvironment?._id === environmentId) {
+        return;
+      }
+
       setCurrentEnvironment(selectEnvironment(environments, environmentId));
 
       /*
@@ -99,7 +103,7 @@ export function EnvironmentProvider({ children }: { children: React.ReactNode })
         navigate(redirectUrl);
       }
     },
-    [queryClient, navigate, setCurrentEnvironment, environments]
+    [queryClient, navigate, setCurrentEnvironment, currentEnvironment, environments]
   );
 
   const switchToProductionEnvironment = useCallback(
@@ -135,7 +139,7 @@ export function EnvironmentProvider({ children }: { children: React.ReactNode })
   );
 
   useEffect(() => {
-    if (environments && environments.length > 0 && !currentEnvironment) {
+    if (environments) {
       switchEnvironment({ environmentId: getEnvironmentId() });
     }
   }, [currentEnvironment, environments, switchEnvironment]);

apps/web/src/ee/billing/utils/hooks/useSubscription.ts

@@ -4,6 +4,7 @@ import { useQuery } from '@tanstack/react-query';
 import { useMemo } from 'react';
 import { differenceInDays, isSameDay } from 'date-fns';
 import { ApiServiceLevelEnum } from '@novu/shared';
+import { useEnvironment } from '../../../../hooks/useEnvironment';
 
 export const useSubscription = () => {
   // TODO: Fix with a useMemo
@@ -14,6 +15,7 @@ export const useSubscription = () => {
     ['billing-subscription', currentOrganization?._id],
     () => api.get('/v1/billing/subscription'),
     {
+      enabled: !!currentOrganization,
       initialData: {
         trialStart: today.toISOString(),
         trialEnd: today.toISOString(),

libs/application-generic/src/services/auth/community.auth.service.ts

@@ -289,30 +289,14 @@ export class CommunityAuthService implements IAuthService {
       ? this.isAuthenticatedForOrganization(payload._id, payload.organizationId)
       : Promise.resolve(true);
 
-    const environmentPromise =
-      payload.organizationId && payload.environmentId
-        ? this.environmentRepository.findByIdAndOrganization(
-            payload.environmentId,
-            payload.organizationId
-          )
-        : Promise.resolve(true);
-
-    const [user, isMember, environment] = await Promise.all([
-      userPromise,
-      isMemberPromise,
-      environmentPromise,
-    ]);
+    const [user, isMember] = await Promise.all([userPromise, isMemberPromise]);
 
     if (!user) throw new UnauthorizedException('User not found');
     if (payload.organizationId && !isMember) {
       throw new UnauthorizedException(
         `User ${payload._id} is not a member of organization ${payload.organizationId}`
       );
     }
-    if (payload.organizationId && payload.environmentId && !environment)
-      throw new UnauthorizedException(
-        `Environment ${payload.environmentId} doesn't belong to organization ${payload.organizationId}`
-      );
 
     return user;
   }