From 7c281743ff9bd98a3dc52780d3b372e08ef40572 Mon Sep 17 00:00:00 2001 From: qian Date: Fri, 29 May 2026 00:57:53 +0800 Subject: [PATCH] =?UTF-8?q?fix(model):=E6=88=90=E5=91=98=E5=8F=AA=E8=83=BD?= =?UTF-8?q?=E7=9C=8B=E8=A7=81=E8=87=AA=E5=B7=B1=E8=83=BD=E8=B0=83=E7=94=A8?= =?UTF-8?q?=E7=9A=84=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/handler/relay/relay.go | 4 +- internal/handler/relay/task.go | 11 +++ internal/logic/open/open_api.go | 28 +++++- internal/logic/relay/provider.go | 103 +++++++++++++++++++- internal/logic/tenant/member_model_scope.go | 70 +++++-------- internal/logic/tenant/model.go | 35 +++++++ relay/common/provider.go | 9 +- relay/handler/chat_handler.go | 5 +- relay/handler/gemini_handler.go | 4 +- relay/handler/realtime_handler.go | 7 ++ 10 files changed, 217 insertions(+), 59 deletions(-) diff --git a/internal/handler/relay/relay.go b/internal/handler/relay/relay.go index 8b052d2..f77c965 100644 --- a/internal/handler/relay/relay.go +++ b/internal/handler/relay/relay.go @@ -317,7 +317,7 @@ func HandleGeminiGenerateContent(r *ghttp.Request) { func HandleModels(r *ghttp.Request) { rc := buildRelayContext(r) // 获取 TenantID - resp, err := handler.HandleModels(r.Context(), rc.TenantID, rc.ApiKeyID, dataProvider) + resp, err := handler.HandleModels(r.Context(), rc.TenantID, rc.ApiKeyID, rc.UserID, dataProvider) if err != nil { r.Response.WriteJsonExit(g.Map{ "error": g.Map{"type": "internal_error", "message": err.Error()}, @@ -371,7 +371,7 @@ func HandleModelDetail(r *ghttp.Request) { func HandleGeminiModels(r *ghttp.Request) { rc := buildRelayContext(r) - resp, err := handler.HandleGeminiModels(r.Context(), rc.TenantID, rc.ApiKeyID, dataProvider) + resp, err := handler.HandleGeminiModels(r.Context(), rc.TenantID, rc.ApiKeyID, rc.UserID, dataProvider) if err != nil { handler.WriteGeminiRelayError(r.Response.Writer, err) return diff --git a/internal/handler/relay/task.go b/internal/handler/relay/task.go index c585a12..a5c7425 100644 --- a/internal/handler/relay/task.go +++ b/internal/handler/relay/task.go @@ -57,6 +57,10 @@ func HandleTaskSubmit(r *ghttp.Request) { statusCode = 403 errType = "permission_denied" errMsg = "当前租户未启用该模型，请联系管理员" + } else if err == relay_common.ErrMemberModelNotAllowed { + statusCode = 403 + errType = "permission_denied" + errMsg = "当前成员无权使用该模型" } else if err == relay_common.ErrChannelUnavailable { statusCode = 503 errType = "server_error" @@ -172,6 +176,13 @@ func selectTaskChannel(r *ghttp.Request, body []byte) (*relay_common.ChannelMeta return nil, err } + // 检查成员模型范围 + if allowed, err := relayDataProvider.CheckMemberModelAccess(ctx, tenantID, userID, req.Model); err != nil { + return nil, err + } else if !allowed { + return nil, relay_common.ErrMemberModelNotAllowed + } + return &relay_common.ChannelMeta{ ChannelID: selection.ChannelID, ChannelType: selection.ChannelType, diff --git a/internal/logic/open/open_api.go b/internal/logic/open/open_api.go index a6525c6..e562032 100644 --- a/internal/logic/open/open_api.go +++ b/internal/logic/open/open_api.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strings" + "time" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/os/gtime" @@ -20,6 +21,9 @@ import ( type sOpen struct{} +// openMemberModelCache 成员模型范围缓存（与 relay provider 中的缓存 key 相同） +var openMemberModelCache = common.NewCache("member_model", 60*time.Second) + func New() *sOpen { return &sOpen{} } @@ -283,20 +287,36 @@ func (s *sOpen) OpenMemberModelsUpdate(ctx context.Context, req *v1.OpenMemberMo } // Insert new scopes - for _, modelID := range req.ModelIDs { - if modelID == 0 { - continue + if len(req.ModelIDs) > 0 { + for _, modelID := range req.ModelIDs { + if modelID == 0 { + continue + } + _, err = dao.TntMemberModelScopes.Ctx(ctx).Data(do.TntMemberModelScopes{ + TenantId: tenantID, + UserId: req.Id, + ModelId: modelID, + }).Insert() + if err != nil { + return nil, err + } } + } else { + // 空列表表示禁止所有模型，插入哨兵记录（model_id = -1） _, err = dao.TntMemberModelScopes.Ctx(ctx).Data(do.TntMemberModelScopes{ TenantId: tenantID, UserId: req.Id, - ModelId: modelID, + ModelId: -1, }).Insert() if err != nil { return nil, err } } + // 清除该成员的模型范围缓存，使下次请求时重新从数据库读取 + cacheKey := fmt.Sprintf("%d:%d", tenantID, req.Id) + openMemberModelCache.Delete(ctx, cacheKey) + return nil, nil } diff --git a/internal/logic/relay/provider.go b/internal/logic/relay/provider.go index 4a69af1..53bd71b 100644 --- a/internal/logic/relay/provider.go +++ b/internal/logic/relay/provider.go @@ -716,7 +716,7 @@ func (p *DataProviderImpl) ResetConsecutiveFailure(ctx context.Context, channelI // GetAvailableModels 实现 DataProvider.GetAvailableModels // 如果 tenantID > 0，返回该租户有权使用的模型列表 // 如果 apiKeyID > 0，进一步按 API Key 的模型范围过滤 -func (p *DataProviderImpl) GetAvailableModels(ctx context.Context, tenantID int64, apiKeyID int64) ([]common.ModelInfo, error) { +func (p *DataProviderImpl) GetAvailableModels(ctx context.Context, tenantID, apiKeyID, userID int64) ([]common.ModelInfo, error) { type modelRow struct { ModelId string `json:"model_id"` ModelName string `json:"model_name"` @@ -821,6 +821,31 @@ func (p *DataProviderImpl) GetAvailableModels(ctx context.Context, tenantID int6 } } + // 按成员模型范围过滤 + if userID > 0 { + allowedNames, err := GetMemberAllowedModelNames(ctx, tenantID, userID) + if err != nil { + return nil, err + } + if allowedNames != nil { + if len(allowedNames) == 0 { + models = nil + } else { + allowed := make(map[string]bool, len(allowedNames)) + for _, n := range allowedNames { + allowed[n] = true + } + filtered := make([]modelRow, 0, len(models)) + for _, m := range models { + if allowed[m.ModelId] { + filtered = append(filtered, m) + } + } + models = filtered + } + } + } + result := make([]common.ModelInfo, 0, len(models)) for _, m := range models { result = append(result, common.ModelInfo{ @@ -1100,8 +1125,57 @@ func (p *DataProviderImpl) InvalidateModelCache(modelName string) { // memberModelCache 成员模型范围缓存（TTL 60s） var memberModelCache = lcommon.NewCache("member_model", 60*time.Second) +// GetMemberAllowedModelNames 获取成员允许使用的 model_id 列表。 +// 返回 nil 表示不限制（向后兼容），返回空切片表示禁止所有模型（哨兵），返回非空切片表示允许的模型集合。 +func GetMemberAllowedModelNames(ctx context.Context, tenantID, userID int64) ([]string, error) { + if userID <= 0 || tenantID <= 0 { + return nil, nil + } + + cacheKey := fmt.Sprintf("%d:%d", tenantID, userID) + + var cachedNames []string + if memberModelCache.GetJSON(ctx, cacheKey, &cachedNames) { + return cachedNames, nil + } + + type scopeRow struct { + ModelID int64 `json:"model_id"` + ModelName string `json:"model_name"` + } + var rows []scopeRow + err := g.DB().Model("tnt_member_model_scopes ms").Ctx(ctx). + LeftJoin("mdl_models m ON ms.model_id = m.id"). + Where("ms.tenant_id", tenantID). + Where("ms.user_id", userID). + Fields("ms.model_id, m.model_id as model_name"). + Scan(&rows) + if err != nil { + return nil, err + } + + if len(rows) == 0 { + return nil, nil + } + + names := make([]string, 0, len(rows)) + for _, r := range rows { + if r.ModelID == -1 { + memberModelCache.Set(ctx, cacheKey, []string{}) + return []string{}, nil + } + if r.ModelName != "" { + names = append(names, r.ModelName) + } + } + + memberModelCache.Set(ctx, cacheKey, names) + return names, nil +} + // CheckMemberModelAccess 检查成员是否有权使用指定模型。 // 无 scope 记录表示不限制（向后兼容）。 +// 哨兵 model_id=-1 表示该成员被限制使用所有模型。 func (p *DataProviderImpl) CheckMemberModelAccess(ctx context.Context, tenantID, userID int64, modelName string) (bool, error) { cacheKey := fmt.Sprintf("%d:%d", tenantID, userID) @@ -1110,6 +1184,12 @@ func (p *DataProviderImpl) CheckMemberModelAccess(ctx context.Context, tenantID, if len(cachedModelNames) == 0 { return true, nil } + // 哨兵值：表示该成员被限制使用所有模型 + for _, name := range cachedModelNames { + if name == "__none__" { + return false, nil + } + } for _, name := range cachedModelNames { if name == modelName { return true, nil @@ -1119,6 +1199,7 @@ func (p *DataProviderImpl) CheckMemberModelAccess(ctx context.Context, tenantID, } type scopeRow struct { + ModelID int64 `json:"model_id"` ModelName string `json:"model_name"` } var rows []scopeRow @@ -1126,7 +1207,7 @@ func (p *DataProviderImpl) CheckMemberModelAccess(ctx context.Context, tenantID, LeftJoin("mdl_models m ON ms.model_id = m.id"). Where("ms.tenant_id", tenantID). Where("ms.user_id", userID). - Fields("m.model_id as model_name"). + Fields("ms.model_id, m.model_id as model_name"). Scan(&rows) if err != nil { return false, err @@ -1134,7 +1215,11 @@ func (p *DataProviderImpl) CheckMemberModelAccess(ctx context.Context, tenantID, names := make([]string, 0, len(rows)) for _, r := range rows { - names = append(names, r.ModelName) + if r.ModelID == -1 { + names = append(names, "__none__") + } else if r.ModelName != "" { + names = append(names, r.ModelName) + } } memberModelCache.Set(ctx, cacheKey, names) @@ -1142,6 +1227,12 @@ func (p *DataProviderImpl) CheckMemberModelAccess(ctx context.Context, tenantID, if len(names) == 0 { return true, nil } + // 哨兵值：表示该成员被限制使用所有模型 + for _, name := range names { + if name == "__none__" { + return false, nil + } + } for _, name := range names { if name == modelName { return true, nil @@ -1150,6 +1241,12 @@ func (p *DataProviderImpl) CheckMemberModelAccess(ctx context.Context, tenantID, return false, nil } +// InvalidateMemberModelCache 实现 DataProvider.InvalidateMemberModelCache +func (p *DataProviderImpl) InvalidateMemberModelCache(ctx context.Context, tenantID, userID int64) { + cacheKey := fmt.Sprintf("%d:%d", tenantID, userID) + memberModelCache.Delete(ctx, cacheKey) +} + // apiKeyModelCache API Key 模型范围缓存（TTL 60s） var apiKeyModelCache = lcommon.NewCache("apikey_model", 60*time.Second) diff --git a/internal/logic/tenant/member_model_scope.go b/internal/logic/tenant/member_model_scope.go index fe298a1..a6e61f2 100644 --- a/internal/logic/tenant/member_model_scope.go +++ b/internal/logic/tenant/member_model_scope.go @@ -2,13 +2,16 @@ package tenant import ( "context" - "github.com/qianfree/team-api/internal/dao" + "fmt" + "time" "github.com/gogf/gf/v2/database/gdb" "github.com/gogf/gf/v2/frame/g" + "github.com/qianfree/team-api/internal/dao" do "github.com/qianfree/team-api/internal/model/do" v1 "github.com/qianfree/team-api/api/tenant/v1" + lcommon "github.com/qianfree/team-api/internal/logic/common" "github.com/qianfree/team-api/internal/middleware" ) @@ -29,11 +32,16 @@ func (s *sTenant) MemberModelScopes(ctx context.Context, req *v1.TenantMemberMod ids := make([]int64, 0, len(rows)) for _, r := range rows { - ids = append(ids, r.ModelID) + if r.ModelID > 0 { + ids = append(ids, r.ModelID) + } } return &v1.TenantMemberModelScopesRes{ModelIDs: ids}, nil } +// memberModelScopeCache 成员模型范围缓存（与 relay provider 中的缓存 key 相同） +var memberModelScopeCache = lcommon.NewCache("member_model", 60*time.Second) + // MemberModelScopesSet sets the available models for a member (full replace). func (s *sTenant) MemberModelScopesSet(ctx context.Context, req *v1.TenantMemberModelScopesSetReq) (*v1.TenantMemberModelScopesSetRes, error) { tenantID := middleware.GetTenantID(ctx) @@ -48,8 +56,8 @@ func (s *sTenant) MemberModelScopesSet(ctx context.Context, req *v1.TenantMember return err } - // Insert new scopes if len(req.ModelIDs) > 0 { + // Insert new scopes data := make([]do.TntMemberModelScopes, len(req.ModelIDs)) for i, mID := range req.ModelIDs { data[i] = do.TntMemberModelScopes{ @@ -62,6 +70,16 @@ func (s *sTenant) MemberModelScopesSet(ctx context.Context, req *v1.TenantMember if err != nil { return err } + } else { + // 空列表表示禁止所有模型，插入哨兵记录（model_id = -1） + _, err = tx.Model("tnt_member_model_scopes").Ctx(ctx).Data(do.TntMemberModelScopes{ + TenantId: tenantID, + UserId: req.Id, + ModelId: -1, + }).Insert() + if err != nil { + return err + } } return nil @@ -70,47 +88,9 @@ func (s *sTenant) MemberModelScopesSet(ctx context.Context, req *v1.TenantMember return nil, err } + // 清除该成员的模型范围缓存，使下次请求时重新从数据库读取 + cacheKey := fmt.Sprintf("%d:%d", tenantID, req.Id) + memberModelScopeCache.Delete(ctx, cacheKey) + return &v1.TenantMemberModelScopesSetRes{}, nil } - -// revokeMemberModels removes specified models from a member and cascades to their API keys. -func revokeMemberModels(ctx context.Context, tenantID, userID int64, modelIDs []int64) error { - if len(modelIDs) == 0 { - return nil - } - - return g.DB().Transaction(ctx, func(ctx context.Context, tx gdb.TX) error { - // Remove from member model scopes - _, err := tx.Model("tnt_member_model_scopes").Ctx(ctx). - Where("tenant_id", tenantID). - Where("user_id", userID). - WhereIn("model_id", modelIDs). - Delete() - if err != nil { - return err - } - - // Cascade: remove model bindings from all keys owned by this user - keyIDs, err := tx.Model("api_keys").Ctx(ctx). - Where("tenant_id", tenantID). - Where("user_id", userID). - Where("status", "active"). - Fields("id"). - Array() - if err != nil { - return err - } - - if len(keyIDs) > 0 { - _, err = tx.Model("api_key_model_scopes").Ctx(ctx). - WhereIn("api_key_id", keyIDs). - WhereIn("model_id", modelIDs). - Delete() - if err != nil { - return err - } - } - - return nil - }) -} diff --git a/internal/logic/tenant/model.go b/internal/logic/tenant/model.go index 1b1200f..30f17ec 100644 --- a/internal/logic/tenant/model.go +++ b/internal/logic/tenant/model.go @@ -258,6 +258,41 @@ func (s *sTenant) ListAvailableModels(ctx context.Context, req *v1.TenantAvailab list = append(list, item) } + // 按成员模型范围过滤 + userID := middleware.GetUserID(ctx) + if userID > 0 { + var memberScopes []struct { + ModelID int64 `json:"model_id"` + ModelName string `json:"model_name"` + } + _ = g.DB().Model("tnt_member_model_scopes ms").Ctx(ctx). + LeftJoin("mdl_models m ON ms.model_id = m.id"). + Where("ms.tenant_id", tenantID). + Where("ms.user_id", userID). + Fields("ms.model_id, m.model_id as model_name"). + Scan(&memberScopes) + + if len(memberScopes) > 0 { + allowed := make(map[string]bool, len(memberScopes)) + for _, s := range memberScopes { + if s.ModelID == -1 { + // 哨兵：禁止所有模型 + return &v1.TenantAvailableModelsRes{List: nil}, nil + } + if s.ModelName != "" { + allowed[s.ModelName] = true + } + } + filtered := make([]v1.TenantAvailableModelItem, 0, len(list)) + for _, item := range list { + if allowed[item.ModelId] { + filtered = append(filtered, item) + } + } + list = filtered + } + } + return &v1.TenantAvailableModelsRes{List: list}, nil } diff --git a/relay/common/provider.go b/relay/common/provider.go index bdabc66..0606938 100644 --- a/relay/common/provider.go +++ b/relay/common/provider.go @@ -18,6 +18,9 @@ var ErrStreamInterrupted = errors.New("stream interrupted by client disconnect") // ErrTenantModelNotEnabled 租户未启用该模型 var ErrTenantModelNotEnabled = errors.New("model not enabled for this tenant") +// ErrMemberModelNotAllowed 成员无权使用该模型 +var ErrMemberModelNotAllowed = errors.New("model not allowed for this member") + // DeprecationInfo 模型弃用信息 type DeprecationInfo struct { Deprecated bool @@ -63,7 +66,8 @@ type DataProvider interface { // GetAvailableModels 获取指定租户可用的模型列表。 // apiKeyID > 0 时进一步按 API Key 的模型范围过滤。 - GetAvailableModels(ctx context.Context, tenantID int64, apiKeyID int64) ([]ModelInfo, error) + // userID > 0 时进一步按成员的模型范围过滤。 + GetAvailableModels(ctx context.Context, tenantID, apiKeyID, userID int64) ([]ModelInfo, error) // GetModelDetail 获取单个模型的详细信息。 // tenantID 用于权限校验（检查租户是否有权使用该模型）。 @@ -87,6 +91,9 @@ type DataProvider interface { // CheckApiKeyModelAccess 检查 API Key 是否有权使用指定模型。 // 无 scope 记录表示不限制（向后兼容）。 CheckApiKeyModelAccess(ctx context.Context, apiKeyID int64, modelName string) (bool, error) + + // InvalidateMemberModelCache 清除指定成员的模型范围缓存。 + InvalidateMemberModelCache(ctx context.Context, tenantID, userID int64) } // ApiKeyInfo API Key 验证结果 diff --git a/relay/handler/chat_handler.go b/relay/handler/chat_handler.go index 8dac0cd..25288a5 100644 --- a/relay/handler/chat_handler.go +++ b/relay/handler/chat_handler.go @@ -41,8 +41,9 @@ func HandleResponses(ctx context.Context, body []byte, path string, headers http // HandleModels 处理 /v1/models 请求（列出可用模型） // tenantID 用于过滤该租户有权使用的模型，apiKeyID 用于进一步按 Key 的模型范围过滤 -func HandleModels(ctx context.Context, tenantID int64, apiKeyID int64, provider common.DataProvider) (*dto.ModelsResponse, error) { - models, err := provider.GetAvailableModels(ctx, tenantID, apiKeyID) +// userID 用于进一步按成员的模型范围过滤 +func HandleModels(ctx context.Context, tenantID, apiKeyID, userID int64, provider common.DataProvider) (*dto.ModelsResponse, error) { + models, err := provider.GetAvailableModels(ctx, tenantID, apiKeyID, userID) if err != nil { return nil, err } diff --git a/relay/handler/gemini_handler.go b/relay/handler/gemini_handler.go index d13ec85..87598ff 100644 --- a/relay/handler/gemini_handler.go +++ b/relay/handler/gemini_handler.go @@ -20,8 +20,8 @@ func HandleGeminiGenerateContent(ctx context.Context, body []byte, path string, } // HandleGeminiModels 处理 GET /v1beta/models 请求（Gemini 原生格式模型列表） -func HandleGeminiModels(ctx context.Context, tenantID int64, apiKeyID int64, provider common.DataProvider) (*dto.GeminiModelsResponse, error) { - models, err := provider.GetAvailableModels(ctx, tenantID, apiKeyID) +func HandleGeminiModels(ctx context.Context, tenantID, apiKeyID, userID int64, provider common.DataProvider) (*dto.GeminiModelsResponse, error) { + models, err := provider.GetAvailableModels(ctx, tenantID, apiKeyID, userID) if err != nil { return nil, err } diff --git a/relay/handler/realtime_handler.go b/relay/handler/realtime_handler.go index fb00e32..a236955 100644 --- a/relay/handler/realtime_handler.go +++ b/relay/handler/realtime_handler.go @@ -71,6 +71,13 @@ func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext, return nil, nil, fmt.Errorf("model not found: %s", modelName) } + // 3.1 检查成员模型范围 + if allowed, err := provider.CheckMemberModelAccess(ctx, rc.TenantID, rc.UserID, modelName); err != nil { + return nil, nil, err + } else if !allowed { + return nil, nil, constant.NewAuthError("model not allowed for this member") + } + // 4. 渠道选择 selection, err := provider.GetChannelForModel(ctx, rc.TenantID, rc.UserID, modelName, nil) if err != nil {