From a9022a57acaab33be14ae23ca8bdd3d36aa5e269 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daphn=C3=A9=20Popin?= <daphne.popin@gmail.com>
Date: Thu, 2 Jan 2025 12:24:46 +0100
Subject: [PATCH] Apply feedback

---
 front/lib/api/user.ts   |  5 +++--
 sdks/js/src/types.ts    | 11 +++++++++--
 types/src/front/user.ts |  9 ++++++++-
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/front/lib/api/user.ts b/front/lib/api/user.ts
index c9ccdf6ee42f..6bb3f7ba6ea0 100644
--- 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";
@@ -142,7 +143,7 @@ export async function fetchRevokedWorkspace(
 export async function getUserWithWorkspaces(
   user: UserResource,
   populateExtensionConfig = false
-): Promise<UserTypeWithWorkspaces> {
+): Promise<UserTypeWithWorkspaces | UserTypeWithExtensionWorkspaces> {
   const { memberships } = await MembershipResource.getActiveMemberships({
     users: [user],
   });
@@ -170,7 +171,7 @@ export async function getUserWithWorkspaces(
         segmentation: w.segmentation || null,
         whiteListedProviders: w.whiteListedProviders,
         defaultEmbeddingProvider: w.defaultEmbeddingProvider,
-        extensionBlacklistedDomains:
+        blacklistedDomains:
           configs.find((c) => c.workspaceId === w.id)?.blacklistedDomains ??
           undefined,
       };
diff --git a/sdks/js/src/types.ts b/sdks/js/src/types.ts
index 8419072e9f50..66eb7304d814 100644
--- 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()).optional(),
 });
 
 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
index 0419d152fd18..2b76aa0826c1 100644
--- 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[];
 };
 
 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;