Merge pull request #320 from oaknational/feat/feature-flag-bootstrap

fix: bootstrap posthog feature flags
oaknational · Oct 31, 2024 · b2eba2a · b2eba2a
2 parents 6d4dacd + c44e1f1
commit b2eba2a
Show file tree

Hide file tree

Showing 11 changed files with 154 additions and 17 deletions.
diff --git a/apps/nextjs/src/app/layout.tsx b/apps/nextjs/src/app/layout.tsx
@@ -23,6 +23,7 @@ import { CookieConsentProvider } from "@/components/ContextProviders/CookieConse
 import FontProvider from "@/components/ContextProviders/FontProvider";
 import { GleapProvider } from "@/components/ContextProviders/GleapProvider";
 import { WebDebuggerPosition } from "@/lib/avo/Avo";
+import { getBootstrappedFeatures } from "@/lib/feature-flags/bootstrap";
 import { SentryIdentify } from "@/lib/sentry/SentryIdentify";
 import { cn } from "@/lib/utils";
 import { TRPCReactProvider } from "@/utils/trpc";
@@ -66,7 +67,9 @@ interface RootLayoutProps {
   children: React.ReactNode;
 }
 
-export default function RootLayout({ children }: Readonly<RootLayoutProps>) {
+export default async function RootLayout({
+  children,
+}: Readonly<RootLayoutProps>) {
   const nonce = headers().get("x-nonce");
   if (!nonce) {
     // Our middleware path matching excludes static paths like /_next/static/...
@@ -75,6 +78,8 @@ export default function RootLayout({ children }: Readonly<RootLayoutProps>) {
     return redirect("/not-found");
   }
 
+  const bootstrappedFeatures = await getBootstrappedFeatures(headers());
+
   return (
     <html lang="en" suppressHydrationWarning className={lexend.variable}>
       <ClerkProvider>
@@ -108,6 +113,7 @@ export default function RootLayout({ children }: Readonly<RootLayoutProps>) {
                           }),
                         },
                       }}
+                      bootstrappedFeatures={bootstrappedFeatures}
                     >
                       <GleapProvider>{children}</GleapProvider>
                     </AnalyticsProvider>

diff --git a/apps/nextjs/src/components/ContextProviders/AnalyticsProvider.tsx b/apps/nextjs/src/components/ContextProviders/AnalyticsProvider.tsx
@@ -74,6 +74,7 @@ export const analyticsContext = createContext<AnalyticsContext | null>(null);
 export type AnalyticsProviderProps = {
   children?: React.ReactNode;
   avoOptions?: Partial<AvoOptions>;
+  bootstrappedFeatures: Record<string, string | boolean>;
 };
 
 if (
@@ -130,6 +131,7 @@ const posthogClientAiBeta = new PostHog();
 export const AnalyticsProvider: React.FC<AnalyticsProviderProps> = ({
   children,
   avoOptions,
+  bootstrappedFeatures,
 }) => {
   const [hubspotScriptLoaded, setHubspotScriptLoadedFn] = useState(false);
   const setHubspotScriptLoaded = useCallback(() => {
@@ -162,7 +164,7 @@ export const AnalyticsProvider: React.FC<AnalyticsProviderProps> = ({
   );
   const posthogAiBeta = useAnalyticsService({
     service: posthogServiceAiBeta,
-    config: posthogAiBetaConfig,
+    config: { ...posthogAiBetaConfig, bootstrappedFeatures },
     consentState: posthogConsentAiBeta,
   });
 

diff --git a/apps/nextjs/src/lib/feature-flags/bootstrap.ts b/apps/nextjs/src/lib/feature-flags/bootstrap.ts
@@ -0,0 +1,46 @@
+"use server";
+
+import { auth } from "@clerk/nextjs/server";
+import { posthogAiBetaServerClient } from "@oakai/core/src/analytics/posthogAiBetaServerClient";
+import { aiLogger } from "@oakai/logger";
+import cookie from "cookie";
+import type { ReadonlyHeaders } from "next/dist/server/web/spec-extension/adapters/headers";
+import invariant from "tiny-invariant";
+
+const log = aiLogger("analytics:feature-flags");
+
+/**
+ * We use posthog feature flags to toggle functionality without deploying code changes.
+ * Fething feature flags on the frontend hasn't been reliable for us:
+ * - you have to wait for a round trip to the posthog API
+ * - we currently have a bug where denying cookie consent prevents feature flags from loading
+ *
+ * Instead, we can bootstrap feature flags by evaluating them on the server.
+ * https://posthog.com/docs/feature-flags/bootstrapping
+ */
+
+function getDistinctIdFromCookie(headers: ReadonlyHeaders) {
+  const cookieHeader = headers.get("cookie");
+  invariant(cookieHeader, "No cookie header");
+  const cookies = cookie.parse(cookieHeader) as Record<string, string>;
+  const phCookieKey = `ph_${process.env.NEXT_PUBLIC_POSTHOG_API_KEY}_posthog`;
+  const phCookie = cookies[phCookieKey];
+  if (!phCookie) {
+    return null;
+  }
+  return (JSON.parse(phCookie) as { distinct_id: string })["distinct_id"];
+}
+
+export async function getBootstrappedFeatures(headers: ReadonlyHeaders) {
+  const { userId } = auth();
+
+  const distinctId = userId ?? getDistinctIdFromCookie(headers) ?? "0";
+  log.info("Evaluating feature flags for", distinctId);
+  const features = await posthogAiBetaServerClient.getAllFlags(distinctId, {
+    // Only bootstrap flags which don't depend on user properties
+    // These are typically flags representing new features
+    onlyEvaluateLocally: true,
+  });
+  log.info("Bootstrapping feature flags", features);
+  return features;
+}
diff --git a/apps/nextjs/src/lib/posthog/posthog.ts b/apps/nextjs/src/lib/posthog/posthog.ts
@@ -12,18 +12,22 @@ export type PosthogConfig = {
   apiKey: string;
   apiHost: string;
   uiHost: string;
+  bootstrappedFeatures?: Record<string, string | boolean>;
 };
 
 export const posthogToAnalyticsService = (
   client: PostHog,
 ): AnalyticsService<PosthogConfig, "posthog"> => ({
   name: "posthog",
-  init: ({ apiKey, apiHost, uiHost }) =>
+  init: ({ apiKey, apiHost, uiHost, bootstrappedFeatures }) =>
     new Promise((resolve) => {
       client.init(apiKey, {
         api_host: apiHost,
         ui_host: uiHost,
         persistence: "localStorage+cookie",
+        bootstrap: {
+          featureFlags: bootstrappedFeatures,
+        },
         loaded: (posthog) => {
           // Enable debug mode in development
           if (process.env.NEXT_PUBLIC_POSTHOG_DEBUG === "true") {

diff --git a/apps/nextjs/src/middleware.ts b/apps/nextjs/src/middleware.ts
@@ -1,3 +1,5 @@
+import { auth } from "@clerk/nextjs/dist/types/server";
+import { posthogAiBetaServerClient } from "@oakai/core/src/analytics/posthogAiBetaServerClient";
 import type { NextMiddlewareResult } from "next/dist/server/web/types";
 import type { NextFetchEvent, NextMiddleware, NextRequest } from "next/server";
 import { NextResponse } from "next/server";

diff --git a/packages/api/src/router/appSessions.ts b/packages/api/src/router/appSessions.ts
@@ -1,6 +1,7 @@
 import type { SignedInAuthObject } from "@clerk/backend/internal";
 import { clerkClient } from "@clerk/nextjs/server";
 import { demoUsers } from "@oakai/core";
+import { posthogAiBetaServerClient } from "@oakai/core/src/analytics/posthogAiBetaServerClient";
 import { rateLimits } from "@oakai/core/src/utils/rateLimiting/rateLimit";
 import { RateLimitExceededError } from "@oakai/core/src/utils/rateLimiting/userBasedRateLimiter";
 import type { Prisma, PrismaClientWithAccelerate } from "@oakai/db";

diff --git a/packages/core/src/analytics/posthogAiBetaServerClient.ts b/packages/core/src/analytics/posthogAiBetaServerClient.ts
diff --git a/packages/core/src/analytics/posthogAiBetaServerClient/featureFlagEvaluation.ts b/packages/core/src/analytics/posthogAiBetaServerClient/featureFlagEvaluation.ts
@@ -0,0 +1,57 @@
+import { aiLogger } from "@oakai/logger";
+import { kv } from "@vercel/kv";
+import type { PostHog } from "posthog-node";
+
+const KV_KEY = "posthog-feature-flag-local-evaluation";
+
+const log = aiLogger("analytics:feature-flags");
+
+const setKv = async (response: Response) => {
+  const value = await response.text();
+  await kv.set(KV_KEY, value, { ex: 60 });
+};
+
+const getKv = async () => {
+  const value = await kv.get(KV_KEY);
+  if (!value) {
+    return null;
+  }
+  return {
+    status: 200,
+    json: () => Promise.resolve(value),
+    text: () => Promise.resolve(JSON.stringify(value)),
+  };
+};
+
+export const cachedFetch: PostHog["fetch"] = async (url, options) => {
+  if (url.includes("api/feature_flag/local_evaluation")) {
+    const kvCachedResponse = await getKv();
+    if (kvCachedResponse) {
+      log.info("evaluations fetched from KV");
+      return kvCachedResponse;
+    }
+    const result = await fetch(url, options);
+
+    if (result.ok) {
+      const cachedResult = result.clone();
+      await setKv(cachedResult);
+      log.info("evaluations cached to KV");
+    } else {
+      log.error("failed to load evaluations", { status: result.status });
+    }
+
+    return result;
+  }
+
+  if (url.includes("/decide")) {
+    log.warn("WARN: feature flag loaded through API");
+  }
+
+  return await fetch(url, options);
+};
+
+const ONE_DAY = 24 * 60 * 60 * 1000;
+const ONE_MINUTE = 60 * 1000;
+export const featureFlagsPollingInterval =
+  // prevent polling timeout from stacking when HMR replaces posthogAiBetaServerClient
+  process.env.NODE_ENV === "development" ? ONE_DAY : ONE_MINUTE;
diff --git a/packages/core/src/analytics/posthogAiBetaServerClient/index.ts b/packages/core/src/analytics/posthogAiBetaServerClient/index.ts
@@ -0,0 +1,30 @@
+import { PostHog } from "posthog-node";
+import invariant from "tiny-invariant";
+
+import {
+  cachedFetch,
+  featureFlagsPollingInterval,
+} from "./featureFlagEvaluation";
+
+const host = process.env.NEXT_PUBLIC_POSTHOG_HOST as string;
+const apiKey = process.env.NEXT_PUBLIC_POSTHOG_API_KEY || "*";
+const personalApiKey = process.env.POSTHOG_PERSONAL_KEY_FLAGS;
+invariant(personalApiKey, "POSTHOG_PERSONAL_KEY_FLAGS is required");
+
+const enableLocalEvaluation = process.env.NODE_ENV !== "test";
+
+/**
+ * This is the posthog nodejs client configured to send events to the
+ * posthog AI BETA instance.
+ */
+export const posthogAiBetaServerClient = new PostHog(apiKey, {
+  host,
+
+  // We evaluate user feature flags on the server to prevent round-trips to posthog.
+  // See https://posthog.com/docs/feature-flags/local-evaluation
+  // As we use edge functions, we can't hold the flag definitions in memory.
+  // Instead we cache them in KV through a custom fetch implementation.
+  fetch: cachedFetch,
+  featureFlagsPollingInterval,
+  personalApiKey: enableLocalEvaluation ? personalApiKey : undefined,
+});
diff --git a/packages/logger/index.ts b/packages/logger/index.ts
@@ -31,6 +31,7 @@ type ChildKey =
   | "aila:rag"
   | "aila:testing"
   | "analytics"
+  | "analytics:feature-flags"
   | "app"
   | "auth"
   | "chat"

diff --git a/turbo.json b/turbo.json
@@ -131,6 +131,7 @@
     "STRICT_CSP",
     "TELEMETRY_ENABLED",
     "UPSTASH_*",
-    "WOLFRAM_CLIENT_SECRET"
+    "WOLFRAM_CLIENT_SECRET",
+    "POSTHOG_PERSONAL_KEY_FLAGS"
   ]
 }