diff --git a/apps/sim/lib/mothership/chat/payload.test.ts b/apps/sim/lib/mothership/chat/payload.test.ts index acf565d95db..323288ca6b0 100644 --- a/apps/sim/lib/mothership/chat/payload.test.ts +++ b/apps/sim/lib/mothership/chat/payload.test.ts @@ -5,6 +5,7 @@ import { billingSubscriptionMock, billingSubscriptionMockFns, } from '@sim/testing/mocks/billing-subscription.mock' +import { featureFlagsMock, featureFlagsMockFns } from '@sim/testing/mocks/feature-flags.mock' import { integrationsAvailabilityMock, integrationsAvailabilityMockFns, @@ -51,6 +52,7 @@ vi.mock('@/lib/mothership/chat/workspace-inventory', () => ({ truncated: [], })), })) +vi.mock('@/lib/core/config/feature-flags', () => featureFlagsMock) vi.mock('@/lib/dashboards/application/availability', () => ({ readDashboardAvailability: { execute: mockDashboardAvailability }, })) @@ -375,19 +377,34 @@ describe('buildCopilotRequestPayload', () => { } ) - it('never grants the workspace-only dashboards entitlement to an organization chat', async () => { - mockDashboardAvailability.mockResolvedValue(true) - const payload = await buildCopilotRequestPayload({ - message: 'Show my dashboard', - userId: 'actor', - userMessageId: 'message-1', - organizationId: 'org-1', - principal: { kind: 'session' as const, userId: 'actor' }, - mode: 'agent', - model: '', - }) - expect(payload.entitlements).toEqual([]) - }) + it.each([ + ['dashboards', 'dashboards', true], + ['dashboards', 'dashboards', false], + ['tests', 'workflow-tests', true], + ['tests', 'workflow-tests', false], + ['changelog', 'changelog', true], + ['changelog', 'changelog', false], + ] as const)( + 'grants an organization Build chat the %s entitlement from the %s rollout: %s', + async (entitlement, rolloutFlag, enabled) => { + featureFlagsMockFns.mockIsFeatureEnabled.mockImplementation( + async (flag) => enabled && flag === rolloutFlag + ) + const payload = await buildCopilotRequestPayload({ + message: 'Build it', + userId: 'actor', + userMessageId: 'message-1', + organizationId: 'org-1', + principal: { kind: 'session' as const, userId: 'actor' }, + mode: 'agent', + model: '', + }) + expect(payload.entitlements).toEqual(enabled ? [entitlement] : []) + expect(featureFlagsMockFns.mockIsFeatureEnabled).toHaveBeenCalledWith(rolloutFlag, { + orgId: 'org-1', + }) + } + ) beforeEach(() => { mockTrackChatUpload.mockResolvedValue({ displayName: 'payroll.xlsx' }) diff --git a/apps/sim/lib/mothership/entitlements.ts b/apps/sim/lib/mothership/entitlements.ts index 47070290d7c..c16ed5c3414 100644 --- a/apps/sim/lib/mothership/entitlements.ts +++ b/apps/sim/lib/mothership/entitlements.ts @@ -1,8 +1,11 @@ import type { Principal } from '@sim/auth/principal' import { readChangelogAvailability } from '@/lib/changelog/application/availability' +import { isChangelogEnabled } from '@/lib/changelog/feature-flag' import { readDashboardAvailability } from '@/lib/dashboards/application/availability' +import { isDashboardsEnabled } from '@/lib/dashboards/feature-flag' import { ENTITLEMENTS, type Entitlement } from '@/lib/mothership/generated/protocol' import { readWorkflowTestAvailability } from '@/lib/workflow-tests/application/availability' +import { isWorkflowTestsEnabled } from '@/lib/workflow-tests/feature-flag' /** The owner of one chat turn: exactly one of a workspace or an organization. */ export interface EntitlementOwner { @@ -22,9 +25,10 @@ interface OrganizationOwner { } /** - * Each entitlement declares the chat scopes it exists in. A scope it does not - * declare is never granted, so a workspace-only capability cannot leak into an - * organization chat that has no workspace to run it against. + * Each entitlement declares the chat scopes it exists in; a scope it does not declare is never + * granted. An organization chat in Build or Plan works in the organization's workspaces, so a + * workspace capability exists there whenever the organization has it rolled out. The operation + * behind it still checks the workspace it runs against. */ interface EntitlementEvaluator { workspace?: (owner: WorkspaceOwner) => Promise @@ -49,14 +53,17 @@ const EVALUATORS: Record = { [ENTITLEMENTS.dashboards]: { workspace: ({ principal, workspaceId }) => readDashboardAvailability.execute({ principal, input: { workspaceId } }), + organization: ({ organizationId }) => isDashboardsEnabled(organizationId), }, [ENTITLEMENTS.tests]: { workspace: ({ principal, workspaceId }) => readWorkflowTestAvailability.execute({ principal, input: { workspaceId } }), + organization: ({ organizationId }) => isWorkflowTestsEnabled(organizationId), }, [ENTITLEMENTS.changelog]: { workspace: ({ principal, workspaceId }) => readChangelogAvailability.execute({ principal, input: { workspaceId } }), + organization: ({ organizationId }) => isChangelogEnabled(organizationId), }, }