tenantpilot/lib/auth/utils.ts

import { db } from "@/lib/db/index";
import { users } from "@/lib/db/schema/auth";
import { DrizzleAdapter } from "@auth/drizzle-adapter";
import { getServerSession } from "next-auth/next";
import { type Adapter } from "next-auth/adapters";
import { redirect } from "next/navigation";
import { z } from "zod";
import { eq } from "drizzle-orm";
import AzureADProvider from "next-auth/providers/azure-ad";

export type AuthSession = {
  session: {
    user: {
      id: string;
      name?: string;
      email?: string;
      tenantId?: string;
    };
  } | null;
};

const envSchema = z.object({
  AZURE_AD_CLIENT_ID: z.string().optional(),
  AZURE_AD_CLIENT_SECRET: z.string().optional(),
});

export const env = envSchema.parse(process.env);

export const authOptions = {
  adapter: DrizzleAdapter(db) as Adapter,
  session: {
    strategy: "jwt" as const, // CRITICAL: Use JWT strategy to access token in session callback
  },
  callbacks: {
    jwt: async ({ token, account, profile }: { token: any; account?: any; profile?: any }) => {
      // Store access token
      if (account) {
        token.accessToken = account.access_token;
      }
      // Extract tenantId from Azure AD tid claim
      if (profile && 'tid' in profile) {
        token.tenantId = profile.tid as string;
        
        // Update tenant_id in database if we have a user id
        if (token.sub) {
          try {
            await db
              .update(users)
              .set({ tenantId: profile.tid as string })
              .where(eq(users.id, token.sub));
          } catch (error) {
            console.error('Failed to update tenant_id:', error);
          }
        }
      }
      return token;
    },
    session: ({ session, token }: { session: any; token: any }) => {
      // Copy user id from token
      if (token?.sub) {
        session.user.id = token.sub;
      }
      // Copy access token
      if (token?.accessToken) {
        session.accessToken = token.accessToken as string;
      }
      // Copy tenantId from token to session
      if (token?.tenantId) {
        session.user.tenantId = token.tenantId as string;
      }
      return session;
    },
  },
  providers: [
    AzureADProvider({
      clientId: env.AZURE_AD_CLIENT_ID ?? "",
      clientSecret: env.AZURE_AD_CLIENT_SECRET ?? "",
      tenantId: "common", // Multi-Tenancy Support
      authorization: {
        params: {
          scope: "openid profile email offline_access User.Read",
        },
      },
    }),
  ],
};


export const getUserAuth = async () => {
  const session = await getServerSession(authOptions);
  return { session } as AuthSession;
};

export const checkAuth = async () => {
  const { session } = await getUserAuth();
  if (!session) redirect("/api/auth/signin");
};