add ai ase review to license claims. Control license before feature usage. Rewrite HasAiCaseReviewEnabled

Author: Pascal-Delange

dto/license_dto.go

@@ -21,6 +21,7 @@ type LicenseEntitlements struct {
 	TestRun        bool `json:"test_run"`
 	Sanctions      bool `json:"sanctions"`
 	CaseAutoAssign bool `json:"auto_assignment"` //nolint:tagliatelle
+	CaseAiAssist   bool `json:"case_ai_assist"`
 }
 
 func AdaptLicenseEntitlements(licenseEntitlements models.LicenseEntitlements) LicenseEntitlements {
@@ -35,6 +36,7 @@ func AdaptLicenseEntitlements(licenseEntitlements models.LicenseEntitlements) Li
 		TestRun:        licenseEntitlements.TestRun,
 		Sanctions:      licenseEntitlements.Sanctions,
 		CaseAutoAssign: licenseEntitlements.CaseAutoAssign,
+		CaseAiAssist:   licenseEntitlements.CaseAiAssist,
 	}
 }
 

dto/organization_feature_access_dto.go

@@ -16,7 +16,8 @@ type APIOrganizationFeatureAccess struct {
 	NameRecognition string `json:"name_recognition"`
 	CaseAutoAssign  string `json:"case_auto_assign"`
 
-	// user-scoped - temporarly at least
+	// user-scoped
+	// Currently only used to control display of the AI assist button in the UI - DO NOT use for anything else as it will be removed
 	AiAssist string `json:"ai_assist"`
 }
 

models/license.go

@@ -61,6 +61,7 @@ type LicenseEntitlements struct {
 	TestRun        bool
 	Sanctions      bool
 	CaseAutoAssign bool
+	CaseAiAssist   bool
 }
 
 type LicenseValidation struct {
@@ -83,6 +84,7 @@ func NewFullLicense() LicenseValidation {
 			TestRun:        true,
 			Sanctions:      true,
 			CaseAutoAssign: true,
+			CaseAiAssist:   true,
 		},
 	}
 }

models/organization_feature_access.go

@@ -16,10 +16,12 @@ type OrganizationFeatureAccess struct {
 	Sanctions       FeatureAccess
 	NameRecognition FeatureAccess
 	CaseAutoAssign  FeatureAccess
+	CaseAiAssist    FeatureAccess
 	CreatedAt       time.Time
 	UpdatedAt       time.Time
 
-	// user-scoped, temporarily at least
+	// user-scoped
+	// Currently only used to control display of the AI assist button in the UI - DO NOT use for anything else as it will be removed
 	AiAssist FeatureAccess
 }
 
@@ -29,6 +31,7 @@ type DbStoredOrganizationFeatureAccess struct {
 	TestRun        FeatureAccess
 	Sanctions      FeatureAccess
 	CaseAutoAssign FeatureAccess
+	CaseAiAssist   FeatureAccess
 	CreatedAt      time.Time
 	UpdatedAt      time.Time
 }
@@ -37,6 +40,8 @@ type UpdateOrganizationFeatureAccessInput struct {
 	OrganizationId string
 	TestRun        *FeatureAccess
 	Sanctions      *FeatureAccess
+	CaseAiAssist   *FeatureAccess
+	CaseAutoAssign *FeatureAccess
 }
 
 type FeaturesConfiguration struct {
@@ -58,6 +63,7 @@ func (f DbStoredOrganizationFeatureAccess) MergeWithLicenseEntitlement(
 		Sanctions:       f.Sanctions,
 		NameRecognition: f.Sanctions,
 		CaseAutoAssign:  f.CaseAutoAssign,
+		CaseAiAssist:    f.CaseAiAssist,
 		CreatedAt:       f.CreatedAt,
 		UpdatedAt:       f.UpdatedAt,
 	}
@@ -90,6 +96,9 @@ func (f DbStoredOrganizationFeatureAccess) MergeWithLicenseEntitlement(
 	if !l.CaseAutoAssign {
 		o.CaseAutoAssign = Restricted
 	}
+	if !l.CaseAiAssist {
+		o.CaseAiAssist = Restricted
+	}
 
 	// remove the feature accesses that are not allowed by the configuration
 	if o.Analytics.IsAllowed() && !c.Analytics {

models/user.go

@@ -5,14 +5,16 @@ import "time"
 type UserId string
 
 type User struct {
-	UserId          UserId
-	Email           string
-	Role            Role
-	OrganizationId  string
-	PartnerId       *string
-	FirstName       string
-	LastName        string
-	DeletedAt       *time.Time
+	UserId         UserId
+	Email          string
+	Role           Role
+	OrganizationId string
+	PartnerId      *string
+	FirstName      string
+	LastName       string
+	DeletedAt      *time.Time
+
+	// Currently only used to control display of the AI assist button in the UI - DO NOT use for anything else as it will be removed
 	AiAssistEnabled bool
 	Picture         string
 }

repositories/dbmodels/db_license.go

@@ -27,6 +27,7 @@ type DBLicense struct {
 	TestRun              bool               `db:"test_run"`
 	Sanctions            bool               `db:"sanctions"`
 	AutoAssignment       bool               `db:"auto_assignment"`
+	CaseAiAssist         bool               `db:"case_ai_assist"`
 }
 
 const TABLE_LICENSES = "licenses"
@@ -53,6 +54,7 @@ func AdaptLicense(db DBLicense) (models.License, error) {
 			TestRun:        db.TestRun,
 			Sanctions:      db.Sanctions,
 			CaseAutoAssign: db.AutoAssignment,
+			CaseAiAssist:   db.CaseAiAssist,
 		},
 	}, nil
 }

repositories/license_repository.go

@@ -79,6 +79,7 @@ func (repo *MarbleDbRepository) CreateLicense(ctx context.Context, exec Executor
 				"test_run",
 				"sanctions",
 				"auto_assignment",
+				"case_ai_assist",
 			).
 			Values(
 				license.Id,
@@ -96,6 +97,7 @@ func (repo *MarbleDbRepository) CreateLicense(ctx context.Context, exec Executor
 				license.LicenseEntitlements.TestRun,
 				license.LicenseEntitlements.Sanctions,
 				license.LicenseEntitlements.CaseAutoAssign,
+				license.LicenseEntitlements.CaseAiAssist,
 			),
 	)
 	return err
@@ -136,6 +138,7 @@ func (repo *MarbleDbRepository) UpdateLicense(ctx context.Context, exec Executor
 		query = query.Set("test_run", licenseEntitlements.TestRun)
 		query = query.Set("sanctions", licenseEntitlements.Sanctions)
 		query = query.Set("auto_assignment", licenseEntitlements.CaseAutoAssign)
+		query = query.Set("case_ai_assist", licenseEntitlements.CaseAiAssist)
 	}
 
 	err := ExecBuilder(

repositories/migrations/20251120132400_licences_ai_assist.sql

@@ -0,0 +1,31 @@
+-- +goose Up
+-- +goose StatementBegin
+ALTER TABLE licenses
+ADD COLUMN case_ai_assist BOOL NOT NULL DEFAULT FALSE;
+
+ALTER TABLE organization_feature_access
+ADD COLUMN case_ai_assist text NOT NULL DEFAULT 'allowed';
+
+UPDATE organization_feature_access
+SET
+    case_ai_assist = 'allowed'
+WHERE
+    org_id IN (
+        SELECT
+            id
+        FROM
+            organizations
+        WHERE
+            ai_case_review_enabled = true
+    );
+
+-- +goose StatementEnd
+-- +goose Down
+ALTER TABLE organization_feature_access
+DROP COLUMN case_ai_assist;
+
+ALTER TABLE licenses
+DROP COLUMN case_ai_assist;
+
+-- +goose StatementBegin
+-- +goose StatementEnd

usecases/ai_agent/ai_agent_usecase.go

@@ -137,6 +137,11 @@ type caseReviewTaskEnqueuer interface {
 	) error
 }
 
+type featureAccessReader interface {
+	GetOrganizationFeatureAccess(ctx context.Context, organizationId string, userId *models.UserId) (
+		models.OrganizationFeatureAccess, error)
+}
+
 type AiAgentUsecase struct {
 	enforceSecurityCase         security.EnforceSecurityCase
 	enforceSecurityOrganization security.EnforceSecurityOrganization
@@ -153,6 +158,7 @@ type AiAgentUsecase struct {
 	caseReviewFileRepository    caseReviewWorkerRepository
 	blobRepository              repositories.BlobRepository
 	caseReviewTaskEnqueuer      caseReviewTaskEnqueuer
+	featureAccessReader         featureAccessReader
 	config                      infra.AIAgentConfiguration
 	caseManagerBucketUrl        string
 
@@ -177,6 +183,7 @@ func NewAiAgentUsecase(
 	blobRepository repositories.BlobRepository,
 	caseReviewTaskEnqueuer caseReviewTaskEnqueuer,
 	transactionFactory executor_factory.TransactionFactory,
+	featureAccessReader featureAccessReader,
 	config infra.AIAgentConfiguration,
 	caseManagerBucketUrl string,
 ) AiAgentUsecase {
@@ -196,6 +203,7 @@ func NewAiAgentUsecase(
 		blobRepository:              blobRepository,
 		caseReviewTaskEnqueuer:      caseReviewTaskEnqueuer,
 		transactionFactory:          transactionFactory,
+		featureAccessReader:         featureAccessReader,
 		config:                      config,
 		caseManagerBucketUrl:        caseManagerBucketUrl,
 	}

usecases/ai_agent/ai_case_review.go

@@ -1051,17 +1051,11 @@ func (uc *AiAgentUsecase) UpdateAiCaseReviewFeedback(
 }
 
 func (uc *AiAgentUsecase) HasAiCaseReviewEnabled(ctx context.Context, orgId string) (bool, error) {
-	// Add check on license here (is not yet handled in license)
-
-	// Check if the organization has AI case review enabled, fetch the organization and check the flag
-	org, err := uc.repository.GetOrganizationById(ctx, uc.executorFactory.NewExecutor(), orgId)
+	featureAccess, err := uc.featureAccessReader.GetOrganizationFeatureAccess(ctx, orgId, nil)
 	if err != nil {
 		return false, err
 	}
-	if !org.AiCaseReviewEnabled {
-		return false, nil
-	}
-	return true, nil
+	return featureAccess.CaseAiAssist.IsAllowed(), nil
 }
 
 var templateFuncMap = template.FuncMap{