Public API/SDK update for /me route to handle extension blacklisted d…

…omains. (#9688)

* Public API/SDK update for /me route to handle extension blacklisted domains

* Apply feedback

* Apply feedback

* Revert change typing

extension/app/src/lib/auth.ts

@@ -117,7 +117,10 @@ const fetchMe = async (
   token: string
 ): Promise<Result<{ user: UserTypeWithWorkspaces }, AuthError>> => {
   const response = await fetch(`${process.env.DUST_DOMAIN}/api/v1/me`, {
-    headers: { Authorization: `Bearer ${token}` },
+    headers: {
+      Authorization: `Bearer ${token}`,
+      "X-Request-Origin": "extension",
+    },
   });
   const me = await response.json();
   if (!response.ok) {

front/lib/api/auth_wrappers.ts

@@ -428,7 +428,11 @@ export function withAuth0TokenAuthentication<T>(
         });
       }
 
-      const userWithWorkspaces = await getUserWithWorkspaces(user);
+      const isFromExtension = req.headers["x-request-origin"] === "extension";
+      const userWithWorkspaces = await getUserWithWorkspaces(
+        user,
+        isFromExtension
+      );
 
       return handler(req, res, userWithWorkspaces);
     }

front/lib/api/user.ts

@@ -1,14 +1,15 @@
 import type {
   Result,
-  RoleType,
   UserMetadataType,
   UserType,
+  UserTypeWithExtensionWorkspaces,
   UserTypeWithWorkspaces,
 } from "@dust-tt/types";
 import { Err, Ok } from "@dust-tt/types";
 
 import type { Authenticator } from "@app/lib/auth";
 import { Workspace } from "@app/lib/models/workspace";
+import { ExtensionConfigurationResource } from "@app/lib/resources/extension";
 import { UserMetadataModel } from "@app/lib/resources/storage/models/user";
 import { UserResource } from "@app/lib/resources/user_resource";
 import logger from "@app/logger/logger";
@@ -139,42 +140,44 @@ export async function fetchRevokedWorkspace(
   return new Ok(workspace);
 }
 
-export async function getUserWithWorkspaces(
-  user: UserResource
-): Promise<UserTypeWithWorkspaces> {
+export async function getUserWithWorkspaces<T extends boolean>(
+  user: UserResource,
+  populateExtensionConfig: T = false as T
+): Promise<
+  T extends true ? UserTypeWithExtensionWorkspaces : UserTypeWithWorkspaces
+> {
   const { memberships } = await MembershipResource.getActiveMemberships({
     users: [user],
   });
+  const workspaceIds = memberships.map((m) => m.workspaceId);
   const workspaces = await Workspace.findAll({
     where: {
-      id: memberships.map((m) => m.workspaceId),
+      id: workspaceIds,
     },
   });
 
+  const configs = populateExtensionConfig
+    ? await ExtensionConfigurationResource.internalFetchForWorkspaces(
+        workspaceIds
+      )
+    : [];
+
   return {
     ...user.toJSON(),
     workspaces: workspaces.map((w) => {
-      const m = memberships.find((m) => m.workspaceId === w.id);
-      let role = "none" as RoleType;
-      if (m) {
-        switch (m.role) {
-          case "admin":
-          case "builder":
-          case "user":
-            role = m.role;
-            break;
-          default:
-            role = "none";
-        }
-      }
       return {
         id: w.id,
         sId: w.sId,
         name: w.name,
-        role,
+        role: memberships.find((m) => m.workspaceId === w.id)?.role ?? "none",
         segmentation: w.segmentation || null,
         whiteListedProviders: w.whiteListedProviders,
         defaultEmbeddingProvider: w.defaultEmbeddingProvider,
+        ...(populateExtensionConfig && {
+          blacklistedDomains:
+            configs.find((c) => c.workspaceId === w.id)?.blacklistedDomains ??
+            null,
+        }),
       };
     }),
   };

front/lib/resources/extension.ts

@@ -115,6 +115,20 @@ export class ExtensionConfigurationResource extends BaseResource<ExtensionConfig
     return config ? new this(ExtensionConfigurationModel, config.get()) : null;
   }
 
+  static async internalFetchForWorkspaces(
+    workspaceIds: ModelId[]
+  ): Promise<ExtensionConfigurationResource[]> {
+    const configs = await this.model.findAll({
+      where: {
+        workspaceId: workspaceIds,
+      },
+    });
+
+    return configs.map(
+      (config) => new this(ExtensionConfigurationModel, config.get())
+    );
+  }
+
   async updateBlacklistedDomains(
     auth: Authenticator,
     {

sdks/js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dust-tt/client",
-  "version": "1.0.16",
+  "version": "1.0.17",
   "description": "Client for Dust API",
   "repository": {
     "type": "git",

sdks/js/src/types.ts

@@ -693,6 +693,10 @@ const WorkspaceSchema = LightWorkspaceSchema.extend({
   ssoEnforced: z.boolean().optional(),
 });
 
+const ExtensionWorkspaceSchema = WorkspaceSchema.extend({
+  blacklistedDomains: z.array(z.string()).nullable(),
+});
+
 const UserProviderSchema = FlexibleEnumSchema<
   "auth0" | "github" | "google" | "okta" | "samlp" | "waad"
 >().nullable();
@@ -2199,7 +2203,11 @@ export type FileUploadedRequestResponseType = z.infer<
 >;
 
 export const MeResponseSchema = z.object({
-  user: UserSchema.and(z.object({ workspaces: LightWorkspaceSchema.array() })),
+  user: UserSchema.and(
+    z.object({
+      workspaces: WorkspaceSchema.array().or(ExtensionWorkspaceSchema.array()),
+    })
+  ),
 });
 
 export type MeResponseType = z.infer<typeof MeResponseSchema>;

types/src/front/user.ts

@@ -46,6 +46,10 @@ export type WorkspaceType = LightWorkspaceType & {
   ssoEnforced?: boolean;
 };
 
+export type ExtensionWorkspaceType = WorkspaceType & {
+  blacklistedDomains: string[] | null;
+};
+
 export type UserProviderType =
   | "auth0"
   | "github"
@@ -69,7 +73,11 @@ export type UserType = {
 };
 
 export type UserTypeWithWorkspaces = UserType & {
-  workspaces: LightWorkspaceType[];
+  workspaces: WorkspaceType[];
+};
+
+export type UserTypeWithExtensionWorkspaces = UserType & {
+  workspaces: ExtensionWorkspaceType[];
 };
 
 export type UserMetadataType = {