From c4d4733469eb54ba0e12726be0d5fbbd850300a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daphn=C3=A9=20Popin?= Date: Thu, 2 Jan 2025 16:39:43 +0100 Subject: [PATCH] Public API/SDK update for /me route to handle extension blacklisted domains. (#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 | 5 +++- front/lib/api/auth_wrappers.ts | 6 ++++- front/lib/api/user.ts | 41 +++++++++++++++++--------------- front/lib/resources/extension.ts | 14 +++++++++++ sdks/js/package-lock.json | 4 ++-- sdks/js/package.json | 2 +- sdks/js/src/types.ts | 10 +++++++- types/src/front/user.ts | 10 +++++++- 8 files changed, 66 insertions(+), 26 deletions(-) diff --git a/extension/app/src/lib/auth.ts b/extension/app/src/lib/auth.ts index e2b9a9375a60..3a9271ed51f8 100644 --- a/extension/app/src/lib/auth.ts +++ b/extension/app/src/lib/auth.ts @@ -117,7 +117,10 @@ const fetchMe = async ( token: string ): Promise> => { 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) { diff --git a/front/lib/api/auth_wrappers.ts b/front/lib/api/auth_wrappers.ts index 73076f79a24f..7f76fd345cb8 100644 --- a/front/lib/api/auth_wrappers.ts +++ b/front/lib/api/auth_wrappers.ts @@ -428,7 +428,11 @@ export function withAuth0TokenAuthentication( }); } - const userWithWorkspaces = await getUserWithWorkspaces(user); + const isFromExtension = req.headers["x-request-origin"] === "extension"; + const userWithWorkspaces = await getUserWithWorkspaces( + user, + isFromExtension + ); return handler(req, res, userWithWorkspaces); } diff --git a/front/lib/api/user.ts b/front/lib/api/user.ts index 4aa9b3690c03..e08ee079c000 100644 --- a/front/lib/api/user.ts +++ b/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 { +export async function getUserWithWorkspaces( + 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, + }), }; }), }; diff --git a/front/lib/resources/extension.ts b/front/lib/resources/extension.ts index 48c0b3e919e1..dd26ba040d42 100644 --- a/front/lib/resources/extension.ts +++ b/front/lib/resources/extension.ts @@ -115,6 +115,20 @@ export class ExtensionConfigurationResource extends BaseResource { + const configs = await this.model.findAll({ + where: { + workspaceId: workspaceIds, + }, + }); + + return configs.map( + (config) => new this(ExtensionConfigurationModel, config.get()) + ); + } + async updateBlacklistedDomains( auth: Authenticator, { diff --git a/sdks/js/package-lock.json b/sdks/js/package-lock.json index 61ad5bab9873..3cef1d9101f6 100644 --- a/sdks/js/package-lock.json +++ b/sdks/js/package-lock.json @@ -1,12 +1,12 @@ { "name": "@dust-tt/client", - "version": "1.0.16", + "version": "1.0.17", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@dust-tt/client", - "version": "1.0.16", + "version": "1.0.17", "license": "ISC", "dependencies": { "eventsource-parser": "^1.1.1", diff --git a/sdks/js/package.json b/sdks/js/package.json index 8415a38d46e6..c97eb41a1298 100644 --- a/sdks/js/package.json +++ b/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", diff --git a/sdks/js/src/types.ts b/sdks/js/src/types.ts index e313464dabce..a2af49357faa 100644 --- a/sdks/js/src/types.ts +++ b/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; diff --git a/types/src/front/user.ts b/types/src/front/user.ts index 656f1a9e6745..26abe7d82bf1 100644 --- a/types/src/front/user.ts +++ b/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 = {