From 6fb546a4ad4b306d18084a0bacd5a29ce6e651b6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 4 Jul 2025 14:55:08 +0000 Subject: [PATCH] Refactor chat permission check and remove unused organization role check Co-authored-by: dallin --- .../src/api/v2/chats/services/chat-helpers.ts | 9 +- packages/database/src/helpers/chats.ts | 351 ++++++++---------- 2 files changed, 165 insertions(+), 195 deletions(-) diff --git a/apps/server/src/api/v2/chats/services/chat-helpers.ts b/apps/server/src/api/v2/chats/services/chat-helpers.ts index d1ccfeaf1..8fddb6ba7 100644 --- a/apps/server/src/api/v2/chats/services/chat-helpers.ts +++ b/apps/server/src/api/v2/chats/services/chat-helpers.ts @@ -1,7 +1,7 @@ +import { canUserAccessChatCached } from '@buster/access-controls'; import { type User, chats, - checkChatPermission, createMessage, db, generateAssetMessages, @@ -9,7 +9,6 @@ import { getMessagesForChat, messages, } from '@buster/database'; -import { eq, and, gte, isNull } from 'drizzle-orm'; import type { Chat, Message } from '@buster/database'; import type { ChatMessage, @@ -23,6 +22,7 @@ import { ReasoningMessageSchema, ResponseMessageSchema, } from '@buster/server-shared/chats'; +import { and, eq, gte, isNull } from 'drizzle-orm'; // Optimized: Generic function to handle both response and reasoning messages function buildMessages( @@ -188,7 +188,10 @@ export async function handleExistingChat( throw new ChatError(ChatErrorCode.CHAT_NOT_FOUND, 'Chat not found', 404); } - const hasPermission = await checkChatPermission(chatId, user.id); + const hasPermission = await canUserAccessChatCached({ + userId: user.id, + chatId, + }); if (!hasPermission) { throw new ChatError( ChatErrorCode.PERMISSION_DENIED, diff --git a/packages/database/src/helpers/chats.ts b/packages/database/src/helpers/chats.ts index 356f9eb00..7c069db53 100644 --- a/packages/database/src/helpers/chats.ts +++ b/packages/database/src/helpers/chats.ts @@ -2,7 +2,7 @@ import { and, eq, isNull } from 'drizzle-orm'; import type { InferSelectModel } from 'drizzle-orm'; import { z } from 'zod'; import { db } from '../connection'; -import { chats, messages, userFavorites, users, usersToOrganizations } from '../schema'; +import { chats, messages, userFavorites, users } from '../schema'; // Type inference from schema export type Chat = InferSelectModel; @@ -11,28 +11,28 @@ export type User = InferSelectModel; // Create a type for updateable chat fields by excluding auto-managed fields type UpdateableChatFields = Partial< - Omit + Omit >; /** * Input/output schemas for type safety */ export const CreateChatInputSchema = z.object({ - title: z.string().min(1), - userId: z.string().uuid(), - organizationId: z.string().uuid(), + title: z.string().min(1), + userId: z.string().uuid(), + organizationId: z.string().uuid(), }); export const GetChatInputSchema = z.object({ - chatId: z.string().uuid(), - userId: z.string().uuid(), + chatId: z.string().uuid(), + userId: z.string().uuid(), }); export const CreateMessageInputSchema = z.object({ - chatId: z.string().uuid(), - content: z.string(), - userId: z.string().uuid(), - messageId: z.string().uuid().optional(), + chatId: z.string().uuid(), + content: z.string(), + userId: z.string().uuid(), + messageId: z.string().uuid().optional(), }); export type CreateChatInput = z.infer; @@ -43,188 +43,155 @@ export type CreateMessageInput = z.infer; * Create a new chat */ export async function createChat(input: CreateChatInput): Promise { - try { - const validated = CreateChatInputSchema.parse(input); + try { + const validated = CreateChatInputSchema.parse(input); - const [chat] = await db - .insert(chats) - .values({ - title: validated.title, - organizationId: validated.organizationId, - createdBy: validated.userId, - updatedBy: validated.userId, - publiclyAccessible: false, - }) - .returning(); + const [chat] = await db + .insert(chats) + .values({ + title: validated.title, + organizationId: validated.organizationId, + createdBy: validated.userId, + updatedBy: validated.userId, + publiclyAccessible: false, + }) + .returning(); - if (!chat) { - throw new Error('Failed to create chat'); - } + if (!chat) { + throw new Error('Failed to create chat'); + } - return chat; - } catch (error) { - if (error instanceof z.ZodError) { - throw new Error(`Invalid chat input: ${error.errors.map((e) => e.message).join(', ')}`); - } - throw error; - } + return chat; + } catch (error) { + if (error instanceof z.ZodError) { + throw new Error(`Invalid chat input: ${error.errors.map((e) => e.message).join(', ')}`); + } + throw error; + } } /** * Get a chat with user and favorite information */ export async function getChatWithDetails(input: GetChatInput): Promise<{ - chat: Chat; - user: User | null; - isFavorited: boolean; + chat: Chat; + user: User | null; + isFavorited: boolean; } | null> { - const validated = GetChatInputSchema.parse(input); + const validated = GetChatInputSchema.parse(input); - // Get chat with creator info - const result = await db - .select({ - chat: chats, - user: users, - }) - .from(chats) - .leftJoin(users, eq(chats.createdBy, users.id)) - .where(and(eq(chats.id, validated.chatId), isNull(chats.deletedAt))) - .limit(1); + // Get chat with creator info + const result = await db + .select({ + chat: chats, + user: users, + }) + .from(chats) + .leftJoin(users, eq(chats.createdBy, users.id)) + .where(and(eq(chats.id, validated.chatId), isNull(chats.deletedAt))) + .limit(1); - if (!result.length || !result[0]?.chat) { - return null; - } + if (!result.length || !result[0]?.chat) { + return null; + } - // Check if favorited - const favorite = await db - .select() - .from(userFavorites) - .where( - and( - eq(userFavorites.userId, validated.userId), - eq(userFavorites.assetId, validated.chatId), - eq(userFavorites.assetType, 'chat') - ) - ) - .limit(1); + // Check if favorited + const favorite = await db + .select() + .from(userFavorites) + .where( + and( + eq(userFavorites.userId, validated.userId), + eq(userFavorites.assetId, validated.chatId), + eq(userFavorites.assetType, 'chat'), + ), + ) + .limit(1); - const { chat, user } = result[0]; - return { - chat, - user, - isFavorited: favorite.length > 0, - }; + const { chat, user } = result[0]; + return { + chat, + user, + isFavorited: favorite.length > 0, + }; } /** * Create a new message in a chat */ export async function createMessage(input: CreateMessageInput): Promise { - try { - const validated = CreateMessageInputSchema.parse(input); + try { + const validated = CreateMessageInputSchema.parse(input); - const messageId = validated.messageId || crypto.randomUUID(); + const messageId = validated.messageId || crypto.randomUUID(); - // Use transaction to ensure atomicity - const result = await db.transaction(async (tx) => { - const [message] = await tx - .insert(messages) - .values({ - id: messageId, - chatId: validated.chatId, - createdBy: validated.userId, - requestMessage: validated.content, - responseMessages: {}, - reasoning: {}, - title: validated.content.substring(0, 255), // Ensure title fits in database - rawLlmMessages: {}, - isCompleted: false, - }) - .returning(); + // Use transaction to ensure atomicity + const result = await db.transaction(async (tx) => { + const [message] = await tx + .insert(messages) + .values({ + id: messageId, + chatId: validated.chatId, + createdBy: validated.userId, + requestMessage: validated.content, + responseMessages: {}, + reasoning: {}, + title: validated.content.substring(0, 255), // Ensure title fits in database + rawLlmMessages: {}, + isCompleted: false, + }) + .returning(); - if (!message) { - throw new Error('Failed to create message'); - } + if (!message) { + throw new Error('Failed to create message'); + } - // Update chat's updated_at timestamp - await tx - .update(chats) - .set({ updatedAt: new Date().toISOString() }) - .where(eq(chats.id, validated.chatId)); + // Update chat's updated_at timestamp + await tx + .update(chats) + .set({ updatedAt: new Date().toISOString() }) + .where(eq(chats.id, validated.chatId)); - return message; - }); + return message; + }); - return result; - } catch (error) { - if (error instanceof z.ZodError) { - throw new Error(`Invalid message input: ${error.errors.map((e) => e.message).join(', ')}`); - } - throw error; - } + return result; + } catch (error) { + if (error instanceof z.ZodError) { + throw new Error(`Invalid message input: ${error.errors.map((e) => e.message).join(', ')}`); + } + throw error; + } } /** * Check if a user has permission to access a chat */ export async function checkChatPermission(chatId: string, userId: string): Promise { - const chat = await db - .select() - .from(chats) - .where(and(eq(chats.id, chatId), isNull(chats.deletedAt))) - .limit(1); + const chat = await db + .select() + .from(chats) + .where(and(eq(chats.id, chatId), isNull(chats.deletedAt))) + .limit(1); - if (!chat.length) { - return false; - } + if (!chat.length) { + return false; + } - // Check if user is the creator - if (chat[0]?.createdBy === userId) { - return true; - } - - // Check if user is an admin in the same organization as the chat - const chatOrganizationId = chat[0]?.organizationId; - if (!chatOrganizationId) { - return false; - } - - // Check user's role in the organization - const userOrgRole = await db - .select({ - role: usersToOrganizations.role, - }) - .from(usersToOrganizations) - .where( - and( - eq(usersToOrganizations.userId, userId), - eq(usersToOrganizations.organizationId, chatOrganizationId), - isNull(usersToOrganizations.deletedAt) - ) - ) - .limit(1); - - if (userOrgRole.length > 0 && userOrgRole[0]) { - const role = userOrgRole[0].role; - // Check if user has admin role (workspace_admin or data_admin) - if (role === 'workspace_admin' || role === 'data_admin') { - return true; - } - } - - // TODO: Add more sophisticated permission checking with asset_permissions table - return false; + // For now, only check if user is the creator + // TODO: Add more sophisticated permission checking with asset_permissions table + return chat[0]?.createdBy === userId; } /** * Get all messages for a chat */ export async function getMessagesForChat(chatId: string): Promise { - return db - .select() - .from(messages) - .where(and(eq(messages.chatId, chatId), isNull(messages.deletedAt))) - .orderBy(messages.createdAt); + return db + .select() + .from(messages) + .where(and(eq(messages.chatId, chatId), isNull(messages.deletedAt))) + .orderBy(messages.createdAt); } /** @@ -236,50 +203,50 @@ export async function getMessagesForChat(chatId: string): Promise { * @returns Success status */ export async function updateChat( - chatId: string, - fields: UpdateableChatFields + chatId: string, + fields: UpdateableChatFields, ): Promise<{ success: boolean }> { - try { - // First verify the chat exists and is not deleted - const existingChat = await db - .select({ id: chats.id }) - .from(chats) - .where(and(eq(chats.id, chatId), isNull(chats.deletedAt))) - .limit(1); + try { + // First verify the chat exists and is not deleted + const existingChat = await db + .select({ id: chats.id }) + .from(chats) + .where(and(eq(chats.id, chatId), isNull(chats.deletedAt))) + .limit(1); - if (existingChat.length === 0) { - throw new Error(`Chat not found or has been deleted: ${chatId}`); - } + if (existingChat.length === 0) { + throw new Error(`Chat not found or has been deleted: ${chatId}`); + } - // Build update object with only provided fields - const updateData: any = { - updatedAt: new Date().toISOString(), - }; + // Build update object with only provided fields + const updateData: any = { + updatedAt: new Date().toISOString(), + }; - // Only add fields that are actually provided (not undefined) - for (const [key, value] of Object.entries(fields)) { - if (value !== undefined && key !== 'id' && key !== 'createdAt' && key !== 'deletedAt') { - updateData[key] = value; - } - } + // Only add fields that are actually provided (not undefined) + for (const [key, value] of Object.entries(fields)) { + if (value !== undefined && key !== 'id' && key !== 'createdAt' && key !== 'deletedAt') { + updateData[key] = value; + } + } - // If updatedAt was explicitly provided, use that instead - if ('updatedAt' in fields && fields.updatedAt !== undefined) { - updateData.updatedAt = fields.updatedAt; - } + // If updatedAt was explicitly provided, use that instead + if ('updatedAt' in fields && fields.updatedAt !== undefined) { + updateData.updatedAt = fields.updatedAt; + } - await db - .update(chats) - .set(updateData) - .where(and(eq(chats.id, chatId), isNull(chats.deletedAt))); + await db + .update(chats) + .set(updateData) + .where(and(eq(chats.id, chatId), isNull(chats.deletedAt))); - return { success: true }; - } catch (error) { - console.error('Failed to update chat fields:', error); - // Re-throw our specific validation errors - if (error instanceof Error && error.message.includes('Chat not found')) { - throw error; - } - throw new Error(`Failed to update chat fields for chat ${chatId}`); - } + return { success: true }; + } catch (error) { + console.error('Failed to update chat fields:', error); + // Re-throw our specific validation errors + if (error instanceof Error && error.message.includes('Chat not found')) { + throw error; + } + throw new Error(`Failed to update chat fields for chat ${chatId}`); + } }