Apply feedback

dust-tt · Jan 2, 2025 · b5dfd1c · b5dfd1c
1 parent 2f833cb
commit b5dfd1c
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 9 deletions.
diff --git a/front/lib/api/user.ts b/front/lib/api/user.ts
@@ -2,6 +2,7 @@ import type {
   Result,
   UserMetadataType,
   UserType,
+  UserTypeWithExtensionWorkspaces,
   UserTypeWithWorkspaces,
 } from "@dust-tt/types";
 import { Err, Ok } from "@dust-tt/types";
@@ -139,10 +140,12 @@ export async function fetchRevokedWorkspace(
   return new Ok(workspace);
 }
 
-export async function getUserWithWorkspaces(
+export async function getUserWithWorkspaces<T extends boolean>(
   user: UserResource,
-  populateExtensionConfig = false
-): Promise<UserTypeWithWorkspaces> {
+  populateExtensionConfig: T = false as T
+): Promise<
+  T extends true ? UserTypeWithExtensionWorkspaces : UserTypeWithWorkspaces
+> {
   const { memberships } = await MembershipResource.getActiveMemberships({
     users: [user],
   });
@@ -170,9 +173,11 @@ export async function getUserWithWorkspaces(
         segmentation: w.segmentation || null,
         whiteListedProviders: w.whiteListedProviders,
         defaultEmbeddingProvider: w.defaultEmbeddingProvider,
-        extensionBlacklistedDomains:
-          configs.find((c) => c.workspaceId === w.id)?.blacklistedDomains ??
-          undefined,
+        ...(populateExtensionConfig && {
+          blacklistedDomains:
+            configs.find((c) => c.workspaceId === w.id)?.blacklistedDomains ??
+            null,
+        }),
       };
     }),
   };

diff --git a/sdks/js/src/types.ts b/sdks/js/src/types.ts
@@ -691,7 +691,10 @@ export type LightWorkspaceType = z.infer<typeof LightWorkspaceSchema>;
 
 const WorkspaceSchema = LightWorkspaceSchema.extend({
   ssoEnforced: z.boolean().optional(),
-  extensionBlacklistedDomains: z.array(z.string()).optional(),
+});
+
+const ExtensionWorkspaceSchema = WorkspaceSchema.extend({
+  blacklistedDomains: z.array(z.string()).nullable(),
 });
 
 const UserProviderSchema = FlexibleEnumSchema<
@@ -2200,7 +2203,11 @@ export type FileUploadedRequestResponseType = z.infer<
 >;
 
 export const MeResponseSchema = z.object({
-  user: UserSchema.and(z.object({ workspaces: WorkspaceSchema.array() })),
+  user: UserSchema.and(
+    z.object({
+      workspaces: WorkspaceSchema.array().or(ExtensionWorkspaceSchema.array()),
+    })
+  ),
 });
 
 export type MeResponseType = z.infer<typeof MeResponseSchema>;

diff --git a/types/src/front/user.ts b/types/src/front/user.ts
@@ -44,7 +44,10 @@ export type LightWorkspaceType = {
 
 export type WorkspaceType = LightWorkspaceType & {
   ssoEnforced?: boolean;
-  extensionBlacklistedDomains?: string[];
+};
+
+export type ExtensionWorkspaceType = WorkspaceType & {
+  blacklistedDomains: string[] | null;
 };
 
 export type UserProviderType =
@@ -73,6 +76,10 @@ export type UserTypeWithWorkspaces = UserType & {
   workspaces: WorkspaceType[];
 };
 
+export type UserTypeWithExtensionWorkspaces = UserType & {
+  workspaces: ExtensionWorkspaceType[];
+};
+
 export type UserMetadataType = {
   key: string;
   value: string;