feat: implement feature flag management and cleanup

- Added feature flag handling in various components and API routes.
- Introduced cookies for managing feature flags and updated authentication flows.
- Removed deprecated feature flag provider and related components.
- Enhanced sidebar and assessment page to utilize feature flags for conditional rendering.
- Updated middleware to enforce feature-based access control.

Author: vprashrex

app/api/auth/logout/route.ts

@@ -1,6 +1,6 @@
 import { NextRequest, NextResponse } from "next/server";
 import { apiClient } from "@/app/lib/apiClient";
-import { clearRoleCookie } from "@/app/lib/authCookie";
+import { clearFeaturesCookie, clearRoleCookie } from "@/app/lib/authCookie";
 
 export async function POST(request: NextRequest) {
   const { status, data, headers } = await apiClient(
@@ -17,6 +17,7 @@ export async function POST(request: NextRequest) {
   }
 
   clearRoleCookie(res);
+  clearFeaturesCookie(res);
 
   return res;
 }

app/api/features/route.ts

@@ -1,6 +1,9 @@
 import { NextRequest, NextResponse } from "next/server";
 import { apiClient } from "@/app/lib/apiClient";
-import { setRoleCookieFromBody } from "@/app/lib/authCookie";
+import {
+  setFeaturesCookieFromBody,
+  setRoleCookieFromBody,
+} from "@/app/lib/authCookie";
 
 export async function GET(request: NextRequest) {
   try {
@@ -9,6 +12,7 @@ export async function GET(request: NextRequest) {
 
     if (status >= 200 && status < 300) {
       setRoleCookieFromBody(res, data);
+      setFeaturesCookieFromBody(res, data);
     }
 
     return res;

app/api/users/me/route.ts

@@ -1,8 +1,11 @@
 "use client";
 
 import { useState, useEffect, useCallback, useRef, Suspense } from "react";
+import { useRouter } from "next/navigation";
 import { colors } from "@/app/lib/colors";
 import { STORAGE_KEY } from "@/app/lib/constants/keystore";
+import { FeatureFlag } from "@/app/lib/constants/featureFlags";
+import { removeFeatureFromClient } from "@/app/lib/featureState";
 import { APIKey } from "@/app/lib/types/credentials";
 import Sidebar from "@/app/components/Sidebar";
 import Loader from "@/app/components/Loader";
@@ -50,6 +53,7 @@ function ShimmerDot({ color }: { color: string }) {
 type IndicatorState = "none" | "processing" | "failed" | "success";
 
 function AssessmentContent() {
+  const router = useRouter();
   const toast = useToast();
   const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
   const [activeTab, setActiveTab] = useState<TabId>("datasets");
@@ -62,6 +66,7 @@ function AssessmentContent() {
   const [selectedKeyId, setSelectedKeyId] = useState("");
   const [evalIndicator, setEvalIndicator] = useState<IndicatorState>("none");
   const dismissedRef = useRef(false);
+  const featureRedirectingRef = useRef(false);
   const [assessmentRefreshToken, setAssessmentRefreshToken] = useState(0);
   const [experimentName, setExperimentName] = useState("");
   const {
@@ -92,12 +97,107 @@ function AssessmentContent() {
 
   const selectedKey = apiKeys.find((k) => k.id === selectedKeyId);
 
+  const redirectIfFeatureDisabled = useCallback(
+    async (
+      response: Response,
+      options?: { notify?: boolean },
+    ): Promise<boolean> => {
+      if (response.status !== 403) return false;
+
+      const errorData = await response
+        .clone()
+        .json()
+        .catch(() => ({}));
+      const message = String(
+        errorData?.error ?? errorData?.message ?? errorData?.detail ?? "",
+      ).toLowerCase();
+
+      if (
+        message.includes("feature") &&
+        message.includes("assessment") &&
+        message.includes("not enabled")
+      ) {
+        if (!featureRedirectingRef.current) {
+          featureRedirectingRef.current = true;
+          if (options?.notify) {
+            toast.error(
+              "Assessment feature is disabled for this organization/project.",
+            );
+          }
+          removeFeatureFromClient(FeatureFlag.ASSESSMENT);
+          // Trigger middleware redirect by navigating to the gated route.
+          window.setTimeout(
+            () => {
+              router.replace("/assessment");
+            },
+            options?.notify ? 300 : 0,
+          );
+        }
+        return true;
+      }
+
+      return false;
+    },
+    [router, toast],
+  );
+
+  useEffect(() => {
+    const originalFetch = window.fetch.bind(window);
+
+    window.fetch = async (...args: Parameters<typeof fetch>) => {
+      const response = await originalFetch(...args);
+      try {
+        const input = args[0];
+        const requestUrl =
+          typeof input === "string"
+            ? input
+            : input instanceof URL
+              ? input.toString()
+              : input.url;
+
+        if (requestUrl.includes("/api/assessment/")) {
+          await redirectIfFeatureDisabled(response);
+        }
+      } catch {
+        // silently ignore
+      }
+      return response;
+    };
+
+    return () => {
+      window.fetch = originalFetch;
+    };
+  }, [redirectIfFeatureDisabled]);
+
+  useEffect(() => {
+    if (!selectedKey?.key) return;
+
+    let cancelled = false;
+
+    (async () => {
+      try {
+        const response = await fetch("/api/assessment/evaluations?limit=1", {
+          headers: { "X-API-KEY": selectedKey.key },
+        });
+        if (cancelled) return;
+        await redirectIfFeatureDisabled(response, { notify: true });
+      } catch {
+        // silently ignore
+      }
+    })();
+
+    return () => {
+      cancelled = true;
+    };
+  }, [redirectIfFeatureDisabled, selectedKey?.key]);
+
   const pollEvalStatus = useCallback(async () => {
     if (!selectedKey) return;
     try {
-      const response = await fetch("/api/assessment/assessments?limit=10", {
+      const response = await fetch("/api/assessment/evaluations?limit=10", {
         headers: { "X-API-KEY": selectedKey.key },
       });
+      if (await redirectIfFeatureDisabled(response)) return;
       if (!response.ok) return;
       const data = await response.json();
       const runs = Array.isArray(data) ? data : data.data || [];
@@ -142,7 +242,7 @@ function AssessmentContent() {
     } catch {
       // silently fail
     }
-  }, [selectedKey]);
+  }, [redirectIfFeatureDisabled, selectedKey]);
 
   useEffect(() => {
     if (!selectedKey) return;
@@ -223,6 +323,8 @@ function AssessmentContent() {
         body: JSON.stringify(payload),
       });
 
+      if (await redirectIfFeatureDisabled(response)) return;
+
       if (!response.ok) {
         const errorData = await response.json().catch(() => ({}));
         throw new Error(

app/assessment/AssessmentPageClient.tsx

@@ -1,19 +1,5 @@
-import { notFound } from "next/navigation";
 import AssessmentPageClient from "@/app/assessment/AssessmentPageClient";
-import FeatureRouteGuard from "@/app/components/FeatureRouteGuard";
-import { FeatureFlag } from "@/app/lib/constants/featureFlags";
-import { getServerFeatureFlags } from "@/app/lib/featureFlags.server";
 
-export default async function AssessmentPage() {
-  const initialFlags = await getServerFeatureFlags();
-
-  if (!initialFlags[FeatureFlag.ASSESSMENT]) {
-    notFound();
-  }
-
-  return (
-    <FeatureRouteGuard featureFlag={FeatureFlag.ASSESSMENT}>
-      <AssessmentPageClient />
-    </FeatureRouteGuard>
-  );
+export default function AssessmentPage() {
+  return <AssessmentPageClient />;
 }

app/assessment/page.tsx

@@ -19,7 +19,6 @@ import {
   ShieldCheckIcon,
   SlidersIcon,
 } from "@/app/components/icons";
-import { useFeatureFlags } from "@/app/lib/FeatureFlagProvider";
 import { MenuItem, SidebarProps } from "@/app/lib/types/nav";
 import { LoginModal } from "@/app/components/auth";
 import { Branding, UserMenuPopover } from "@/app/components/user-menu";
@@ -34,9 +33,9 @@ export default function Sidebar({
   activeRoute = "/evaluations",
 }: SidebarProps) {
   const router = useRouter();
-  const { isEnabled, isLoaded: flagsLoaded } = useFeatureFlags();
   const [hasMounted, setHasMounted] = useState(false);
-  const { currentUser, googleProfile, isAuthenticated, logout } = useAuth();
+  const { currentUser, googleProfile, isAuthenticated, logout, hasFeature } =
+    useAuth();
   const [expandedMenus, setExpandedMenus] = useState<Record<string, boolean>>({
     Evaluations: true,
     Configurations: false,
@@ -132,7 +131,7 @@ export default function Sidebar({
     if (item.superuserOnly && !currentUser?.is_superuser) return false;
     if (item.featureFlag) {
       if (!hasMounted) return false;
-      if (flagsLoaded && !isEnabled(item.featureFlag)) return false;
+      if (!hasFeature(item.featureFlag)) return false;
     }
     return true;
   }).map((item) => ({

app/components/FeatureRouteGuard.tsx

@@ -1,11 +1,6 @@
 import type { Metadata } from "next";
 import { Inter, JetBrains_Mono } from "next/font/google";
 import "./globals.css";
-import { ToastProvider } from "@/app/components/Toast";
-import { AuthProvider } from "@/app/lib/context/AuthContext";
-import { AppProvider } from "@/app/lib/context/AppContext";
-import { FeatureFlagProvider } from "./lib/FeatureFlagProvider";
-import { getServerFeatureFlags } from "./lib/featureFlags.server";
 import { Providers } from "@/app/components/providers";
 import { APP_NAME } from "@/app/lib/constants";
 
@@ -26,25 +21,16 @@ export const metadata: Metadata = {
   description: "",
 };
 
-export default async function RootLayout({
+export default function RootLayout({
   children,
 }: Readonly<{
   children: React.ReactNode;
 }>) {
-  const initialFlags = await getServerFeatureFlags();
-
   return (
     <html lang="en">
       <body
         className={`${inter.variable} ${jetbrainsMono.variable} antialiased`}
       >
-        <FeatureFlagProvider initialFlags={initialFlags}>
-          <ToastProvider>
-            <AuthProvider>
-              <AppProvider>{children}</AppProvider>
-            </AuthProvider>
-          </ToastProvider>
-        </FeatureFlagProvider>
         <Providers>{children}</Providers>
       </body>
     </html>