Merge branch 'LAION-AI:main' into messageNavigation

website/src/pages/api/auth/[...nextauth].ts

@@ -59,6 +59,17 @@ if (boolean(process.env.DEBUG_LOGIN) || process.env.NODE_ENV === "development")
   );
 }
 
+// Create a map of provider types to a set of admin user identifiers based on
+// the environment variables.  We assume the list is separated by ',' and each
+// entry is separated by ':'.
+const adminUserMap = process.env.ADMIN_USERS.split(",").reduce((result, entry) => {
+  const [authType, id] = entry.split(":");
+  const s = result.get(authType) || new Set();
+  s.add(id);
+  result.set(authType, s);
+  return result;
+}, new Map());
+
 export const authOptions: AuthOptions = {
   // Ensure we can store user data in a database.
   adapter: PrismaAdapter(prisma),
@@ -68,6 +79,56 @@ export const authOptions: AuthOptions = {
     verifyRequest: "/auth/verify",
     // error: "/auth/error", -Will be used later
   },
+  callbacks: {
+    /**
+     * Ensure we propagate the user's role when creating the session from the
+     * token.
+     */
+    async session({ session, token }) {
+      session.user.role = token.role;
+      return session;
+    },
+    /**
+     * When creating a token, fetch the user's role and inject it in the token.
+     * This let's use forward the role to the session object.
+     */
+    async jwt({ token }) {
+      const { role } = await prisma.user.findUnique({
+        where: { id: token.sub },
+        select: { role: true },
+      });
+      token.role = role;
+      return token;
+    },
+  },
+  events: {
+    /**
+     * Update the user's role after they have successfully signed in
+     */
+    async signIn({ user, account }) {
+      // Get the admin list for the user's auth type.
+      const adminForAccountType = adminUserMap.get(account.provider);
+
+      // Return early if there's no admin list.
+      if (!adminForAccountType) {
+        return;
+      }
+
+      // TODO(#236): Reduce the number of times we update the role field.
+
+      // Update the database if the user is an admin.
+      if (adminForAccountType.has(account.providerAccountId)) {
+        await prisma.user.update({
+          data: {
+            role: "admin",
+          },
+          where: {
+            id: user.id,
+          },
+        });
+      }
+    },
+  },
   session: {
     strategy: "jwt",
   },

website/src/pages/api/new_task/[task_type].ts

@@ -1,4 +1,5 @@
 import { getToken } from "next-auth/jwt";
+import { oasstApiClient } from "src/lib/oasst_api_client";
 import prisma from "src/lib/prismadb";
 
 /**
@@ -21,24 +22,7 @@ const handler = async (req, res) => {
   }
 
   // Fetch the new task.
-  //
-  // This needs to be refactored into an easier to use library.
-  const taskRes = await fetch(`${process.env.FASTAPI_URL}/api/v1/tasks/`, {
-    method: "POST",
-    headers: {
-      "X-API-Key": process.env.FASTAPI_KEY,
-      "Content-Type": "application/json",
-    },
-    body: JSON.stringify({
-      type: task_type,
-      user: {
-        id: token.sub,
-        display_name: token.name || token.email,
-        auth_method: "local",
-      },
-    }),
-  });
-  const task = await taskRes.json();
+  const task = await oasstApiClient.fetchTask(task_type, token);
 
   // Store the task and link it to the user..
   const registeredTask = await prisma.registeredTask.create({
@@ -53,16 +37,7 @@ const handler = async (req, res) => {
   });
 
   // Update the backend with our Task ID
-  await fetch(`${process.env.FASTAPI_URL}/api/v1/tasks/${task.id}/ack`, {
-    method: "POST",
-    headers: {
-      "X-API-Key": process.env.FASTAPI_KEY,
-      "Content-Type": "application/json",
-    },
-    body: JSON.stringify({
-      message_id: registeredTask.id,
-    }),
-  });
+  await oasstApiClient.ackTask(task.id, registeredTask.id);
 
   // Send the results to the client.
   res.status(200).json(registeredTask);

website/src/pages/api/update_task.ts

@@ -1,4 +1,5 @@
 import { getToken } from "next-auth/jwt";
+import { oasstApiClient } from "src/lib/oasst_api_client";
 import prisma from "src/lib/prismadb";
 
 /**
@@ -34,27 +35,7 @@ const handler = async (req, res) => {
     },
   });
 
-  // Send the interaction to the Task Backend.  This automatically fetches the
-  // next task in the sequence (or the done task).
-  const interactionRes = await fetch(`${process.env.FASTAPI_URL}/api/v1/tasks/interaction`, {
-    method: "POST",
-    headers: {
-      "X-API-Key": process.env.FASTAPI_KEY,
-      "Content-Type": "application/json",
-    },
-    body: JSON.stringify({
-      type: update_type,
-      user: {
-        id: token.sub,
-        display_name: token.name || token.email,
-        auth_method: "local",
-      },
-      message_id: id,
-      user_message_id: interaction.id,
-      ...content,
-    }),
-  });
-  const newTask = await interactionRes.json();
+  const newTask = await oasstApiClient.interactTask(update_type, id, interaction.id, content, token);
 
   // Stores the new task with our database.
   const newRegisteredTask = await prisma.registeredTask.create({