fix(model): block deletion when model is referenced by KB or agent

Reject model delete with HTTP 400 when the current tenant still has
knowledge bases or custom agents bound to the model ID.

internal/application/repository/custom_agent.go

@@ -60,3 +60,16 @@ func (r *customAgentRepository) UpdateAgent(ctx context.Context, agent *types.Cu
 func (r *customAgentRepository) DeleteAgent(ctx context.Context, id string, tenantID uint64) error {
 	return r.db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&types.CustomAgent{}).Error
 }
+
+// CountByModelID counts active agents whose config references modelID.
+func (r *customAgentRepository) CountByModelID(
+	ctx context.Context, tenantID uint64, modelID string,
+) (int64, error) {
+	var count int64
+	query := r.db.WithContext(ctx).
+		Model(&types.CustomAgent{}).
+		Where("tenant_id = ?", tenantID)
+	query = scopeCustomAgentsByModelID(query, modelID)
+	err := query.Count(&count).Error
+	return count, err
+}

internal/application/repository/knowledgebase.go

@@ -203,3 +203,17 @@ func (r *knowledgeBaseRepository) CountByVectorStoreID(
 		Count(&count).Error
 	return count, err
 }
+
+// CountByModelID counts active knowledge bases that reference modelID in any
+// model-binding column (scalar fields or JSON config blobs).
+func (r *knowledgeBaseRepository) CountByModelID(
+	ctx context.Context, tenantID uint64, modelID string,
+) (int64, error) {
+	var count int64
+	query := r.db.WithContext(ctx).
+		Model(&types.KnowledgeBase{}).
+		Where("tenant_id = ?", tenantID)
+	query = scopeKnowledgeBasesByModelID(query, modelID)
+	err := query.Count(&count).Error
+	return count, err
+}

internal/application/repository/model_usage.go

@@ -0,0 +1,49 @@
+package repository
+
+import (
+	"gorm.io/gorm"
+)
+
+// scopeKnowledgeBasesByModelID filters knowledge_bases rows that reference
+// modelID in any model-binding field.
+func scopeKnowledgeBasesByModelID(db *gorm.DB, modelID string) *gorm.DB {
+	if db.Dialector.Name() == "postgres" {
+		return db.Where(
+			"embedding_model_id = ? OR summary_model_id = ? OR "+
+				"image_processing_config->>'model_id' = ? OR "+
+				"vlm_config->>'model_id' = ? OR "+
+				"asr_config->>'model_id' = ? OR "+
+				"wiki_config->>'synthesis_model_id' = ?",
+			modelID, modelID, modelID, modelID, modelID, modelID,
+		)
+	}
+	return db.Where(
+		"embedding_model_id = ? OR summary_model_id = ? OR "+
+			"json_extract(image_processing_config, '$.model_id') = ? OR "+
+			"json_extract(vlm_config, '$.model_id') = ? OR "+
+			"json_extract(asr_config, '$.model_id') = ? OR "+
+			"json_extract(wiki_config, '$.synthesis_model_id') = ?",
+		modelID, modelID, modelID, modelID, modelID, modelID,
+	)
+}
+
+// scopeCustomAgentsByModelID filters custom_agents rows whose config JSON
+// references modelID in any model-binding field.
+func scopeCustomAgentsByModelID(db *gorm.DB, modelID string) *gorm.DB {
+	if db.Dialector.Name() == "postgres" {
+		return db.Where(
+			"config->>'model_id' = ? OR config->>'rerank_model_id' = ? OR "+
+				"config->>'vlm_model_id' = ? OR config->>'asr_model_id' = ? OR "+
+				"config->>'query_understand_model_id' = ?",
+			modelID, modelID, modelID, modelID, modelID,
+		)
+	}
+	return db.Where(
+		"json_extract(config, '$.model_id') = ? OR "+
+			"json_extract(config, '$.rerank_model_id') = ? OR "+
+			"json_extract(config, '$.vlm_model_id') = ? OR "+
+			"json_extract(config, '$.asr_model_id') = ? OR "+
+			"json_extract(config, '$.query_understand_model_id') = ?",
+		modelID, modelID, modelID, modelID, modelID,
+	)
+}

internal/application/repository/model_usage_test.go

@@ -0,0 +1,113 @@
+package repository
+
+import (
+	"context"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+const customAgentsTestDDL = `
+CREATE TABLE IF NOT EXISTS custom_agents (
+    id VARCHAR(36) NOT NULL,
+    name VARCHAR(255) NOT NULL,
+    description TEXT,
+    avatar VARCHAR(64),
+    is_builtin BOOLEAN NOT NULL DEFAULT 0,
+    tenant_id INTEGER NOT NULL,
+    created_by VARCHAR(36),
+    config TEXT NOT NULL DEFAULT '{}',
+    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    deleted_at DATETIME,
+    PRIMARY KEY (id, tenant_id)
+);
+`
+
+func setupModelUsageTestDB(t *testing.T) *gorm.DB {
+	t.Helper()
+	db := setupKBTestDB(t)
+	require.NoError(t, db.Exec(customAgentsTestDDL).Error)
+	return db
+}
+
+func TestCountByModelID_KnowledgeBase(t *testing.T) {
+	ctx := context.Background()
+	db := setupModelUsageTestDB(t)
+	repo := NewKnowledgeBaseRepository(db)
+	modelID := "embed-model-1"
+
+	kb := makeKB(nil)
+	kb.EmbeddingModelID = modelID
+	require.NoError(t, db.Create(kb).Error)
+
+	count, err := repo.CountByModelID(ctx, 1, modelID)
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), count)
+
+	count, err = repo.CountByModelID(ctx, 1, "other-model")
+	require.NoError(t, err)
+	assert.Equal(t, int64(0), count)
+
+	kb2 := makeKB(nil)
+	kb2.ID = uuid.New().String()
+	kb2.VLMConfig = types.VLMConfig{Enabled: true, ModelID: modelID}
+	require.NoError(t, db.Create(kb2).Error)
+
+	count, err = repo.CountByModelID(ctx, 1, modelID)
+	require.NoError(t, err)
+	assert.Equal(t, int64(2), count)
+
+	require.NoError(t, db.Delete(kb2).Error)
+	count, err = repo.CountByModelID(ctx, 1, modelID)
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), count)
+}
+
+func TestCountByModelID_CustomAgent(t *testing.T) {
+	ctx := context.Background()
+	db := setupModelUsageTestDB(t)
+	repo := NewCustomAgentRepository(db)
+	modelID := "chat-model-1"
+
+	agent := &types.CustomAgent{
+		ID:       uuid.New().String(),
+		Name:     "test-agent",
+		TenantID: 1,
+		Config: types.CustomAgentConfig{
+			ModelID: modelID,
+		},
+	}
+	require.NoError(t, repo.CreateAgent(ctx, agent))
+
+	count, err := repo.CountByModelID(ctx, 1, modelID)
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), count)
+
+	agent2 := &types.CustomAgent{
+		ID:       uuid.New().String(),
+		Name:     "rerank-agent",
+		TenantID: 1,
+		Config: types.CustomAgentConfig{
+			RerankModelID: modelID,
+		},
+	}
+	require.NoError(t, repo.CreateAgent(ctx, agent2))
+
+	count, err = repo.CountByModelID(ctx, 1, modelID)
+	require.NoError(t, err)
+	assert.Equal(t, int64(2), count)
+
+	count, err = repo.CountByModelID(ctx, 2, modelID)
+	require.NoError(t, err)
+	assert.Equal(t, int64(0), count)
+
+	require.NoError(t, repo.DeleteAgent(ctx, agent2.ID, 1))
+	count, err = repo.CountByModelID(ctx, 1, modelID)
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), count)
+}

internal/application/service/knowledgebase_pr3_test.go

@@ -108,6 +108,9 @@ func (r *fakeKBRepo) TogglePinKnowledgeBase(_ context.Context, _ string, _ uint6
 func (r *fakeKBRepo) CountByVectorStoreID(_ context.Context, _ *gorm.DB, _ uint64, _ string) (int64, error) {
 	return 0, nil
 }
+func (r *fakeKBRepo) CountByModelID(_ context.Context, _ uint64, _ string) (int64, error) {
+	return 0, nil
+}
 func (r *fakeKBRepo) SetUserKBPin(_ context.Context, _ uint64, _ string, _ string, _ bool) (*time.Time, error) {
 	return nil, nil
 }

internal/application/service/model.go

@@ -3,7 +3,9 @@ package service
 import (
 	"context"
 	"errors"
+	"fmt"
 
+	apperrors "github.com/Tencent/WeKnora/internal/errors"
 	"github.com/Tencent/WeKnora/internal/logger"
 	"github.com/Tencent/WeKnora/internal/models/asr"
 	"github.com/Tencent/WeKnora/internal/models/chat"
@@ -23,6 +25,8 @@ var ErrModelNotFound = errors.New("model not found")
 // modelService implements the model service interface
 type modelService struct {
 	repo          interfaces.ModelRepository
+	kbRepo        interfaces.KnowledgeBaseRepository
+	agentRepo     interfaces.CustomAgentRepository
 	ollamaService *ollama.OllamaService
 	pooler        embedding.EmbedderPooler
 	tenantService interfaces.TenantService
@@ -30,12 +34,16 @@ type modelService struct {
 
 // NewModelService creates a new model service instance
 func NewModelService(repo interfaces.ModelRepository,
+	kbRepo interfaces.KnowledgeBaseRepository,
+	agentRepo interfaces.CustomAgentRepository,
 	ollamaService *ollama.OllamaService,
 	pooler embedding.EmbedderPooler,
 	tenantService interfaces.TenantService,
 ) interfaces.ModelService {
 	return &modelService{
 		repo:          repo,
+		kbRepo:        kbRepo,
+		agentRepo:     agentRepo,
 		ollamaService: ollamaService,
 		pooler:        pooler,
 		tenantService: tenantService,
@@ -341,9 +349,31 @@ func (s *modelService) DeleteModel(ctx context.Context, id string) error {
 		})
 		return err
 	}
-	if existingModel != nil && existingModel.IsBuiltin {
+	if existingModel == nil {
+		return ErrModelNotFound
+	}
+	if existingModel.IsBuiltin {
 		logger.Warnf(ctx, "Attempted to delete builtin model: %s", id)
-		return errors.New("builtin models cannot be deleted")
+		return apperrors.NewBadRequestError("builtin models cannot be deleted")
+	}
+
+	kbCount, err := s.kbRepo.CountByModelID(ctx, tenantID, id)
+	if err != nil {
+		logger.ErrorWithFields(ctx, err, map[string]interface{}{
+			"model_id": id,
+		})
+		return err
+	}
+	agentCount, err := s.agentRepo.CountByModelID(ctx, tenantID, id)
+	if err != nil {
+		logger.ErrorWithFields(ctx, err, map[string]interface{}{
+			"model_id": id,
+		})
+		return err
+	}
+	if kbCount > 0 || agentCount > 0 {
+		logger.Warnf(ctx, "Model %s is in use: kb=%d agent=%d", id, kbCount, agentCount)
+		return apperrors.NewBadRequestError(formatModelInUseMessage(kbCount, agentCount))
 	}
 
 	// Delete model from repository
@@ -582,3 +612,26 @@ func (s *modelService) GetASRModel(ctx context.Context, modelId string) (asr.ASR
 
 	return sttModel, nil
 }
+
+func formatModelInUseMessage(kbCount, agentCount int64) string {
+	switch {
+	case kbCount > 0 && agentCount > 0:
+		return fmt.Sprintf(
+			"model is used by %d knowledge base(s) and %d agent(s); "+
+				"reconfigure or remove those references before deleting",
+			kbCount, agentCount,
+		)
+	case kbCount > 0:
+		return fmt.Sprintf(
+			"model is used by %d knowledge base(s); "+
+				"reconfigure or remove those references before deleting",
+			kbCount,
+		)
+	default:
+		return fmt.Sprintf(
+			"model is used by %d agent(s); "+
+				"reconfigure or remove those references before deleting",
+			agentCount,
+		)
+	}
+}

internal/application/service/model_delete_test.go

@@ -0,0 +1,175 @@
+package service
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	apperrors "github.com/Tencent/WeKnora/internal/errors"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+type stubKBRepoForModelDelete struct {
+	count int64
+}
+
+func (s *stubKBRepoForModelDelete) CreateKnowledgeBase(context.Context, *types.KnowledgeBase) error {
+	return nil
+}
+func (s *stubKBRepoForModelDelete) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) {
+	return nil, nil
+}
+func (s *stubKBRepoForModelDelete) GetKnowledgeBaseByIDAndTenant(context.Context, string, uint64) (*types.KnowledgeBase, error) {
+	return nil, nil
+}
+func (s *stubKBRepoForModelDelete) GetKnowledgeBaseByIDs(context.Context, []string) ([]*types.KnowledgeBase, error) {
+	return nil, nil
+}
+func (s *stubKBRepoForModelDelete) ListKnowledgeBases(context.Context) ([]*types.KnowledgeBase, error) {
+	return nil, nil
+}
+func (s *stubKBRepoForModelDelete) ListKnowledgeBasesByTenantID(context.Context, uint64) ([]*types.KnowledgeBase, error) {
+	return nil, nil
+}
+func (s *stubKBRepoForModelDelete) UpdateKnowledgeBase(context.Context, *types.KnowledgeBase) error {
+	return nil
+}
+func (s *stubKBRepoForModelDelete) DeleteKnowledgeBase(context.Context, string) error { return nil }
+func (s *stubKBRepoForModelDelete) CountByVectorStoreID(context.Context, *gorm.DB, uint64, string) (int64, error) {
+	return 0, nil
+}
+func (s *stubKBRepoForModelDelete) CountByModelID(context.Context, uint64, string) (int64, error) {
+	return s.count, nil
+}
+func (s *stubKBRepoForModelDelete) SetUserKBPin(context.Context, uint64, string, string, bool) (*time.Time, error) {
+	return nil, nil
+}
+func (s *stubKBRepoForModelDelete) ListUserKBPinIDs(context.Context, uint64, string) (map[string]time.Time, error) {
+	return nil, nil
+}
+
+type stubAgentRepoForModelDelete struct {
+	count int64
+}
+
+func (s *stubAgentRepoForModelDelete) CreateAgent(context.Context, *types.CustomAgent) error {
+	return nil
+}
+func (s *stubAgentRepoForModelDelete) GetAgentByID(context.Context, string, uint64) (*types.CustomAgent, error) {
+	return nil, nil
+}
+func (s *stubAgentRepoForModelDelete) ListAgentsByTenantID(context.Context, uint64) ([]*types.CustomAgent, error) {
+	return nil, nil
+}
+func (s *stubAgentRepoForModelDelete) UpdateAgent(context.Context, *types.CustomAgent) error {
+	return nil
+}
+func (s *stubAgentRepoForModelDelete) DeleteAgent(context.Context, string, uint64) error { return nil }
+func (s *stubAgentRepoForModelDelete) CountByModelID(context.Context, uint64, string) (int64, error) {
+	return s.count, nil
+}
+
+type stubModelRepoForDelete struct {
+	model  *types.Model
+	delete func(id string) error
+}
+
+func (s *stubModelRepoForDelete) Create(context.Context, *types.Model) error { return nil }
+func (s *stubModelRepoForDelete) GetByID(_ context.Context, _ uint64, id string) (*types.Model, error) {
+	if s.model != nil && s.model.ID == id {
+		return s.model, nil
+	}
+	return nil, nil
+}
+func (s *stubModelRepoForDelete) List(context.Context, uint64, types.ModelType, types.ModelSource) ([]*types.Model, error) {
+	return nil, nil
+}
+func (s *stubModelRepoForDelete) Update(context.Context, *types.Model) error { return nil }
+func (s *stubModelRepoForDelete) Delete(_ context.Context, _ uint64, id string) error {
+	if s.delete != nil {
+		return s.delete(id)
+	}
+	return nil
+}
+func (s *stubModelRepoForDelete) ClearDefaultByType(context.Context, uint, types.ModelType, string) error {
+	return nil
+}
+
+func TestDeleteModel_RejectsWhenReferenced(t *testing.T) {
+	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
+	modelID := "model-in-use"
+
+	svc := NewModelService(
+		&stubModelRepoForDelete{model: &types.Model{ID: modelID, TenantID: 1}},
+		&stubKBRepoForModelDelete{count: 1},
+		&stubAgentRepoForModelDelete{count: 0},
+		nil, nil, nil,
+	)
+
+	err := svc.DeleteModel(ctx, modelID)
+	require.Error(t, err)
+	appErr, ok := apperrors.IsAppError(err)
+	require.True(t, ok)
+	assert.Equal(t, apperrors.ErrBadRequest, appErr.Code)
+	assert.Contains(t, appErr.Message, "knowledge base")
+}
+
+func TestDeleteModel_RejectsWhenUsedByAgent(t *testing.T) {
+	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
+	modelID := "agent-model"
+
+	svc := NewModelService(
+		&stubModelRepoForDelete{model: &types.Model{ID: modelID, TenantID: 1}},
+		&stubKBRepoForModelDelete{count: 0},
+		&stubAgentRepoForModelDelete{count: 2},
+		nil, nil, nil,
+	)
+
+	err := svc.DeleteModel(ctx, modelID)
+	require.Error(t, err)
+	appErr, ok := apperrors.IsAppError(err)
+	require.True(t, ok)
+	assert.Contains(t, appErr.Message, "2 agent(s)")
+}
+
+func TestDeleteModel_SucceedsWhenUnreferenced(t *testing.T) {
+	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
+	modelID := "free-model"
+	deleted := false
+
+	svc := NewModelService(
+		&stubModelRepoForDelete{
+			model: &types.Model{ID: modelID, TenantID: 1},
+			delete: func(id string) error {
+				assert.Equal(t, modelID, id)
+				deleted = true
+				return nil
+			},
+		},
+		&stubKBRepoForModelDelete{},
+		&stubAgentRepoForModelDelete{},
+		nil, nil, nil,
+	)
+
+	require.NoError(t, svc.DeleteModel(ctx, modelID))
+	assert.True(t, deleted)
+}
+
+func TestFormatModelInUseMessage(t *testing.T) {
+	t.Parallel()
+	assert.Equal(t,
+		"model is used by 1 knowledge base(s); reconfigure or remove those references before deleting",
+		formatModelInUseMessage(1, 0),
+	)
+	assert.Equal(t,
+		"model is used by 2 agent(s); reconfigure or remove those references before deleting",
+		formatModelInUseMessage(0, 2),
+	)
+	assert.Equal(t,
+		"model is used by 1 knowledge base(s) and 1 agent(s); reconfigure or remove those references before deleting",
+		formatModelInUseMessage(1, 1),
+	)
+}

internal/application/service/vectorstore_test.go

@@ -921,6 +921,9 @@ func (r *realKBRepo) CountByVectorStoreID(ctx context.Context, db *gorm.DB, tena
 		Count(&count).Error
 	return count, err
 }
+func (r *realKBRepo) CountByModelID(_ context.Context, _ uint64, _ string) (int64, error) {
+	return 0, nil
+}
 
 // The remaining methods are not called by the tested code paths; declare them
 // so realKBRepo satisfies interfaces.KnowledgeBaseRepository.

internal/handler/model.go

@@ -346,6 +346,10 @@ func (h *ModelHandler) DeleteModel(c *gin.Context) {
 			c.Error(errors.NewNotFoundError("Model not found"))
 			return
 		}
+		if appErr, ok := errors.IsAppError(err); ok {
+			c.Error(appErr)
+			return
+		}
 		logger.ErrorWithFields(ctx, err, nil)
 		c.Error(errors.NewInternalServerError(err.Error()))
 		return

internal/types/interfaces/custom_agent.go

@@ -126,4 +126,8 @@ type CustomAgentRepository interface {
 	// Returns:
 	//   - Possible errors such as record not existing, database errors, etc.
 	DeleteAgent(ctx context.Context, id string, tenantID uint64) error
+
+	// CountByModelID counts active agents in the tenant whose config references
+	// the given model ID (chat, rerank, VLM, ASR, query-understand, etc.).
+	CountByModelID(ctx context.Context, tenantID uint64, modelID string) (int64, error)
 }

internal/types/interfaces/knowledgebase.go

@@ -216,6 +216,10 @@ type KnowledgeBaseRepository interface {
 	// scope on KnowledgeBase; implementations MUST NOT add an explicit
 	// `deleted_at IS NULL` predicate (avoids divergence with the auto-scope).
 	CountByVectorStoreID(ctx context.Context, db *gorm.DB, tenantID uint64, storeID string) (int64, error)
+
+	// CountByModelID counts active KBs in the tenant that reference the given
+	// model ID in any model-binding field (embedding, summary, VLM, ASR, etc.).
+	CountByModelID(ctx context.Context, tenantID uint64, modelID string) (int64, error)
 	// SetUserKBPin inserts or removes a row in user_kb_pins for the given
 	// (tenant, user, kb) triple. Returns the resulting pinned_at (nil when
 	// pinned=false) and an error. The tenant_id is captured to support