diff --git a/apps/studio/lib/ai/ai-details.test.ts b/apps/studio/lib/ai/ai-details.test.ts index 56e4b61b3a2..343d36ea6eb 100644 --- a/apps/studio/lib/ai/ai-details.test.ts +++ b/apps/studio/lib/ai/ai-details.test.ts @@ -91,6 +91,7 @@ describe('getAIDetails', () => { hasAccessToAdvanceModel: false, hasHipaaAddon: false, orgId: 1, + orgSlug: ORG_SLUG, planId: 'pro', region: 'us-east-1', isSensitive: false, @@ -189,6 +190,7 @@ describe('getAIDetails', () => { expect(result.aiOptInLevel).toBe('disabled') expect(result.hasAccessToAdvanceModel).toBe(false) expect(result.orgId).toBeUndefined() + expect(result.orgSlug).toBeUndefined() expect(result.planId).toBeUndefined() }) @@ -216,6 +218,7 @@ describe('getAIDetails', () => { expect(result.aiOptInLevel).toBe('disabled') expect(result.orgId).toBeUndefined() + expect(result.orgSlug).toBeUndefined() }) it('falls back to the most restrictive posture when project detail is unavailable', async () => { diff --git a/apps/studio/lib/ai/ai-details.ts b/apps/studio/lib/ai/ai-details.ts index be61660368d..9a4fe56875c 100644 --- a/apps/studio/lib/ai/ai-details.ts +++ b/apps/studio/lib/ai/ai-details.ts @@ -11,6 +11,7 @@ export type AIDetails = { hasAccessToAdvanceModel: boolean hasHipaaAddon: boolean | undefined orgId: number | undefined + orgSlug: string | undefined planId: string | undefined region: string | undefined isSensitive: boolean | null | undefined @@ -56,6 +57,7 @@ export const getAIDetails = async ({ // Undefined rather than false so isTracingAllowed fails closed hasHipaaAddon: undefined, orgId: undefined, + orgSlug: undefined, planId: undefined, region, isSensitive, @@ -72,6 +74,7 @@ export const getAIDetails = async ({ hasAccessToAdvanceModel: advanceModelAccess.hasAccess, hasHipaaAddon, orgId: selectedOrg.id, + orgSlug: selectedOrg.slug, planId: selectedOrg.plan.id, region, isSensitive, diff --git a/apps/studio/lib/ai/generate-assistant-response.ts b/apps/studio/lib/ai/generate-assistant-response.ts index a46b4c6f89e..ccb02cc2866 100644 --- a/apps/studio/lib/ai/generate-assistant-response.ts +++ b/apps/studio/lib/ai/generate-assistant-response.ts @@ -39,6 +39,7 @@ export async function generateAssistantResponse({ supportMode, userId, orgId, + orgSlug, planId, includesLogsSnippets, isExplorerEnabled, @@ -60,6 +61,7 @@ export async function generateAssistantResponse({ supportMode?: boolean userId?: string orgId?: number + orgSlug?: string planId?: string /** Whether any user message in the conversation attached a logs (ClickHouse) query. */ includesLogsSnippets?: boolean @@ -171,6 +173,7 @@ export async function generateAssistantResponse({ aiOptInLevel, userId, orgId, + orgSlug, planId, requestedModel, gitBranch: process.env.VERCEL_GIT_COMMIT_REF, diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index 1b056af7c64..bd16a812d6c 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -96,7 +96,7 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: Jw messages: rawMessages, projectRef, connectionString, - orgSlug, + orgSlug: rawOrgSlug, chatId, chatName, model: rawRequestedModel, @@ -126,6 +126,7 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: Jw let projectIsSensitive: boolean | null | undefined let projectRegion: string | undefined let orgId: number | undefined + let orgSlug: string | undefined let planId: string | undefined if (!IS_PLATFORM) { @@ -133,14 +134,15 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: Jw hasAccessToAdvanceModel = true } - if (IS_PLATFORM && orgSlug && authorization && projectRef) { + if (IS_PLATFORM && rawOrgSlug && authorization && projectRef) { try { - const aiDetails = await getAIDetails({ orgSlug, projectRef, authorization }) + const aiDetails = await getAIDetails({ orgSlug: rawOrgSlug, projectRef, authorization }) aiOptInLevel = aiDetails.aiOptInLevel hasAccessToAdvanceModel = aiDetails.hasAccessToAdvanceModel orgHasHipaaAddon = aiDetails.hasHipaaAddon orgId = aiDetails.orgId + orgSlug = aiDetails.orgSlug planId = aiDetails.planId projectIsSensitive = aiDetails.isSensitive projectRegion = aiDetails.region @@ -234,6 +236,7 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: Jw supportMode, userId, orgId, + orgSlug, planId, includesLogsSnippets, isExplorerEnabled: explorerEnabled,