diff --git a/constant/context_key.go b/constant/context_key.go index ccb8010f9476..eabcb829d5f8 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -68,6 +68,10 @@ const ( ContextKeyLanguage ContextKey = "language" ContextKeyIsStream ContextKey = "is_stream" + // ContextKeySuppressUpstreamResponseLog prevents user-controlled upstream + // response bodies from being written to debug and relay error logs. + ContextKeySuppressUpstreamResponseLog ContextKey = "suppress_upstream_response_log" + // ContextKeyAuditLogged marks that the current request has already recorded // a manage/operation audit log inside the handler. When set, the admin-audit // fallback in authHelper (finishAdminAudit) skips its record to avoid diff --git a/controller/channel-test.go b/controller/channel-test.go index fffc59d24d5a..76d0a00265e2 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -41,6 +41,45 @@ type testResult struct { newAPIError *types.NewAPIError } +type channelTestOptions struct { + UseSSRFProtectedClient bool + SkipConsumeLog bool + SkipPricingValidation bool + GroupOverride string +} + +type channelTestResponseRecorder struct { + *httptest.ResponseRecorder + maxBytes int + exceeded bool +} + +func newChannelTestResponseRecorder(maxBytes int) *channelTestResponseRecorder { + return &channelTestResponseRecorder{ + ResponseRecorder: httptest.NewRecorder(), + maxBytes: maxBytes, + } +} + +func (recorder *channelTestResponseRecorder) Write(data []byte) (int, error) { + remaining := recorder.maxBytes - recorder.Body.Len() + if remaining > 0 { + writeBytes := len(data) + if writeBytes > remaining { + writeBytes = remaining + } + _, _ = recorder.ResponseRecorder.Write(data[:writeBytes]) + } + if len(data) > remaining { + recorder.exceeded = true + } + return len(data), nil +} + +func (recorder *channelTestResponseRecorder) WriteString(data string) (int, error) { + return recorder.Write([]byte(data)) +} + func normalizeChannelTestEndpoint(channel *model.Channel, endpointType string) string { normalized := strings.TrimSpace(endpointType) if normalized != "" { @@ -70,9 +109,16 @@ func resolveChannelTestUserID(c *gin.Context) (int, error) { } func testChannel(ctx context.Context, channel *model.Channel, testUserID int, testModel string, endpointType string, isStream bool) testResult { + return testChannelWithOptions(ctx, channel, testUserID, testModel, endpointType, isStream, channelTestOptions{}) +} + +func testChannelWithOptions(ctx context.Context, channel *model.Channel, testUserID int, testModel string, endpointType string, isStream bool, options channelTestOptions) testResult { if ctx == nil { ctx = context.Background() } + if options.UseSSRFProtectedClient { + ctx = context.WithValue(ctx, constant.ContextKeySuppressUpstreamResponseLog, true) + } tik := time.Now() var unsupportedTestChannelTypes = []int{ constant.ChannelTypeMidjourney, @@ -89,7 +135,7 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te localErr: fmt.Errorf("%s channel test is not supported", channelTypeName), } } - w := httptest.NewRecorder() + w := newChannelTestResponseRecorder(int(service.StrictSSRFProtectedResponseBodyLimitBytes)) c, _ := gin.CreateTestContext(w) testModel = strings.TrimSpace(testModel) @@ -165,7 +211,11 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te c.Set("channel", channel.Type) c.Set("base_url", channel.GetBaseURL()) group, _ := model.GetUserGroup(testUserID, false) - c.Set("group", group) + if strings.TrimSpace(options.GroupOverride) != "" { + group = strings.TrimSpace(options.GroupOverride) + common.SetContextKey(c, constant.ContextKeyUserGroup, group) + } + common.SetContextKey(c, constant.ContextKeyUsingGroup, group) newAPIError := middleware.SetupContextForSelectedChannel(c, channel, testModel) if newAPIError != nil { @@ -239,6 +289,7 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te } info.IsChannelTest = true + info.UseSSRFProtectedClient = options.UseSSRFProtectedClient info.InitChannelMeta(c) err = attachTestBillingRequestInput(info, request) @@ -286,12 +337,15 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te //logInfo.ApiKey = "" common.SysLog(fmt.Sprintf("testing channel %d with model %s , info %+v ", channel.Id, testModel, info.ToString())) - priceData, err := helper.ModelPriceHelper(c, info, 0, request.GetTokenCountMeta()) - if err != nil { - return testResult{ - context: c, - localErr: err, - newAPIError: types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest)), + priceData := hosttypes.PriceData{} + if !options.SkipPricingValidation { + priceData, err = helper.ModelPriceHelper(c, info, 0, request.GetTokenCountMeta()) + if err != nil { + return testResult{ + context: c, + localErr: err, + newAPIError: types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest)), + } } } @@ -437,7 +491,11 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te if resp != nil { httpResp = resp.(*http.Response) if httpResp.StatusCode != http.StatusOK { - err := service.RelayErrorHandler(c.Request.Context(), httpResp, true) + err := service.RelayErrorHandler(c.Request.Context(), httpResp, !options.UseSSRFProtectedClient) + logErr := error(err) + if options.UseSSRFProtectedClient { + logErr = sanitizeChannelCredentialError(err, channel.Key, channel.GetBaseURL()) + } common.SysError(fmt.Sprintf( "channel test bad response: channel_id=%d name=%s type=%d model=%s endpoint_type=%s status=%d err=%v", channel.Id, @@ -446,7 +504,7 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te testModel, endpointType, httpResp.StatusCode, - err, + logErr, )) return testResult{ context: c, @@ -463,6 +521,14 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te newAPIError: respErr, } } + if w.exceeded { + err := fmt.Errorf("channel test response exceeds %d bytes", service.StrictSSRFProtectedResponseBodyLimitBytes) + return testResult{ + context: c, + localErr: err, + newAPIError: types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError), + } + } usage, usageErr := coerceTestUsage(usageA, isStream, info.GetEstimatePromptTokens()) if usageErr != nil { return testResult{ @@ -494,20 +560,24 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te milliseconds := tok.Sub(tik).Milliseconds() consumedTime := float64(milliseconds) / 1000.0 other := buildTestLogOther(c, info, priceData, usage, tieredResult) - model.RecordConsumeLog(c, testUserID, model.RecordConsumeLogParams{ - ChannelId: channel.Id, - PromptTokens: usage.PromptTokens, - CompletionTokens: usage.CompletionTokens, - ModelName: info.OriginModelName, - TokenName: "模型测试", - Quota: quota, - Content: "模型测试", - UseTimeSeconds: int(consumedTime), - IsStream: info.IsStream, - Group: info.UsingGroup, - Other: other, - }) - common.SysLog(fmt.Sprintf("testing channel #%d, response: \n%s", channel.Id, string(respBody))) + if !options.SkipConsumeLog { + model.RecordConsumeLog(c, testUserID, model.RecordConsumeLogParams{ + ChannelId: channel.Id, + PromptTokens: usage.PromptTokens, + CompletionTokens: usage.CompletionTokens, + ModelName: info.OriginModelName, + TokenName: "模型测试", + Quota: quota, + Content: "模型测试", + UseTimeSeconds: int(consumedTime), + IsStream: info.IsStream, + Group: info.UsingGroup, + Other: other, + }) + } + if !options.UseSSRFProtectedClient { + common.SysLog(fmt.Sprintf("testing channel #%d, response: \n%s", channel.Id, string(respBody))) + } return testResult{ context: c, localErr: nil, @@ -1024,6 +1094,7 @@ func runChannelTestTask(ctx context.Context, mode string, notify bool, report fu } func selectChannelsForAutomaticTest(channels []*model.Channel, mode string) []*model.Channel { + channels = model.FilterNonContributionChannels(channels) selected := make([]*model.Channel, 0, len(channels)) for _, channel := range channels { if channel.Status == common.ChannelStatusManuallyDisabled { diff --git a/controller/channel.go b/controller/channel.go index 3a1e58328923..5aaf27c7a4c2 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -165,6 +165,10 @@ func GetAllChannels(c *gin.Context) { } } + if err := model.PopulateContributionChannelFlags(channelData); err != nil { + common.ApiError(c, err) + return + } for _, datum := range channelData { clearChannelInfo(datum) } @@ -378,6 +382,10 @@ func SearchChannels(c *gin.Context) { pagedData := channelData[startIdx:endIdx] + if err := model.PopulateContributionChannelFlags(pagedData); err != nil { + common.ApiError(c, err) + return + } for _, datum := range pagedData { clearChannelInfo(datum) } @@ -406,6 +414,10 @@ func GetChannel(c *gin.Context) { return } if channel != nil { + if err := model.PopulateContributionChannelFlags([]*model.Channel{channel}); err != nil { + common.ApiError(c, err) + return + } clearChannelInfo(channel) } c.JSON(http.StatusOK, gin.H{ diff --git a/controller/channel_authz.go b/controller/channel_authz.go index f85ffef92769..c76dd999b5c4 100644 --- a/controller/channel_authz.go +++ b/controller/channel_authz.go @@ -88,6 +88,7 @@ var channelReadOnlyFields = map[string]struct{}{ "balance": {}, "balance_updated_time": {}, "used_quota": {}, + "is_contribution": {}, } func clearChannelReadOnlyFields(channel *PatchChannel, requestData map[string]any) { diff --git a/controller/channel_contribution.go b/controller/channel_contribution.go new file mode 100644 index 000000000000..3c8798b0a0e9 --- /dev/null +++ b/controller/channel_contribution.go @@ -0,0 +1,997 @@ +package controller + +import ( + "crypto/sha256" + "errors" + "fmt" + "net/http" + "net/url" + "sort" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/operation_setting" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +const ( + channelContributionMaxModels = 100 + channelContributionTestResultTTLSeconds = int64(30 * 60) + channelContributionProbeTimeoutSeconds = 30 +) + +type channelContributionInput struct { + Name string `json:"name"` + Type int `json:"type"` + BaseURL string `json:"base_url"` + APIEndpoint string `json:"api_endpoint"` + APIKey *string `json:"api_key"` + Key *string `json:"key"` + Group string `json:"group"` + Models []string `json:"models"` + ModelMapping map[string]string `json:"model_mapping"` +} + +type channelContributionSubmitInput struct { + TestRunId int64 `json:"test_run_id"` + AgreementAccepted bool `json:"agreement_accepted"` + AgreementVersion string `json:"agreement_version"` +} + +type channelContributionAdminReviewInput struct { + TestRunId int64 `json:"test_run_id"` + Reason string `json:"reason"` +} + +type channelContributionSettingsInput struct { + Tag *string `json:"tag"` + AllowedGroups []string `json:"allowed_groups"` + AllowedChannelTypes []int `json:"allowed_channel_types"` + Priority *int64 `json:"priority"` + Weight *uint `json:"weight"` + UnavailableDeleteHours *int `json:"unavailable_delete_hours"` + HealthCheckIntervalMinutes *int `json:"health_check_interval_minutes"` + RewardBps *int `json:"reward_bps"` + AgreementVersion *string `json:"agreement_version"` + AgreementContent *string `json:"agreement_content"` +} + +type channelContributionChannelTypeOption struct { + Value int `json:"value"` + Label string `json:"label"` +} + +type channelContributionSettingsResponse struct { + operation_setting.ChannelContributionSetting + SupportedChannelTypes []channelContributionChannelTypeOption `json:"supported_channel_types"` +} + +type channelContributionRevisionResponse struct { + Id int `json:"id"` + RevisionNumber int `json:"revision_number"` + Name string `json:"name"` + Type int `json:"type"` + BaseURL string `json:"base_url"` + HasAPIKey bool `json:"has_api_key"` + Group string `json:"group"` + Models []string `json:"models"` + ModelMapping map[string]string `json:"model_mapping"` + Status model.ChannelContributionRevisionStatus `json:"status"` + PriceConfigured bool `json:"price_configured"` + UnpricedModels []string `json:"unpriced_models"` + AgreementVersion string `json:"agreement_version"` + AgreementHash string `json:"agreement_hash"` + AgreementAcceptedAt int64 `json:"agreement_accepted_at"` + SubmittedAt int64 `json:"submitted_at"` + ReviewerId int `json:"reviewer_id"` + ReviewerUsername string `json:"reviewer_username"` + ReviewedAt int64 `json:"reviewed_at"` + ReviewReason string `json:"review_reason"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` +} + +type channelContributionModelHealthResponse struct { + Id int64 `json:"id"` + ContributionId int `json:"contribution_id"` + RevisionId int `json:"revision_id"` + ChannelId int `json:"channel_id"` + Model string `json:"model"` + Healthy bool `json:"healthy"` + FailureSince int64 `json:"failure_since"` + LastCheckedAt int64 `json:"last_checked_at"` + LastSuccessAt int64 `json:"last_success_at"` + LastFailureAt int64 `json:"last_failure_at"` + LastError string `json:"last_error"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` +} + +type channelContributionResponse struct { + Id int `json:"id"` + UserId int `json:"user_id"` + Username string `json:"username"` + Status model.ChannelContributionStatus `json:"status"` + RevisionStatus model.ChannelContributionRevisionStatus `json:"revision_status"` + ChannelId *int `json:"channel_id"` + CurrentRevisionId *int `json:"current_revision_id"` + PendingRevisionId *int `json:"pending_revision_id"` + ApprovedRevisionId *int `json:"approved_revision_id"` + CurrentRevision *channelContributionRevisionResponse `json:"current_revision"` + PendingRevision *channelContributionRevisionResponse `json:"pending_revision"` + ApprovedRevision *channelContributionRevisionResponse `json:"approved_revision"` + LatestTestRun *channelContributionTestRunResponse `json:"latest_test_run"` + SubmittedAt int64 `json:"submitted_at"` + ReviewerId int `json:"reviewer_id"` + ReviewerUsername string `json:"reviewer_username"` + ReviewedAt int64 `json:"reviewed_at"` + ReviewReason string `json:"review_reason"` + UnavailableSince int64 `json:"unavailable_since"` + ModelHealth []channelContributionModelHealthResponse `json:"model_health,omitempty"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` +} + +func channelContributionAgreementHash(content string) string { + sum := sha256.Sum256([]byte(content)) + return fmt.Sprintf("%x", sum[:]) +} + +func channelContributionModels(raw string) []string { + parts := strings.Split(raw, ",") + models := make([]string, 0, len(parts)) + for _, item := range parts { + item = strings.TrimSpace(item) + if item != "" { + models = append(models, item) + } + } + return models +} + +func channelContributionMapping(raw string) map[string]string { + mapping := map[string]string{} + if strings.TrimSpace(raw) != "" { + _ = common.UnmarshalJsonStr(raw, &mapping) + } + return mapping +} + +func channelContributionPriceStatus(models []string) (bool, []string) { + unpriced := make([]string, 0) + for _, modelName := range models { + if !helper.HasModelBillingConfig(modelName) { + unpriced = append(unpriced, modelName) + } + } + return len(models) > 0 && len(unpriced) == 0, unpriced +} + +func channelContributionTypeOptions(channelTypes []int) []channelContributionChannelTypeOption { + options := make([]channelContributionChannelTypeOption, 0, len(channelTypes)) + for _, channelType := range channelTypes { + options = append(options, channelContributionChannelTypeOption{ + Value: channelType, + Label: constant.GetChannelTypeName(channelType), + }) + } + return options +} + +func buildChannelContributionSettingsResponse(setting operation_setting.ChannelContributionSetting) channelContributionSettingsResponse { + return channelContributionSettingsResponse{ + ChannelContributionSetting: setting, + SupportedChannelTypes: channelContributionTypeOptions(operation_setting.GetSupportedChannelContributionTypes()), + } +} + +func normalizeChannelContributionInput(input channelContributionInput, previous *model.ChannelContributionRevision) (*model.ChannelContributionRevision, error) { + setting := operation_setting.GetChannelContributionSetting() + revision := &model.ChannelContributionRevision{} + + revision.Name = strings.TrimSpace(input.Name) + if revision.Name == "" || len(revision.Name) > 128 { + return nil, errors.New("name must contain 1 to 128 characters") + } + + revision.Type = input.Type + if revision.Type == 0 && previous != nil { + revision.Type = previous.Type + } + if revision.Type == 0 { + revision.Type = constant.ChannelTypeOpenAI + } + if err := model.ValidateContributionChannelType(revision.Type); err != nil { + return nil, err + } + if !setting.IsChannelTypeAllowed(revision.Type) { + return nil, errors.New("channel type is not enabled for contribution") + } + + revision.BaseURL = strings.TrimSpace(input.BaseURL) + if revision.BaseURL == "" { + revision.BaseURL = strings.TrimSpace(input.APIEndpoint) + } + revision.BaseURL = strings.TrimRight(revision.BaseURL, "/") + parsedURL, err := url.ParseRequestURI(revision.BaseURL) + if err != nil || parsedURL.Host == "" || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") || parsedURL.User != nil || parsedURL.RawQuery != "" || parsedURL.Fragment != "" { + return nil, errors.New("base_url must be an absolute HTTP or HTTPS URL without credentials, query, or fragment") + } + if err := service.ValidateStrictSSRFProtectedFetchURL(revision.BaseURL); err != nil { + return nil, fmt.Errorf("base_url is not allowed: %w", err) + } + + if input.APIKey != nil && input.Key != nil && strings.TrimSpace(*input.APIKey) != strings.TrimSpace(*input.Key) { + return nil, errors.New("api_key and key must match when both are provided") + } + providedKey := input.APIKey + if providedKey == nil { + providedKey = input.Key + } + if providedKey != nil && strings.TrimSpace(*providedKey) != "" { + revision.Key = strings.TrimSpace(*providedKey) + } else if previous != nil { + revision.Key = previous.Key + } + if revision.Key == "" { + return nil, errors.New("api_key is required") + } + if len(revision.Key) > 16_384 || strings.ContainsAny(revision.Key, "\r\n") { + return nil, errors.New("api_key must be a single key no longer than 16384 characters") + } + + revision.Group = strings.TrimSpace(input.Group) + if revision.Group == "" && previous != nil { + revision.Group = previous.Group + } + if revision.Group == "" && len(setting.AllowedGroups) > 0 { + revision.Group = strings.TrimSpace(setting.AllowedGroups[0]) + } + if !setting.IsGroupAllowed(revision.Group) { + return nil, errors.New("group is not enabled for contribution") + } + + seenModels := make(map[string]struct{}, len(input.Models)) + models := make([]string, 0, len(input.Models)) + for _, rawModel := range input.Models { + modelName := strings.TrimSpace(rawModel) + if modelName == "" || len(modelName) > 255 || strings.Contains(modelName, ",") { + return nil, errors.New("models contains an invalid model name") + } + if _, exists := seenModels[modelName]; exists { + continue + } + seenModels[modelName] = struct{}{} + models = append(models, modelName) + } + if len(models) > channelContributionMaxModels { + return nil, fmt.Errorf("at most %d models may be contributed", channelContributionMaxModels) + } + + mapping := make(map[string]string, len(input.ModelMapping)) + for modelName, upstreamName := range input.ModelMapping { + modelName = strings.TrimSpace(modelName) + upstreamName = strings.TrimSpace(upstreamName) + if _, exists := seenModels[modelName]; !exists { + return nil, fmt.Errorf("model mapping source %q is not in models", modelName) + } + if upstreamName == "" || len(upstreamName) > 255 { + return nil, fmt.Errorf("model mapping target for %q is invalid", modelName) + } + mapping[modelName] = upstreamName + } + mappingJSON, err := common.Marshal(mapping) + if err != nil { + return nil, err + } + revision.Models = strings.Join(models, ",") + revision.ModelMapping = string(mappingJSON) + + if len(models) > 0 { + revisionForProbe := *revision + if _, err := resolveChannelContributionProbeSpecs(&revisionForProbe); err != nil { + return nil, err + } + } + + revision.ConfigHash, err = model.ComputeChannelContributionConfigHash(revision) + if err != nil { + return nil, err + } + return revision, nil +} + +func buildChannelContributionRevisionResponse(revision *model.ChannelContributionRevision) *channelContributionRevisionResponse { + if revision == nil { + return nil + } + models := channelContributionModels(revision.Models) + priceConfigured, unpriced := channelContributionPriceStatus(models) + return &channelContributionRevisionResponse{ + Id: revision.Id, + RevisionNumber: revision.RevisionNumber, + Name: revision.Name, + Type: revision.Type, + BaseURL: revision.BaseURL, + HasAPIKey: strings.TrimSpace(revision.Key) != "", + Group: revision.Group, + Models: models, + ModelMapping: channelContributionMapping(revision.ModelMapping), + Status: revision.Status, + PriceConfigured: priceConfigured, + UnpricedModels: unpriced, + AgreementVersion: revision.AgreementVersion, + AgreementHash: revision.AgreementHash, + AgreementAcceptedAt: revision.AgreementAcceptedAt, + SubmittedAt: revision.SubmittedAt, + ReviewerId: revision.ReviewerId, + ReviewerUsername: revision.ReviewerUsername, + ReviewedAt: revision.ReviewedAt, + ReviewReason: revision.ReviewReason, + CreatedAt: revision.CreatedAt, + UpdatedAt: revision.UpdatedAt, + } +} + +func loadChannelContributionRevision(id *int) (*model.ChannelContributionRevision, error) { + if id == nil || *id <= 0 { + return nil, nil + } + revision, err := model.GetChannelContributionRevisionById(*id) + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + return revision, err +} + +func buildChannelContributionResponse(contribution *model.ChannelContribution, includeTestResults bool) (*channelContributionResponse, error) { + current, err := loadChannelContributionRevision(contribution.CurrentRevisionId) + if err != nil { + return nil, err + } + pending, err := loadChannelContributionRevision(contribution.PendingRevisionId) + if err != nil { + return nil, err + } + approved, err := loadChannelContributionRevision(contribution.ApprovedRevisionId) + if err != nil { + return nil, err + } + + var latestRunResponse *channelContributionTestRunResponse + if current != nil { + latestRun, runErr := model.GetLatestChannelContributionTestRun(current.Id, current.ConfigHash) + if runErr == nil { + latestRunResponse, err = buildChannelContributionTestRunResponse(latestRun, includeTestResults) + if err != nil { + return nil, err + } + } else if !errors.Is(runErr, gorm.ErrRecordNotFound) { + return nil, runErr + } + } + + var healthResponse []channelContributionModelHealthResponse + if includeTestResults { + healthRows, healthErr := model.GetChannelContributionModelHealth(contribution.Id) + if healthErr != nil { + return nil, healthErr + } + healthResponse = make([]channelContributionModelHealthResponse, 0, len(healthRows)) + for _, health := range healthRows { + healthResponse = append(healthResponse, channelContributionModelHealthResponse{ + Id: health.Id, + ContributionId: health.ContributionId, + RevisionId: health.RevisionId, + ChannelId: health.ChannelId, + Model: health.Model, + Healthy: health.Healthy, + FailureSince: health.FailureSince, + LastCheckedAt: health.LastCheckedAt, + LastSuccessAt: health.LastSuccessAt, + LastFailureAt: health.LastFailureAt, + LastError: health.LastError, + CreatedAt: health.CreatedAt, + UpdatedAt: health.UpdatedAt, + }) + } + } + + revisionStatus := model.ChannelContributionRevisionStatus("") + if pending != nil { + revisionStatus = pending.Status + } else if current != nil { + revisionStatus = current.Status + } + return &channelContributionResponse{ + Id: contribution.Id, + UserId: contribution.UserId, + Username: contribution.Username, + Status: contribution.Status, + RevisionStatus: revisionStatus, + ChannelId: contribution.ChannelId, + CurrentRevisionId: contribution.CurrentRevisionId, + PendingRevisionId: contribution.PendingRevisionId, + ApprovedRevisionId: contribution.ApprovedRevisionId, + CurrentRevision: buildChannelContributionRevisionResponse(current), + PendingRevision: buildChannelContributionRevisionResponse(pending), + ApprovedRevision: buildChannelContributionRevisionResponse(approved), + LatestTestRun: latestRunResponse, + SubmittedAt: contribution.SubmittedAt, + ReviewerId: contribution.ReviewerId, + ReviewerUsername: contribution.ReviewerUsername, + ReviewedAt: contribution.ReviewedAt, + ReviewReason: contribution.ReviewReason, + UnavailableSince: contribution.UnavailableSince, + ModelHealth: healthResponse, + CreatedAt: contribution.CreatedAt, + UpdatedAt: contribution.UpdatedAt, + }, nil +} + +func channelContributionId(c *gin.Context) (int, error) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil || id <= 0 { + return 0, errors.New("invalid contribution id") + } + return id, nil +} + +func GetChannelContributionConfig(c *gin.Context) { + setting := operation_setting.GetChannelContributionSetting() + types := channelContributionTypeOptions(setting.AllowedChannelTypes) + common.ApiSuccess(c, gin.H{ + "enabled": true, + "allowed_groups": setting.AllowedGroups, + "allowed_channel_types": types, + "max_models": channelContributionMaxModels, + "test_result_ttl_seconds": channelContributionTestResultTTLSeconds, + "probe_timeout_seconds": channelContributionProbeTimeoutSeconds, + "unavailable_delete_hours": setting.UnavailableDeleteHours, + "health_check_interval_minutes": setting.HealthCheckIntervalMinutes, + "reward_bps": setting.RewardBps, + "agreement_version": setting.AgreementVersion, + "agreement_content": setting.AgreementContent, + "agreement_hash": channelContributionAgreementHash(setting.AgreementContent), + }) +} + +func ListUserChannelContributions(c *gin.Context) { + pageInfo := common.GetPageQuery(c) + contributions, total, err := model.ListUserChannelContributions(c.GetInt("id"), pageInfo.GetStartIdx(), pageInfo.GetPageSize()) + if err != nil { + common.ApiError(c, err) + return + } + items := make([]*channelContributionResponse, 0, len(contributions)) + for _, contribution := range contributions { + item, err := buildChannelContributionResponse(contribution, false) + if err != nil { + common.ApiError(c, err) + return + } + items = append(items, item) + } + pageInfo.SetTotal(int(total)) + pageInfo.SetItems(items) + common.ApiSuccess(c, pageInfo) +} + +func CreateChannelContribution(c *gin.Context) { + var input channelContributionInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid request body"}) + return + } + revision, err := normalizeChannelContributionInput(input, nil) + if err != nil { + common.ApiError(c, err) + return + } + contribution := &model.ChannelContribution{ + UserId: c.GetInt("id"), + Username: c.GetString("username"), + Status: model.ChannelContributionStatusDraft, + } + if err := model.CreateChannelContributionWithRevision(contribution, revision); err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionResponse(contribution, false) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func GetUserChannelContribution(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + contribution, err := model.GetUserChannelContributionById(id, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionResponse(contribution, true) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func UpdateUserChannelContribution(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + contribution, err := model.GetUserChannelContributionById(id, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + current, err := loadChannelContributionRevision(contribution.CurrentRevisionId) + if err != nil || current == nil { + if err == nil { + err = errors.New("current revision is missing") + } + common.ApiError(c, err) + return + } + var input channelContributionInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid request body"}) + return + } + revision, err := normalizeChannelContributionInput(input, current) + if err != nil { + common.ApiError(c, err) + return + } + if err := model.CreateChannelContributionRevision(id, c.GetInt("id"), revision); err != nil { + common.ApiError(c, err) + return + } + contribution, err = model.GetUserChannelContributionById(id, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionResponse(contribution, false) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func FetchChannelContributionModels(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + contribution, err := model.GetUserChannelContributionById(id, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + revision, err := loadChannelContributionRevision(contribution.CurrentRevisionId) + if err != nil || revision == nil { + if err == nil { + err = errors.New("current revision is missing") + } + common.ApiError(c, err) + return + } + baseURL := revision.BaseURL + mapping := revision.ModelMapping + channel := &model.Channel{ + Type: revision.Type, + Key: revision.Key, + Name: revision.Name, + BaseURL: &baseURL, + Group: revision.Group, + ModelMapping: &mapping, + } + strictClient := *service.GetStrictSSRFProtectedHTTPClient() + strictClient.Timeout = channelContributionProbeTimeoutSeconds * time.Second + models, err := fetchChannelUpstreamModelIDsWithOptions(channel, fetchChannelModelsOptions{ + UseSSRFProtectedClient: true, + HTTPClient: &strictClient, + }) + if err != nil { + common.ApiError(c, sanitizeChannelCredentialError(err, revision.Key, revision.BaseURL)) + return + } + normalized := make([]string, 0, len(models)) + seen := make(map[string]struct{}, len(models)) + for _, modelName := range models { + modelName = strings.TrimSpace(modelName) + if modelName == "" || len(modelName) > 255 || strings.Contains(modelName, ",") { + continue + } + if _, exists := seen[modelName]; exists { + continue + } + seen[modelName] = struct{}{} + normalized = append(normalized, modelName) + } + sort.Strings(normalized) + if len(normalized) > channelContributionMaxModels { + normalized = normalized[:channelContributionMaxModels] + } + common.ApiSuccess(c, gin.H{"models": normalized}) +} + +func SubmitUserChannelContribution(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + var input channelContributionSubmitInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid request body"}) + return + } + if !input.AgreementAccepted { + common.ApiErrorMsg(c, "channel contribution agreement must be accepted") + return + } + setting := operation_setting.GetChannelContributionSetting() + if input.AgreementVersion != setting.AgreementVersion { + common.ApiErrorMsg(c, "channel contribution agreement version has changed") + return + } + contribution, err := model.GetUserChannelContributionById(id, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + revision, err := loadChannelContributionRevision(contribution.CurrentRevisionId) + if err != nil || revision == nil { + if err == nil { + err = errors.New("current revision is missing") + } + common.ApiError(c, err) + return + } + computedConfigHash, err := model.ComputeChannelContributionConfigHash(revision) + if err != nil { + common.ApiError(c, err) + return + } + if computedConfigHash != revision.ConfigHash { + common.ApiErrorMsg(c, "channel contribution configuration changed; save and test it again") + return + } + if err := validateChannelContributionSubmissionRun(input.TestRunId, contribution, revision, model.ChannelContributionTestActorUser); err != nil { + common.ApiError(c, err) + return + } + priceReady, unpriced := channelContributionPriceStatus(channelContributionModels(revision.Models)) + if !priceReady { + common.ApiError(c, fmt.Errorf("models without configured price: %s", strings.Join(unpriced, ", "))) + return + } + now := common.GetTimestamp() + agreementContent := setting.AgreementContent + if err := model.SubmitChannelContribution( + id, + c.GetInt("id"), + revision.Id, + revision.ConfigHash, + setting.AgreementVersion, + agreementContent, + channelContributionAgreementHash(agreementContent), + now, + ); err != nil { + common.ApiError(c, err) + return + } + contribution, err = model.GetUserChannelContributionById(id, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionResponse(contribution, false) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func WithdrawUserChannelContribution(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + if err := model.WithdrawChannelContribution(id, c.GetInt("id")); err != nil { + common.ApiError(c, err) + return + } + contribution, err := model.GetUserChannelContributionById(id, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionResponse(contribution, false) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func ListAdminChannelContributions(c *gin.Context) { + pageInfo := common.GetPageQuery(c) + status := model.ChannelContributionStatus(strings.TrimSpace(c.Query("status"))) + if status != "" && !model.IsValidChannelContributionStatus(status) { + common.ApiErrorMsg(c, "invalid contribution status") + return + } + contributions, total, err := model.ListChannelContributions(status, pageInfo.GetStartIdx(), pageInfo.GetPageSize()) + if err != nil { + common.ApiError(c, err) + return + } + items := make([]*channelContributionResponse, 0, len(contributions)) + for _, contribution := range contributions { + item, err := buildChannelContributionResponse(contribution, false) + if err != nil { + common.ApiError(c, err) + return + } + items = append(items, item) + } + pageInfo.SetTotal(int(total)) + pageInfo.SetItems(items) + common.ApiSuccess(c, pageInfo) +} + +func GetAdminChannelContribution(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + contribution, err := model.GetChannelContributionById(id) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionResponse(contribution, true) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func ApproveAdminChannelContribution(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + var input channelContributionAdminReviewInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid request body"}) + return + } + contribution, err := model.GetChannelContributionById(id) + if err != nil { + common.ApiError(c, err) + return + } + revision, err := loadChannelContributionRevision(contribution.PendingRevisionId) + if err != nil || revision == nil { + if err == nil { + err = errors.New("pending revision is missing") + } + common.ApiError(c, err) + return + } + computedConfigHash, err := model.ComputeChannelContributionConfigHash(revision) + if err != nil { + common.ApiError(c, err) + return + } + if computedConfigHash != revision.ConfigHash { + common.ApiErrorMsg(c, "channel contribution configuration changed; submit and test it again") + return + } + if err := validateChannelContributionSubmissionRun(input.TestRunId, contribution, revision, model.ChannelContributionTestActorAdmin); err != nil { + common.ApiError(c, err) + return + } + priceReady, unpriced := channelContributionPriceStatus(channelContributionModels(revision.Models)) + if !priceReady { + common.ApiError(c, fmt.Errorf("models without configured price: %s", strings.Join(unpriced, ", "))) + return + } + setting := operation_setting.GetChannelContributionSetting() + approved, _, err := model.ApproveChannelContribution(id, revision.Id, model.ChannelContributionApproval{ + ReviewerId: c.GetInt("id"), + ReviewerUsername: c.GetString("username"), + Tag: setting.Tag, + Priority: setting.Priority, + Weight: setting.Weight, + }) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionResponse(approved, false) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func RejectAdminChannelContribution(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + var input channelContributionAdminReviewInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid request body"}) + return + } + reason := strings.TrimSpace(input.Reason) + if reason == "" || len(reason) > 500 { + common.ApiErrorMsg(c, "reason must contain 1 to 500 characters") + return + } + contribution, err := model.GetChannelContributionById(id) + if err != nil { + common.ApiError(c, err) + return + } + if contribution.PendingRevisionId == nil { + common.ApiErrorMsg(c, "pending revision is missing") + return + } + if err := model.RejectChannelContribution(id, *contribution.PendingRevisionId, c.GetInt("id"), c.GetString("username"), reason); err != nil { + common.ApiError(c, err) + return + } + contribution, err = model.GetChannelContributionById(id) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionResponse(contribution, false) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func DeleteAdminChannelContribution(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + reason := strings.TrimSpace(c.Query("reason")) + if len(reason) > 500 { + common.ApiErrorMsg(c, "reason must not exceed 500 characters") + return + } + if err := model.DeleteChannelContribution(id, c.GetInt("id"), c.GetString("username"), reason); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, nil) +} + +func GetAdminChannelContributionSettings(c *gin.Context) { + setting := *operation_setting.GetChannelContributionSetting() + common.ApiSuccess(c, buildChannelContributionSettingsResponse(setting)) +} + +func UpdateAdminChannelContributionSettings(c *gin.Context) { + var input channelContributionSettingsInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid request body"}) + return + } + current := *operation_setting.GetChannelContributionSetting() + updated := current + if input.Tag != nil { + updated.Tag = strings.TrimSpace(*input.Tag) + } + if input.AllowedGroups != nil { + updated.AllowedGroups = input.AllowedGroups + } + if input.AllowedChannelTypes != nil { + updated.AllowedChannelTypes = input.AllowedChannelTypes + } + if input.Priority != nil { + updated.Priority = *input.Priority + } + if input.Weight != nil { + updated.Weight = *input.Weight + } + if input.UnavailableDeleteHours != nil { + updated.UnavailableDeleteHours = *input.UnavailableDeleteHours + } + if input.HealthCheckIntervalMinutes != nil { + updated.HealthCheckIntervalMinutes = *input.HealthCheckIntervalMinutes + } + if input.RewardBps != nil { + updated.RewardBps = *input.RewardBps + } + if input.AgreementVersion != nil { + updated.AgreementVersion = strings.TrimSpace(*input.AgreementVersion) + } + if input.AgreementContent != nil { + updated.AgreementContent = *input.AgreementContent + } + versionChanged := updated.AgreementVersion != current.AgreementVersion + contentChanged := updated.AgreementContent != current.AgreementContent + if versionChanged != contentChanged { + common.ApiErrorMsg(c, "agreement_version and agreement_content must change together") + return + } + + groupsJSON, err := common.Marshal(updated.AllowedGroups) + if err != nil { + common.ApiError(c, err) + return + } + typesJSON, err := common.Marshal(updated.AllowedChannelTypes) + if err != nil { + common.ApiError(c, err) + return + } + values := map[string]string{ + operation_setting.ChannelContributionSettingPrefix + "tag": updated.Tag, + operation_setting.ChannelContributionSettingPrefix + "allowed_groups": string(groupsJSON), + operation_setting.ChannelContributionSettingPrefix + "allowed_channel_types": string(typesJSON), + operation_setting.ChannelContributionSettingPrefix + "priority": strconv.FormatInt(updated.Priority, 10), + operation_setting.ChannelContributionSettingPrefix + "weight": strconv.FormatUint(uint64(updated.Weight), 10), + operation_setting.ChannelContributionSettingPrefix + "unavailable_delete_hours": strconv.Itoa(updated.UnavailableDeleteHours), + operation_setting.ChannelContributionSettingPrefix + "health_check_interval_minutes": strconv.Itoa(updated.HealthCheckIntervalMinutes), + operation_setting.ChannelContributionSettingPrefix + "reward_bps": strconv.Itoa(updated.RewardBps), + operation_setting.ChannelContributionSettingPrefix + "agreement_version": updated.AgreementVersion, + operation_setting.ChannelContributionSettingPrefix + "agreement_content": updated.AgreementContent, + } + keys := make([]string, 0, len(values)) + for key := range values { + keys = append(keys, key) + } + sort.Strings(keys) + for _, key := range keys { + if err := operation_setting.ValidateChannelContributionOption(key, values[key]); err != nil { + common.ApiError(c, err) + return + } + } + if err := model.UpdateOptionsBulk(values); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, buildChannelContributionSettingsResponse(*operation_setting.GetChannelContributionSetting())) +} diff --git a/controller/channel_contribution_error.go b/controller/channel_contribution_error.go new file mode 100644 index 000000000000..f8c85386fc11 --- /dev/null +++ b/controller/channel_contribution_error.go @@ -0,0 +1,17 @@ +package controller + +import "unicode/utf8" + +func truncateChannelContributionError(message string, maxBytes int) string { + if maxBytes <= 0 { + return "" + } + if len(message) <= maxBytes { + return message + } + limit := maxBytes + for limit > 0 && !utf8.RuneStart(message[limit]) { + limit-- + } + return message[:limit] +} diff --git a/controller/channel_contribution_error_test.go b/controller/channel_contribution_error_test.go new file mode 100644 index 000000000000..cff993b09313 --- /dev/null +++ b/controller/channel_contribution_error_test.go @@ -0,0 +1,17 @@ +package controller + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/stretchr/testify/assert" +) + +func TestTruncateChannelContributionErrorPreservesUTF8(t *testing.T) { + message := strings.Repeat("错误", 1_100) + truncated := truncateChannelContributionError(message, 2_000) + assert.LessOrEqual(t, len(truncated), 2_000) + assert.True(t, utf8.ValidString(truncated)) + assert.Equal(t, message, truncateChannelContributionError(message, len(message))) +} diff --git a/controller/channel_contribution_health.go b/controller/channel_contribution_health.go new file mode 100644 index 000000000000..289e2c20d58a --- /dev/null +++ b/controller/channel_contribution_health.go @@ -0,0 +1,389 @@ +package controller + +import ( + "context" + "errors" + "fmt" + "net/url" + "strings" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/operation_setting" +) + +const ( + channelContributionHealthBatchSize = 100 + channelContributionHealthWorkerCount = 4 + channelContributionHealthRequestTimeout = 30 * time.Second +) + +type channelContributionHealthHandler struct{} + +type channelContributionHealthSummary struct { + Contributions int `json:"contributions"` + Models int `json:"models"` + Succeeded int `json:"succeeded"` + Failed int `json:"failed"` + Paused int `json:"paused"` + Deleted int `json:"deleted"` +} + +type channelContributionHealthEntry struct { + Candidate model.ChannelContributionHealthCandidate + Channel *model.Channel + Revision *model.ChannelContributionRevision + Specs []channelContributionProbeSpec + Observations map[string]model.ChannelContributionModelObservation +} + +type channelContributionHealthWork struct { + EntryIndex int + Channel *model.Channel + Group string + Spec channelContributionProbeSpec +} + +type channelContributionHealthWorkResult struct { + EntryIndex int + Observation model.ChannelContributionModelObservation +} + +func RegisterChannelContributionHealthTask() { + service.RegisterSystemTaskHandler(channelContributionHealthHandler{}) +} + +func (channelContributionHealthHandler) Type() string { + return model.SystemTaskTypeChannelContributionHealth +} + +func (channelContributionHealthHandler) Enabled() bool { + return true +} + +func (channelContributionHealthHandler) Interval() time.Duration { + minutes := operation_setting.GetChannelContributionSetting().HealthCheckIntervalMinutes + if minutes <= 0 { + minutes = 10 + } + return time.Duration(minutes) * time.Minute +} + +func (channelContributionHealthHandler) NewPayload() any { return nil } + +func (channelContributionHealthHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) { + summary, err := runChannelContributionHealthTask(ctx, service.NewSystemTaskProgressReporter(task, runnerID)) + if err != nil { + finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, err) + return + } + finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil) +} + +func runChannelContributionHealthTask(ctx context.Context, report func(processed, total int)) (summary channelContributionHealthSummary, runErr error) { + testUserId, err := resolveChannelTestUserID(nil) + if err != nil { + return summary, err + } + setting := operation_setting.GetChannelContributionSetting() + deleteAfterSeconds := int64(setting.UnavailableDeleteHours) * int64(time.Hour/time.Second) + afterContributionId := 0 + processed := 0 + cacheChanged := false + defer func() { + if !cacheChanged { + return + } + model.InitChannelCache() + }() + + for { + if err := ctx.Err(); err != nil { + return summary, err + } + candidates, err := model.ListContributionChannelsForHealthAfter(afterContributionId, channelContributionHealthBatchSize) + if err != nil { + return summary, err + } + if len(candidates) == 0 { + break + } + afterContributionId = candidates[len(candidates)-1].ContributionId + + channelIds := make([]int, 0, len(candidates)) + revisionIds := make([]int, 0, len(candidates)) + for _, candidate := range candidates { + channelIds = append(channelIds, candidate.ChannelId) + revisionIds = append(revisionIds, candidate.RevisionId) + } + var channels []*model.Channel + if err := model.DB.Where("id IN ?", channelIds).Find(&channels).Error; err != nil { + return summary, err + } + channelById := make(map[int]*model.Channel, len(channels)) + for _, channel := range channels { + channelById[channel.Id] = channel + } + var revisions []*model.ChannelContributionRevision + if err := model.DB.Where("id IN ?", revisionIds).Find(&revisions).Error; err != nil { + return summary, err + } + revisionById := make(map[int]*model.ChannelContributionRevision, len(revisions)) + for _, revision := range revisions { + revisionById[revision.Id] = revision + } + + entries := make([]channelContributionHealthEntry, 0, len(candidates)) + for _, candidate := range candidates { + entry := channelContributionHealthEntry{ + Candidate: candidate, + Channel: channelById[candidate.ChannelId], + Revision: revisionById[candidate.RevisionId], + Observations: make(map[string]model.ChannelContributionModelObservation), + } + if entry.Revision == nil || entry.Revision.ConfigHash != candidate.ConfigHash { + entries = append(entries, entry) + continue + } + if entry.Channel == nil || entry.Channel.Status == common.ChannelStatusManuallyDisabled { + entries = append(entries, entry) + continue + } + + specs, specErr := resolveChannelContributionProbeSpecs(entry.Revision) + channelModels := channelContributionModels(entry.Channel.Models) + if specErr != nil { + for _, modelName := range channelModels { + entry.Observations[modelName] = model.ChannelContributionModelObservation{ + Model: modelName, + Error: specErr.Error(), + } + } + entries = append(entries, entry) + continue + } + specByModel := make(map[string]channelContributionProbeSpec, len(specs)) + for _, spec := range specs { + specByModel[spec.Model] = spec + } + entry.Specs = make([]channelContributionProbeSpec, 0, len(channelModels)) + for _, modelName := range channelModels { + spec, ok := specByModel[modelName] + if !ok { + entry.Observations[modelName] = model.ChannelContributionModelObservation{ + Model: modelName, + Error: "model is absent from the approved contribution revision", + } + continue + } + entry.Specs = append(entry.Specs, spec) + } + entries = append(entries, entry) + } + + work := buildChannelContributionHealthWork(entries) + results := executeChannelContributionHealthWork(ctx, testUserId, work) + if err := ctx.Err(); err != nil { + return summary, err + } + for _, result := range results { + entry := &entries[result.EntryIndex] + entry.Observations[result.Observation.Model] = result.Observation + } + + totalHint := processed + len(entries) + for index := range entries { + entry := &entries[index] + processed++ + summary.Contributions++ + if report != nil { + report(processed, totalHint) + } + if entry.Revision == nil || entry.Revision.ConfigHash != entry.Candidate.ConfigHash { + continue + } + + observations := make([]model.ChannelContributionModelObservation, 0, len(entry.Observations)) + for _, modelName := range channelContributionModels(entry.channelModels()) { + observation, ok := entry.Observations[modelName] + if ok { + observations = append(observations, observation) + } + } + if entry.Channel != nil && entry.Channel.Status != common.ChannelStatusManuallyDisabled { + summary.Models += len(observations) + for _, observation := range observations { + if observation.Healthy { + summary.Succeeded++ + } else { + summary.Failed++ + } + } + } + + cycle, err := model.ApplyChannelContributionHealthCycle( + entry.Candidate.ContributionId, + entry.Candidate.ChannelId, + entry.Candidate.RevisionId, + entry.Candidate.ConfigHash, + observations, + common.GetTimestamp(), + deleteAfterSeconds, + ) + if errors.Is(err, model.ErrStaleChannelContributionHealthProbe) { + continue + } + if err != nil { + return summary, fmt.Errorf("apply contribution health channel=%d: %w", entry.Candidate.ChannelId, err) + } + if cycle.Paused { + summary.Paused++ + } + if cycle.Deleted { + summary.Deleted++ + } + cacheChanged = cacheChanged || cycle.StateChanged + } + if len(candidates) < channelContributionHealthBatchSize { + break + } + } + + if report != nil { + report(processed, processed) + } + return summary, nil +} + +func (entry *channelContributionHealthEntry) channelModels() string { + if entry.Channel != nil { + return entry.Channel.Models + } + if entry.Revision != nil { + return entry.Revision.Models + } + return "" +} + +func buildChannelContributionHealthWork(entries []channelContributionHealthEntry) []channelContributionHealthWork { + maxModels := 0 + total := 0 + for _, entry := range entries { + total += len(entry.Specs) + if len(entry.Specs) > maxModels { + maxModels = len(entry.Specs) + } + } + work := make([]channelContributionHealthWork, 0, total) + for modelIndex := 0; modelIndex < maxModels; modelIndex++ { + for entryIndex := range entries { + entry := &entries[entryIndex] + if entry.Channel == nil || modelIndex >= len(entry.Specs) { + continue + } + group := entry.Channel.Group + if entry.Revision != nil { + group = entry.Revision.Group + } + work = append(work, channelContributionHealthWork{ + EntryIndex: entryIndex, + Channel: entry.Channel, + Group: group, + Spec: entry.Specs[modelIndex], + }) + } + } + return work +} + +func executeChannelContributionHealthWork(ctx context.Context, testUserId int, work []channelContributionHealthWork) []channelContributionHealthWorkResult { + if len(work) == 0 { + return nil + } + jobs := make(chan channelContributionHealthWork, len(work)) + results := make(chan channelContributionHealthWorkResult, len(work)) + workerCount := channelContributionHealthWorkerCount + if len(work) < workerCount { + workerCount = len(work) + } + var workers sync.WaitGroup + workers.Add(workerCount) + for workerIndex := 0; workerIndex < workerCount; workerIndex++ { + go func() { + defer workers.Done() + for item := range jobs { + results <- channelContributionHealthWorkResult{ + EntryIndex: item.EntryIndex, + Observation: probeChannelContributionModel(ctx, testUserId, item.Channel, item.Group, item.Spec), + } + } + }() + } + for _, item := range work { + jobs <- item + } + close(jobs) + workers.Wait() + close(results) + + collected := make([]channelContributionHealthWorkResult, 0, len(work)) + for result := range results { + collected = append(collected, result) + } + return collected +} + +func probeChannelContributionModel(ctx context.Context, testUserId int, channel *model.Channel, group string, spec channelContributionProbeSpec) model.ChannelContributionModelObservation { + observation := model.ChannelContributionModelObservation{Model: spec.Model} + failures := make([]string, 0, len(spec.Streams)) + for _, stream := range spec.Streams { + probeCtx, cancel := context.WithTimeout(ctx, channelContributionHealthRequestTimeout) + result := testChannelWithOptions( + probeCtx, + channel, + testUserId, + spec.Model, + string(spec.EndpointType), + stream, + channelTestOptions{ + UseSSRFProtectedClient: true, + SkipConsumeLog: true, + SkipPricingValidation: true, + GroupOverride: group, + }, + ) + cancel() + if result.localErr == nil && result.newAPIError == nil { + continue + } + mode := "non-stream" + if stream { + mode = "stream" + } + failures = append(failures, mode+": "+contributionHealthError(channel, result)) + } + observation.Healthy = len(failures) == 0 + observation.Error = strings.Join(failures, "; ") + return observation +} + +func contributionHealthError(channel *model.Channel, result testResult) string { + message := "channel health test failed" + if result.newAPIError != nil { + message = result.newAPIError.Error() + } else if result.localErr != nil { + message = result.localErr.Error() + } + if channel != nil { + if baseURL := strings.TrimSpace(channel.GetBaseURL()); baseURL != "" { + for _, candidate := range []string{baseURL, url.QueryEscape(baseURL), url.PathEscape(baseURL)} { + message = strings.ReplaceAll(message, candidate, "[UPSTREAM]") + } + } + message = sanitizeChannelCredentialError(errors.New(message), channel.Key, "").Error() + } + return truncateChannelContributionError(message, 500) +} diff --git a/controller/channel_contribution_health_test.go b/controller/channel_contribution_health_test.go new file mode 100644 index 000000000000..8a8b9b89bf0d --- /dev/null +++ b/controller/channel_contribution_health_test.go @@ -0,0 +1,70 @@ +package controller + +import ( + "errors" + "net/url" + "strings" + "testing" + "unicode/utf8" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" +) + +func TestBuildChannelContributionHealthWorkUsesModelRoundRobinOrder(t *testing.T) { + entries := []channelContributionHealthEntry{ + { + Channel: &model.Channel{Id: 1}, + Specs: []channelContributionProbeSpec{ + {Model: "a-1", EndpointType: constant.EndpointTypeOpenAI}, + {Model: "a-2", EndpointType: constant.EndpointTypeOpenAI}, + {Model: "a-3", EndpointType: constant.EndpointTypeOpenAI}, + }, + }, + { + Channel: &model.Channel{Id: 2, Group: "channel-group"}, + Revision: &model.ChannelContributionRevision{Group: "revision-group"}, + Specs: []channelContributionProbeSpec{ + {Model: "b-1", EndpointType: constant.EndpointTypeOpenAI}, + {Model: "b-2", EndpointType: constant.EndpointTypeOpenAI}, + }, + }, + } + + work := buildChannelContributionHealthWork(entries) + models := make([]string, 0, len(work)) + channelIDs := make([]int, 0, len(work)) + groups := make([]string, 0, len(work)) + for _, item := range work { + models = append(models, item.Spec.Model) + channelIDs = append(channelIDs, item.Channel.Id) + groups = append(groups, item.Group) + } + assert.Equal(t, []string{"a-1", "b-1", "a-2", "b-2", "a-3"}, models) + assert.Equal(t, []int{1, 2, 1, 2, 1}, channelIDs) + assert.Equal(t, []string{"", "revision-group", "", "revision-group", ""}, groups) +} + +func TestContributionHealthErrorRedactsSecretsAndPreservesUTF8(t *testing.T) { + baseURL := "https://upstream.example/v1" + channel := &model.Channel{ + Key: "super-secret-key", + BaseURL: &baseURL, + } + result := testResult{localErr: errors.New( + baseURL + " rejected super-secret-key " + url.QueryEscape(channel.Key) + + " Authorization: Bearer reflected-token: " + strings.Repeat("错误", 300), + )} + + message := contributionHealthError(channel, result) + + assert.NotContains(t, message, baseURL) + assert.NotContains(t, message, channel.Key) + assert.NotContains(t, message, url.QueryEscape(channel.Key)) + assert.NotContains(t, message, "reflected-token") + assert.Contains(t, message, "[UPSTREAM]") + assert.Contains(t, message, "[REDACTED]") + assert.LessOrEqual(t, len(message), 500) + assert.True(t, utf8.ValidString(message)) +} diff --git a/controller/channel_contribution_pricing_test.go b/controller/channel_contribution_pricing_test.go new file mode 100644 index 000000000000..7893cb1f2b2f --- /dev/null +++ b/controller/channel_contribution_pricing_test.go @@ -0,0 +1,68 @@ +package controller + +import ( + "testing" + + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestBuildChannelContributionTestRunResponseUsesCurrentSystemPricing(t *testing.T) { + previousDB := model.DB + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + t.Cleanup(func() { + model.DB = previousDB + }) + require.NoError(t, db.AutoMigrate( + &model.ChannelContributionRevision{}, + &model.ChannelContributionTestRun{}, + &model.ChannelContributionTestResult{}, + )) + + previousPrices := ratio_setting.ModelPrice2JSONString() + previousRatios := ratio_setting.ModelRatio2JSONString() + t.Cleanup(func() { + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(previousPrices)) + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(previousRatios)) + }) + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"central-priced-model":1}`)) + + revision := &model.ChannelContributionRevision{ + ContributionId: 1, + Models: "central-priced-model", + } + require.NoError(t, db.Create(revision).Error) + run := &model.ChannelContributionTestRun{ + ContributionId: 1, + RevisionId: revision.Id, + PricingReady: false, + } + require.NoError(t, db.Create(run).Error) + require.NoError(t, db.Create(&model.ChannelContributionTestResult{ + TestRunId: run.Id, + RevisionId: revision.Id, + Model: "central-priced-model", + Success: true, + }).Error) + + response, err := buildChannelContributionTestRunResponse(run, true) + require.NoError(t, err) + assert.True(t, response.PricingReady) + require.Len(t, response.Results, 1) + assert.True(t, response.Results[0].PriceConfigured) + + run.PricingReady = true + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{}`)) + response, err = buildChannelContributionTestRunResponse(run, true) + require.NoError(t, err) + assert.False(t, response.PricingReady) + require.Len(t, response.Results, 1) + assert.False(t, response.Results[0].PriceConfigured) +} diff --git a/controller/channel_contribution_probe.go b/controller/channel_contribution_probe.go new file mode 100644 index 000000000000..841c3966d4de --- /dev/null +++ b/controller/channel_contribution_probe.go @@ -0,0 +1,110 @@ +package controller + +import ( + "errors" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" +) + +type channelContributionProbeSpec struct { + Model string + EndpointType constant.EndpointType + Streams []bool +} + +func resolveChannelContributionProbeSpecs(revision *model.ChannelContributionRevision) ([]channelContributionProbeSpec, error) { + if revision == nil { + return nil, errors.New("contribution revision is required") + } + models := channelContributionModels(revision.Models) + if len(models) == 0 { + return nil, errors.New("at least one model is required before testing") + } + if len(models) > channelContributionMaxModels { + return nil, errors.New("contribution contains too many models") + } + + // Populate endpoint metadata before reading the per-model endpoint map. + model.GetPricing() + specs := make([]channelContributionProbeSpec, 0, len(models)) + for _, modelName := range models { + if common.IsImageGenerationModel(modelName) { + return nil, errors.New("image generation models are not supported for channel contribution") + } + endpointTypes := model.GetModelSupportEndpointTypes(modelName) + if err := validateChannelContributionEndpointTypes(endpointTypes); err != nil { + return nil, err + } + + endpointType := selectChannelContributionEndpoint(revision.Type, modelName, endpointTypes) + streams := []bool{false, true} + if endpointType == constant.EndpointTypeEmbeddings || endpointType == constant.EndpointTypeJinaRerank { + streams = []bool{false} + } + specs = append(specs, channelContributionProbeSpec{ + Model: modelName, + EndpointType: endpointType, + Streams: streams, + }) + } + return specs, nil +} + +func validateChannelContributionEndpointTypes(endpointTypes []constant.EndpointType) error { + for _, endpointType := range endpointTypes { + switch endpointType { + case constant.EndpointTypeImageGeneration, constant.EndpointTypeOpenAIVideo: + return errors.New("asynchronous image and video models are not supported for channel contribution") + } + } + return nil +} + +func selectChannelContributionEndpoint(channelType int, modelName string, endpointTypes []constant.EndpointType) constant.EndpointType { + for _, preferred := range []constant.EndpointType{ + constant.EndpointTypeEmbeddings, + constant.EndpointTypeJinaRerank, + constant.EndpointTypeOpenAIResponse, + } { + for _, endpointType := range endpointTypes { + if endpointType == preferred { + return endpointType + } + } + } + + lowerModel := strings.ToLower(modelName) + if strings.Contains(lowerModel, "rerank") { + return constant.EndpointTypeJinaRerank + } + if strings.Contains(lowerModel, "embedding") || strings.Contains(lowerModel, "embed") || + strings.HasPrefix(lowerModel, "m3e") || strings.Contains(lowerModel, "bge-") { + return constant.EndpointTypeEmbeddings + } + if common.IsOpenAIResponseOnlyModel(modelName) { + return constant.EndpointTypeOpenAIResponse + } + + switch channelType { + case constant.ChannelTypeAnthropic: + return constant.EndpointTypeAnthropic + case constant.ChannelTypeGemini: + return constant.EndpointTypeGemini + case constant.ChannelTypeSub2API, constant.ChannelTypeNewAPI: + for _, preferred := range []constant.EndpointType{ + constant.EndpointTypeAnthropic, + constant.EndpointTypeGemini, + constant.EndpointTypeOpenAI, + } { + for _, endpointType := range endpointTypes { + if endpointType == preferred { + return endpointType + } + } + } + } + return constant.EndpointTypeOpenAI +} diff --git a/controller/channel_contribution_probe_test.go b/controller/channel_contribution_probe_test.go new file mode 100644 index 000000000000..e0aa8c35fbf0 --- /dev/null +++ b/controller/channel_contribution_probe_test.go @@ -0,0 +1,93 @@ +package controller + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSelectChannelContributionEndpointRespectsChannelProtocol(t *testing.T) { + tests := []struct { + name string + channelType int + modelName string + endpointTypes []constant.EndpointType + expected constant.EndpointType + }{ + { + name: "openai compatible claude model remains openai", + channelType: constant.ChannelTypeOpenAI, + modelName: "claude-3-7-sonnet", + endpointTypes: []constant.EndpointType{constant.EndpointTypeAnthropic}, + expected: constant.EndpointTypeOpenAI, + }, + { + name: "openai compatible gemini model remains openai", + channelType: constant.ChannelTypeOpenRouter, + modelName: "gemini-2.5-pro", + endpointTypes: []constant.EndpointType{constant.EndpointTypeGemini}, + expected: constant.EndpointTypeOpenAI, + }, + { + name: "native anthropic channel", + channelType: constant.ChannelTypeAnthropic, + modelName: "claude-3-7-sonnet", + expected: constant.EndpointTypeAnthropic, + }, + { + name: "native gemini channel", + channelType: constant.ChannelTypeGemini, + modelName: "gemini-2.5-pro", + expected: constant.EndpointTypeGemini, + }, + { + name: "embedding capability overrides chat protocol", + channelType: constant.ChannelTypeOpenAI, + modelName: "text-embedding-3-small", + endpointTypes: []constant.EndpointType{constant.EndpointTypeEmbeddings}, + expected: constant.EndpointTypeEmbeddings, + }, + { + name: "multi protocol new api follows metadata", + channelType: constant.ChannelTypeNewAPI, + modelName: "claude-3-7-sonnet", + endpointTypes: []constant.EndpointType{constant.EndpointTypeAnthropic, constant.EndpointTypeOpenAI}, + expected: constant.EndpointTypeAnthropic, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + assert.Equal(t, test.expected, selectChannelContributionEndpoint(test.channelType, test.modelName, test.endpointTypes)) + }) + } +} + +func TestValidateChannelContributionEndpointTypesRejectsAsyncMedia(t *testing.T) { + for _, endpointType := range []constant.EndpointType{ + constant.EndpointTypeImageGeneration, + constant.EndpointTypeOpenAIVideo, + } { + t.Run(string(endpointType), func(t *testing.T) { + require.Error(t, validateChannelContributionEndpointTypes([]constant.EndpointType{endpointType})) + }) + } + require.NoError(t, validateChannelContributionEndpointTypes([]constant.EndpointType{ + constant.EndpointTypeOpenAI, + constant.EndpointTypeOpenAIResponse, + constant.EndpointTypeEmbeddings, + })) +} + +func TestChannelTestResponseRecorderCapsBufferedOutput(t *testing.T) { + recorder := newChannelTestResponseRecorder(8) + written, err := recorder.Write([]byte(strings.Repeat("x", 32))) + + require.NoError(t, err) + assert.Equal(t, 32, written) + assert.True(t, recorder.exceeded) + assert.Equal(t, "xxxxxxxx", recorder.Body.String()) +} diff --git a/controller/channel_contribution_reward.go b/controller/channel_contribution_reward.go new file mode 100644 index 000000000000..845bd64f0bd0 --- /dev/null +++ b/controller/channel_contribution_reward.go @@ -0,0 +1,99 @@ +package controller + +import ( + "net/http" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + + "github.com/gin-gonic/gin" +) + +type channelContributionRewardTransferInput struct { + Amount int `json:"amount"` +} + +type channelContributionRewardEntryResponse struct { + Id int64 `json:"id"` + UserId int `json:"user_id"` + ContributionId int `json:"contribution_id"` + ChannelId int `json:"channel_id"` + RequestId string `json:"request_id"` + EntryType string `json:"entry_type"` + Amount int64 `json:"amount"` + BalanceAfter int64 `json:"balance_after"` + SourceQuota int `json:"source_quota"` + RewardBps int `json:"reward_bps"` + CreatedAt int64 `json:"created_at"` +} + +func buildChannelContributionRewardEntries(entries []*model.ChannelContributionRewardLedger) []channelContributionRewardEntryResponse { + response := make([]channelContributionRewardEntryResponse, 0, len(entries)) + for _, entry := range entries { + if entry == nil { + continue + } + response = append(response, channelContributionRewardEntryResponse{ + Id: entry.Id, + UserId: entry.UserId, + ContributionId: entry.ContributionId, + ChannelId: entry.ChannelId, + RequestId: entry.RequestId, + EntryType: entry.EntryType, + Amount: entry.Amount, + BalanceAfter: entry.BalanceAfter, + SourceQuota: entry.SourceQuota, + RewardBps: entry.RewardBps, + CreatedAt: entry.CreatedAt, + }) + } + return response +} + +func GetChannelContributionRewards(c *gin.Context) { + userId := c.GetInt("id") + pageInfo := common.GetPageQuery(c) + account, err := model.GetChannelContributionRewardAccount(userId) + if err != nil { + common.ApiError(c, err) + return + } + entries, total, err := model.ListChannelContributionRewardLedger(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize()) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{ + "account": account, + "items": buildChannelContributionRewardEntries(entries), + "total": total, + }) +} + +func ListChannelContributionRewardTransfers(c *gin.Context) { + pageInfo := common.GetPageQuery(c) + entries, total, err := model.ListChannelContributionRewardTransfers(c.GetInt("id"), pageInfo.GetStartIdx(), pageInfo.GetPageSize()) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{ + "items": buildChannelContributionRewardEntries(entries), + "total": total, + }) +} + +func TransferChannelContributionReward(c *gin.Context) { + var input channelContributionRewardTransferInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid request body"}) + return + } + entry, err := model.TransferChannelContributionReward(c.GetInt("id"), input.Amount) + if err != nil { + common.ApiError(c, err) + return + } + entries := buildChannelContributionRewardEntries([]*model.ChannelContributionRewardLedger{entry}) + common.ApiSuccess(c, entries[0]) +} diff --git a/controller/channel_contribution_test_run.go b/controller/channel_contribution_test_run.go new file mode 100644 index 000000000000..55f498917200 --- /dev/null +++ b/controller/channel_contribution_test_run.go @@ -0,0 +1,512 @@ +package controller + +import ( + "context" + "errors" + "fmt" + "strconv" + "strings" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/service" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +type channelContributionTestResultResponse struct { + Id int64 `json:"id"` + Model string `json:"model"` + EndpointType string `json:"endpoint_type"` + Stream bool `json:"stream"` + Mode string `json:"mode"` + Success bool `json:"success"` + PriceConfigured bool `json:"price_configured"` + LatencyMs int64 `json:"latency_ms"` + Error string `json:"error"` + CreatedAt int64 `json:"created_at"` +} + +type channelContributionTestRunResponse struct { + Id int64 `json:"id"` + ContributionId int `json:"contribution_id"` + RevisionId int `json:"revision_id"` + RevisionNumber int `json:"revision_number"` + ActorId int `json:"actor_id"` + ActorType string `json:"actor_type"` + Status model.ChannelContributionTestRunStatus `json:"status"` + PricingReady bool `json:"pricing_ready"` + Total int `json:"total"` + Passed int `json:"passed"` + Failed int `json:"failed"` + Error string `json:"error"` + StartedAt int64 `json:"started_at"` + CompletedAt int64 `json:"completed_at"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` + Results []channelContributionTestResultResponse `json:"results,omitempty"` +} + +func buildChannelContributionTestRunResponse(run *model.ChannelContributionTestRun, includeResults bool) (*channelContributionTestRunResponse, error) { + if run == nil { + return nil, nil + } + revision, err := model.GetChannelContributionRevision(run.ContributionId, run.RevisionId) + if err != nil { + return nil, err + } + pricingReady, _ := channelContributionPriceStatus(channelContributionModels(revision.Models)) + response := &channelContributionTestRunResponse{ + Id: run.Id, + ContributionId: run.ContributionId, + RevisionId: run.RevisionId, + RevisionNumber: revision.RevisionNumber, + ActorId: run.ActorId, + ActorType: run.ActorType, + Status: run.Status, + PricingReady: pricingReady, + Total: run.Total, + Passed: run.Passed, + Failed: run.Failed, + Error: run.Error, + StartedAt: run.StartedAt, + CompletedAt: run.CompletedAt, + CreatedAt: run.CreatedAt, + UpdatedAt: run.UpdatedAt, + } + if !includeResults { + return response, nil + } + results, err := model.ListChannelContributionTestResults(run.Id) + if err != nil { + return nil, err + } + response.Results = make([]channelContributionTestResultResponse, 0, len(results)) + for _, result := range results { + mode := "non_stream" + if result.Stream { + mode = "stream" + } + response.Results = append(response.Results, channelContributionTestResultResponse{ + Id: result.Id, + Model: result.Model, + EndpointType: result.EndpointType, + Stream: result.Stream, + Mode: mode, + Success: result.Success, + PriceConfigured: helper.HasModelBillingConfig(result.Model), + LatencyMs: result.LatencyMs, + Error: result.Error, + CreatedAt: result.CreatedAt, + }) + } + return response, nil +} + +func validateChannelContributionSubmissionRun(runId int64, contribution *model.ChannelContribution, revision *model.ChannelContributionRevision, actorType string) error { + if runId <= 0 { + return errors.New("test_run_id is required") + } + run, err := model.GetChannelContributionTestRunForContribution(runId, contribution.Id) + if err != nil { + return err + } + if run.RevisionId != revision.Id || run.ConfigHash != revision.ConfigHash { + return errors.New("test run does not match the submitted revision") + } + if run.ActorType != actorType { + return fmt.Errorf("a %s test run is required", actorType) + } + if run.Status != model.ChannelContributionTestRunStatusSucceeded || run.Failed != 0 || run.Passed != run.Total || run.Total <= 0 { + return errors.New("test run has not passed all required probes") + } + specs, err := resolveChannelContributionProbeSpecs(revision) + if err != nil { + return err + } + expected := 0 + for _, spec := range specs { + expected += len(spec.Streams) + } + if run.Total != expected { + return errors.New("test run does not cover every required model mode") + } + now := common.GetTimestamp() + if run.CompletedAt <= 0 || run.CompletedAt > now+60 || now-run.CompletedAt > channelContributionTestResultTTLSeconds { + return errors.New("test run has expired; run all model tests again") + } + return nil +} + +func enqueueChannelContributionTestRun(contribution *model.ChannelContribution, revision *model.ChannelContributionRevision, actorId int, actorType string) (*model.ChannelContributionTestRun, error) { + if contribution == nil || revision == nil { + return nil, errors.New("contribution and revision are required") + } + if _, err := resolveChannelContributionProbeSpecs(revision); err != nil { + return nil, err + } + run := &model.ChannelContributionTestRun{ + ContributionId: contribution.Id, + RevisionId: revision.Id, + ConfigHash: revision.ConfigHash, + ActorId: actorId, + ActorType: actorType, + Status: model.ChannelContributionTestRunStatusQueued, + } + if err := model.CreateChannelContributionTestRun(run); err != nil { + return nil, err + } + if _, _, err := service.EnqueueSystemTask(model.SystemTaskTypeChannelContributionTest, nil); err != nil { + common.SysError(fmt.Sprintf("failed to wake channel contribution test runner: run_id=%d err=%v", run.Id, err)) + } + return run, nil +} + +func CreateUserChannelContributionTestRun(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + contribution, err := model.GetUserChannelContributionById(id, c.GetInt("id")) + if err != nil { + common.ApiError(c, err) + return + } + revision, err := loadChannelContributionRevision(contribution.CurrentRevisionId) + if err != nil || revision == nil { + if err == nil { + err = errors.New("current revision is missing") + } + common.ApiError(c, err) + return + } + if revision.Status == model.ChannelContributionRevisionStatusPending || revision.Status == model.ChannelContributionRevisionStatusApproved || revision.Status == model.ChannelContributionRevisionStatusSuperseded { + common.ApiErrorMsg(c, "current revision is not editable or testable by the contributor") + return + } + run, err := enqueueChannelContributionTestRun(contribution, revision, c.GetInt("id"), model.ChannelContributionTestActorUser) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionTestRunResponse(run, false) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func GetUserChannelContributionTestRun(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + if _, err := model.GetUserChannelContributionById(id, c.GetInt("id")); err != nil { + common.ApiError(c, err) + return + } + runId, err := strconv.ParseInt(c.Param("runId"), 10, 64) + if err != nil || runId <= 0 { + common.ApiErrorMsg(c, "invalid test run id") + return + } + run, err := model.GetChannelContributionTestRunForContribution(runId, id) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionTestRunResponse(run, true) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func CreateAdminChannelContributionTestRun(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + contribution, err := model.GetChannelContributionById(id) + if err != nil { + common.ApiError(c, err) + return + } + revision, err := loadChannelContributionRevision(contribution.PendingRevisionId) + if err != nil || revision == nil { + if err == nil { + err = errors.New("pending revision is missing") + } + common.ApiError(c, err) + return + } + run, err := enqueueChannelContributionTestRun(contribution, revision, c.GetInt("id"), model.ChannelContributionTestActorAdmin) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionTestRunResponse(run, false) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +func GetAdminChannelContributionTestRun(c *gin.Context) { + id, err := channelContributionId(c) + if err != nil { + common.ApiError(c, err) + return + } + runId, err := strconv.ParseInt(c.Param("runId"), 10, 64) + if err != nil || runId <= 0 { + common.ApiErrorMsg(c, "invalid test run id") + return + } + run, err := model.GetChannelContributionTestRunForContribution(runId, id) + if err != nil { + common.ApiError(c, err) + return + } + response, err := buildChannelContributionTestRunResponse(run, true) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, response) +} + +type channelContributionTestTaskSummary struct { + Runs int `json:"runs"` + Succeeded int `json:"succeeded"` + Failed int `json:"failed"` +} + +type channelContributionTestHandler struct{} + +func (channelContributionTestHandler) Type() string { + return model.SystemTaskTypeChannelContributionTest +} + +func (channelContributionTestHandler) Enabled() bool { + return model.HasUnfinishedChannelContributionTestRuns() +} + +func (channelContributionTestHandler) Interval() time.Duration { return 15 * time.Second } + +func (channelContributionTestHandler) NewPayload() any { return nil } + +func (channelContributionTestHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) { + summary := channelContributionTestTaskSummary{} + if requeued, err := model.RequeueRunningChannelContributionTestRuns(); err != nil { + finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, err) + return + } else if requeued > 0 { + common.SysLog(fmt.Sprintf("requeued %d interrupted channel contribution test runs", requeued)) + } + for ctx.Err() == nil { + run, err := model.ClaimNextQueuedChannelContributionTestRun() + if errors.Is(err, gorm.ErrRecordNotFound) { + break + } + if err != nil { + finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, err) + return + } + summary.Runs++ + if err := executeChannelContributionTestRun(ctx, run); err != nil { + summary.Failed++ + common.SysError(fmt.Sprintf("channel contribution test run failed: run_id=%d err=%v", run.Id, err)) + } else { + reloaded, reloadErr := model.GetChannelContributionTestRun(run.Id) + if reloadErr != nil || reloaded.Status != model.ChannelContributionTestRunStatusSucceeded { + summary.Failed++ + } else { + summary.Succeeded++ + } + } + } + finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil) +} + +type channelContributionProbeJob struct { + Index int + Model string + EndpointType string + Stream bool +} + +func executeChannelContributionProbe( + ctx context.Context, + run *model.ChannelContributionTestRun, + revision *model.ChannelContributionRevision, + channel *model.Channel, + job channelContributionProbeJob, +) model.ChannelContributionTestResult { + probeCtx, cancel := context.WithTimeout(ctx, channelContributionProbeTimeoutSeconds*time.Second) + defer cancel() + started := time.Now() + result := testChannelWithOptions( + probeCtx, + channel, + run.ActorId, + job.Model, + job.EndpointType, + job.Stream, + channelTestOptions{ + UseSSRFProtectedClient: true, + SkipConsumeLog: true, + SkipPricingValidation: true, + GroupOverride: revision.Group, + }, + ) + probeError := "" + if result.localErr != nil { + probeError = result.localErr.Error() + } else if result.newAPIError != nil { + probeError = result.newAPIError.Error() + } + if probeError != "" { + probeError = sanitizeChannelCredentialError(errors.New(probeError), revision.Key, revision.BaseURL).Error() + probeError = truncateChannelContributionError(probeError, 2_000) + } + return model.ChannelContributionTestResult{ + Model: job.Model, + EndpointType: job.EndpointType, + Stream: job.Stream, + Success: probeError == "", + LatencyMs: time.Since(started).Milliseconds(), + Error: probeError, + } +} + +func executeChannelContributionTestRun(ctx context.Context, run *model.ChannelContributionTestRun) error { + revision, err := model.GetChannelContributionRevision(run.ContributionId, run.RevisionId) + if err != nil { + return finishFailedChannelContributionRun(run, nil, false, err) + } + if revision.ConfigHash != run.ConfigHash { + return finishFailedChannelContributionRun(run, nil, false, errors.New("revision config hash changed")) + } + specs, err := resolveChannelContributionProbeSpecs(revision) + if err != nil { + return finishFailedChannelContributionRun(run, nil, false, err) + } + priceReady, _ := channelContributionPriceStatus(channelContributionModels(revision.Models)) + baseURL := revision.BaseURL + mapping := revision.ModelMapping + channel := &model.Channel{ + Type: revision.Type, + Key: revision.Key, + Status: common.ChannelStatusEnabled, + Name: revision.Name, + BaseURL: &baseURL, + Models: revision.Models, + Group: revision.Group, + ModelMapping: &mapping, + } + + jobs := make([]channelContributionProbeJob, 0, len(specs)*2) + for _, spec := range specs { + for _, stream := range spec.Streams { + jobs = append(jobs, channelContributionProbeJob{ + Index: len(jobs), + Model: spec.Model, + EndpointType: string(spec.EndpointType), + Stream: stream, + }) + } + } + results := make([]model.ChannelContributionTestResult, len(jobs)) + completed := make([]bool, len(jobs)) + jobQueue := make(chan channelContributionProbeJob, len(jobs)) + for _, job := range jobs { + jobQueue <- job + } + close(jobQueue) + workerCount := 4 + if len(jobs) < workerCount { + workerCount = len(jobs) + } + var workers sync.WaitGroup + workers.Add(workerCount) + for range workerCount { + go func() { + defer workers.Done() + for { + if ctx.Err() != nil { + return + } + select { + case <-ctx.Done(): + return + case job, ok := <-jobQueue: + if !ok { + return + } + results[job.Index] = executeChannelContributionProbe(ctx, run, revision, channel, job) + completed[job.Index] = true + } + } + }() + } + workers.Wait() + + allPassed := ctx.Err() == nil + cancellationError := "" + if ctx.Err() != nil { + cancellationError = sanitizeChannelCredentialError(ctx.Err(), revision.Key, revision.BaseURL).Error() + } + for index, job := range jobs { + if !completed[index] { + results[index] = model.ChannelContributionTestResult{ + Model: job.Model, + EndpointType: job.EndpointType, + Stream: job.Stream, + Success: false, + Error: cancellationError, + } + } + if !results[index].Success { + allPassed = false + } + } + status := model.ChannelContributionTestRunStatusSucceeded + runError := "" + if !allPassed { + status = model.ChannelContributionTestRunStatusFailed + if cancellationError != "" { + runError = cancellationError + } else { + runError = "one or more model probes failed" + } + } + return model.FinishChannelContributionTestRun(run.Id, status, priceReady, results, runError) +} + +func finishFailedChannelContributionRun(run *model.ChannelContributionTestRun, results []model.ChannelContributionTestResult, priceReady bool, runErr error) error { + message := runErr.Error() + message = truncateChannelContributionError(message, 2_000) + if err := model.FinishChannelContributionTestRun(run.Id, model.ChannelContributionTestRunStatusFailed, priceReady, results, message); err != nil { + return fmt.Errorf("%v; failed to persist terminal state: %w", runErr, err) + } + return runErr +} + +func channelContributionTestRunPathId(c *gin.Context) (int64, error) { + runId, err := strconv.ParseInt(strings.TrimSpace(c.Param("runId")), 10, 64) + if err != nil || runId <= 0 { + return 0, errors.New("invalid test run id") + } + return runId, nil +} diff --git a/controller/channel_test_internal_test.go b/controller/channel_test_internal_test.go index 904e08da21de..47dae411f784 100644 --- a/controller/channel_test_internal_test.go +++ b/controller/channel_test_internal_test.go @@ -56,6 +56,31 @@ func TestValidateChannelProxy(t *testing.T) { } } +func TestChannelTestResponseRecorderLimitsStringWrites(t *testing.T) { + recorder := newChannelTestResponseRecorder(4) + written, err := recorder.WriteString("streamed") + + require.NoError(t, err) + assert.Equal(t, len("streamed"), written) + assert.Equal(t, "stre", recorder.Body.String()) + assert.True(t, recorder.exceeded) +} + +func TestFetchModelsResponseBodyRejectsOversizedPayload(t *testing.T) { + payload := bytes.Repeat([]byte("x"), int(service.StrictSSRFProtectedResponseBodyLimitBytes)+1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write(payload) + })) + t.Cleanup(server.Close) + + body, err := getFetchModelsResponseBody(http.MethodGet, server.URL, nil, nil, fetchChannelModelsOptions{ + HTTPClient: server.Client(), + }) + + require.ErrorContains(t, err, "model list response exceeds") + assert.Nil(t, body) +} + func TestValidateChannelRequiresNewAPIBaseURL(t *testing.T) { tests := []struct { name string diff --git a/controller/channel_upstream_update.go b/controller/channel_upstream_update.go index 71ab0e53fafe..1df8d5e0f448 100644 --- a/controller/channel_upstream_update.go +++ b/controller/channel_upstream_update.go @@ -281,7 +281,9 @@ func parseOpenAIModelIDs(body []byte) ([]string, error) { return ids, nil } -func sanitizeFetchModelsError(err error, key string) error { +var authorizationCredentialPattern = regexp.MustCompile(`(?i)(authorization(?:\s*|["']\s*)[:=]\s*["']?)(?:bearer\s+)?[^"'\s,}\r\n]+`) + +func sanitizeChannelCredentialError(err error, key string, baseURL string) error { if err == nil { return nil } @@ -295,16 +297,28 @@ func sanitizeFetchModelsError(err error, key string) error { } message := err.Error() - key = strings.TrimSpace(key) - if key != "" { - message = strings.ReplaceAll(message, key, "[REDACTED]") - message = strings.ReplaceAll(message, url.QueryEscape(key), "[REDACTED]") - message = strings.ReplaceAll(message, url.PathEscape(key), "[REDACTED]") + for _, candidate := range []string{strings.TrimSpace(key), strings.TrimSpace(baseURL)} { + if candidate == "" { + continue + } + message = strings.ReplaceAll(message, candidate, "[REDACTED]") + message = strings.ReplaceAll(message, url.QueryEscape(candidate), "[REDACTED]") + message = strings.ReplaceAll(message, url.PathEscape(candidate), "[REDACTED]") } + message = authorizationCredentialPattern.ReplaceAllString(message, `${1}[REDACTED]`) return errors.New(message) } -func getFetchModelsResponseBody(method string, requestURL string, channel *model.Channel, headers http.Header) ([]byte, error) { +func sanitizeFetchModelsError(err error, key string) error { + return sanitizeChannelCredentialError(err, key, "") +} + +type fetchChannelModelsOptions struct { + UseSSRFProtectedClient bool + HTTPClient *http.Client +} + +func getFetchModelsResponseBody(method string, requestURL string, channel *model.Channel, headers http.Header, options fetchChannelModelsOptions) ([]byte, error) { request, err := http.NewRequest(method, requestURL, nil) if err != nil { return nil, err @@ -317,9 +331,17 @@ func getFetchModelsResponseBody(method string, requestURL string, channel *model request.Host = headers.Get(name) } } - client, err := service.NewProxyHttpClient(channel.GetSetting().Proxy) - if err != nil { - return nil, err + var client *http.Client + if options.HTTPClient != nil { + client = options.HTTPClient + } else if options.UseSSRFProtectedClient { + client = service.GetStrictSSRFProtectedHTTPClient() + } else { + var err error + client, err = service.NewProxyHttpClient(channel.GetSetting().Proxy) + if err != nil { + return nil, err + } } response, err := client.Do(request) if err != nil { @@ -329,10 +351,21 @@ func getFetchModelsResponseBody(method string, requestURL string, channel *model if response.StatusCode != http.StatusOK { return nil, fmt.Errorf("status code: %d", response.StatusCode) } - return io.ReadAll(response.Body) + body, err := io.ReadAll(io.LimitReader(response.Body, service.StrictSSRFProtectedResponseBodyLimitBytes+1)) + if err != nil { + return nil, err + } + if int64(len(body)) > service.StrictSSRFProtectedResponseBodyLimitBytes { + return nil, fmt.Errorf("model list response exceeds %d bytes", service.StrictSSRFProtectedResponseBodyLimitBytes) + } + return body, nil } func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) { + return fetchChannelUpstreamModelIDsWithOptions(channel, fetchChannelModelsOptions{}) +} + +func fetchChannelUpstreamModelIDsWithOptions(channel *model.Channel, options fetchChannelModelsOptions) ([]string, error) { baseURL := constant.ChannelBaseURLs[channel.Type] if channel.GetBaseURL() != "" { baseURL = channel.GetBaseURL() @@ -340,7 +373,15 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) { if channel.Type == constant.ChannelTypeOllama { key := strings.TrimSpace(strings.Split(channel.Key, "\n")[0]) - models, err := ollama.FetchOllamaModels(baseURL, key) + var models []ollama.OllamaModel + var err error + if options.HTTPClient != nil { + models, err = ollama.FetchOllamaModelsWithClient(options.HTTPClient, baseURL, key) + } else if options.UseSSRFProtectedClient { + models, err = ollama.FetchOllamaModelsWithClient(service.GetStrictSSRFProtectedHTTPClient(), baseURL, key) + } else { + models, err = ollama.FetchOllamaModels(baseURL, key) + } if err != nil { return nil, err } @@ -355,7 +396,15 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) { return nil, fmt.Errorf("获取渠道密钥失败: %w", apiErr) } key = strings.TrimSpace(key) - models, err := gemini.FetchGeminiModels(baseURL, key, channel.GetSetting().Proxy) + var models []string + var err error + if options.HTTPClient != nil { + models, err = gemini.FetchGeminiModelsWithClient(options.HTTPClient, baseURL, key) + } else if options.UseSSRFProtectedClient { + models, err = gemini.FetchGeminiModelsWithClient(service.GetStrictSSRFProtectedHTTPClient(), baseURL, key) + } else { + models, err = gemini.FetchGeminiModels(baseURL, key, channel.GetSetting().Proxy) + } if err != nil { return nil, err } @@ -363,7 +412,7 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) { } if channel.Type == constant.ChannelTypeAdvancedCustom { - return fetchAdvancedCustomUpstreamModelIDs(channel, baseURL) + return fetchAdvancedCustomUpstreamModelIDs(channel, baseURL, options) } if channel.Type == constant.ChannelTypeCodex { @@ -407,7 +456,7 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) { return nil, sanitizeFetchModelsError(err, key) } - body, err := getFetchModelsResponseBody(http.MethodGet, url, channel, headers) + body, err := getFetchModelsResponseBody(http.MethodGet, url, channel, headers, options) if err != nil { return nil, sanitizeFetchModelsError(err, key) } @@ -425,7 +474,7 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) { return normalizeModelNames(ids), nil } -func fetchAdvancedCustomUpstreamModelIDs(channel *model.Channel, baseURL string) ([]string, error) { +func fetchAdvancedCustomUpstreamModelIDs(channel *model.Channel, baseURL string, options fetchChannelModelsOptions) ([]string, error) { key, _, apiErr := channel.GetNextEnabledKey() if apiErr != nil { return nil, fmt.Errorf("获取渠道密钥失败: %w", apiErr) @@ -453,7 +502,7 @@ func fetchAdvancedCustomUpstreamModelIDs(channel *model.Channel, baseURL string) return nil, sanitizeFetchModelsError(err, key) } - body, err := getFetchModelsResponseBody(http.MethodGet, url, channel, headers) + body, err := getFetchModelsResponseBody(http.MethodGet, url, channel, headers, options) if err != nil { return nil, sanitizeFetchModelsError(err, key) } @@ -694,6 +743,7 @@ scanLoop: break } lastID = channels[len(channels)-1].Id + channels = model.FilterNonContributionChannels(channels) for _, channel := range channels { if channel == nil { diff --git a/controller/model_list_test.go b/controller/model_list_test.go index 812207b8fd44..38f4e47fc99a 100644 --- a/controller/model_list_test.go +++ b/controller/model_list_test.go @@ -49,7 +49,17 @@ func setupModelListControllerTestDB(t *testing.T) *gorm.DB { model.DB = db model.LOG_DB = db - require.NoError(t, db.AutoMigrate(&model.User{}, &model.Channel{}, &model.Ability{}, &model.Model{}, &model.Vendor{})) + require.NoError(t, db.AutoMigrate( + &model.User{}, + &model.Channel{}, + &model.Ability{}, + &model.Model{}, + &model.Vendor{}, + &model.Option{}, + &model.ChannelContribution{}, + &model.ChannelContributionRevision{}, + &model.ChannelContributionModelHealth{}, + )) t.Cleanup(func() { sqlDB, err := db.DB() diff --git a/controller/relay.go b/controller/relay.go index 8dccfe76dddd..72bb85d951a9 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -125,6 +125,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed) return } + service.SnapshotChannelContributionReward(c, relayInfo) needSensitiveCheck := setting.ShouldCheckPromptSensitive() needCountToken := constant.CountToken diff --git a/controller/system_task_handlers.go b/controller/system_task_handlers.go index c31059d148da..90db1df02098 100644 --- a/controller/system_task_handlers.go +++ b/controller/system_task_handlers.go @@ -20,6 +20,8 @@ import ( func RegisterScheduledSystemTasks() { service.RegisterSystemTaskHandler(channelTestHandler{}) service.RegisterSystemTaskHandler(modelUpdateHandler{}) + service.RegisterSystemTaskHandler(channelContributionTestHandler{}) + RegisterChannelContributionHealthTask() service.RegisterSystemTaskHandler(midjourneyPollHandler{}) service.RegisterSystemTaskHandler(asyncTaskPollHandler{}) } diff --git a/docs/images/channel-contribution-admin.png b/docs/images/channel-contribution-admin.png new file mode 100644 index 000000000000..2ce4d43b70ac Binary files /dev/null and b/docs/images/channel-contribution-admin.png differ diff --git a/docs/images/channel-contribution-desktop.png b/docs/images/channel-contribution-desktop.png new file mode 100644 index 000000000000..a53769fa4a9b Binary files /dev/null and b/docs/images/channel-contribution-desktop.png differ diff --git a/docs/images/channel-contribution-mobile.png b/docs/images/channel-contribution-mobile.png new file mode 100644 index 000000000000..f85170d6d3ef Binary files /dev/null and b/docs/images/channel-contribution-mobile.png differ diff --git a/logger/logger.go b/logger/logger.go index 867d88322430..da3aa7322053 100644 --- a/logger/logger.go +++ b/logger/logger.go @@ -11,6 +11,7 @@ import ( "time" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/bytedance/gopkg/util/gopool" @@ -86,6 +87,9 @@ func LogError(ctx context.Context, msg string) { } func LogDebug(ctx context.Context, msg string, args ...any) { + if ctx != nil && ctx.Value(constant.ContextKeySuppressUpstreamResponseLog) == true { + return + } if common.DebugEnabled { if len(args) > 0 { msg = fmt.Sprintf(msg, args...) diff --git a/model/ability.go b/model/ability.go index d950a6adbfc4..4709bbb34180 100644 --- a/model/ability.go +++ b/model/ability.go @@ -121,6 +121,7 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha if err != nil { return nil, err } + abilities = applyContributionHealthToAbilities(abilities) abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, model) channel := Channel{} if len(abilities) > 0 { @@ -194,6 +195,22 @@ func filterAbilitiesByRequestPathAndModel(abilities []Ability, requestPath strin } func (channel *Channel) AddAbilities(tx *gorm.DB) error { + useDB := tx + if useDB == nil { + useDB = DB + } + unhealthyModels := map[string]struct{}{} + if channel.Id > 0 { + unhealthyByChannel, err := contributionUnhealthyModelSet(useDB, []int{channel.Id}) + if err != nil { + return err + } + unhealthyModels = unhealthyByChannel[channel.Id] + } + return createChannelAbilitiesTx(useDB, channel, channel.Status == common.ChannelStatusEnabled, unhealthyModels) +} + +func createChannelAbilitiesTx(tx *gorm.DB, channel *Channel, abilityEnabled bool, unhealthyModels map[string]struct{}) error { models_ := strings.Split(channel.Models, ",") groups_ := strings.Split(channel.Group, ",") abilitySet := make(map[string]struct{}) @@ -205,11 +222,12 @@ func (channel *Channel) AddAbilities(tx *gorm.DB) error { continue } abilitySet[key] = struct{}{} + _, unhealthy := unhealthyModels[model] ability := Ability{ Group: group, Model: model, ChannelId: channel.Id, - Enabled: channel.Status == common.ChannelStatusEnabled, + Enabled: abilityEnabled && !unhealthy, Priority: channel.Priority, Weight: uint(channel.GetWeight()), Tag: channel.Tag, @@ -220,13 +238,8 @@ func (channel *Channel) AddAbilities(tx *gorm.DB) error { if len(abilities) == 0 { return nil } - // choose DB or provided tx - useDB := DB - if tx != nil { - useDB = tx - } for _, chunk := range lo.Chunk(abilities, 50) { - err := useDB.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error + err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error if err != nil { return err } @@ -256,8 +269,7 @@ func (channel *Channel) UpdateAbilities(tx *gorm.DB) error { }() } - // First delete all abilities of this channel - err := tx.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error + unhealthyByChannel, err := contributionUnhealthyModelSet(tx, []int{channel.Id}) if err != nil { if isNewTx { tx.Rollback() @@ -265,41 +277,19 @@ func (channel *Channel) UpdateAbilities(tx *gorm.DB) error { return err } - // Then add new abilities - models_ := strings.Split(channel.Models, ",") - groups_ := strings.Split(channel.Group, ",") - abilitySet := make(map[string]struct{}) - abilities := make([]Ability, 0, len(models_)) - for _, model := range models_ { - for _, group := range groups_ { - key := group + "|" + model - if _, exists := abilitySet[key]; exists { - continue - } - abilitySet[key] = struct{}{} - ability := Ability{ - Group: group, - Model: model, - ChannelId: channel.Id, - Enabled: channel.Status == common.ChannelStatusEnabled, - Priority: channel.Priority, - Weight: uint(channel.GetWeight()), - Tag: channel.Tag, - } - abilities = append(abilities, ability) + err = tx.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error + if err != nil { + if isNewTx { + tx.Rollback() } + return err } - if len(abilities) > 0 { - for _, chunk := range lo.Chunk(abilities, 50) { - err = tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error - if err != nil { - if isNewTx { - tx.Rollback() - } - return err - } + if err := createChannelAbilitiesTx(tx, channel, channel.Status == common.ChannelStatusEnabled, unhealthyByChannel[channel.Id]); err != nil { + if isNewTx { + tx.Rollback() } + return err } // 如果是新创建的事务，需要提交 @@ -311,11 +301,41 @@ func (channel *Channel) UpdateAbilities(tx *gorm.DB) error { } func UpdateAbilityStatus(channelId int, status bool) error { - return DB.Model(&Ability{}).Where("channel_id = ?", channelId).Select("enabled").Update("enabled", status).Error + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Model(&Ability{}).Where("channel_id = ?", channelId).Select("enabled").Update("enabled", status).Error; err != nil { + return err + } + if !status { + return nil + } + unhealthy, err := contributionUnhealthyModelSet(tx, []int{channelId}) + if err != nil { + return err + } + return applyContributionUnhealthyAbilitiesTx(tx, unhealthy) + }) } func UpdateAbilityStatusByTag(tag string, status bool) error { - return DB.Model(&Ability{}).Where("tag = ?", tag).Select("enabled").Update("enabled", status).Error + return DB.Transaction(func(tx *gorm.DB) error { + var channelIds []int + if status { + if err := tx.Model(&Ability{}).Where("tag = ?", tag).Distinct("channel_id").Pluck("channel_id", &channelIds).Error; err != nil { + return err + } + } + if err := tx.Model(&Ability{}).Where("tag = ?", tag).Select("enabled").Update("enabled", status).Error; err != nil { + return err + } + if !status { + return nil + } + unhealthy, err := contributionUnhealthyModelSet(tx, channelIds) + if err != nil { + return err + } + return applyContributionUnhealthyAbilitiesTx(tx, unhealthy) + }) } func UpdateAbilityByTag(tag string, newTag *string, priority *int64, weight *uint) error { @@ -341,50 +361,42 @@ func FixAbility() (int, int, error) { } defer fixLock.Unlock() - // truncate abilities table - if common.UsingMainDatabase(common.DatabaseTypeSQLite) { - err := DB.Exec("DELETE FROM abilities").Error - if err != nil { - common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error())) - return 0, 0, err - } - } else { - err := DB.Exec("TRUNCATE TABLE abilities").Error - if err != nil { - common.SysLog(fmt.Sprintf("Truncate abilities failed: %s", err.Error())) - return 0, 0, err - } - } - var channels []*Channel - // Find all channels - err := DB.Model(&Channel{}).Find(&channels).Error - if err != nil { - return 0, 0, err - } - if len(channels) == 0 { - return 0, 0, nil - } successCount := 0 failCount := 0 - for _, chunk := range lo.Chunk(channels, 50) { - ids := lo.Map(chunk, func(c *Channel, _ int) int { return c.Id }) - // Delete all abilities of this channel - err = DB.Where("channel_id IN ?", ids).Delete(&Ability{}).Error + err := DB.Transaction(func(tx *gorm.DB) error { + var channelIds []int + if err := tx.Model(&Channel{}).Order("id ASC").Pluck("id", &channelIds).Error; err != nil { + return err + } + if _, err := lockActiveChannelContributionsTx(tx, channelIds); err != nil { + return err + } + var channels []*Channel + if len(channelIds) > 0 { + if err := lockForUpdate(tx).Where("id IN ?", channelIds).Order("id ASC").Find(&channels).Error; err != nil { + return err + } + } + unhealthyByChannel, err := contributionUnhealthyModelSet(tx, channelIds) if err != nil { + return err + } + if err := tx.Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&Ability{}).Error; err != nil { common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error())) - failCount += len(chunk) - continue + return err } - // Then add new abilities - for _, channel := range chunk { - err = channel.AddAbilities(nil) - if err != nil { + for _, channel := range channels { + if err := createChannelAbilitiesTx(tx, channel, channel.Status == common.ChannelStatusEnabled, unhealthyByChannel[channel.Id]); err != nil { common.SysLog(fmt.Sprintf("Add abilities for channel %d failed: %s", channel.Id, err.Error())) failCount++ - } else { - successCount++ + return err } + successCount++ } + return nil + }) + if err != nil { + return 0, failCount, err } InitChannelCache() return successCount, failCount, nil diff --git a/model/channel.go b/model/channel.go index 0f8cdb101ec8..f188812f0b15 100644 --- a/model/channel.go +++ b/model/channel.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "math/rand" + "reflect" "strings" "sync" @@ -56,7 +57,8 @@ type Channel struct { OtherSettings string `json:"settings" gorm:"column:settings"` // 其他设置，存储azure版本等不需要检索的信息，详见dto.ChannelOtherSettings // cache info - Keys []string `json:"-" gorm:"-"` + Keys []string `json:"-" gorm:"-"` + IsContribution bool `json:"is_contribution" gorm:"-"` } type ChannelInfo struct { @@ -350,6 +352,10 @@ func (channel *Channel) Save() error { // Keeping this allowlist here prevents a stale channel snapshot from // overwriting credentials, accounting counters, or channel configuration. func (channel *Channel) saveStatusState() error { + return channel.saveStatusStateWithDB(DB) +} + +func (channel *Channel) saveStatusStateWithDB(db *gorm.DB) error { if channel.Id == 0 { return errors.New("channel ID is 0") } @@ -360,7 +366,7 @@ func (channel *Channel) saveStatusState() error { if channel.ChannelInfo.IsMultiKey { updates["channel_info"] = channel.ChannelInfo } - return DB.Model(&Channel{}).Where("id = ?", channel.Id).Updates(updates).Error + return db.Model(&Channel{}).Where("id = ?", channel.Id).Updates(updates).Error } func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool, sortOptions ...ChannelSortOptions) ([]*Channel, error) { @@ -473,6 +479,10 @@ func BatchDeleteChannels(ids []int) (int64, error) { } var deletedCount int64 for _, chunk := range lo.Chunk(ids, 200) { + if err := markChannelContributionsDeletedTx(tx, chunk, common.GetTimestamp()); err != nil { + tx.Rollback() + return 0, err + } result := tx.Where("id in (?)", chunk).Delete(&Channel{}) if result.Error != nil { tx.Rollback() @@ -540,6 +550,87 @@ func (channel *Channel) Insert() error { } func (channel *Channel) Update() error { + var activeContribution *ChannelContribution + if err := DB.Transaction(func(tx *gorm.DB) error { + contribution, err := lockActiveChannelContributionTx(tx, channel.Id) + if err != nil { + return err + } + activeContribution = contribution + if contribution == nil { + return nil + } + var existing Channel + if err := lockForUpdate(tx).Where("id = ?", channel.Id).First(&existing).Error; err != nil { + return err + } + if channel.Name == "" { + channel.Name = existing.Name + } + if channel.Type == 0 { + channel.Type = existing.Type + } + if channel.BaseURL == nil { + channel.BaseURL = existing.BaseURL + } + if channel.Key == "" { + channel.Key = existing.Key + } + if channel.Group == "" { + channel.Group = existing.Group + } + if channel.Models == "" { + channel.Models = existing.Models + } + if channel.ModelMapping == nil { + channel.ModelMapping = existing.ModelMapping + } + if channelContributionReviewedFieldsChanged(&existing, channel) { + return ErrChannelContributionRequiresReview + } + beforeStatus := existing.Status + updates := make(map[string]any, 4) + if channel.Status != 0 { + updates["status"] = channel.Status + } + if channel.Tag != nil { + updates["tag"] = *channel.Tag + } + if channel.Priority != nil { + updates["priority"] = *channel.Priority + } + if channel.Weight != nil { + updates["weight"] = *channel.Weight + } + if len(updates) > 0 { + if err := tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(updates).Error; err != nil { + return err + } + } + if err := tx.First(channel, "id = ?", channel.Id).Error; err != nil { + return err + } + if channel.Status != beforeStatus { + now := common.GetTimestamp() + if channel.Status == common.ChannelStatusManuallyDisabled { + if err := setLockedContributionHealthPausedTx(tx, contribution, channel.Id, true, now); err != nil { + return err + } + } else if beforeStatus == common.ChannelStatusManuallyDisabled { + if err := setLockedContributionHealthPausedTx(tx, contribution, channel.Id, false, now); err != nil { + return err + } + } + } + return channel.UpdateAbilities(tx) + }); err != nil { + return err + } + if activeContribution != nil { + CacheUpdateChannel(channel) + return nil + } + // If this is a multi-key channel, recalculate MultiKeySize based on the current key list to avoid inconsistency after editing keys if channel.ChannelInfo.IsMultiKey { var keyStr string @@ -588,6 +679,54 @@ func (channel *Channel) Update() error { return err } +// UpdateChannelAtomically applies an update to the latest channel row while +// preserving the review boundary for contributed channel configuration. +func UpdateChannelAtomically(channelId int, apply func(*Channel) error) (*Channel, error) { + if channelId <= 0 || apply == nil { + return nil, errors.New("invalid channel update") + } + channel := &Channel{} + err := DB.Transaction(func(tx *gorm.DB) error { + contribution, err := lockActiveChannelContributionTx(tx, channelId) + if err != nil { + return err + } + if err := lockForUpdate(tx).Where("id = ?", channelId).First(channel).Error; err != nil { + return err + } + before := *channel + if err := apply(channel); err != nil { + return err + } + channel.Id = channelId + channel.Keys = nil + if contribution != nil && channelContributionReviewedFieldsChanged(&before, channel) { + return ErrChannelContributionRequiresReview + } + if err := tx.Model(&Channel{}).Where("id = ?", channelId).Select("*").Omit("id").Updates(channel).Error; err != nil { + return err + } + if contribution != nil && channel.Status != before.Status { + now := common.GetTimestamp() + if channel.Status == common.ChannelStatusManuallyDisabled { + if err := setLockedContributionHealthPausedTx(tx, contribution, channelId, true, now); err != nil { + return err + } + } else if before.Status == common.ChannelStatusManuallyDisabled { + if err := setLockedContributionHealthPausedTx(tx, contribution, channelId, false, now); err != nil { + return err + } + } + } + return channel.UpdateAbilities(tx) + }) + if err != nil { + return nil, err + } + CacheUpdateChannel(channel) + return channel, nil +} + func (channel *Channel) UpdateResponseTime(responseTime int64) { err := DB.Model(channel).Select("response_time", "test_time").Updates(Channel{ TestTime: common.GetTimestamp(), @@ -609,17 +748,45 @@ func (channel *Channel) UpdateBalance(balance float64) { } func (channel *Channel) Delete() error { - var err error - err = DB.Delete(channel).Error - if err != nil { - return err - } - err = channel.DeleteAbilities() - return err + return DB.Transaction(func(tx *gorm.DB) error { + if err := markChannelContributionsDeletedTx(tx, []int{channel.Id}, common.GetTimestamp()); err != nil { + return err + } + if err := tx.Delete(channel).Error; err != nil { + return err + } + return tx.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error + }) } var channelStatusLock sync.Mutex +var ErrChannelContributionRequiresReview = errors.New("channel contribution configuration changes require a new reviewed revision") + +func channelContributionReviewedFieldsChanged(before *Channel, after *Channel) bool { + if before == nil || after == nil { + return true + } + type reviewedFields struct { + Name string + Type int + BaseURL *string + Key string + Group string + Models string + ModelMapping *string + } + left := reviewedFields{ + Name: before.Name, Type: before.Type, BaseURL: before.BaseURL, Key: before.Key, + Group: before.Group, Models: before.Models, ModelMapping: before.ModelMapping, + } + right := reviewedFields{ + Name: after.Name, Type: after.Type, BaseURL: after.BaseURL, Key: after.Key, + Group: after.Group, Models: after.Models, ModelMapping: after.ModelMapping, + } + return !reflect.DeepEqual(left, right) +} + // channelPollingLocks stores locks for each channel.id to ensure thread-safe polling var channelPollingLocks sync.Map @@ -720,6 +887,15 @@ func hasEnabledMultiKey(keys []string, statusList map[int]int) bool { } func UpdateChannelStatus(channelId int, usingKey string, status int, reason string) bool { + changed, err := UpdateChannelStatusWithError(channelId, usingKey, status, reason) + if err != nil { + common.SysLog(fmt.Sprintf("failed to update channel status: channel_id=%d, status=%d, error=%v", channelId, status, err)) + return false + } + return changed +} + +func UpdateChannelStatusWithError(channelId int, usingKey string, status int, reason string) (bool, error) { if common.MemoryCacheEnabled { channelStatusLock.Lock() defer channelStatusLock.Unlock() @@ -732,87 +908,132 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri pollingLock.Lock() defer pollingLock.Unlock() - if common.MemoryCacheEnabled { - channelCache, _ := CacheGetChannel(channelId) - if channelCache == nil { - return false - } - if channelCache.ChannelInfo.IsMultiKey { - beforeStatus := channelCache.Status - // 如果是多Key模式，更新缓存中的状态 - handlerMultiKeyUpdate(channelCache, usingKey, status, reason) - if beforeStatus != channelCache.Status { - CacheUpdateChannelStatus(channelId, channelCache.Status) - } - //CacheUpdateChannel(channelCache) - //return true - } else { - // 如果缓存渠道存在，且状态已是目标状态，直接返回 - if channelCache.Status == status { - return false - } - CacheUpdateChannelStatus(channelId, status) - } - } - - shouldUpdateAbilities := false - defer func() { - if shouldUpdateAbilities { - err := UpdateAbilityStatus(channelId, status == common.ChannelStatusEnabled) - if err != nil { - common.SysLog(fmt.Sprintf("failed to update ability status: channel_id=%d, error=%v", channelId, err)) - } + channel := &Channel{} + changed := false + err := DB.Transaction(func(tx *gorm.DB) error { + contribution, err := lockActiveChannelContributionTx(tx, channelId) + if err != nil { + return err } - }() - channel, err := GetChannelById(channelId, true) - if err != nil { - return false - } else { - if channel.Status == status { - return false + if err := lockForUpdate(tx).Where("id = ?", channelId).First(channel).Error; err != nil { + return err } - + beforeStatus := channel.Status if channel.ChannelInfo.IsMultiKey { - beforeStatus := channel.Status + beforeInfo, err := common.Marshal(channel.ChannelInfo) + if err != nil { + return err + } handlerMultiKeyUpdate(channel, usingKey, status, reason) - if beforeStatus != channel.Status { - shouldUpdateAbilities = true + afterInfo, err := common.Marshal(channel.ChannelInfo) + if err != nil { + return err } + changed = beforeStatus != channel.Status || string(beforeInfo) != string(afterInfo) } else { + if channel.Status == status { + return nil + } info := channel.GetOtherInfo() info["status_reason"] = reason info["status_time"] = common.GetTimestamp() channel.SetOtherInfo(info) channel.Status = status - shouldUpdateAbilities = true + changed = true } - err = channel.saveStatusState() - if err != nil { - common.SysLog(fmt.Sprintf("failed to update channel status: channel_id=%d, status=%d, error=%v", channel.Id, status, err)) - return false + if !changed { + return nil + } + if err := channel.saveStatusStateWithDB(tx); err != nil { + return err + } + if channel.Status != beforeStatus { + now := common.GetTimestamp() + if channel.Status == common.ChannelStatusManuallyDisabled { + if err := setLockedContributionHealthPausedTx(tx, contribution, channelId, true, now); err != nil { + return err + } + } else if beforeStatus == common.ChannelStatusManuallyDisabled { + if err := setLockedContributionHealthPausedTx(tx, contribution, channelId, false, now); err != nil { + return err + } + } + abilityEnabled := channel.Status == common.ChannelStatusEnabled + if err := tx.Model(&Ability{}).Where("channel_id = ?", channelId).Update("enabled", abilityEnabled).Error; err != nil { + return err + } + if abilityEnabled { + if err := reapplyContributionHealthToAbilitiesTx(tx, []int{channelId}); err != nil { + return err + } + } } + return nil + }) + if err != nil { + return false, err } - return true + if changed { + CacheUpdateChannel(channel) + } + return changed, nil } func EnableChannelByTag(tag string) error { - err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusEnabled).Error - if err != nil { - return err + err := updateChannelStatusByTag(tag, common.ChannelStatusEnabled) + if err == nil { + InitChannelCache() } - err = UpdateAbilityStatusByTag(tag, true) return err } func DisableChannelByTag(tag string) error { - err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusManuallyDisabled).Error - if err != nil { - return err + err := updateChannelStatusByTag(tag, common.ChannelStatusManuallyDisabled) + if err == nil { + InitChannelCache() } - err = UpdateAbilityStatusByTag(tag, false) return err } +func updateChannelStatusByTag(tag string, status int) error { + return DB.Transaction(func(tx *gorm.DB) error { + var channels []Channel + if err := tx.Where("tag = ?", tag).Order("id ASC").Find(&channels).Error; err != nil { + return err + } + channelIds := make([]int, 0, len(channels)) + for index := range channels { + channelIds = append(channelIds, channels[index].Id) + } + contributions, err := lockActiveChannelContributionsTx(tx, channelIds) + if err != nil { + return err + } + now := common.GetTimestamp() + for index := range channels { + channel := &channels[index] + beforeStatus := channel.Status + if err := tx.Model(&Channel{}).Where("id = ?", channel.Id).Update("status", status).Error; err != nil { + return err + } + paused := status == common.ChannelStatusManuallyDisabled + if paused || beforeStatus == common.ChannelStatusManuallyDisabled { + if err := setLockedContributionHealthPausedTx(tx, contributions[channel.Id], channel.Id, paused, now); err != nil { + return err + } + } + abilityEnabled := status == common.ChannelStatusEnabled + if err := tx.Model(&Ability{}).Where("channel_id = ?", channel.Id).Update("enabled", abilityEnabled).Error; err != nil { + return err + } + } + if status == common.ChannelStatusEnabled { + return reapplyContributionHealthToAbilitiesTx(tx, channelIds) + } + return nil + }) +} + func EditChannelByTag(tag string, newTag *string, modelMapping *string, models *string, group *string, priority *int64, weight *uint, paramOverride *string, headerOverride *string) error { updateData := Channel{} shouldReCreateAbilities := false @@ -846,27 +1067,53 @@ func EditChannelByTag(tag string, newTag *string, modelMapping *string, models * updateData.HeaderOverride = headerOverride } - err := DB.Model(&Channel{}).Where("tag = ?", tag).Updates(updateData).Error - if err != nil { - return err - } - if shouldReCreateAbilities { - channels, err := GetChannelsByTag(updatedTag, false, false) - if err == nil { + hasSensitiveUpdates := modelMapping != nil || + (models != nil && *models != "") || + (group != nil && *group != "") || + paramOverride != nil || + headerOverride != nil + + return DB.Transaction(func(tx *gorm.DB) error { + if hasSensitiveUpdates && tx.Migrator().HasTable(&ChannelContribution{}) { + var channelIds []int + if err := tx.Model(&Channel{}).Where("tag = ?", tag).Pluck("id", &channelIds).Error; err != nil { + return err + } + contributions, err := lockActiveChannelContributionsTx(tx, channelIds) + if err != nil { + return err + } + if len(contributions) > 0 { + return ErrChannelContributionRequiresReview + } + } + if err := tx.Model(&Channel{}).Where("tag = ?", tag).Updates(updateData).Error; err != nil { + return err + } + if shouldReCreateAbilities { + var channels []*Channel + if err := tx.Where("tag = ?", updatedTag).Find(&channels).Error; err != nil { + return err + } for _, channel := range channels { - err = channel.UpdateAbilities(nil) - if err != nil { - common.SysLog(fmt.Sprintf("failed to update abilities: channel_id=%d, tag=%s, error=%v", channel.Id, channel.GetTag(), err)) + if err := channel.UpdateAbilities(tx); err != nil { + return fmt.Errorf("failed to update abilities: channel_id=%d, tag=%s: %w", channel.Id, channel.GetTag(), err) } } + return nil } - } else { - err := UpdateAbilityByTag(tag, newTag, priority, weight) - if err != nil { - return err + ability := Ability{} + if newTag != nil { + ability.Tag = newTag } - } - return nil + if priority != nil { + ability.Priority = priority + } + if weight != nil { + ability.Weight = *weight + } + return tx.Model(&Ability{}).Where("tag = ?", tag).Updates(ability).Error + }) } func UpdateChannelUsedQuota(id int, quota int) { @@ -885,13 +1132,37 @@ func updateChannelUsedQuota(id int, quota int) { } func DeleteChannelByStatus(status int64) (int64, error) { - result := DB.Where("status = ?", status).Delete(&Channel{}) - return result.RowsAffected, result.Error + return deleteChannelsByStatuses([]int64{status}) } func DeleteDisabledChannel() (int64, error) { - result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{}) - return result.RowsAffected, result.Error + return deleteChannelsByStatuses([]int64{common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled}) +} + +func deleteChannelsByStatuses(statuses []int64) (int64, error) { + var deletedCount int64 + err := DB.Transaction(func(tx *gorm.DB) error { + var ids []int + if err := tx.Model(&Channel{}).Where("status IN ?", statuses).Order("id ASC").Pluck("id", &ids).Error; err != nil { + return err + } + if len(ids) == 0 { + return nil + } + if err := markChannelContributionsDeletedTx(tx, ids, common.GetTimestamp()); err != nil { + return err + } + result := tx.Where("id IN ? AND status IN ?", ids, statuses).Delete(&Channel{}) + if result.Error != nil { + return result.Error + } + if err := tx.Where("channel_id IN ?", ids).Delete(&Ability{}).Error; err != nil { + return err + } + deletedCount = result.RowsAffected + return nil + }) + return deletedCount, err } func GetPaginatedTags(offset int, limit int) ([]*string, error) { diff --git a/model/channel_cache.go b/model/channel_cache.go index 86c594384d50..9839595f0057 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -5,7 +5,6 @@ import ( "fmt" "math/rand" "sort" - "strings" "sync" "time" @@ -22,6 +21,7 @@ var channelsIDM map[int]*Channel // all channels include dis // path-aware selection avoids re-parsing JSON per request. Refreshed on full sync. var channel2advancedCustomConfig map[int]*dto.AdvancedCustomConfig var channelSyncLock sync.RWMutex +var channelCacheGeneration uint64 func InitChannelCache() { if !common.MemoryCacheEnabled { @@ -40,30 +40,21 @@ func InitChannelCache() { } } } - var abilities []*Ability - DB.Find(&abilities) - groups := make(map[string]bool) - for _, ability := range abilities { - groups[ability.Group] = true - } newGroup2model2channels := make(map[string]map[string][]int) - for group := range groups { - newGroup2model2channels[group] = make(map[string][]int) + var abilities []Ability + if err := DB.Where("enabled = ?", true).Find(&abilities).Error; err != nil { + common.SysError(fmt.Sprintf("failed to load enabled abilities for channel cache: %v", err)) + return } - for _, channel := range channels { - if channel.Status != common.ChannelStatusEnabled { + for _, ability := range abilities { + channel := newChannelId2channel[ability.ChannelId] + if channel == nil || channel.Status != common.ChannelStatusEnabled { continue // skip disabled channels } - groups := strings.Split(channel.Group, ",") - for _, group := range groups { - models := strings.Split(channel.Models, ",") - for _, model := range models { - if _, ok := newGroup2model2channels[group][model]; !ok { - newGroup2model2channels[group][model] = make([]int, 0) - } - newGroup2model2channels[group][model] = append(newGroup2model2channels[group][model], channel.Id) - } + if newGroup2model2channels[ability.Group] == nil { + newGroup2model2channels[ability.Group] = make(map[string][]int) } + newGroup2model2channels[ability.Group][ability.Model] = append(newGroup2model2channels[ability.Group][ability.Model], channel.Id) } // sort by priority @@ -94,6 +85,7 @@ func InitChannelCache() { } channelsIDM = newChannelId2channel channel2advancedCustomConfig = newChannel2advancedCustomConfig + channelCacheGeneration++ channelSyncLock.Unlock() // Lock ordering: InvalidatePricingCache acquires updatePricingLock, and // GetPricing (holding updatePricingLock) nests channelSyncLock.RLock via @@ -272,36 +264,80 @@ func CacheUpdateChannelStatus(id int, status int) { if !common.MemoryCacheEnabled { return } + var abilities []Ability + if status == common.ChannelStatusEnabled { + if err := DB.Where("channel_id = ? AND enabled = ?", id, true).Find(&abilities).Error; err != nil { + common.SysError(fmt.Sprintf("failed to refresh channel routing abilities: channel_id=%d err=%v", id, err)) + return + } + } channelSyncLock.Lock() - defer channelSyncLock.Unlock() if channel, ok := channelsIDM[id]; ok { channel.Status = status + refreshChannelRoutingCacheLocked(channel, abilities) } - if status != common.ChannelStatusEnabled { - // delete the channel from group2model2channels - for group, model2channels := range group2model2channels { - for model, channels := range model2channels { - for i, channelId := range channels { - if channelId == id { - // remove the channel from the slice - group2model2channels[group][model] = append(channels[:i], channels[i+1:]...) - break - } + channelCacheGeneration++ + channelSyncLock.Unlock() + InvalidatePricingCache() +} + +func refreshChannelRoutingCacheLocked(channel *Channel, abilities []Ability) { + if channel == nil { + return + } + if group2model2channels == nil { + group2model2channels = make(map[string]map[string][]int) + } + for group, model2channels := range group2model2channels { + for modelName, channelIds := range model2channels { + filtered := channelIds[:0] + for _, channelId := range channelIds { + if channelId != channel.Id { + filtered = append(filtered, channelId) } } + group2model2channels[group][modelName] = filtered } } + if channel.Status != common.ChannelStatusEnabled { + return + } + + for _, ability := range abilities { + if group2model2channels[ability.Group] == nil { + group2model2channels[ability.Group] = make(map[string][]int) + } + channelIds := append(group2model2channels[ability.Group][ability.Model], channel.Id) + sort.Slice(channelIds, func(i, j int) bool { + left := channelsIDM[channelIds[i]] + right := channelsIDM[channelIds[j]] + if left == nil { + return false + } + if right == nil { + return true + } + return left.GetPriority() > right.GetPriority() + }) + group2model2channels[ability.Group][ability.Model] = channelIds + } } func CacheUpdateChannel(channel *Channel) { if !common.MemoryCacheEnabled { return } - channelSyncLock.Lock() if channel == nil { - channelSyncLock.Unlock() return } + var abilities []Ability + if channel.Status == common.ChannelStatusEnabled { + if err := DB.Where("channel_id = ? AND enabled = ?", channel.Id, true).Find(&abilities).Error; err != nil { + common.SysError(fmt.Sprintf("failed to refresh channel routing abilities: channel_id=%d err=%v", channel.Id, err)) + return + } + } + channelSyncLock.Lock() if channelsIDM == nil { channelsIDM = make(map[int]*Channel) @@ -310,6 +346,7 @@ func CacheUpdateChannel(channel *Channel) { logger.LogDebug(nil, "CacheUpdateChannel before: id=%d, name=%s, status=%d, polling_index=%d", channel.Id, channel.Name, channel.Status, oldChannel.ChannelInfo.MultiKeyPollingIndex) } channelsIDM[channel.Id] = channel + channelCacheGeneration++ if channel2advancedCustomConfig == nil { channel2advancedCustomConfig = make(map[int]*dto.AdvancedCustomConfig) } @@ -319,6 +356,7 @@ func CacheUpdateChannel(channel *Channel) { channel2advancedCustomConfig[channel.Id] = config } } + refreshChannelRoutingCacheLocked(channel, abilities) logger.LogDebug(nil, "CacheUpdateChannel after: id=%d, name=%s, status=%d, polling_index=%d", channel.Id, channel.Name, channel.Status, channel.ChannelInfo.MultiKeyPollingIndex) // Lock ordering: do NOT hold channelSyncLock while calling // InvalidatePricingCache. GetPricing acquires updatePricingLock first and then diff --git a/model/channel_contribution.go b/model/channel_contribution.go new file mode 100644 index 000000000000..016ec6c44c26 --- /dev/null +++ b/model/channel_contribution.go @@ -0,0 +1,907 @@ +package model + +import ( + "errors" + "fmt" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + + "gorm.io/gorm" +) + +type ChannelContributionStatus string + +const ( + ChannelContributionStatusDraft ChannelContributionStatus = "draft" + ChannelContributionStatusPending ChannelContributionStatus = "pending" + ChannelContributionStatusApproved ChannelContributionStatus = "approved" + ChannelContributionStatusRejected ChannelContributionStatus = "rejected" + ChannelContributionStatusUnavailable ChannelContributionStatus = "unavailable" + ChannelContributionStatusDeleted ChannelContributionStatus = "deleted" +) + +type ChannelContributionRevisionStatus string + +const ( + ChannelContributionRevisionStatusDraft ChannelContributionRevisionStatus = "draft" + ChannelContributionRevisionStatusPending ChannelContributionRevisionStatus = "pending" + ChannelContributionRevisionStatusApproved ChannelContributionRevisionStatus = "approved" + ChannelContributionRevisionStatusRejected ChannelContributionRevisionStatus = "rejected" + ChannelContributionRevisionStatusWithdrawn ChannelContributionRevisionStatus = "withdrawn" + ChannelContributionRevisionStatusSuperseded ChannelContributionRevisionStatus = "superseded" +) + +type ChannelContributionTestRunStatus string + +const ( + ChannelContributionTestRunStatusQueued ChannelContributionTestRunStatus = "queued" + ChannelContributionTestRunStatusRunning ChannelContributionTestRunStatus = "running" + ChannelContributionTestRunStatusSucceeded ChannelContributionTestRunStatus = "succeeded" + ChannelContributionTestRunStatusFailed ChannelContributionTestRunStatus = "failed" +) + +const ( + ChannelContributionTestActorUser = "user" + ChannelContributionTestActorAdmin = "admin" +) + +type ChannelContribution struct { + Id int `json:"id"` + UserId int `json:"user_id" gorm:"index;not null"` + Username string `json:"username" gorm:"type:varchar(64);not null"` + Status ChannelContributionStatus `json:"status" gorm:"type:varchar(32);index;not null"` + ChannelId *int `json:"channel_id" gorm:"index"` + CurrentRevisionId *int `json:"current_revision_id" gorm:"index"` + PendingRevisionId *int `json:"pending_revision_id" gorm:"index"` + ApprovedRevisionId *int `json:"approved_revision_id" gorm:"index"` + SubmittedAt int64 `json:"submitted_at" gorm:"bigint;index;not null"` + ReviewerId int `json:"reviewer_id" gorm:"index;not null"` + ReviewerUsername string `json:"reviewer_username" gorm:"type:varchar(64);not null"` + ReviewedAt int64 `json:"reviewed_at" gorm:"bigint;not null"` + ReviewReason string `json:"review_reason" gorm:"type:varchar(500);not null"` + UnavailableSince int64 `json:"unavailable_since" gorm:"bigint;index;not null"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index;not null"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint;index;not null"` +} + +type ChannelContributionRevision struct { + Id int `json:"id"` + ContributionId int `json:"contribution_id" gorm:"uniqueIndex:uk_channel_contribution_revision,priority:1;index;not null"` + RevisionNumber int `json:"revision_number" gorm:"uniqueIndex:uk_channel_contribution_revision,priority:2;not null"` + Name string `json:"name" gorm:"type:varchar(128);not null"` + Type int `json:"type" gorm:"not null"` + BaseURL string `json:"base_url" gorm:"column:base_url;type:text;not null"` + Key string `json:"-" gorm:"not null"` + Group string `json:"group" gorm:"type:varchar(64);not null"` + Models string `json:"models" gorm:"type:text;not null"` + ModelMapping string `json:"model_mapping" gorm:"type:text;not null"` + ConfigHash string `json:"-" gorm:"type:varchar(64);index;not null"` + Status ChannelContributionRevisionStatus `json:"status" gorm:"type:varchar(32);index;not null"` + AgreementVersion string `json:"agreement_version" gorm:"type:varchar(64);not null"` + AgreementContent string `json:"agreement_content" gorm:"type:text;not null"` + AgreementHash string `json:"agreement_hash" gorm:"type:varchar(64);not null"` + AgreementAcceptedAt int64 `json:"agreement_accepted_at" gorm:"bigint;not null"` + SubmittedAt int64 `json:"submitted_at" gorm:"bigint;index;not null"` + ReviewerId int `json:"reviewer_id" gorm:"index;not null"` + ReviewerUsername string `json:"reviewer_username" gorm:"type:varchar(64);not null"` + ReviewedAt int64 `json:"reviewed_at" gorm:"bigint;not null"` + ReviewReason string `json:"review_reason" gorm:"type:varchar(500);not null"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index;not null"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint;index;not null"` +} + +type ChannelContributionTestRun struct { + Id int64 `json:"id"` + ContributionId int `json:"contribution_id" gorm:"index;not null"` + RevisionId int `json:"revision_id" gorm:"index;not null"` + ConfigHash string `json:"-" gorm:"type:varchar(64);index;not null"` + ActorId int `json:"actor_id" gorm:"index;not null"` + ActorType string `json:"actor_type" gorm:"type:varchar(16);not null"` + ActiveUserId *int `json:"-" gorm:"uniqueIndex:uk_channel_contribution_active_user"` + Status ChannelContributionTestRunStatus `json:"status" gorm:"type:varchar(32);index;not null"` + PricingReady bool `json:"pricing_ready" gorm:"not null"` + Total int `json:"total" gorm:"not null"` + Passed int `json:"passed" gorm:"not null"` + Failed int `json:"failed" gorm:"not null"` + Error string `json:"error" gorm:"type:text;not null"` + StartedAt int64 `json:"started_at" gorm:"bigint;not null"` + CompletedAt int64 `json:"completed_at" gorm:"bigint;index;not null"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index;not null"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint;index;not null"` +} + +type ChannelContributionTestResult struct { + Id int64 `json:"id"` + TestRunId int64 `json:"test_run_id" gorm:"index;not null"` + RevisionId int `json:"revision_id" gorm:"index;not null"` + Model string `json:"model" gorm:"type:varchar(255);index;not null"` + EndpointType string `json:"endpoint_type" gorm:"type:varchar(64);not null"` + Stream bool `json:"stream" gorm:"not null"` + Success bool `json:"success" gorm:"index;not null"` + LatencyMs int64 `json:"latency_ms" gorm:"bigint;not null"` + Error string `json:"error" gorm:"type:text;not null"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index;not null"` +} + +type ChannelContributionApproval struct { + ReviewerId int + ReviewerUsername string + Tag string + Priority int64 + Weight uint +} + +func ComputeChannelContributionConfigHash(revision *ChannelContributionRevision) (string, error) { + if revision == nil { + return "", errors.New("channel contribution revision is required") + } + payload := struct { + Name string `json:"name"` + Type int `json:"type"` + BaseURL string `json:"base_url"` + Key string `json:"key"` + Group string `json:"group"` + Models string `json:"models"` + ModelMapping string `json:"model_mapping"` + }{ + Name: revision.Name, + Type: revision.Type, + BaseURL: revision.BaseURL, + Key: revision.Key, + Group: revision.Group, + Models: revision.Models, + ModelMapping: revision.ModelMapping, + } + encoded, err := common.Marshal(payload) + if err != nil { + return "", err + } + key := []byte("channel-contribution-config-v1:" + common.CryptoSecret) + return common.GenerateHMACWithKey(key, string(encoded)), nil +} + +func (contribution *ChannelContribution) BeforeCreate(_ *gorm.DB) error { + now := common.GetTimestamp() + if contribution.CreatedAt == 0 { + contribution.CreatedAt = now + } + if contribution.UpdatedAt == 0 { + contribution.UpdatedAt = now + } + if contribution.Status == "" { + contribution.Status = ChannelContributionStatusDraft + } + return nil +} + +func (revision *ChannelContributionRevision) BeforeCreate(_ *gorm.DB) error { + now := common.GetTimestamp() + if revision.CreatedAt == 0 { + revision.CreatedAt = now + } + if revision.UpdatedAt == 0 { + revision.UpdatedAt = now + } + if revision.Status == "" { + revision.Status = ChannelContributionRevisionStatusDraft + } + if strings.TrimSpace(revision.ModelMapping) == "" { + revision.ModelMapping = "{}" + } + return nil +} + +func (run *ChannelContributionTestRun) BeforeCreate(_ *gorm.DB) error { + now := common.GetTimestamp() + if run.CreatedAt == 0 { + run.CreatedAt = now + } + if run.UpdatedAt == 0 { + run.UpdatedAt = now + } + if run.Status == "" { + run.Status = ChannelContributionTestRunStatusQueued + } + return nil +} + +func CreateChannelContribution(contribution *ChannelContribution) error { + return DB.Create(contribution).Error +} + +func CreateChannelContributionWithRevision(contribution *ChannelContribution, revision *ChannelContributionRevision) error { + if contribution == nil || revision == nil { + return errors.New("contribution and revision are required") + } + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Create(contribution).Error; err != nil { + return err + } + revision.ContributionId = contribution.Id + revision.RevisionNumber = 1 + revision.Status = ChannelContributionRevisionStatusDraft + if err := tx.Create(revision).Error; err != nil { + return err + } + contribution.CurrentRevisionId = &revision.Id + return tx.Model(&ChannelContribution{}). + Where("id = ?", contribution.Id). + Updates(map[string]any{ + "current_revision_id": revision.Id, + "updated_at": common.GetTimestamp(), + }).Error + }) +} + +func CreateChannelContributionRevision(contributionId int, userId int, revision *ChannelContributionRevision) error { + if revision == nil { + return errors.New("revision is required") + } + return DB.Transaction(func(tx *gorm.DB) error { + var contribution ChannelContribution + if err := lockForUpdate(tx).Where("id = ? AND user_id = ?", contributionId, userId).First(&contribution).Error; err != nil { + return err + } + if contribution.PendingRevisionId != nil || contribution.Status == ChannelContributionStatusDeleted { + return errors.New("contribution is not editable") + } + + var latestNumber int + if err := tx.Model(&ChannelContributionRevision{}). + Where("contribution_id = ?", contributionId). + Select("COALESCE(MAX(revision_number), 0)"). + Scan(&latestNumber).Error; err != nil { + return err + } + revision.ContributionId = contributionId + revision.RevisionNumber = latestNumber + 1 + revision.Status = ChannelContributionRevisionStatusDraft + if err := tx.Create(revision).Error; err != nil { + return err + } + + status := contribution.Status + if contribution.ApprovedRevisionId == nil { + status = ChannelContributionStatusDraft + } + return tx.Model(&ChannelContribution{}). + Where("id = ?", contributionId). + Updates(map[string]any{ + "current_revision_id": revision.Id, + "status": status, + "updated_at": common.GetTimestamp(), + }).Error + }) +} + +func GetChannelContributionById(id int) (*ChannelContribution, error) { + var contribution ChannelContribution + if err := DB.First(&contribution, "id = ?", id).Error; err != nil { + return nil, err + } + return &contribution, nil +} + +func GetUserChannelContributionById(id int, userId int) (*ChannelContribution, error) { + var contribution ChannelContribution + if err := DB.Where("id = ? AND user_id = ?", id, userId).First(&contribution).Error; err != nil { + return nil, err + } + return &contribution, nil +} + +func GetChannelContributionRevisionById(id int) (*ChannelContributionRevision, error) { + var revision ChannelContributionRevision + if err := DB.First(&revision, "id = ?", id).Error; err != nil { + return nil, err + } + return &revision, nil +} + +func GetChannelContributionRevision(contributionId int, revisionId int) (*ChannelContributionRevision, error) { + var revision ChannelContributionRevision + if err := DB.Where("id = ? AND contribution_id = ?", revisionId, contributionId).First(&revision).Error; err != nil { + return nil, err + } + return &revision, nil +} + +func ListChannelContributionRevisions(contributionId int) ([]*ChannelContributionRevision, error) { + var revisions []*ChannelContributionRevision + err := DB.Where("contribution_id = ?", contributionId).Order("revision_number desc").Find(&revisions).Error + return revisions, err +} + +func ListUserChannelContributions(userId int, offset int, limit int) ([]*ChannelContribution, int64, error) { + var contributions []*ChannelContribution + var total int64 + query := DB.Model(&ChannelContribution{}).Where("user_id = ?", userId) + if err := query.Count(&total).Error; err != nil { + return nil, 0, err + } + if err := query.Order("id desc").Offset(offset).Limit(limit).Find(&contributions).Error; err != nil { + return nil, 0, err + } + return contributions, total, nil +} + +func ListChannelContributions(status ChannelContributionStatus, offset int, limit int) ([]*ChannelContribution, int64, error) { + var contributions []*ChannelContribution + var total int64 + query := DB.Model(&ChannelContribution{}) + if status == ChannelContributionStatusPending { + query = query.Where("pending_revision_id IS NOT NULL") + } else if status != "" { + query = query.Where("status = ?", status) + } + if err := query.Count(&total).Error; err != nil { + return nil, 0, err + } + if err := query.Order("id desc").Offset(offset).Limit(limit).Find(&contributions).Error; err != nil { + return nil, 0, err + } + return contributions, total, nil +} + +func CreateChannelContributionTestRun(run *ChannelContributionTestRun) error { + if run == nil { + return errors.New("test run is required") + } + if run.ActorId <= 0 { + return errors.New("test run actor is required") + } + return DB.Transaction(func(tx *gorm.DB) error { + var contribution ChannelContribution + if err := lockForUpdate(tx). + Select("id", "user_id"). + Where("id = ?", run.ContributionId). + First(&contribution).Error; err != nil { + return err + } + var revisionCount int64 + if err := tx.Model(&ChannelContributionRevision{}). + Where("id = ? AND contribution_id = ? AND config_hash = ?", run.RevisionId, run.ContributionId, run.ConfigHash). + Count(&revisionCount).Error; err != nil { + return err + } + if revisionCount != 1 { + return errors.New("test run revision does not match contribution configuration") + } + run.ActiveUserId = &run.ActorId + return tx.Create(run).Error + }) +} + +func GetChannelContributionTestRun(id int64) (*ChannelContributionTestRun, error) { + var run ChannelContributionTestRun + if err := DB.First(&run, "id = ?", id).Error; err != nil { + return nil, err + } + return &run, nil +} + +func GetChannelContributionTestRunForContribution(id int64, contributionId int) (*ChannelContributionTestRun, error) { + var run ChannelContributionTestRun + if err := DB.Where("id = ? AND contribution_id = ?", id, contributionId).First(&run).Error; err != nil { + return nil, err + } + return &run, nil +} + +func GetLatestChannelContributionTestRun(revisionId int, configHash string) (*ChannelContributionTestRun, error) { + var run ChannelContributionTestRun + err := DB.Where("revision_id = ? AND config_hash = ?", revisionId, configHash). + Order("id desc").First(&run).Error + if err != nil { + return nil, err + } + return &run, nil +} + +func GetLatestSuccessfulChannelContributionTestRun(revisionId int, configHash string) (*ChannelContributionTestRun, error) { + var run ChannelContributionTestRun + err := DB.Where("revision_id = ? AND config_hash = ? AND status = ?", revisionId, configHash, ChannelContributionTestRunStatusSucceeded). + Order("completed_at desc, id desc").First(&run).Error + if err != nil { + return nil, err + } + return &run, nil +} + +func ListChannelContributionTestResults(runId int64) ([]*ChannelContributionTestResult, error) { + var results []*ChannelContributionTestResult + err := DB.Where("test_run_id = ?", runId).Order("id asc").Find(&results).Error + return results, err +} + +func HasUnfinishedChannelContributionTestRuns() bool { + var count int64 + if err := DB.Model(&ChannelContributionTestRun{}). + Where("status IN ?", []ChannelContributionTestRunStatus{ + ChannelContributionTestRunStatusQueued, + ChannelContributionTestRunStatusRunning, + }). + Limit(1). + Count(&count).Error; err != nil { + return false + } + return count > 0 +} + +func ClaimNextQueuedChannelContributionTestRun() (*ChannelContributionTestRun, error) { + var claimed ChannelContributionTestRun + err := DB.Transaction(func(tx *gorm.DB) error { + var run ChannelContributionTestRun + if err := lockForUpdate(tx). + Where("status = ?", ChannelContributionTestRunStatusQueued). + Order("id asc"). + First(&run).Error; err != nil { + return err + } + now := common.GetTimestamp() + result := tx.Model(&ChannelContributionTestRun{}). + Where("id = ? AND status = ?", run.Id, ChannelContributionTestRunStatusQueued). + Updates(map[string]any{ + "status": ChannelContributionTestRunStatusRunning, + "started_at": now, + "updated_at": now, + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return gorm.ErrRecordNotFound + } + run.Status = ChannelContributionTestRunStatusRunning + run.StartedAt = now + run.UpdatedAt = now + claimed = run + return nil + }) + if err != nil { + return nil, err + } + return &claimed, nil +} + +func FinishChannelContributionTestRun(runId int64, status ChannelContributionTestRunStatus, pricingReady bool, results []ChannelContributionTestResult, runError string) error { + if status != ChannelContributionTestRunStatusSucceeded && status != ChannelContributionTestRunStatusFailed { + return errors.New("test run terminal status is invalid") + } + return DB.Transaction(func(tx *gorm.DB) error { + var run ChannelContributionTestRun + if err := lockForUpdate(tx).Where("id = ?", runId).First(&run).Error; err != nil { + return err + } + if run.Status != ChannelContributionTestRunStatusRunning { + return errors.New("test run is not running") + } + now := common.GetTimestamp() + passed := 0 + failed := 0 + for index := range results { + results[index].TestRunId = run.Id + results[index].RevisionId = run.RevisionId + if results[index].CreatedAt == 0 { + results[index].CreatedAt = now + } + if results[index].Success { + passed++ + } else { + failed++ + } + } + if len(results) > 0 { + if err := tx.Create(&results).Error; err != nil { + return err + } + } + return tx.Model(&ChannelContributionTestRun{}). + Where("id = ? AND status = ?", run.Id, ChannelContributionTestRunStatusRunning). + Updates(map[string]any{ + "status": status, + "active_user_id": nil, + "pricing_ready": pricingReady, + "total": len(results), + "passed": passed, + "failed": failed, + "error": runError, + "completed_at": now, + "updated_at": now, + }).Error + }) +} + +func RequeueRunningChannelContributionTestRuns() (int64, error) { + result := DB.Model(&ChannelContributionTestRun{}). + Where("status = ?", ChannelContributionTestRunStatusRunning). + Updates(map[string]any{ + "status": ChannelContributionTestRunStatusQueued, + "started_at": int64(0), + "updated_at": common.GetTimestamp(), + }) + return result.RowsAffected, result.Error +} + +func SubmitChannelContribution(id int, userId int, revisionId int, configHash string, agreementVersion string, agreementContent string, agreementHash string, acceptedAt int64) error { + return DB.Transaction(func(tx *gorm.DB) error { + var contribution ChannelContribution + if err := lockForUpdate(tx).Where("id = ? AND user_id = ?", id, userId).First(&contribution).Error; err != nil { + return err + } + if contribution.PendingRevisionId != nil || contribution.Status == ChannelContributionStatusDeleted { + return errors.New("contribution is not ready for submission") + } + if contribution.CurrentRevisionId == nil || *contribution.CurrentRevisionId != revisionId { + return errors.New("contribution revision changed before submission") + } + var revision ChannelContributionRevision + if err := lockForUpdate(tx).Where("id = ? AND contribution_id = ?", revisionId, id).First(&revision).Error; err != nil { + return err + } + if revision.ConfigHash != configHash { + return errors.New("contribution configuration changed before submission") + } + computedConfigHash, err := ComputeChannelContributionConfigHash(&revision) + if err != nil { + return err + } + if computedConfigHash != revision.ConfigHash { + return errors.New("contribution configuration fingerprint is stale") + } + switch revision.Status { + case ChannelContributionRevisionStatusDraft, ChannelContributionRevisionStatusRejected, ChannelContributionRevisionStatusWithdrawn: + default: + return errors.New("contribution revision is not ready for submission") + } + if err := tx.Model(&ChannelContributionRevision{}). + Where("id = ?", revision.Id). + Updates(map[string]any{ + "status": ChannelContributionRevisionStatusPending, + "agreement_version": agreementVersion, + "agreement_content": agreementContent, + "agreement_hash": agreementHash, + "agreement_accepted_at": acceptedAt, + "submitted_at": acceptedAt, + "reviewer_id": 0, + "reviewer_username": "", + "reviewed_at": int64(0), + "review_reason": "", + "updated_at": acceptedAt, + }).Error; err != nil { + return err + } + status := contribution.Status + if contribution.ApprovedRevisionId == nil { + status = ChannelContributionStatusPending + } + return tx.Model(&ChannelContribution{}). + Where("id = ?", id). + Updates(map[string]any{ + "status": status, + "pending_revision_id": revision.Id, + "submitted_at": acceptedAt, + "reviewer_id": 0, + "reviewer_username": "", + "reviewed_at": int64(0), + "review_reason": "", + "updated_at": acceptedAt, + }).Error + }) +} + +func WithdrawChannelContribution(id int, userId int) error { + deletedChannel := false + err := DB.Transaction(func(tx *gorm.DB) error { + var contribution ChannelContribution + if err := lockForUpdate(tx).Where("id = ? AND user_id = ?", id, userId).First(&contribution).Error; err != nil { + return err + } + if contribution.Status == ChannelContributionStatusDeleted { + return nil + } + now := common.GetTimestamp() + if contribution.PendingRevisionId != nil { + if err := tx.Model(&ChannelContributionRevision{}). + Where("id = ? AND status = ?", *contribution.PendingRevisionId, ChannelContributionRevisionStatusPending). + Updates(map[string]any{ + "status": ChannelContributionRevisionStatusWithdrawn, + "updated_at": now, + }).Error; err != nil { + return err + } + } + if contribution.ChannelId != nil { + if err := tx.Where("channel_id = ?", *contribution.ChannelId).Delete(&Ability{}).Error; err != nil { + return err + } + if err := tx.Where("id = ?", *contribution.ChannelId).Delete(&Channel{}).Error; err != nil { + return err + } + deletedChannel = true + } + if err := tx.Model(&ChannelContributionRevision{}). + Where("contribution_id = ?", id). + Updates(map[string]any{"key": "", "updated_at": now}).Error; err != nil { + return err + } + return tx.Model(&ChannelContribution{}). + Where("id = ?", id). + Updates(map[string]any{ + "status": ChannelContributionStatusDeleted, + "pending_revision_id": nil, + "updated_at": now, + }).Error + }) + if err == nil && deletedChannel { + InitChannelCache() + } + return err +} + +func ApproveChannelContribution(id int, revisionId int, approval ChannelContributionApproval) (*ChannelContribution, *Channel, error) { + var approvedContribution ChannelContribution + var approvedChannel Channel + err := DB.Transaction(func(tx *gorm.DB) error { + var contribution ChannelContribution + if err := lockForUpdate(tx).Where("id = ?", id).First(&contribution).Error; err != nil { + return err + } + if contribution.PendingRevisionId == nil || *contribution.PendingRevisionId != revisionId { + return errors.New("contribution is not pending this revision") + } + var revision ChannelContributionRevision + if err := lockForUpdate(tx). + Where("id = ? AND contribution_id = ? AND status = ?", revisionId, id, ChannelContributionRevisionStatusPending). + First(&revision).Error; err != nil { + return err + } + computedConfigHash, err := ComputeChannelContributionConfigHash(&revision) + if err != nil { + return err + } + if computedConfigHash != revision.ConfigHash { + return errors.New("contribution configuration fingerprint is stale") + } + + baseURL := strings.TrimSpace(revision.BaseURL) + mapping := revision.ModelMapping + tag := strings.TrimSpace(approval.Tag) + priority := approval.Priority + weight := approval.Weight + remark := fmt.Sprintf("贡献者：%d %s", contribution.UserId, contribution.Username) + + if contribution.ChannelId == nil { + channel := Channel{ + Type: revision.Type, + Key: revision.Key, + Status: common.ChannelStatusEnabled, + Name: revision.Name, + Weight: &weight, + CreatedTime: common.GetTimestamp(), + BaseURL: &baseURL, + Models: revision.Models, + Group: revision.Group, + ModelMapping: &mapping, + Priority: &priority, + Tag: &tag, + Remark: &remark, + } + if err := tx.Create(&channel).Error; err != nil { + return err + } + contribution.ChannelId = &channel.Id + if err := ResetChannelContributionHealthForRevision(tx, &contribution, revision.Id, revision.ConfigHash); err != nil { + return err + } + if err := channel.AddAbilities(tx); err != nil { + return err + } + approvedChannel = channel + } else { + if err := lockForUpdate(tx).Where("id = ?", *contribution.ChannelId).First(&approvedChannel).Error; err != nil { + return err + } + approvedChannel.Type = revision.Type + approvedChannel.Key = revision.Key + if approvedChannel.Status != common.ChannelStatusManuallyDisabled { + approvedChannel.Status = common.ChannelStatusEnabled + } + approvedChannel.Name = revision.Name + approvedChannel.BaseURL = &baseURL + approvedChannel.Models = revision.Models + approvedChannel.Group = revision.Group + approvedChannel.ModelMapping = &mapping + approvedChannel.Remark = &remark + if err := tx.Model(&Channel{}). + Where("id = ?", approvedChannel.Id). + Select("type", "key", "status", "name", "base_url", "models", "group", "model_mapping", "remark"). + Updates(&approvedChannel).Error; err != nil { + return err + } + if err := ResetChannelContributionHealthForRevision(tx, &contribution, revision.Id, revision.ConfigHash); err != nil { + return err + } + if err := approvedChannel.UpdateAbilities(tx); err != nil { + return err + } + } + + now := common.GetTimestamp() + if contribution.ApprovedRevisionId != nil && *contribution.ApprovedRevisionId != revision.Id { + if err := tx.Model(&ChannelContributionRevision{}). + Where("id = ? AND status = ?", *contribution.ApprovedRevisionId, ChannelContributionRevisionStatusApproved). + Updates(map[string]any{ + "status": ChannelContributionRevisionStatusSuperseded, + "updated_at": now, + }).Error; err != nil { + return err + } + } + if err := tx.Model(&ChannelContributionRevision{}). + Where("id = ?", revision.Id). + Updates(map[string]any{ + "status": ChannelContributionRevisionStatusApproved, + "reviewer_id": approval.ReviewerId, + "reviewer_username": approval.ReviewerUsername, + "reviewed_at": now, + "review_reason": "", + "updated_at": now, + }).Error; err != nil { + return err + } + if err := tx.Model(&ChannelContribution{}). + Where("id = ?", id). + Updates(map[string]any{ + "status": ChannelContributionStatusApproved, + "channel_id": approvedChannel.Id, + "current_revision_id": revision.Id, + "pending_revision_id": nil, + "approved_revision_id": revision.Id, + "reviewer_id": approval.ReviewerId, + "reviewer_username": approval.ReviewerUsername, + "reviewed_at": now, + "review_reason": "", + "unavailable_since": int64(0), + "updated_at": now, + }).Error; err != nil { + return err + } + contribution.Status = ChannelContributionStatusApproved + contribution.CurrentRevisionId = &revision.Id + contribution.PendingRevisionId = nil + contribution.ApprovedRevisionId = &revision.Id + contribution.ReviewerId = approval.ReviewerId + contribution.ReviewerUsername = approval.ReviewerUsername + contribution.ReviewedAt = now + contribution.ReviewReason = "" + contribution.UnavailableSince = 0 + contribution.UpdatedAt = now + approvedContribution = contribution + return nil + }) + if err != nil { + return nil, nil, err + } + InitChannelCache() + return &approvedContribution, &approvedChannel, nil +} + +func RejectChannelContribution(id int, revisionId int, reviewerId int, reviewerUsername string, reason string) error { + return DB.Transaction(func(tx *gorm.DB) error { + var contribution ChannelContribution + if err := lockForUpdate(tx).Where("id = ?", id).First(&contribution).Error; err != nil { + return err + } + if contribution.PendingRevisionId == nil || *contribution.PendingRevisionId != revisionId { + return errors.New("contribution is not pending this revision") + } + now := common.GetTimestamp() + if err := tx.Model(&ChannelContributionRevision{}). + Where("id = ? AND status = ?", revisionId, ChannelContributionRevisionStatusPending). + Updates(map[string]any{ + "status": ChannelContributionRevisionStatusRejected, + "reviewer_id": reviewerId, + "reviewer_username": reviewerUsername, + "reviewed_at": now, + "review_reason": reason, + "updated_at": now, + }).Error; err != nil { + return err + } + status := contribution.Status + if contribution.ApprovedRevisionId == nil { + status = ChannelContributionStatusRejected + } + return tx.Model(&ChannelContribution{}). + Where("id = ?", id). + Updates(map[string]any{ + "status": status, + "pending_revision_id": nil, + "reviewer_id": reviewerId, + "reviewer_username": reviewerUsername, + "reviewed_at": now, + "review_reason": reason, + "updated_at": now, + }).Error + }) +} + +func DeleteChannelContribution(id int, reviewerId int, reviewerUsername string, reason string) error { + deletedChannel := false + err := DB.Transaction(func(tx *gorm.DB) error { + var contribution ChannelContribution + if err := lockForUpdate(tx).Where("id = ?", id).First(&contribution).Error; err != nil { + return err + } + if contribution.Status == ChannelContributionStatusDeleted { + return nil + } + if contribution.ChannelId != nil { + if err := tx.Where("channel_id = ?", *contribution.ChannelId).Delete(&Ability{}).Error; err != nil { + return err + } + if err := tx.Where("id = ?", *contribution.ChannelId).Delete(&Channel{}).Error; err != nil { + return err + } + deletedChannel = true + } + now := common.GetTimestamp() + if contribution.PendingRevisionId != nil { + if err := tx.Model(&ChannelContributionRevision{}). + Where("id = ? AND status = ?", *contribution.PendingRevisionId, ChannelContributionRevisionStatusPending). + Updates(map[string]any{"status": ChannelContributionRevisionStatusWithdrawn, "updated_at": now}).Error; err != nil { + return err + } + } + if err := tx.Model(&ChannelContributionRevision{}). + Where("contribution_id = ?", id). + Updates(map[string]any{"key": "", "updated_at": now}).Error; err != nil { + return err + } + return tx.Model(&ChannelContribution{}). + Where("id = ?", id). + Updates(map[string]any{ + "status": ChannelContributionStatusDeleted, + "pending_revision_id": nil, + "reviewer_id": reviewerId, + "reviewer_username": reviewerUsername, + "reviewed_at": now, + "review_reason": reason, + "updated_at": now, + }).Error + }) + if err == nil && deletedChannel { + InitChannelCache() + } + return err +} + +func IsValidChannelContributionStatus(status ChannelContributionStatus) bool { + switch status { + case ChannelContributionStatusDraft, + ChannelContributionStatusPending, + ChannelContributionStatusApproved, + ChannelContributionStatusRejected, + ChannelContributionStatusUnavailable, + ChannelContributionStatusDeleted: + return true + default: + return false + } +} + +func ValidateContributionChannelType(channelType int) error { + if channelType <= 0 { + return errors.New("channel type is required") + } + if channelType >= len(constant.ChannelBaseURLs) { + return fmt.Errorf("unsupported channel type %d", channelType) + } + if _, ok := constant.ChannelTypeNames[channelType]; !ok { + return fmt.Errorf("unsupported channel type %d", channelType) + } + return nil +} diff --git a/model/channel_contribution_health.go b/model/channel_contribution_health.go new file mode 100644 index 000000000000..9c41e0b0b097 --- /dev/null +++ b/model/channel_contribution_health.go @@ -0,0 +1,765 @@ +package model + +import ( + "errors" + "fmt" + "strings" + "unicode/utf8" + + "github.com/QuantumNous/new-api/common" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +type ChannelContributionModelHealth struct { + Id int64 `json:"id" gorm:"primaryKey"` + ContributionId int `json:"contribution_id" gorm:"uniqueIndex:uk_contribution_model_health,priority:1;index;not null"` + RevisionId int `json:"revision_id" gorm:"index;not null"` + ConfigHash string `json:"-" gorm:"type:varchar(64);index;not null"` + ChannelId int `json:"channel_id" gorm:"index;not null"` + Model string `json:"model" gorm:"type:varchar(255);uniqueIndex:uk_contribution_model_health,priority:2;not null"` + Healthy bool `json:"healthy" gorm:"index;not null"` + FailureSince int64 `json:"failure_since" gorm:"bigint;index;not null"` + LastCheckedAt int64 `json:"last_checked_at" gorm:"bigint;index;not null"` + LastSuccessAt int64 `json:"last_success_at" gorm:"bigint;not null"` + LastFailureAt int64 `json:"last_failure_at" gorm:"bigint;not null"` + LastError string `json:"last_error" gorm:"type:varchar(500);not null"` + CreatedAt int64 `json:"created_at" gorm:"bigint;not null"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint;index;not null"` +} + +type ChannelContributionHealthState struct { + ContributionId int `json:"contribution_id" gorm:"primaryKey;autoIncrement:false"` + ChannelId int `json:"channel_id" gorm:"uniqueIndex;not null"` + FailureSince int64 `json:"failure_since" gorm:"bigint;index;not null"` + PausedAt int64 `json:"paused_at" gorm:"bigint;index;not null"` + LastCheckedAt int64 `json:"last_checked_at" gorm:"bigint;index;not null"` + CreatedAt int64 `json:"created_at" gorm:"bigint;not null"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint;index;not null"` +} + +type ChannelContributionHealthCandidate struct { + ContributionId int `json:"contribution_id"` + ChannelId int `json:"channel_id"` + UserId int `json:"user_id"` + RevisionId int `json:"revision_id"` + ConfigHash string `json:"-"` + Status ChannelContributionStatus `json:"status"` +} + +type ChannelContributionModelObservation struct { + Model string `json:"model"` + Healthy bool `json:"healthy"` + Error string `json:"error,omitempty"` +} + +type ChannelContributionHealthCycleResult struct { + AllFailed bool `json:"all_failed"` + Paused bool `json:"paused"` + Deleted bool `json:"deleted"` + StateChanged bool `json:"state_changed"` + UnhealthyModels []string `json:"unhealthy_models"` +} + +var ErrStaleChannelContributionHealthProbe = errors.New("stale channel contribution health probe") + +func (health *ChannelContributionModelHealth) BeforeCreate(_ *gorm.DB) error { + now := common.GetTimestamp() + if health.CreatedAt == 0 { + health.CreatedAt = now + } + if health.UpdatedAt == 0 { + health.UpdatedAt = now + } + return nil +} + +func (state *ChannelContributionHealthState) BeforeCreate(_ *gorm.DB) error { + now := common.GetTimestamp() + if state.CreatedAt == 0 { + state.CreatedAt = now + } + if state.UpdatedAt == 0 { + state.UpdatedAt = now + } + return nil +} + +func ListContributionChannelsForHealthAfter(afterContributionId int, limit int) ([]ChannelContributionHealthCandidate, error) { + if limit <= 0 { + limit = 100 + } + if limit > 500 { + limit = 500 + } + var candidates []ChannelContributionHealthCandidate + err := DB.Table("channel_contributions AS contributions"). + Select("contributions.id AS contribution_id, contributions.user_id, contributions.status, contributions.channel_id, revisions.id AS revision_id, revisions.config_hash"). + Joins("JOIN channel_contribution_revisions AS revisions ON revisions.id = contributions.approved_revision_id"). + Where("contributions.id > ? AND contributions.channel_id IS NOT NULL AND contributions.status IN ? AND revisions.status = ?", afterContributionId, []ChannelContributionStatus{ + ChannelContributionStatusApproved, + ChannelContributionStatusUnavailable, + }, ChannelContributionRevisionStatusApproved). + Order("contributions.id asc"). + Limit(limit). + Scan(&candidates).Error + if err != nil { + return nil, err + } + return candidates, nil +} + +func GetChannelContributionModelHealth(contributionId int) ([]ChannelContributionModelHealth, error) { + var rows []ChannelContributionModelHealth + err := DB.Where("contribution_id = ?", contributionId).Order("model asc").Find(&rows).Error + return rows, err +} + +func IsContributionChannel(channelId int) bool { + if channelId <= 0 { + return false + } + var count int64 + if err := DB.Model(&ChannelContribution{}). + Where("channel_id = ? AND status IN ?", channelId, []ChannelContributionStatus{ + ChannelContributionStatusApproved, + ChannelContributionStatusUnavailable, + }).Count(&count).Error; err != nil { + return false + } + return count > 0 +} + +func PopulateContributionChannelFlags(channels []*Channel) error { + channelIds := make([]int, 0, len(channels)) + seen := make(map[int]struct{}, len(channels)) + for _, channel := range channels { + if channel == nil { + continue + } + channel.IsContribution = false + if channel.Id <= 0 { + continue + } + if _, exists := seen[channel.Id]; exists { + continue + } + seen[channel.Id] = struct{}{} + channelIds = append(channelIds, channel.Id) + } + if len(channelIds) == 0 { + return nil + } + + var contributionChannelIds []int + if err := DB.Model(&ChannelContribution{}). + Where("channel_id IN ? AND status IN ?", channelIds, []ChannelContributionStatus{ + ChannelContributionStatusApproved, + ChannelContributionStatusUnavailable, + }). + Distinct(). + Pluck("channel_id", &contributionChannelIds).Error; err != nil { + return err + } + contributionSet := make(map[int]struct{}, len(contributionChannelIds)) + for _, channelId := range contributionChannelIds { + contributionSet[channelId] = struct{}{} + } + for _, channel := range channels { + if channel == nil { + continue + } + _, channel.IsContribution = contributionSet[channel.Id] + } + return nil +} + +func lockActiveChannelContributionTx(tx *gorm.DB, channelId int) (*ChannelContribution, error) { + contributions, err := lockActiveChannelContributionsTx(tx, []int{channelId}) + if err != nil { + return nil, err + } + return contributions[channelId], nil +} + +func lockActiveChannelContributionsTx(tx *gorm.DB, channelIds []int) (map[int]*ChannelContribution, error) { + result := make(map[int]*ChannelContribution) + if tx == nil || len(channelIds) == 0 || !tx.Migrator().HasTable(&ChannelContribution{}) { + return result, nil + } + var contributions []ChannelContribution + if err := lockForUpdate(tx). + Where("channel_id IN ? AND status IN ?", channelIds, []ChannelContributionStatus{ + ChannelContributionStatusApproved, + ChannelContributionStatusUnavailable, + }). + Order("id ASC"). + Find(&contributions).Error; err != nil { + return nil, err + } + for index := range contributions { + contribution := &contributions[index] + if contribution.ChannelId != nil { + result[*contribution.ChannelId] = contribution + } + } + return result, nil +} + +func FilterNonContributionChannels(channels []*Channel) []*Channel { + if len(channels) == 0 { + return channels + } + ids := make([]int, 0, len(channels)) + for _, channel := range channels { + if channel != nil && channel.Id > 0 { + ids = append(ids, channel.Id) + } + } + var contributionChannelIds []int + if err := DB.Model(&ChannelContribution{}). + Where("channel_id IN ? AND status IN ?", ids, []ChannelContributionStatus{ + ChannelContributionStatusApproved, + ChannelContributionStatusUnavailable, + }).Pluck("channel_id", &contributionChannelIds).Error; err != nil { + return channels + } + excluded := make(map[int]struct{}, len(contributionChannelIds)) + for _, id := range contributionChannelIds { + excluded[id] = struct{}{} + } + filtered := make([]*Channel, 0, len(channels)-len(excluded)) + for _, channel := range channels { + if channel == nil { + continue + } + if _, exists := excluded[channel.Id]; !exists { + filtered = append(filtered, channel) + } + } + return filtered +} + +func setContributionHealthPausedTx(tx *gorm.DB, channelId int, paused bool, now int64) error { + contribution, err := lockActiveChannelContributionTx(tx, channelId) + if err != nil { + return err + } + return setLockedContributionHealthPausedTx(tx, contribution, channelId, paused, now) +} + +func setLockedContributionHealthPausedTx(tx *gorm.DB, contribution *ChannelContribution, channelId int, paused bool, now int64) error { + if contribution == nil || !tx.Migrator().HasTable(&ChannelContributionHealthState{}) { + return nil + } + if contribution.ChannelId == nil || *contribution.ChannelId != channelId { + return errors.New("locked channel contribution does not match channel") + } + state := ChannelContributionHealthState{ContributionId: contribution.Id, ChannelId: channelId} + if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&state).Error; err != nil { + return err + } + if err := lockForUpdate(tx).Where("contribution_id = ?", contribution.Id).First(&state).Error; err != nil { + return err + } + if paused { + if state.PausedAt != 0 { + return nil + } + return tx.Model(&ChannelContributionHealthState{}). + Where("contribution_id = ?", contribution.Id). + Updates(map[string]any{"paused_at": now, "updated_at": now}).Error + } + if state.PausedAt == 0 { + return nil + } + pauseDuration := now - state.PausedAt + if pauseDuration < 0 { + pauseDuration = 0 + } + updates := map[string]any{"paused_at": int64(0), "updated_at": now} + if state.FailureSince > 0 && pauseDuration > 0 { + updates["failure_since"] = state.FailureSince + pauseDuration + if err := tx.Model(&ChannelContribution{}). + Where("id = ? AND unavailable_since > 0", contribution.Id). + Update("unavailable_since", gorm.Expr("unavailable_since + ?", pauseDuration)).Error; err != nil { + return err + } + } + if pauseDuration > 0 && tx.Migrator().HasTable(&ChannelContributionModelHealth{}) { + if err := tx.Model(&ChannelContributionModelHealth{}). + Where("contribution_id = ? AND failure_since > 0", contribution.Id). + Update("failure_since", gorm.Expr("failure_since + ?", pauseDuration)).Error; err != nil { + return err + } + } + return tx.Model(&ChannelContributionHealthState{}). + Where("contribution_id = ?", contribution.Id). + Updates(updates).Error +} + +func PauseContributionHealthForChannel(channelId int, now int64) error { + return DB.Transaction(func(tx *gorm.DB) error { + return setContributionHealthPausedTx(tx, channelId, true, now) + }) +} + +func ResumeContributionHealthForChannel(channelId int, now int64) error { + return DB.Transaction(func(tx *gorm.DB) error { + return setContributionHealthPausedTx(tx, channelId, false, now) + }) +} + +func ApplyChannelContributionHealthCycle( + contributionId int, + channelId int, + revisionId int, + configHash string, + observations []ChannelContributionModelObservation, + now int64, + deleteAfterSeconds int64, +) (ChannelContributionHealthCycleResult, error) { + result := ChannelContributionHealthCycleResult{} + err := DB.Transaction(func(tx *gorm.DB) error { + var contribution ChannelContribution + if err := lockForUpdate(tx). + Where("id = ? AND channel_id = ? AND approved_revision_id = ? AND status IN ?", contributionId, channelId, revisionId, []ChannelContributionStatus{ + ChannelContributionStatusApproved, + ChannelContributionStatusUnavailable, + }).First(&contribution).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return ErrStaleChannelContributionHealthProbe + } + return err + } + var revision ChannelContributionRevision + if err := tx.Select("id", "config_hash", "status"). + Where("id = ? AND contribution_id = ? AND config_hash = ? AND status = ?", revisionId, contributionId, configHash, ChannelContributionRevisionStatusApproved). + First(&revision).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return ErrStaleChannelContributionHealthProbe + } + return err + } + var channel Channel + if err := lockForUpdate(tx).Where("id = ?", channelId).First(&channel).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + if err := markLockedChannelContributionsDeletedTx(tx, []ChannelContribution{contribution}, now); err != nil { + return err + } + result.Deleted = true + result.StateChanged = true + return nil + } + return err + } + + if channel.Status == common.ChannelStatusManuallyDisabled { + if err := setLockedContributionHealthPausedTx(tx, &contribution, channelId, true, now); err != nil { + return err + } + result.Paused = true + return nil + } + if err := setLockedContributionHealthPausedTx(tx, &contribution, channelId, false, now); err != nil { + return err + } + + models := normalizeContributionHealthModels(channel.Models) + observationByModel := make(map[string]ChannelContributionModelObservation, len(observations)) + for _, observation := range observations { + modelName := strings.TrimSpace(observation.Model) + if modelName == "" { + continue + } + observation.Model = modelName + observationByModel[modelName] = observation + } + if len(models) == 0 || len(observationByModel) != len(models) { + return errors.New("health observations must cover every configured model") + } + + allFailed := true + modelHealthChanged := false + unhealthyModels := make([]string, 0) + for _, modelName := range models { + observation, exists := observationByModel[modelName] + if !exists { + return fmt.Errorf("missing health observation for model %s", modelName) + } + if observation.Healthy { + allFailed = false + } else { + unhealthyModels = append(unhealthyModels, modelName) + } + changed, err := upsertContributionModelHealthTx(tx, contributionId, channelId, revisionId, configHash, observation, now) + if err != nil { + return err + } + modelHealthChanged = modelHealthChanged || changed + } + if err := tx.Where("contribution_id = ? AND model NOT IN ?", contributionId, models). + Delete(&ChannelContributionModelHealth{}).Error; err != nil { + return err + } + + state := ChannelContributionHealthState{ContributionId: contributionId, ChannelId: channelId} + if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&state).Error; err != nil { + return err + } + if err := lockForUpdate(tx).Where("contribution_id = ?", contributionId).First(&state).Error; err != nil { + return err + } + + if allFailed { + if state.FailureSince == 0 { + state.FailureSince = now + result.StateChanged = true + } + if contribution.Status != ChannelContributionStatusUnavailable || contribution.UnavailableSince != state.FailureSince { + result.StateChanged = true + } + if err := tx.Model(&ChannelContribution{}).Where("id = ?", contributionId).Updates(map[string]any{ + "status": ChannelContributionStatusUnavailable, + "unavailable_since": state.FailureSince, + "updated_at": now, + }).Error; err != nil { + return err + } + if channel.Status == common.ChannelStatusEnabled { + channel.Status = common.ChannelStatusAutoDisabled + result.StateChanged = true + if err := tx.Model(&Channel{}).Where("id = ?", channelId).Update("status", channel.Status).Error; err != nil { + return err + } + } + abilityResult := tx.Model(&Ability{}). + Where("channel_id = ? AND enabled = ?", channelId, true). + Update("enabled", false) + if abilityResult.Error != nil { + return abilityResult.Error + } + if abilityResult.RowsAffected > 0 { + result.StateChanged = true + } + if deleteAfterSeconds > 0 && now-state.FailureSince >= deleteAfterSeconds { + if err := tx.Where("channel_id = ?", channelId).Delete(&Ability{}).Error; err != nil { + return err + } + if err := tx.Where("id = ?", channelId).Delete(&Channel{}).Error; err != nil { + return err + } + if err := markLockedChannelContributionsDeletedTx(tx, []ChannelContribution{contribution}, now); err != nil { + return err + } + result.Deleted = true + result.StateChanged = true + } + } else { + if state.FailureSince != 0 || contribution.Status != ChannelContributionStatusApproved || contribution.UnavailableSince != 0 { + result.StateChanged = true + } + state.FailureSince = 0 + if err := tx.Model(&ChannelContribution{}).Where("id = ?", contributionId).Updates(map[string]any{ + "status": ChannelContributionStatusApproved, + "unavailable_since": int64(0), + "updated_at": now, + }).Error; err != nil { + return err + } + if channel.Status == common.ChannelStatusAutoDisabled { + channel.Status = common.ChannelStatusEnabled + result.StateChanged = true + if err := tx.Model(&Channel{}).Where("id = ?", channelId).Update("status", channel.Status).Error; err != nil { + return err + } + } + for _, modelName := range models { + observation := observationByModel[modelName] + abilityResult := tx.Model(&Ability{}). + Where("channel_id = ? AND model = ? AND enabled <> ?", channelId, modelName, observation.Healthy). + Update("enabled", observation.Healthy) + if abilityResult.Error != nil { + return abilityResult.Error + } + if abilityResult.RowsAffected > 0 { + result.StateChanged = true + } + } + } + if !result.Deleted { + if err := tx.Model(&ChannelContributionHealthState{}). + Where("contribution_id = ?", contributionId). + Updates(map[string]any{ + "failure_since": state.FailureSince, + "paused_at": int64(0), + "last_checked_at": now, + "updated_at": now, + }).Error; err != nil { + return err + } + } + result.AllFailed = allFailed + result.StateChanged = result.StateChanged || modelHealthChanged + result.UnhealthyModels = unhealthyModels + return nil + }) + return result, err +} + +func upsertContributionModelHealthTx( + tx *gorm.DB, + contributionId int, + channelId int, + revisionId int, + configHash string, + observation ChannelContributionModelObservation, + now int64, +) (bool, error) { + var current ChannelContributionModelHealth + err := tx.Where("contribution_id = ? AND model = ?", contributionId, observation.Model).First(¤t).Error + if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { + return false, err + } + healthChanged := current.Id != 0 && current.Healthy != observation.Healthy + if errors.Is(err, gorm.ErrRecordNotFound) { + healthChanged = !observation.Healthy + current = ChannelContributionModelHealth{ + ContributionId: contributionId, + ChannelId: channelId, + Model: observation.Model, + CreatedAt: now, + } + } + current.RevisionId = revisionId + current.ConfigHash = configHash + current.Healthy = observation.Healthy + current.LastCheckedAt = now + current.UpdatedAt = now + if observation.Healthy { + current.FailureSince = 0 + current.LastSuccessAt = now + current.LastError = "" + } else { + if current.FailureSince == 0 { + current.FailureSince = now + } + current.LastFailureAt = now + current.LastError = truncateContributionHealthError(observation.Error) + } + if current.Id == 0 { + return healthChanged, tx.Create(¤t).Error + } + return healthChanged, tx.Save(¤t).Error +} + +func normalizeContributionHealthModels(raw string) []string { + seen := make(map[string]struct{}) + models := make([]string, 0) + for _, item := range strings.Split(raw, ",") { + modelName := strings.TrimSpace(item) + if modelName == "" { + continue + } + if _, exists := seen[modelName]; exists { + continue + } + seen[modelName] = struct{}{} + models = append(models, modelName) + } + return models +} + +func truncateContributionHealthError(message string) string { + message = strings.TrimSpace(message) + if len(message) <= 500 { + return message + } + end := 500 + for end > 0 && !utf8.ValidString(message[:end]) { + end-- + } + return message[:end] +} + +func contributionUnhealthyModelSet(db *gorm.DB, channelIds []int) (map[int]map[string]struct{}, error) { + result := make(map[int]map[string]struct{}) + if len(channelIds) == 0 { + return result, nil + } + if db == nil { + db = DB + } + if !db.Migrator().HasTable(&ChannelContribution{}) || + !db.Migrator().HasTable(&ChannelContributionRevision{}) || + !db.Migrator().HasTable(&ChannelContributionModelHealth{}) { + return result, nil + } + var rows []ChannelContributionModelHealth + if err := db.Table("channel_contribution_model_healths AS health"). + Select("health.channel_id, health.model"). + Joins("JOIN channel_contributions AS contributions ON contributions.id = health.contribution_id AND contributions.channel_id = health.channel_id"). + Joins("JOIN channel_contribution_revisions AS revisions ON revisions.id = contributions.approved_revision_id AND revisions.id = health.revision_id AND revisions.config_hash = health.config_hash"). + Where("health.channel_id IN ? AND health.healthy = ? AND contributions.status IN ? AND revisions.status = ?", channelIds, false, []ChannelContributionStatus{ + ChannelContributionStatusApproved, + ChannelContributionStatusUnavailable, + }, ChannelContributionRevisionStatusApproved). + Scan(&rows).Error; err != nil { + return nil, err + } + for _, row := range rows { + if result[row.ChannelId] == nil { + result[row.ChannelId] = make(map[string]struct{}) + } + result[row.ChannelId][row.Model] = struct{}{} + } + return result, nil +} + +func applyContributionHealthToAbilities(abilities []Ability) []Ability { + if len(abilities) == 0 { + return abilities + } + channelIds := make([]int, 0, len(abilities)) + seen := make(map[int]struct{}) + for _, ability := range abilities { + if _, exists := seen[ability.ChannelId]; exists { + continue + } + seen[ability.ChannelId] = struct{}{} + channelIds = append(channelIds, ability.ChannelId) + } + unhealthy, err := contributionUnhealthyModelSet(DB, channelIds) + if err != nil { + common.SysError(fmt.Sprintf("failed to apply contribution health overlay: %v", err)) + return abilities + } + filtered := abilities[:0] + for _, ability := range abilities { + if models := unhealthy[ability.ChannelId]; models != nil { + if _, disabled := models[ability.Model]; disabled { + continue + } + } + filtered = append(filtered, ability) + } + return filtered +} + +func reapplyContributionHealthToAbilitiesTx(tx *gorm.DB, channelIds []int) error { + unhealthy, err := contributionUnhealthyModelSet(tx, channelIds) + if err != nil { + return err + } + return applyContributionUnhealthyAbilitiesTx(tx, unhealthy) +} + +func applyContributionUnhealthyAbilitiesTx(tx *gorm.DB, unhealthy map[int]map[string]struct{}) error { + for channelId, models := range unhealthy { + modelNames := make([]string, 0, len(models)) + for modelName := range models { + modelNames = append(modelNames, modelName) + } + if len(modelNames) == 0 { + continue + } + if err := tx.Model(&Ability{}). + Where("channel_id = ? AND model IN ?", channelId, modelNames). + Update("enabled", false).Error; err != nil { + return err + } + } + return nil +} + +func markChannelContributionsDeletedTx(tx *gorm.DB, channelIds []int, now int64) error { + if len(channelIds) == 0 { + return nil + } + if !tx.Migrator().HasTable(&ChannelContribution{}) { + return nil + } + var contributions []ChannelContribution + if err := lockForUpdate(tx). + Select("id", "channel_id", "status"). + Where("channel_id IN ?", channelIds). + Order("id ASC"). + Find(&contributions).Error; err != nil { + return err + } + if len(contributions) == 0 { + return nil + } + return markLockedChannelContributionsDeletedTx(tx, contributions, now) +} + +func markLockedChannelContributionsDeletedTx(tx *gorm.DB, contributions []ChannelContribution, now int64) error { + if len(contributions) == 0 { + return nil + } + contributionIds := make([]int, 0, len(contributions)) + for _, contribution := range contributions { + contributionIds = append(contributionIds, contribution.Id) + } + if tx.Migrator().HasTable(&ChannelContributionRevision{}) { + if err := tx.Model(&ChannelContributionRevision{}). + Where("contribution_id IN ? AND status = ?", contributionIds, ChannelContributionRevisionStatusPending). + Updates(map[string]any{ + "status": ChannelContributionRevisionStatusWithdrawn, + "updated_at": now, + }).Error; err != nil { + return err + } + if err := tx.Model(&ChannelContributionRevision{}). + Where("contribution_id IN ?", contributionIds). + Updates(map[string]any{ + "key": "", + "updated_at": now, + }).Error; err != nil { + return err + } + } + return tx.Model(&ChannelContribution{}). + Where("id IN ?", contributionIds). + Updates(map[string]any{ + "status": ChannelContributionStatusDeleted, + "pending_revision_id": nil, + "updated_at": now, + }).Error +} + +// ResetChannelContributionHealthForRevision removes health state inherited +// from the previously approved revision. Call it while the contribution and +// channel rows are already locked, before rebuilding abilities for the new +// revision. +func ResetChannelContributionHealthForRevision( + tx *gorm.DB, + contribution *ChannelContribution, + revisionId int, + configHash string, +) error { + if tx == nil || contribution == nil || contribution.Id <= 0 || revisionId <= 0 || strings.TrimSpace(configHash) == "" { + return errors.New("invalid channel contribution health reset") + } + var revisionCount int64 + if err := tx.Model(&ChannelContributionRevision{}). + Where("id = ? AND contribution_id = ? AND config_hash = ?", revisionId, contribution.Id, configHash). + Count(&revisionCount).Error; err != nil { + return err + } + if revisionCount != 1 { + return ErrStaleChannelContributionHealthProbe + } + if tx.Migrator().HasTable(&ChannelContributionModelHealth{}) { + if err := tx.Where("contribution_id = ?", contribution.Id).Delete(&ChannelContributionModelHealth{}).Error; err != nil { + return err + } + } + if tx.Migrator().HasTable(&ChannelContributionHealthState{}) { + if err := tx.Where("contribution_id = ?", contribution.Id).Delete(&ChannelContributionHealthState{}).Error; err != nil { + return err + } + } + return nil +} diff --git a/model/channel_contribution_health_reward_test.go b/model/channel_contribution_health_reward_test.go new file mode 100644 index 000000000000..0dcb8cf9fcd7 --- /dev/null +++ b/model/channel_contribution_health_reward_test.go @@ -0,0 +1,600 @@ +package model + +import ( + "errors" + "math" + "strings" + "sync" + "testing" + "unicode/utf8" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func prepareChannelContributionFeatureTables(t *testing.T) { + t.Helper() + require.NoError(t, DB.AutoMigrate( + &ChannelContributionRevision{}, + &ChannelContributionModelHealth{}, + &ChannelContributionHealthState{}, + &ChannelContributionRewardAccount{}, + &ChannelContributionRewardLedger{}, + )) + clear := func() { + DB.Exec("DELETE FROM channel_contribution_model_healths") + DB.Exec("DELETE FROM channel_contribution_health_states") + DB.Exec("DELETE FROM channel_contribution_reward_ledgers") + DB.Exec("DELETE FROM channel_contribution_reward_accounts") + DB.Exec("DELETE FROM channel_contribution_revisions") + DB.Exec("DELETE FROM channel_contributions") + DB.Exec("DELETE FROM abilities") + DB.Exec("DELETE FROM channels") + DB.Exec("DELETE FROM users") + } + clear() + t.Cleanup(clear) +} + +func seedApprovedContributionHealthFixture(t *testing.T, modelNames ...string) (*ChannelContribution, *ChannelContributionRevision, *Channel) { + t.Helper() + priority := int64(100) + weight := uint(0) + tag := "donate" + channel := &Channel{ + Type: constant.ChannelTypeOpenAI, + Key: "sk-health-test", + Status: common.ChannelStatusEnabled, + Name: "contributed channel", + Weight: &weight, + Models: strings.Join(modelNames, ","), + Group: "default", + Priority: &priority, + Tag: &tag, + CreatedTime: 1, + } + require.NoError(t, channel.Insert()) + + channelID := channel.Id + contribution := &ChannelContribution{ + UserId: 42, + Username: "contributor", + Status: ChannelContributionStatusApproved, + ChannelId: &channelID, + } + require.NoError(t, DB.Create(contribution).Error) + revision := &ChannelContributionRevision{ + ContributionId: contribution.Id, + RevisionNumber: 1, + Name: channel.Name, + Type: channel.Type, + BaseURL: "https://example.com", + Key: channel.Key, + Group: channel.Group, + Models: channel.Models, + ModelMapping: "{}", + ConfigHash: "config-v1", + Status: ChannelContributionRevisionStatusApproved, + } + require.NoError(t, DB.Create(revision).Error) + revisionID := revision.Id + require.NoError(t, DB.Model(&ChannelContribution{}). + Where("id = ?", contribution.Id). + Updates(map[string]any{ + "current_revision_id": revisionID, + "approved_revision_id": revisionID, + }).Error) + contribution.CurrentRevisionId = &revisionID + contribution.ApprovedRevisionId = &revisionID + return contribution, revision, channel +} + +func loadContributionAbilities(t *testing.T, channelID int) map[string]Ability { + t.Helper() + var abilities []Ability + require.NoError(t, DB.Where("channel_id = ?", channelID).Find(&abilities).Error) + result := make(map[string]Ability, len(abilities)) + for _, ability := range abilities { + result[ability.Model] = ability + } + return result +} + +func TestApplyChannelContributionHealthCycleDisablesOnlyFailedModels(t *testing.T) { + prepareChannelContributionFeatureTables(t) + contribution, revision, channel := seedApprovedContributionHealthFixture(t, "model-a", "model-b") + + result, err := ApplyChannelContributionHealthCycle( + contribution.Id, + channel.Id, + revision.Id, + revision.ConfigHash, + []ChannelContributionModelObservation{ + {Model: "model-a", Error: "upstream unavailable"}, + {Model: "model-b", Healthy: true}, + }, + 1_000, + 48*60*60, + ) + require.NoError(t, err) + assert.False(t, result.AllFailed) + assert.True(t, result.StateChanged) + assert.Equal(t, []string{"model-a"}, result.UnhealthyModels) + + abilities := loadContributionAbilities(t, channel.Id) + assert.False(t, abilities["model-a"].Enabled) + assert.True(t, abilities["model-b"].Enabled) + require.NoError(t, UpdateAbilityStatus(channel.Id, true)) + assert.False(t, loadContributionAbilities(t, channel.Id)["model-a"].Enabled) + _, _, err = FixAbility() + require.NoError(t, err) + assert.False(t, loadContributionAbilities(t, channel.Id)["model-a"].Enabled) + reloadedChannel, err := GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, common.ChannelStatusEnabled, reloadedChannel.Status) + reloadedContribution, err := GetChannelContributionById(contribution.Id) + require.NoError(t, err) + assert.Equal(t, ChannelContributionStatusApproved, reloadedContribution.Status) +} + +func TestApplyChannelContributionHealthCyclePausesManualDisableAndDeletesAfterActiveFailureWindow(t *testing.T) { + prepareChannelContributionFeatureTables(t) + contribution, revision, channel := seedApprovedContributionHealthFixture(t, "model-a", "model-b") + pendingRevision := &ChannelContributionRevision{ + ContributionId: contribution.Id, + RevisionNumber: 2, + Name: revision.Name, + Type: revision.Type, + BaseURL: revision.BaseURL, + Key: "sk-pending", + Group: revision.Group, + Models: revision.Models, + ModelMapping: revision.ModelMapping, + ConfigHash: "config-pending", + Status: ChannelContributionRevisionStatusPending, + } + require.NoError(t, DB.Create(pendingRevision).Error) + require.NoError(t, DB.Model(&ChannelContribution{}). + Where("id = ?", contribution.Id). + Update("pending_revision_id", pendingRevision.Id).Error) + failed := []ChannelContributionModelObservation{ + {Model: "model-a", Error: "failed"}, + {Model: "model-b", Error: "failed"}, + } + + result, err := ApplyChannelContributionHealthCycle(contribution.Id, channel.Id, revision.Id, revision.ConfigHash, failed, 1_000, 100) + require.NoError(t, err) + assert.True(t, result.AllFailed) + assert.False(t, result.Deleted) + + changed, err := UpdateChannelStatusWithError(channel.Id, "", common.ChannelStatusManuallyDisabled, "maintenance") + require.NoError(t, err) + assert.True(t, changed) + result, err = ApplyChannelContributionHealthCycle(contribution.Id, channel.Id, revision.Id, revision.ConfigHash, nil, 1_500, 100) + require.NoError(t, err) + assert.True(t, result.Paused) + assert.False(t, result.Deleted) + var state ChannelContributionHealthState + require.NoError(t, DB.Where("contribution_id = ?", contribution.Id).First(&state).Error) + assert.NotZero(t, state.PausedAt) + + changed, err = UpdateChannelStatusWithError(channel.Id, "", common.ChannelStatusEnabled, "maintenance complete") + require.NoError(t, err) + assert.True(t, changed) + require.NoError(t, DB.Where("contribution_id = ?", contribution.Id).First(&state).Error) + assert.Zero(t, state.PausedAt) + require.NoError(t, DB.Model(&ChannelContributionHealthState{}).Where("contribution_id = ?", contribution.Id).Update("failure_since", int64(1_000)).Error) + require.NoError(t, DB.Model(&ChannelContribution{}).Where("id = ?", contribution.Id).Update("unavailable_since", int64(1_000)).Error) + require.NoError(t, PauseContributionHealthForChannel(channel.Id, 1_000)) + require.NoError(t, ResumeContributionHealthForChannel(channel.Id, 2_000)) + require.NoError(t, DB.Where("contribution_id = ?", contribution.Id).First(&state).Error) + assert.Equal(t, int64(2_000), state.FailureSince) + + result, err = ApplyChannelContributionHealthCycle(contribution.Id, channel.Id, revision.Id, revision.ConfigHash, failed, 2_099, 100) + require.NoError(t, err) + assert.False(t, result.Deleted) + result, err = ApplyChannelContributionHealthCycle(contribution.Id, channel.Id, revision.Id, revision.ConfigHash, failed, 2_100, 100) + require.NoError(t, err) + assert.True(t, result.Deleted) + + reloadedContribution, err := GetChannelContributionById(contribution.Id) + require.NoError(t, err) + assert.Equal(t, ChannelContributionStatusDeleted, reloadedContribution.Status) + assert.Nil(t, reloadedContribution.PendingRevisionId) + var channelCount int64 + require.NoError(t, DB.Model(&Channel{}).Where("id = ?", channel.Id).Count(&channelCount).Error) + assert.Zero(t, channelCount) + var revisionKey string + require.NoError(t, DB.Model(&ChannelContributionRevision{}).Where("id = ?", revision.Id).Pluck("key", &revisionKey).Error) + assert.Empty(t, revisionKey) + var reloadedPending ChannelContributionRevision + require.NoError(t, DB.First(&reloadedPending, pendingRevision.Id).Error) + assert.Empty(t, reloadedPending.Key) + assert.Equal(t, ChannelContributionRevisionStatusWithdrawn, reloadedPending.Status) +} + +func TestApplyChannelContributionHealthCycleRejectsStaleRevisionAndConfig(t *testing.T) { + prepareChannelContributionFeatureTables(t) + contribution, revision, channel := seedApprovedContributionHealthFixture(t, "model-a") + observations := []ChannelContributionModelObservation{{Model: "model-a", Error: "failed"}} + + _, err := ApplyChannelContributionHealthCycle(contribution.Id, channel.Id, revision.Id, "old-config", observations, 1_000, 100) + require.ErrorIs(t, err, ErrStaleChannelContributionHealthProbe) + assert.True(t, loadContributionAbilities(t, channel.Id)["model-a"].Enabled) + + nextRevision := &ChannelContributionRevision{ + ContributionId: contribution.Id, + RevisionNumber: 2, + Name: revision.Name, + Type: revision.Type, + BaseURL: revision.BaseURL, + Key: "sk-next", + Group: revision.Group, + Models: revision.Models, + ModelMapping: "{}", + ConfigHash: "config-v2", + Status: ChannelContributionRevisionStatusApproved, + } + require.NoError(t, DB.Create(nextRevision).Error) + require.NoError(t, DB.Model(&ChannelContribution{}). + Where("id = ?", contribution.Id). + Update("approved_revision_id", nextRevision.Id).Error) + _, err = ApplyChannelContributionHealthCycle(contribution.Id, channel.Id, revision.Id, revision.ConfigHash, observations, 1_001, 100) + require.ErrorIs(t, err, ErrStaleChannelContributionHealthProbe) + assert.True(t, loadContributionAbilities(t, channel.Id)["model-a"].Enabled) + var healthCount int64 + require.NoError(t, DB.Model(&ChannelContributionModelHealth{}).Where("contribution_id = ?", contribution.Id).Count(&healthCount).Error) + assert.Zero(t, healthCount) +} + +func TestResetChannelContributionHealthForRevisionClearsOldOverlayBeforeAbilityRebuild(t *testing.T) { + prepareChannelContributionFeatureTables(t) + contribution, revision, channel := seedApprovedContributionHealthFixture(t, "model-a", "model-b") + _, err := ApplyChannelContributionHealthCycle( + contribution.Id, + channel.Id, + revision.Id, + revision.ConfigHash, + []ChannelContributionModelObservation{{Model: "model-a", Error: "failed"}, {Model: "model-b", Healthy: true}}, + 1_000, + 100, + ) + require.NoError(t, err) + assert.False(t, loadContributionAbilities(t, channel.Id)["model-a"].Enabled) + + nextRevision := &ChannelContributionRevision{ + ContributionId: contribution.Id, + RevisionNumber: 2, + Name: revision.Name, + Type: revision.Type, + BaseURL: revision.BaseURL, + Key: "sk-next", + Group: revision.Group, + Models: revision.Models, + ModelMapping: "{}", + ConfigHash: "config-v2", + Status: ChannelContributionRevisionStatusPending, + } + require.NoError(t, DB.Create(nextRevision).Error) + require.NoError(t, DB.Transaction(func(tx *gorm.DB) error { + var lockedContribution ChannelContribution + if err := lockForUpdate(tx).Where("id = ?", contribution.Id).First(&lockedContribution).Error; err != nil { + return err + } + var lockedChannel Channel + if err := lockForUpdate(tx).Where("id = ?", channel.Id).First(&lockedChannel).Error; err != nil { + return err + } + if err := ResetChannelContributionHealthForRevision(tx, &lockedContribution, nextRevision.Id, nextRevision.ConfigHash); err != nil { + return err + } + return lockedChannel.UpdateAbilities(tx) + })) + abilities := loadContributionAbilities(t, channel.Id) + assert.True(t, abilities["model-a"].Enabled) + assert.True(t, abilities["model-b"].Enabled) +} + +func TestHistoricalContributionHealthDoesNotDisableReusedOrdinaryChannelID(t *testing.T) { + prepareChannelContributionFeatureTables(t) + contribution, revision, contributed := seedApprovedContributionHealthFixture(t, "model-a") + _, err := ApplyChannelContributionHealthCycle( + contribution.Id, + contributed.Id, + revision.Id, + revision.ConfigHash, + []ChannelContributionModelObservation{{Model: "model-a", Error: "failed"}}, + 1_000, + 100, + ) + require.NoError(t, err) + assert.False(t, loadContributionAbilities(t, contributed.Id)["model-a"].Enabled) + + reusedID := contributed.Id + require.NoError(t, DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Where("channel_id = ?", reusedID).Delete(&Ability{}).Error; err != nil { + return err + } + if err := tx.Where("id = ?", reusedID).Delete(&Channel{}).Error; err != nil { + return err + } + return tx.Model(&ChannelContribution{}). + Where("id = ?", contribution.Id). + Update("status", ChannelContributionStatusDeleted).Error + })) + + ordinary := &Channel{ + Id: reusedID, + Type: constant.ChannelTypeOpenAI, + Key: "sk-ordinary", + Status: common.ChannelStatusEnabled, + Name: "ordinary reused channel", + Models: "model-a", + Group: "default", + CreatedTime: 2, + } + require.NoError(t, ordinary.Insert()) + assert.True(t, loadContributionAbilities(t, ordinary.Id)["model-a"].Enabled) +} + +func TestChannelContributionHealthErrorPreservesUTF8WhenTruncated(t *testing.T) { + prepareChannelContributionFeatureTables(t) + contribution, revision, channel := seedApprovedContributionHealthFixture(t, "model-a") + message := strings.Repeat("中", 200) + + _, err := ApplyChannelContributionHealthCycle( + contribution.Id, + channel.Id, + revision.Id, + revision.ConfigHash, + []ChannelContributionModelObservation{{Model: "model-a", Error: message}}, + 1_000, + 100, + ) + require.NoError(t, err) + + var health ChannelContributionModelHealth + require.NoError(t, DB.Where("contribution_id = ? AND model = ?", contribution.Id, "model-a").First(&health).Error) + assert.LessOrEqual(t, len(health.LastError), 500) + assert.True(t, utf8.ValidString(health.LastError)) + assert.Equal(t, strings.Repeat("中", 166), health.LastError) +} + +func TestInitChannelCacheUsesOnlyEnabledAbilities(t *testing.T) { + prepareChannelContributionFeatureTables(t) + _, _, channel := seedApprovedContributionHealthFixture(t, "model-a", "model-b") + require.NoError(t, DB.Model(&Ability{}). + Where("channel_id = ? AND model = ?", channel.Id, "model-a"). + Update("enabled", false).Error) + + originalMemoryCacheEnabled := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + channelSyncLock.Lock() + originalGroupRoutes := group2model2channels + originalChannels := channelsIDM + originalAdvancedCustom := channel2advancedCustomConfig + originalGeneration := channelCacheGeneration + channelSyncLock.Unlock() + t.Cleanup(func() { + common.MemoryCacheEnabled = originalMemoryCacheEnabled + channelSyncLock.Lock() + group2model2channels = originalGroupRoutes + channelsIDM = originalChannels + channel2advancedCustomConfig = originalAdvancedCustom + channelCacheGeneration = originalGeneration + channelSyncLock.Unlock() + }) + + InitChannelCache() + channelSyncLock.RLock() + failedRoutes := append([]int(nil), group2model2channels["default"]["model-a"]...) + healthyRoutes := append([]int(nil), group2model2channels["default"]["model-b"]...) + channelSyncLock.RUnlock() + assert.Empty(t, failedRoutes) + assert.Equal(t, []int{channel.Id}, healthyRoutes) +} + +func TestFilterNonContributionChannelsExcludesApprovedAndUnavailableChannels(t *testing.T) { + prepareChannelContributionFeatureTables(t) + contribution, _, contributed := seedApprovedContributionHealthFixture(t, "model-a") + ordinary := &Channel{Name: "ordinary", Type: constant.ChannelTypeOpenAI, Key: "sk-ordinary", Status: common.ChannelStatusEnabled, Models: "model-a", Group: "default"} + require.NoError(t, ordinary.Insert()) + + filtered := FilterNonContributionChannels([]*Channel{contributed, ordinary}) + require.Len(t, filtered, 1) + assert.Equal(t, ordinary.Id, filtered[0].Id) + require.NoError(t, DB.Model(&ChannelContribution{}). + Where("id = ?", contribution.Id). + Update("status", ChannelContributionStatusUnavailable).Error) + filtered = FilterNonContributionChannels([]*Channel{contributed, ordinary}) + require.Len(t, filtered, 1) + assert.Equal(t, ordinary.Id, filtered[0].Id) +} + +func TestPopulateContributionChannelFlagsMarksOnlyActiveContributions(t *testing.T) { + prepareChannelContributionFeatureTables(t) + contribution, _, contributed := seedApprovedContributionHealthFixture(t, "model-a") + ordinary := &Channel{Name: "ordinary", Type: constant.ChannelTypeOpenAI, Key: "sk-ordinary", Status: common.ChannelStatusEnabled, Models: "model-a", Group: "default"} + require.NoError(t, ordinary.Insert()) + + channels := []*Channel{ordinary, contributed} + require.NoError(t, PopulateContributionChannelFlags(channels)) + assert.False(t, ordinary.IsContribution) + assert.True(t, contributed.IsContribution) + + require.NoError(t, DB.Model(&ChannelContribution{}). + Where("id = ?", contribution.Id). + Update("status", ChannelContributionStatusDeleted).Error) + require.NoError(t, PopulateContributionChannelFlags(channels)) + assert.False(t, contributed.IsContribution) +} + +func TestContributionChannelAdminMutationRequiresRevisionForSensitiveFields(t *testing.T) { + prepareChannelContributionFeatureTables(t) + _, _, channel := seedApprovedContributionHealthFixture(t, "model-a") + originalKey := channel.Key + + _, err := UpdateChannelAtomically(channel.Id, func(current *Channel) error { + current.Key = "sk-bypassed-review" + return nil + }) + require.ErrorIs(t, err, ErrChannelContributionRequiresReview) + reloaded, err := GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, originalKey, reloaded.Key) + + newPriority := int64(250) + newWeight := uint(7) + newTag := "reviewed-donation" + updated, err := UpdateChannelAtomically(channel.Id, func(current *Channel) error { + current.Priority = &newPriority + current.Weight = &newWeight + current.Tag = &newTag + return nil + }) + require.NoError(t, err) + assert.Equal(t, originalKey, updated.Key) + assert.Equal(t, newPriority, updated.GetPriority()) + assert.Equal(t, int(newWeight), updated.GetWeight()) + require.NotNil(t, updated.Tag) + assert.Equal(t, newTag, *updated.Tag) + + mapping := `{"model-a":"other-upstream"}` + err = EditChannelByTag(newTag, nil, &mapping, nil, nil, nil, nil, nil, nil) + require.ErrorIs(t, err, ErrChannelContributionRequiresReview) +} + +func TestContributionChannelPartialAdminUpdatePreservesReviewedFields(t *testing.T) { + prepareChannelContributionFeatureTables(t) + _, _, channel := seedApprovedContributionHealthFixture(t, "model-a") + newTag := "partial-update" + + patch := &Channel{Id: channel.Id, Tag: &newTag} + require.NoError(t, patch.Update()) + + reloaded, err := GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, channel.Name, reloaded.Name) + assert.Equal(t, channel.Type, reloaded.Type) + assert.Equal(t, channel.Key, reloaded.Key) + assert.Equal(t, channel.Group, reloaded.Group) + assert.Equal(t, channel.Models, reloaded.Models) + require.NotNil(t, reloaded.Tag) + assert.Equal(t, newTag, *reloaded.Tag) +} + +func TestChannelContributionRewardCreditIsIdempotentAndAuditsSaturation(t *testing.T) { + prepareChannelContributionFeatureTables(t) + clamp := &common.QuotaClamp{ + Op: "QuotaFromFloat", + Kind: common.QuotaClampOverflow, + Original: math.MaxFloat64, + Clamped: common.MaxQuota, + } + credited, err := CreditChannelContributionReward(42, 10, 20, "request-1", common.MaxQuota, 10_000, common.MaxQuota, clamp) + require.NoError(t, err) + assert.True(t, credited) + credited, err = CreditChannelContributionReward(42, 10, 20, "request-1", common.MaxQuota, 10_000, common.MaxQuota, clamp) + require.NoError(t, err) + assert.False(t, credited) + + account, err := GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(common.MaxQuota), account.Balance) + assert.Equal(t, int64(common.MaxQuota), account.LifetimeEarned) + entries, total, err := ListChannelContributionRewardLedger(42, 0, 20) + require.NoError(t, err) + assert.Equal(t, int64(1), total) + require.Len(t, entries, 1) + assert.True(t, entries[0].QuotaSaturated) + assert.Contains(t, entries[0].QuotaSaturation, `"kind":"overflow"`) +} + +func TestTransferChannelContributionRewardMovesBalanceToUserQuotaAtomically(t *testing.T) { + prepareChannelContributionFeatureTables(t) + user := &User{Id: 42, Username: "contributor", Quota: 100, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(user).Error) + credited, err := CreditChannelContributionReward(42, 10, 20, "request-1", 1_000, 500, 50, nil) + require.NoError(t, err) + require.True(t, credited) + + entry, err := TransferChannelContributionReward(42, 30) + require.NoError(t, err) + assert.Equal(t, int64(-30), entry.Amount) + assert.Equal(t, int64(20), entry.BalanceAfter) + account, err := GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(20), account.Balance) + assert.Equal(t, int64(30), account.LifetimeTransferred) + var reloadedUser User + require.NoError(t, DB.First(&reloadedUser, 42).Error) + assert.Equal(t, 130, reloadedUser.Quota) + transfers, total, err := ListChannelContributionRewardTransfers(42, 0, 20) + require.NoError(t, err) + assert.Equal(t, int64(1), total) + require.Len(t, transfers, 1) + assert.Equal(t, ChannelContributionRewardEntryTransfer, transfers[0].EntryType) + + _, err = TransferChannelContributionReward(42, 21) + require.ErrorIs(t, err, ErrChannelContributionRewardInsufficientBalance) + account, err = GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(20), account.Balance) +} + +func TestTransferChannelContributionRewardPreventsConcurrentOverspend(t *testing.T) { + prepareChannelContributionFeatureTables(t) + user := &User{Id: 42, Username: "contributor", Quota: 100, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(user).Error) + credited, err := CreditChannelContributionReward(42, 10, 20, "request-1", 1_000, 500, 50, nil) + require.NoError(t, err) + require.True(t, credited) + + errorsCh := make(chan error, 2) + var transfers sync.WaitGroup + transfers.Add(2) + for range 2 { + go func() { + defer transfers.Done() + _, transferErr := TransferChannelContributionReward(42, 30) + errorsCh <- transferErr + }() + } + transfers.Wait() + close(errorsCh) + succeeded := 0 + insufficient := 0 + for transferErr := range errorsCh { + if transferErr == nil { + succeeded++ + } else if errors.Is(transferErr, ErrChannelContributionRewardInsufficientBalance) { + insufficient++ + } + } + assert.Equal(t, 1, succeeded) + assert.Equal(t, 1, insufficient) + account, err := GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(20), account.Balance) + var reloadedUser User + require.NoError(t, DB.First(&reloadedUser, 42).Error) + assert.Equal(t, 130, reloadedUser.Quota) +} + +func TestTransferChannelContributionRewardRejectsUserQuotaOverflow(t *testing.T) { + prepareChannelContributionFeatureTables(t) + user := &User{Id: 42, Username: "contributor", Quota: common.MaxQuota - 10, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(user).Error) + credited, err := CreditChannelContributionReward(42, 10, 20, "request-1", 1_000, 500, 50, nil) + require.NoError(t, err) + require.True(t, credited) + + _, err = TransferChannelContributionReward(42, 20) + require.EqualError(t, err, "user quota would exceed the supported limit") + account, err := GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(50), account.Balance) +} diff --git a/model/channel_contribution_reward.go b/model/channel_contribution_reward.go new file mode 100644 index 000000000000..65fe903d3ba9 --- /dev/null +++ b/model/channel_contribution_reward.go @@ -0,0 +1,295 @@ +package model + +import ( + "errors" + "fmt" + "math" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +const ( + ChannelContributionRewardEntryEarn = "earn" + ChannelContributionRewardEntryTransfer = "transfer" +) + +var ErrChannelContributionRewardInsufficientBalance = errors.New("channel contribution reward balance is insufficient") + +type ChannelContributionRewardAccount struct { + UserId int `json:"user_id" gorm:"primaryKey;autoIncrement:false"` + Balance int64 `json:"balance" gorm:"type:bigint;not null"` + LifetimeEarned int64 `json:"lifetime_earned" gorm:"type:bigint;not null"` + LifetimeTransferred int64 `json:"lifetime_transferred" gorm:"type:bigint;not null"` + CreatedAt int64 `json:"created_at" gorm:"type:bigint;not null"` + UpdatedAt int64 `json:"updated_at" gorm:"type:bigint;index;not null"` +} + +type ChannelContributionRewardLedger struct { + Id int64 `json:"id" gorm:"primaryKey"` + UserId int `json:"user_id" gorm:"index;not null"` + ContributionId int `json:"contribution_id" gorm:"index;not null"` + ChannelId int `json:"channel_id" gorm:"uniqueIndex:uk_contribution_reward_request,priority:1;index;not null"` + RequestId string `json:"request_id" gorm:"type:varchar(128);uniqueIndex:uk_contribution_reward_request,priority:2;not null"` + EntryType string `json:"entry_type" gorm:"type:varchar(32);not null"` + Amount int64 `json:"amount" gorm:"type:bigint;not null"` + BalanceAfter int64 `json:"balance_after" gorm:"type:bigint;not null"` + SourceQuota int `json:"source_quota" gorm:"not null"` + RewardBps int `json:"reward_bps" gorm:"not null"` + QuotaSaturated bool `json:"quota_saturated" gorm:"not null"` + QuotaSaturation string `json:"quota_saturation" gorm:"type:text;not null"` + CreatedAt int64 `json:"created_at" gorm:"type:bigint;index;not null"` +} + +type ChannelContributionRewardTarget struct { + ContributionId int + UserId int +} + +func (account *ChannelContributionRewardAccount) BeforeCreate(_ *gorm.DB) error { + now := common.GetTimestamp() + if account.CreatedAt == 0 { + account.CreatedAt = now + } + if account.UpdatedAt == 0 { + account.UpdatedAt = now + } + return nil +} + +func GetActiveChannelContributionRewardTarget(channelId int) (*ChannelContributionRewardTarget, error) { + if channelId <= 0 { + return nil, nil + } + var contribution ChannelContribution + err := DB.Select("id", "user_id"). + Where("channel_id = ? AND status = ?", channelId, ChannelContributionStatusApproved). + First(&contribution).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &ChannelContributionRewardTarget{ + ContributionId: contribution.Id, + UserId: contribution.UserId, + }, nil +} + +func GetChannelContributionRewardAccount(userId int) (*ChannelContributionRewardAccount, error) { + var account ChannelContributionRewardAccount + err := DB.Where("user_id = ?", userId).First(&account).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return &ChannelContributionRewardAccount{UserId: userId}, nil + } + if err != nil { + return nil, err + } + return &account, nil +} + +func ListChannelContributionRewardLedger(userId int, offset int, limit int) ([]*ChannelContributionRewardLedger, int64, error) { + if limit <= 0 { + limit = 20 + } + if limit > 100 { + limit = 100 + } + var total int64 + query := DB.Model(&ChannelContributionRewardLedger{}).Where("user_id = ?", userId) + if err := query.Count(&total).Error; err != nil { + return nil, 0, err + } + var entries []*ChannelContributionRewardLedger + if err := query.Order("id desc").Offset(offset).Limit(limit).Find(&entries).Error; err != nil { + return nil, 0, err + } + return entries, total, nil +} + +func ListChannelContributionRewardTransfers(userId int, offset int, limit int) ([]*ChannelContributionRewardLedger, int64, error) { + if limit <= 0 { + limit = 20 + } + if limit > 100 { + limit = 100 + } + var total int64 + query := DB.Model(&ChannelContributionRewardLedger{}). + Where("user_id = ? AND entry_type = ?", userId, ChannelContributionRewardEntryTransfer) + if err := query.Count(&total).Error; err != nil { + return nil, 0, err + } + var entries []*ChannelContributionRewardLedger + if err := query.Order("id desc").Offset(offset).Limit(limit).Find(&entries).Error; err != nil { + return nil, 0, err + } + return entries, total, nil +} + +func CreditChannelContributionReward( + userId int, + contributionId int, + channelId int, + requestId string, + sourceQuota int, + rewardBps int, + amount int, + quotaClamp *common.QuotaClamp, +) (bool, error) { + if userId <= 0 || contributionId <= 0 || channelId <= 0 || requestId == "" { + return false, errors.New("invalid channel contribution reward identity") + } + if sourceQuota <= 0 || rewardBps <= 0 || amount <= 0 { + return false, nil + } + + quotaSaturation := "" + if quotaClamp != nil { + encoded, err := common.Marshal(quotaClamp.AuditMap()) + if err != nil { + return false, fmt.Errorf("marshal channel contribution reward saturation: %w", err) + } + quotaSaturation = string(encoded) + } + + credited := false + err := DB.Transaction(func(tx *gorm.DB) error { + now := common.GetTimestamp() + entry := ChannelContributionRewardLedger{ + UserId: userId, + ContributionId: contributionId, + ChannelId: channelId, + RequestId: requestId, + EntryType: ChannelContributionRewardEntryEarn, + Amount: int64(amount), + SourceQuota: sourceQuota, + RewardBps: rewardBps, + QuotaSaturated: quotaClamp != nil, + QuotaSaturation: quotaSaturation, + CreatedAt: now, + } + result := tx.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "channel_id"}, {Name: "request_id"}}, + DoNothing: true, + }).Create(&entry) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return nil + } + + account := ChannelContributionRewardAccount{UserId: userId} + if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&account).Error; err != nil { + return err + } + if err := lockForUpdate(tx).Where("user_id = ?", userId).First(&account).Error; err != nil { + return err + } + if account.Balance > math.MaxInt64-int64(amount) || account.LifetimeEarned > math.MaxInt64-int64(amount) { + return errors.New("channel contribution reward balance overflow") + } + account.Balance += int64(amount) + account.LifetimeEarned += int64(amount) + account.UpdatedAt = now + if err := tx.Model(&ChannelContributionRewardAccount{}). + Where("user_id = ?", userId). + Updates(map[string]any{ + "balance": account.Balance, + "lifetime_earned": account.LifetimeEarned, + "updated_at": now, + }).Error; err != nil { + return err + } + if err := tx.Model(&ChannelContributionRewardLedger{}). + Where("id = ?", entry.Id). + Update("balance_after", account.Balance).Error; err != nil { + return err + } + credited = true + return nil + }) + return credited, err +} + +func TransferChannelContributionReward(userId int, amount int) (*ChannelContributionRewardLedger, error) { + if userId <= 0 || amount <= 0 { + return nil, errors.New("transfer amount must be positive") + } + if amount > common.MaxQuota { + return nil, fmt.Errorf("transfer amount exceeds quota limit: %d", amount) + } + + var entry ChannelContributionRewardLedger + err := DB.Transaction(func(tx *gorm.DB) error { + var account ChannelContributionRewardAccount + if err := lockForUpdate(tx).Where("user_id = ?", userId).First(&account).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return ErrChannelContributionRewardInsufficientBalance + } + return err + } + if account.Balance < int64(amount) { + return ErrChannelContributionRewardInsufficientBalance + } + var user User + if err := lockForUpdate(tx).Select("id", "quota").Where("id = ?", userId).First(&user).Error; err != nil { + return err + } + newQuota := int64(user.Quota) + int64(amount) + if newQuota > int64(common.MaxQuota) { + return errors.New("user quota would exceed the supported limit") + } + + now := common.GetTimestamp() + newBalance := account.Balance - int64(amount) + if account.LifetimeTransferred > math.MaxInt64-int64(amount) { + return errors.New("channel contribution transfer counter overflow") + } + result := tx.Model(&ChannelContributionRewardAccount{}). + Where("user_id = ? AND balance >= ?", userId, amount). + Updates(map[string]any{ + "balance": newBalance, + "lifetime_transferred": account.LifetimeTransferred + int64(amount), + "updated_at": now, + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return ErrChannelContributionRewardInsufficientBalance + } + if err := tx.Model(&User{}).Where("id = ?", userId).Update("quota", int(newQuota)).Error; err != nil { + return err + } + transferId, err := common.GenerateRandomCharsKey(24) + if err != nil { + return err + } + entry = ChannelContributionRewardLedger{ + UserId: userId, + RequestId: "transfer_" + transferId, + EntryType: ChannelContributionRewardEntryTransfer, + Amount: -int64(amount), + BalanceAfter: newBalance, + CreatedAt: now, + } + if err := tx.Create(&entry).Error; err != nil { + return err + } + return nil + }) + if err != nil { + return nil, err + } + if err := cacheIncrUserQuota(userId, int64(amount)); err != nil { + common.SysError(fmt.Sprintf("failed to update user quota cache after contribution reward transfer: user_id=%d err=%v", userId, err)) + } + RecordLog(userId, LogTypeTopup, fmt.Sprintf("渠道贡献奖励划转 %s", logger.LogQuota(amount))) + return &entry, nil +} diff --git a/model/channel_contribution_test.go b/model/channel_contribution_test.go new file mode 100644 index 000000000000..c1436b90e167 --- /dev/null +++ b/model/channel_contribution_test.go @@ -0,0 +1,295 @@ +package model + +import ( + "errors" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func newChannelContributionDraft(t *testing.T, userId int, name string, models string) (*ChannelContribution, *ChannelContributionRevision) { + t.Helper() + revision := &ChannelContributionRevision{ + Name: name, + Type: constant.ChannelTypeOpenAI, + BaseURL: "https://example.com", + Key: "sk-test", + Group: "default", + Models: models, + ModelMapping: `{"gpt-test":"upstream-gpt-test"}`, + } + configHash, err := ComputeChannelContributionConfigHash(revision) + require.NoError(t, err) + revision.ConfigHash = configHash + contribution := &ChannelContribution{ + UserId: userId, + Username: "contributor", + Status: ChannelContributionStatusDraft, + } + require.NoError(t, CreateChannelContributionWithRevision(contribution, revision)) + return contribution, revision +} + +func newPendingChannelContribution(t *testing.T) (*ChannelContribution, *ChannelContributionRevision) { + t.Helper() + contribution, revision := newChannelContributionDraft(t, 42, "shared upstream", "gpt-test,gpt-test-mini") + require.NoError(t, SubmitChannelContribution( + contribution.Id, + contribution.UserId, + revision.Id, + revision.ConfigHash, + "v1", + "agreement", + "agreement-hash", + 100, + )) + contribution, err := GetChannelContributionById(contribution.Id) + require.NoError(t, err) + revision, err = GetChannelContributionRevision(contribution.Id, revision.Id) + require.NoError(t, err) + return contribution, revision +} + +func approveChannelContributionFixture(t *testing.T) (*ChannelContribution, *ChannelContributionRevision, *Channel) { + t.Helper() + contribution, revision := newPendingChannelContribution(t) + approved, channel, err := ApproveChannelContribution(contribution.Id, revision.Id, ChannelContributionApproval{ + ReviewerId: 7, + ReviewerUsername: "reviewer", + Tag: "donate", + Priority: 100, + Weight: 0, + }) + require.NoError(t, err) + return approved, revision, channel +} + +func TestChannelContributionRevisionKeepsLateTestResultsIsolated(t *testing.T) { + truncateTables(t) + contribution, firstRevision := newChannelContributionDraft(t, 42, "first", "gpt-test") + run := &ChannelContributionTestRun{ + ContributionId: contribution.Id, + RevisionId: firstRevision.Id, + ConfigHash: firstRevision.ConfigHash, + ActorId: contribution.UserId, + ActorType: ChannelContributionTestActorUser, + } + require.NoError(t, CreateChannelContributionTestRun(run)) + claimed, err := ClaimNextQueuedChannelContributionTestRun() + require.NoError(t, err) + assert.Equal(t, run.Id, claimed.Id) + + secondRevision := &ChannelContributionRevision{ + Name: "second", + Type: constant.ChannelTypeOpenAI, + BaseURL: "https://second.example.com", + Key: "sk-second", + Group: "default", + Models: "gpt-test", + ModelMapping: "{}", + } + secondRevision.ConfigHash, err = ComputeChannelContributionConfigHash(secondRevision) + require.NoError(t, err) + require.NoError(t, CreateChannelContributionRevision(contribution.Id, contribution.UserId, secondRevision)) + require.NoError(t, FinishChannelContributionTestRun(run.Id, ChannelContributionTestRunStatusSucceeded, true, []ChannelContributionTestResult{{ + Model: "gpt-test", + EndpointType: string(constant.EndpointTypeOpenAI), + Success: true, + }}, "")) + + firstRun, err := GetLatestSuccessfulChannelContributionTestRun(firstRevision.Id, firstRevision.ConfigHash) + require.NoError(t, err) + assert.Equal(t, run.Id, firstRun.Id) + _, err = GetLatestSuccessfulChannelContributionTestRun(secondRevision.Id, secondRevision.ConfigHash) + assert.True(t, errors.Is(err, gorm.ErrRecordNotFound)) + reloaded, err := GetChannelContributionById(contribution.Id) + require.NoError(t, err) + require.NotNil(t, reloaded.CurrentRevisionId) + assert.Equal(t, secondRevision.Id, *reloaded.CurrentRevisionId) +} + +func TestChannelContributionTestRunAllowsOnlyOneActiveRunPerUser(t *testing.T) { + truncateTables(t) + contribution, revision := newChannelContributionDraft(t, 42, "first", "gpt-test") + first := &ChannelContributionTestRun{ + ContributionId: contribution.Id, + RevisionId: revision.Id, + ConfigHash: revision.ConfigHash, + ActorId: contribution.UserId, + ActorType: ChannelContributionTestActorUser, + } + require.NoError(t, CreateChannelContributionTestRun(first)) + second := &ChannelContributionTestRun{ + ContributionId: contribution.Id, + RevisionId: revision.Id, + ConfigHash: revision.ConfigHash, + ActorId: contribution.UserId, + ActorType: ChannelContributionTestActorUser, + } + require.Error(t, CreateChannelContributionTestRun(second)) +} + +func TestChannelContributionAdminTestRunUniquenessUsesAdminActor(t *testing.T) { + truncateTables(t) + firstContribution, firstRevision := newChannelContributionDraft(t, 42, "first", "gpt-test") + secondContribution, secondRevision := newChannelContributionDraft(t, 43, "second", "gpt-test") + firstAdminRun := &ChannelContributionTestRun{ + ContributionId: firstContribution.Id, + RevisionId: firstRevision.Id, + ConfigHash: firstRevision.ConfigHash, + ActorId: 99, + ActorType: ChannelContributionTestActorAdmin, + } + require.NoError(t, CreateChannelContributionTestRun(firstAdminRun)) + sameAdminOtherContribution := &ChannelContributionTestRun{ + ContributionId: secondContribution.Id, + RevisionId: secondRevision.Id, + ConfigHash: secondRevision.ConfigHash, + ActorId: 99, + ActorType: ChannelContributionTestActorAdmin, + } + require.Error(t, CreateChannelContributionTestRun(sameAdminOtherContribution)) + otherAdminSameContribution := &ChannelContributionTestRun{ + ContributionId: firstContribution.Id, + RevisionId: firstRevision.Id, + ConfigHash: firstRevision.ConfigHash, + ActorId: 100, + ActorType: ChannelContributionTestActorAdmin, + } + require.NoError(t, CreateChannelContributionTestRun(otherAdminSameContribution)) +} + +func TestRequeueRunningChannelContributionTestRunsRecoversInterruptedRun(t *testing.T) { + truncateTables(t) + contribution, revision := newChannelContributionDraft(t, 42, "first", "gpt-test") + run := &ChannelContributionTestRun{ + ContributionId: contribution.Id, + RevisionId: revision.Id, + ConfigHash: revision.ConfigHash, + ActorId: contribution.UserId, + ActorType: ChannelContributionTestActorUser, + } + require.NoError(t, CreateChannelContributionTestRun(run)) + assert.True(t, HasUnfinishedChannelContributionTestRuns()) + _, err := ClaimNextQueuedChannelContributionTestRun() + require.NoError(t, err) + assert.True(t, HasUnfinishedChannelContributionTestRuns()) + requeued, err := RequeueRunningChannelContributionTestRuns() + require.NoError(t, err) + assert.Equal(t, int64(1), requeued) + reloaded, err := GetChannelContributionTestRun(run.Id) + require.NoError(t, err) + assert.Equal(t, ChannelContributionTestRunStatusQueued, reloaded.Status) + assert.Zero(t, reloaded.StartedAt) +} + +func TestApproveChannelContributionCreatesChannelAndAbilities(t *testing.T) { + truncateTables(t) + approved, _, channel := approveChannelContributionFixture(t) + + assert.Equal(t, ChannelContributionStatusApproved, approved.Status) + assert.Equal(t, common.ChannelStatusEnabled, channel.Status) + assert.Equal(t, "sk-test", channel.Key) + assert.Equal(t, "default", channel.Group) + assert.Equal(t, int64(100), channel.GetPriority()) + assert.Equal(t, 0, channel.GetWeight()) + require.NotNil(t, channel.Tag) + assert.Equal(t, "donate", *channel.Tag) + require.NotNil(t, channel.Remark) + assert.Equal(t, "贡献者：42 contributor", *channel.Remark) + + var abilities []Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).Order("model asc").Find(&abilities).Error) + require.Len(t, abilities, 2) + assert.Equal(t, "gpt-test", abilities[0].Model) + assert.Equal(t, "gpt-test-mini", abilities[1].Model) +} + +func TestModificationReviewPreservesActiveChannelAndAdminRouting(t *testing.T) { + truncateTables(t) + approved, _, channel := approveChannelContributionFixture(t) + adminTag := "admin-managed" + adminPriority := int64(7) + adminWeight := uint(9) + require.NoError(t, DB.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]any{ + "status": common.ChannelStatusManuallyDisabled, + "tag": adminTag, + "priority": adminPriority, + "weight": adminWeight, + }).Error) + + nextRevision := &ChannelContributionRevision{ + Name: "updated upstream", + Type: constant.ChannelTypeOpenAI, + BaseURL: "https://updated.example.com", + Key: "sk-updated", + Group: "default", + Models: "gpt-test", + ModelMapping: "{}", + } + var err error + nextRevision.ConfigHash, err = ComputeChannelContributionConfigHash(nextRevision) + require.NoError(t, err) + require.NoError(t, CreateChannelContributionRevision(approved.Id, approved.UserId, nextRevision)) + require.NoError(t, SubmitChannelContribution( + approved.Id, + approved.UserId, + nextRevision.Id, + nextRevision.ConfigHash, + "v2", + "updated agreement", + "updated-agreement-hash", + 200, + )) + + pending, err := GetChannelContributionById(approved.Id) + require.NoError(t, err) + assert.Equal(t, ChannelContributionStatusApproved, pending.Status) + require.NotNil(t, pending.PendingRevisionId) + activeChannel, err := GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, "shared upstream", activeChannel.Name) + assert.Equal(t, "sk-test", activeChannel.Key) + + _, updatedChannel, err := ApproveChannelContribution(approved.Id, nextRevision.Id, ChannelContributionApproval{ + ReviewerId: 8, + ReviewerUsername: "second-reviewer", + Tag: "new-default", + Priority: 999, + Weight: 1, + }) + require.NoError(t, err) + assert.Equal(t, "updated upstream", updatedChannel.Name) + assert.Equal(t, "sk-updated", updatedChannel.Key) + assert.Equal(t, common.ChannelStatusManuallyDisabled, updatedChannel.Status) + require.NotNil(t, updatedChannel.Tag) + assert.Equal(t, adminTag, *updatedChannel.Tag) + assert.Equal(t, adminPriority, updatedChannel.GetPriority()) + assert.Equal(t, int(adminWeight), updatedChannel.GetWeight()) +} + +func TestWithdrawChannelContributionClearsKeysAndDeletesChannel(t *testing.T) { + truncateTables(t) + approved, _, channel := approveChannelContributionFixture(t) + require.NoError(t, WithdrawChannelContribution(approved.Id, approved.UserId)) + + withdrawn, err := GetChannelContributionById(approved.Id) + require.NoError(t, err) + assert.Equal(t, ChannelContributionStatusDeleted, withdrawn.Status) + revisions, err := ListChannelContributionRevisions(approved.Id) + require.NoError(t, err) + require.NotEmpty(t, revisions) + for _, revision := range revisions { + assert.Empty(t, revision.Key) + } + var channelCount int64 + require.NoError(t, DB.Model(&Channel{}).Where("id = ?", channel.Id).Count(&channelCount).Error) + assert.Zero(t, channelCount) + var abilityCount int64 + require.NoError(t, DB.Model(&Ability{}).Where("channel_id = ?", channel.Id).Count(&abilityCount).Error) + assert.Zero(t, abilityCount) +} diff --git a/model/main.go b/model/main.go index 21445593e54e..ca7efc3b4826 100644 --- a/model/main.go +++ b/model/main.go @@ -260,6 +260,14 @@ func migrateDB() error { err := DB.AutoMigrate( &Channel{}, + &ChannelContribution{}, + &ChannelContributionRevision{}, + &ChannelContributionTestRun{}, + &ChannelContributionTestResult{}, + &ChannelContributionModelHealth{}, + &ChannelContributionHealthState{}, + &ChannelContributionRewardAccount{}, + &ChannelContributionRewardLedger{}, &Token{}, &User{}, &UserSession{}, @@ -323,6 +331,14 @@ func migrateDBFast() error { name string }{ {&Channel{}, "Channel"}, + {&ChannelContribution{}, "ChannelContribution"}, + {&ChannelContributionRevision{}, "ChannelContributionRevision"}, + {&ChannelContributionTestRun{}, "ChannelContributionTestRun"}, + {&ChannelContributionTestResult{}, "ChannelContributionTestResult"}, + {&ChannelContributionModelHealth{}, "ChannelContributionModelHealth"}, + {&ChannelContributionHealthState{}, "ChannelContributionHealthState"}, + {&ChannelContributionRewardAccount{}, "ChannelContributionRewardAccount"}, + {&ChannelContributionRewardLedger{}, "ChannelContributionRewardLedger"}, {&Token{}, "Token"}, {&User{}, "User"}, {&UserSession{}, "UserSession"}, diff --git a/model/option.go b/model/option.go index e7fda5231be7..a371a09a0d6d 100644 --- a/model/option.go +++ b/model/option.go @@ -206,6 +206,9 @@ func SyncOptions(frequency int) { } func validateOptionValue(key string, value string) error { + if err := operation_setting.ValidateChannelContributionOption(key, value); err != nil { + return err + } if key == operation_setting.ToolPriceOptionKey { return operation_setting.ValidateToolPricesJSON(value) } diff --git a/model/system_task.go b/model/system_task.go index c811409b487d..182a4b8a0ce7 100644 --- a/model/system_task.go +++ b/model/system_task.go @@ -16,11 +16,13 @@ const ( SystemTaskStatusSucceeded SystemTaskStatus = "succeeded" SystemTaskStatusFailed SystemTaskStatus = "failed" - SystemTaskTypeLogCleanup = "log_cleanup" - SystemTaskTypeChannelTest = "channel_test" - SystemTaskTypeModelUpdate = "model_update" - SystemTaskTypeMidjourneyPoll = "midjourney_poll" - SystemTaskTypeAsyncTaskPoll = "async_task_poll" + SystemTaskTypeLogCleanup = "log_cleanup" + SystemTaskTypeChannelTest = "channel_test" + SystemTaskTypeModelUpdate = "model_update" + SystemTaskTypeMidjourneyPoll = "midjourney_poll" + SystemTaskTypeAsyncTaskPoll = "async_task_poll" + SystemTaskTypeChannelContributionTest = "channel_contribution_test" + SystemTaskTypeChannelContributionHealth = "channel_contribution_health" ) var ErrSystemTaskLockLost = errors.New("system task lock lost") diff --git a/model/task_cas_test.go b/model/task_cas_test.go index a53804d3baab..52a4903c8d45 100644 --- a/model/task_cas_test.go +++ b/model/task_cas_test.go @@ -46,6 +46,14 @@ func TestMain(m *testing.M) { &TwoFABackupCode{}, &Log{}, &Channel{}, + &ChannelContribution{}, + &ChannelContributionRevision{}, + &ChannelContributionTestRun{}, + &ChannelContributionTestResult{}, + &ChannelContributionModelHealth{}, + &ChannelContributionHealthState{}, + &ChannelContributionRewardAccount{}, + &ChannelContributionRewardLedger{}, &QuotaData{}, &Ability{}, &TopUp{}, @@ -79,6 +87,14 @@ func truncateTables(t *testing.T) { DB.Exec("DELETE FROM users") DB.Exec("DELETE FROM logs") DB.Exec("DELETE FROM channels") + DB.Exec("DELETE FROM channel_contribution_model_healths") + DB.Exec("DELETE FROM channel_contribution_health_states") + DB.Exec("DELETE FROM channel_contribution_reward_ledgers") + DB.Exec("DELETE FROM channel_contribution_reward_accounts") + DB.Exec("DELETE FROM channel_contribution_test_results") + DB.Exec("DELETE FROM channel_contribution_test_runs") + DB.Exec("DELETE FROM channel_contribution_revisions") + DB.Exec("DELETE FROM channel_contributions") DB.Exec("DELETE FROM quota_data") DB.Exec("DELETE FROM abilities") DB.Exec("DELETE FROM top_ups") diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index 48241b14a5e9..a71459e815be 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -488,9 +488,15 @@ func keepUpstreamRedirectResponse(_ *http.Request, _ []*http.Request) error { } func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { - client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting) - if err != nil { - return nil, fmt.Errorf("new proxy http client failed: %w", err) + var client *http.Client + if info.UseSSRFProtectedClient { + client = service.GetStrictSSRFProtectedHTTPClient() + } else { + var err error + client, err = service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting) + if err != nil { + return nil, fmt.Errorf("new proxy http client failed: %w", err) + } } // Clients are cached and shared across channels, so override redirect // behavior on a shallow copy instead of mutating the cached client. This diff --git a/relay/channel/api_request_ssrf_test.go b/relay/channel/api_request_ssrf_test.go new file mode 100644 index 000000000000..df385f75e138 --- /dev/null +++ b/relay/channel/api_request_ssrf_test.go @@ -0,0 +1,63 @@ +package channel + +import ( + "net/http" + "net/http/httptest" + "testing" + + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDoRequestUsesSSRFProtectedClientForUserControlledUpstream(t *testing.T) { + fetchSetting := system_setting.GetFetchSetting() + originalSetting := *fetchSetting + defer func() { + *fetchSetting = originalSetting + }() + + fetchSetting.EnableSSRFProtection = true + fetchSetting.AllowPrivateIp = false + fetchSetting.DomainFilterMode = false + fetchSetting.IpFilterMode = false + fetchSetting.DomainList = nil + fetchSetting.IpList = nil + fetchSetting.AllowedPorts = []string{"1-65535"} + fetchSetting.ApplyIPFilterForDomain = true + service.InitHttpClient() + + requestCount := 0 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requestCount++ + w.WriteHeader(http.StatusNoContent) + })) + defer upstream.Close() + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/relay", nil) + + protectedRequest, err := http.NewRequest(http.MethodGet, upstream.URL, http.NoBody) + require.NoError(t, err) + _, err = doRequest(ctx, protectedRequest, &relaycommon.RelayInfo{ + UseSSRFProtectedClient: true, + ChannelMeta: &relaycommon.ChannelMeta{}, + }) + require.Error(t, err) + assert.Zero(t, requestCount) + + standardRequest, err := http.NewRequest(http.MethodGet, upstream.URL, http.NoBody) + require.NoError(t, err) + response, err := doRequest(ctx, standardRequest, &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{}, + }) + require.NoError(t, err) + defer response.Body.Close() + assert.Equal(t, http.StatusNoContent, response.StatusCode) + assert.Equal(t, 1, requestCount) +} diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 84acea73c585..1187a1e6e7e4 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -498,6 +498,13 @@ func FetchGeminiModels(baseURL, apiKey, proxyURL string) ([]string, error) { if err != nil { return nil, fmt.Errorf("创建HTTP客户端失败: %v", err) } + return FetchGeminiModelsWithClient(client, baseURL, apiKey) +} + +func FetchGeminiModelsWithClient(client *http.Client, baseURL, apiKey string) ([]string, error) { + if client == nil { + return nil, errors.New("HTTP client is required") + } allModels := make([]string, 0) nextPageToken := "" diff --git a/relay/channel/ollama/relay-ollama.go b/relay/channel/ollama/relay-ollama.go index e517a1e6aa03..5984e1a9fd3e 100644 --- a/relay/channel/ollama/relay-ollama.go +++ b/relay/channel/ollama/relay-ollama.go @@ -1,6 +1,7 @@ package ollama import ( + "errors" "fmt" "io" "net/http" @@ -332,9 +333,15 @@ func ollamaEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h } func FetchOllamaModels(baseURL, apiKey string) ([]OllamaModel, error) { + return FetchOllamaModelsWithClient(&http.Client{}, baseURL, apiKey) +} + +func FetchOllamaModelsWithClient(client *http.Client, baseURL, apiKey string) ([]OllamaModel, error) { + if client == nil { + return nil, errors.New("HTTP client is required") + } url := fmt.Sprintf("%s/api/tags", baseURL) - client := &http.Client{} request, err := http.NewRequest("GET", url, nil) if err != nil { return nil, fmt.Errorf("创建请求失败: %v", err) diff --git a/relay/channel/ollama/stream.go b/relay/channel/ollama/stream.go index a0d7839f9f6d..2fcb15699465 100644 --- a/relay/channel/ollama/stream.go +++ b/relay/channel/ollama/stream.go @@ -214,7 +214,7 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R } service.CloseResponseBodyGracefully(resp) raw := string(body) - if common.DebugEnabled { + if common.DebugEnabled && c.Request.Context().Value(constant.ContextKeySuppressUpstreamResponseLog) != true { println("ollama non-stream raw resp:", raw) } diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 9a0619eb27f5..ce589b42864a 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -169,7 +169,11 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re shouldSendLastResp := true if err := handleLastResponse(lastStreamData, &responseId, &createAt, &systemFingerprint, &model, &usage, &containStreamUsage, info, &shouldSendLastResp); err != nil { - logger.LogError(c, fmt.Sprintf("error handling last response: %s, lastStreamData: [%s]", err.Error(), lastStreamData)) + if c.Request.Context().Value(constant.ContextKeySuppressUpstreamResponseLog) == true { + logger.LogError(c, "error handling last response: "+err.Error()) + } else { + logger.LogError(c, fmt.Sprintf("error handling last response: %s, lastStreamData: [%s]", err.Error(), lastStreamData)) + } } if info.RelayFormat == types.RelayFormatOpenAI { diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index b0bb19bdca3b..cb5f16f33f05 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -143,6 +143,9 @@ type RelayInfo struct { SubscriptionAmountUsedAfterPreConsume int64 IsClaudeBetaQuery bool // /v1/messages?beta=true IsChannelTest bool // channel test request + UseSSRFProtectedClient bool // user-controlled upstream URL must use the protected outbound client + ContributionRewardBps int + ContributionRewardSnapshotted bool RetryIndex int LastError *types.NewAPIError RuntimeHeadersOverride map[string]interface{} diff --git a/router/api-router.go b/router/api-router.go index 31c595e00db2..00e658a284cf 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -232,6 +232,7 @@ func SetApiRouter(router *gin.Engine) { ratioSyncRoute.POST("/fetch", controller.FetchUpstreamRatios) } registerChannelRoutes(apiRouter) + registerChannelContributionRoutes(apiRouter) registerAuthzRoutes(apiRouter) tokenRoute := apiRouter.Group("/token") tokenRoute.Use(middleware.UserAuth()) diff --git a/router/channel-contribution-router.go b/router/channel-contribution-router.go new file mode 100644 index 000000000000..ca49b78d1a5d --- /dev/null +++ b/router/channel-contribution-router.go @@ -0,0 +1,42 @@ +package router + +import ( + "github.com/QuantumNous/new-api/controller" + "github.com/QuantumNous/new-api/middleware" + + "github.com/gin-gonic/gin" +) + +func registerChannelContributionRoutes(apiRouter *gin.RouterGroup) { + adminRoute := apiRouter.Group("/channel-contributions/admin") + adminRoute.Use(middleware.AdminAuth()) + { + adminRoute.GET("/settings", controller.GetAdminChannelContributionSettings) + adminRoute.PUT("/settings", controller.UpdateAdminChannelContributionSettings) + adminRoute.GET("", controller.ListAdminChannelContributions) + adminRoute.GET("/:id", controller.GetAdminChannelContribution) + adminRoute.POST("/:id/test-runs", controller.CreateAdminChannelContributionTestRun) + adminRoute.GET("/:id/test-runs/:runId", controller.GetAdminChannelContributionTestRun) + adminRoute.POST("/:id/approve", controller.ApproveAdminChannelContribution) + adminRoute.POST("/:id/reject", controller.RejectAdminChannelContribution) + adminRoute.DELETE("/:id", controller.DeleteAdminChannelContribution) + } + + userRoute := apiRouter.Group("/channel-contributions") + userRoute.Use(middleware.UserAuth()) + { + userRoute.GET("/config", controller.GetChannelContributionConfig) + userRoute.GET("/rewards", controller.GetChannelContributionRewards) + userRoute.GET("/reward-transfers", controller.ListChannelContributionRewardTransfers) + userRoute.POST("/reward-transfers", middleware.UserCriticalRateLimit("channel-contribution-reward-transfer"), controller.TransferChannelContributionReward) + userRoute.GET("", controller.ListUserChannelContributions) + userRoute.POST("", controller.CreateChannelContribution) + userRoute.GET("/:id", controller.GetUserChannelContribution) + userRoute.PUT("/:id", controller.UpdateUserChannelContribution) + userRoute.POST("/:id/fetch-models", middleware.UserCriticalRateLimit("channel-contribution-fetch-models"), controller.FetchChannelContributionModels) + userRoute.POST("/:id/test-runs", middleware.UserCriticalRateLimit("channel-contribution-test"), controller.CreateUserChannelContributionTestRun) + userRoute.GET("/:id/test-runs/:runId", controller.GetUserChannelContributionTestRun) + userRoute.POST("/:id/submit", middleware.UserCriticalRateLimit("channel-contribution-submit"), middleware.TurnstileCheck(), controller.SubmitUserChannelContribution) + userRoute.POST("/:id/withdraw", controller.WithdrawUserChannelContribution) + } +} diff --git a/router/channel_contribution_router_test.go b/router/channel_contribution_router_test.go new file mode 100644 index 000000000000..86072c2cfec7 --- /dev/null +++ b/router/channel_contribution_router_test.go @@ -0,0 +1,91 @@ +package router + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupChannelContributionRouterTest(t *testing.T) string { + t.Helper() + previousDB := model.DB + previousType := common.MainDatabaseType() + previousRedis := common.RedisEnabled + previousTurnstile := common.TurnstileCheckEnabled + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.ChannelContribution{})) + model.DB = db + common.SetMainDatabaseType(common.DatabaseTypeSQLite) + common.RedisEnabled = false + token := "channel-contribution-router-token" + user := &model.User{ + Username: "contribution-router-user", + Password: "password-placeholder", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + AccessToken: &token, + AuthVersion: 1, + AffCode: "contribution-router-aff-code", + } + require.NoError(t, db.Create(user).Error) + t.Cleanup(func() { + model.DB = previousDB + common.SetMainDatabaseType(previousType) + common.RedisEnabled = previousRedis + common.TurnstileCheckEnabled = previousTurnstile + }) + return token +} + +func performChannelContributionRouteRequest(router http.Handler, token string, method string, path string, body string) *httptest.ResponseRecorder { + request := httptest.NewRequest(method, path, strings.NewReader(body)) + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + return response +} + +func TestChannelContributionSubmitAloneRequiresTurnstile(t *testing.T) { + token := setupChannelContributionRouterTest(t) + common.TurnstileCheckEnabled = true + gin.SetMode(gin.TestMode) + engine := gin.New() + registerChannelContributionRoutes(engine.Group("/api")) + + submit := performChannelContributionRouteRequest(engine, token, http.MethodPost, "/api/channel-contributions/1/submit", `{}`) + assert.Contains(t, submit.Body.String(), "Turnstile token") + + for _, path := range []string{ + "/api/channel-contributions/1/fetch-models", + "/api/channel-contributions/1/test-runs", + } { + response := performChannelContributionRouteRequest(engine, token, http.MethodPost, path, "") + assert.NotContains(t, response.Body.String(), "Turnstile", path) + assert.Contains(t, response.Body.String(), "record not found", path) + } +} + +func TestChannelContributionSubmitSkipsTurnstileWhenDisabled(t *testing.T) { + token := setupChannelContributionRouterTest(t) + common.TurnstileCheckEnabled = false + gin.SetMode(gin.TestMode) + engine := gin.New() + registerChannelContributionRoutes(engine.Group("/api")) + + response := performChannelContributionRouteRequest(engine, token, http.MethodPost, "/api/channel-contributions/1/submit", `{}`) + assert.NotContains(t, response.Body.String(), "Turnstile") + assert.Contains(t, response.Body.String(), "agreement must be accepted") +} diff --git a/service/billing.go b/service/billing.go index 7ffa537192af..31cf2f097cf7 100644 --- a/service/billing.go +++ b/service/billing.go @@ -18,6 +18,7 @@ const ( // PreConsumeBilling 根据用户计费偏好创建 BillingSession 并执行预扣费。 // 会话存储在 relayInfo.Billing 上，供后续 Settle / Refund 使用。 func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.NewAPIError { + SnapshotChannelContributionReward(c, relayInfo) if relayInfo != nil && relayInfo.QuotaClamp != nil { return types.NewErrorWithStatusCode( relayInfo.QuotaClamp, @@ -83,13 +84,17 @@ func SettleBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, actualQuo checkAndSendQuotaNotify(relayInfo, actualQuota-preConsumed, preConsumed) } } + SettleChannelContributionReward(ctx, relayInfo, actualQuota) return nil } // 回退：无 BillingSession 时使用旧路径 quotaDelta := actualQuota - relayInfo.FinalPreConsumedQuota if quotaDelta != 0 { - return PostConsumeQuota(relayInfo, quotaDelta, relayInfo.FinalPreConsumedQuota, true) + if err := PostConsumeQuota(relayInfo, quotaDelta, relayInfo.FinalPreConsumedQuota, true); err != nil { + return err + } } + SettleChannelContributionReward(ctx, relayInfo, actualQuota) return nil } diff --git a/service/channel_contribution_reward.go b/service/channel_contribution_reward.go new file mode 100644 index 000000000000..1fe0274fe3e6 --- /dev/null +++ b/service/channel_contribution_reward.go @@ -0,0 +1,96 @@ +package service + +import ( + "fmt" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + + "github.com/gin-gonic/gin" +) + +const channelContributionRewardBpsOptionKey = "channel_contribution_setting.reward_bps" + +func currentChannelContributionRewardBps() int { + common.OptionMapRWMutex.RLock() + raw := strings.TrimSpace(common.OptionMap[channelContributionRewardBpsOptionKey]) + common.OptionMapRWMutex.RUnlock() + bps, err := strconv.Atoi(raw) + if err != nil || bps <= 0 { + return 0 + } + if bps > 10_000 { + return 10_000 + } + return bps +} + +// SnapshotChannelContributionReward captures the configured reward rate before +// channel selection and retries. Settlement resolves the final channel owner. +func SnapshotChannelContributionReward(_ *gin.Context, info *relaycommon.RelayInfo) { + if info == nil { + return + } + if info.ContributionRewardSnapshotted { + return + } + info.ContributionRewardBps = currentChannelContributionRewardBps() + info.ContributionRewardSnapshotted = true +} + +// SettleChannelContributionReward credits the contributor once for the final +// charged quota. The model layer enforces channel+request idempotency. +func SettleChannelContributionReward(ctx *gin.Context, info *relaycommon.RelayInfo, chargedQuota int) { + if info == nil || chargedQuota <= 0 || info.IsChannelTest || info.ContributionRewardBps <= 0 || + info.UserId <= 0 || info.ChannelMeta == nil || info.ChannelId <= 0 || info.RequestId == "" { + return + } + target, err := model.GetActiveChannelContributionRewardTarget(info.ChannelId) + if err != nil { + logger.LogWarn(ctx, fmt.Sprintf("failed to resolve final channel contribution reward target: channel_id=%d err=%v", info.ChannelId, err)) + return + } + if target == nil || target.UserId == info.UserId { + return + } + reward, clamp := common.QuotaFromFloatChecked(float64(chargedQuota) * float64(info.ContributionRewardBps) / 10_000) + if clamp != nil { + if info.QuotaClamp == nil { + info.QuotaClamp = clamp + } + logger.LogWarn(ctx, fmt.Sprintf( + "channel contribution reward saturated: request_id=%s channel_id=%d source_quota=%d reward_bps=%d clamp=%v", + info.RequestId, + info.ChannelId, + chargedQuota, + info.ContributionRewardBps, + clamp, + )) + } + if reward <= 0 { + return + } + _, err = model.CreditChannelContributionReward( + target.UserId, + target.ContributionId, + info.ChannelId, + info.RequestId, + chargedQuota, + info.ContributionRewardBps, + reward, + clamp, + ) + if err != nil { + logger.LogWarn(ctx, fmt.Sprintf( + "failed to settle channel contribution reward: request_id=%s channel_id=%d contributor_id=%d err=%v", + info.RequestId, + info.ChannelId, + target.UserId, + err, + )) + } +} diff --git a/service/channel_contribution_reward_test.go b/service/channel_contribution_reward_test.go new file mode 100644 index 000000000000..f09460c5e7a2 --- /dev/null +++ b/service/channel_contribution_reward_test.go @@ -0,0 +1,203 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func prepareChannelContributionRewardServiceTest(t *testing.T) { + t.Helper() + require.NoError(t, model.DB.AutoMigrate( + &model.Channel{}, + &model.Ability{}, + &model.ChannelContribution{}, + &model.ChannelContributionRevision{}, + &model.ChannelContributionModelHealth{}, + &model.ChannelContributionRewardAccount{}, + &model.ChannelContributionRewardLedger{}, + )) + clear := func() { + model.DB.Exec("DELETE FROM channel_contribution_reward_ledgers") + model.DB.Exec("DELETE FROM channel_contribution_reward_accounts") + model.DB.Exec("DELETE FROM channel_contribution_revisions") + model.DB.Exec("DELETE FROM channel_contributions") + model.DB.Exec("DELETE FROM abilities") + model.DB.Exec("DELETE FROM channels") + } + clear() + t.Cleanup(clear) + + common.OptionMapRWMutex.Lock() + wasNil := common.OptionMap == nil + if wasNil { + common.OptionMap = make(map[string]string) + } + previous, existed := common.OptionMap[channelContributionRewardBpsOptionKey] + common.OptionMap[channelContributionRewardBpsOptionKey] = "500" + common.OptionMapRWMutex.Unlock() + t.Cleanup(func() { + common.OptionMapRWMutex.Lock() + if wasNil { + common.OptionMap = nil + } else if existed { + common.OptionMap[channelContributionRewardBpsOptionKey] = previous + } else { + delete(common.OptionMap, channelContributionRewardBpsOptionKey) + } + common.OptionMapRWMutex.Unlock() + }) +} + +func seedRewardTargetChannel(t *testing.T, contributorID int) (*model.Channel, *model.Channel) { + t.Helper() + ordinary := &model.Channel{Name: "ordinary", Type: 1, Key: "ordinary-key", Status: common.ChannelStatusEnabled, Models: "test-model", Group: "default"} + require.NoError(t, ordinary.Insert()) + contributed := &model.Channel{Name: "contributed", Type: 1, Key: "contributed-key", Status: common.ChannelStatusEnabled, Models: "test-model", Group: "default"} + require.NoError(t, contributed.Insert()) + channelID := contributed.Id + contribution := &model.ChannelContribution{ + UserId: contributorID, + Username: "contributor", + Status: model.ChannelContributionStatusApproved, + ChannelId: &channelID, + } + require.NoError(t, model.DB.Create(contribution).Error) + return ordinary, contributed +} + +func TestSnapshotChannelContributionRewardDoesNotRequireSelectedChannel(t *testing.T) { + prepareChannelContributionRewardServiceTest(t) + info := &relaycommon.RelayInfo{} + require.NotPanics(t, func() { + SnapshotChannelContributionReward(nil, info) + }) + assert.Equal(t, 500, info.ContributionRewardBps) + assert.True(t, info.ContributionRewardSnapshotted) + + common.OptionMapRWMutex.Lock() + common.OptionMap[channelContributionRewardBpsOptionKey] = "900" + common.OptionMapRWMutex.Unlock() + SnapshotChannelContributionReward(nil, info) + assert.Equal(t, 500, info.ContributionRewardBps) +} + +func TestSettleChannelContributionRewardUsesFinalChannelAndIsIdempotent(t *testing.T) { + prepareChannelContributionRewardServiceTest(t) + ordinary, contributed := seedRewardTargetChannel(t, 42) + info := &relaycommon.RelayInfo{UserId: 7, RequestId: "request-final-contribution"} + SnapshotChannelContributionReward(nil, info) + + info.ChannelMeta = &relaycommon.ChannelMeta{ChannelId: ordinary.Id} + SettleChannelContributionReward(nil, info, 1_000) + account, err := model.GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Zero(t, account.Balance) + + info.ChannelMeta = &relaycommon.ChannelMeta{ChannelId: contributed.Id} + SettleChannelContributionReward(nil, info, 1_000) + SettleChannelContributionReward(nil, info, 1_000) + account, err = model.GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(50), account.Balance) + entries, total, err := model.ListChannelContributionRewardLedger(42, 0, 20) + require.NoError(t, err) + assert.Equal(t, int64(1), total) + require.Len(t, entries, 1) + assert.Equal(t, 500, entries[0].RewardBps) + assert.Equal(t, 1_000, entries[0].SourceQuota) + + info.RequestId = "request-final-ordinary" + info.ChannelMeta = &relaycommon.ChannelMeta{ChannelId: contributed.Id} + SnapshotChannelContributionReward(nil, info) + info.ChannelMeta = &relaycommon.ChannelMeta{ChannelId: ordinary.Id} + SettleChannelContributionReward(nil, info, 1_000) + account, err = model.GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(50), account.Balance) + deleted, err := model.BatchDeleteChannels([]int{contributed.Id}) + require.NoError(t, err) + assert.Equal(t, int64(1), deleted) + account, err = model.GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(50), account.Balance) + _, total, err = model.ListChannelContributionRewardLedger(42, 0, 20) + require.NoError(t, err) + assert.Equal(t, int64(1), total) +} + +func TestSettleChannelContributionRewardExcludesSelfTestAndFreeRequests(t *testing.T) { + prepareChannelContributionRewardServiceTest(t) + _, contributed := seedRewardTargetChannel(t, 42) + + cases := []struct { + name string + info *relaycommon.RelayInfo + quota int + }{ + { + name: "self use", + info: &relaycommon.RelayInfo{ + UserId: 42, + RequestId: "self-use", + ContributionRewardBps: 500, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: contributed.Id}, + }, + quota: 1_000, + }, + { + name: "channel test", + info: &relaycommon.RelayInfo{ + UserId: 7, + RequestId: "channel-test", + IsChannelTest: true, + ContributionRewardBps: 500, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: contributed.Id}, + }, + quota: 1_000, + }, + { + name: "free request", + info: &relaycommon.RelayInfo{ + UserId: 7, + RequestId: "free-request", + ContributionRewardBps: 500, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: contributed.Id}, + }, + quota: 0, + }, + } + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + SettleChannelContributionReward(nil, test.info, test.quota) + }) + } + account, err := model.GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Zero(t, account.Balance) + _, total, err := model.ListChannelContributionRewardLedger(42, 0, 20) + require.NoError(t, err) + assert.Zero(t, total) +} + +func TestSettleBillingCreditsRewardWhenPreConsumedQuotaMatchesFinalQuota(t *testing.T) { + prepareChannelContributionRewardServiceTest(t) + _, contributed := seedRewardTargetChannel(t, 42) + info := &relaycommon.RelayInfo{ + UserId: 7, + RequestId: "request-settle-billing", + FinalPreConsumedQuota: 1_000, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: contributed.Id}, + ContributionRewardBps: 500, + ContributionRewardSnapshotted: true, + } + + require.NoError(t, SettleBilling(nil, info, 1_000)) + account, err := model.GetChannelContributionRewardAccount(42) + require.NoError(t, err) + assert.Equal(t, int64(50), account.Balance) +} diff --git a/service/contribution_strict_fetch_test.go b/service/contribution_strict_fetch_test.go new file mode 100644 index 000000000000..e204224fadfc --- /dev/null +++ b/service/contribution_strict_fetch_test.go @@ -0,0 +1,24 @@ +package service + +import ( + "context" + "net/http" + "testing" + + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestStrictSSRFClientBlocksLoopbackWhenGlobalProtectionIsDisabled(t *testing.T) { + fetchSetting := system_setting.GetFetchSetting() + previous := fetchSetting.EnableSSRFProtection + fetchSetting.EnableSSRFProtection = false + t.Cleanup(func() { fetchSetting.EnableSSRFProtection = previous }) + + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://127.0.0.1:65535/v1/models", nil) + require.NoError(t, err) + _, err = GetStrictSSRFProtectedHTTPClient().Do(req) + require.Error(t, err) + assert.Contains(t, err.Error(), "private IP address not allowed") +} diff --git a/service/error.go b/service/error.go index f14f1bbab660..ba25b65ca0b1 100644 --- a/service/error.go +++ b/service/error.go @@ -12,6 +12,7 @@ import ( "strings" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" taskdto "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/relaykit/dto" @@ -86,6 +87,7 @@ func ClaudeErrorWrapperLocal(err error, code string, statusCode int) *dto.Claude func RelayErrorHandler(ctx context.Context, resp *http.Response, showBodyWhenFail bool) (newApiErr *types.NewAPIError) { newApiErr = types.InitOpenAIError(types.ErrorCodeBadResponseStatusCode, resp.StatusCode) + suppressResponseLog := ctx != nil && ctx.Value(constant.ContextKeySuppressUpstreamResponseLog) == true responseBody, err := io.ReadAll(resp.Body) if err != nil { @@ -106,8 +108,10 @@ func RelayErrorHandler(ctx context.Context, resp *http.Response, showBodyWhenFai if err != nil { if showBodyWhenFail { newApiErr.Err = buildErrWithBody("") - } else { + } else if !suppressResponseLog { logger.LogError(ctx, fmt.Sprintf("bad response status code %d, body: %s", resp.StatusCode, responseBodyPreview)) + } + if !showBodyWhenFail { newApiErr.Err = fmt.Errorf("bad response status code %d", resp.StatusCode) } return @@ -125,7 +129,7 @@ func RelayErrorHandler(ctx context.Context, resp *http.Response, showBodyWhenFai } } message := errResponse.ToMessage() - if message == "" { + if message == "" && !suppressResponseLog { // The body parsed as JSON but carried no usable error message; log the // raw body so the upstream failure remains diagnosable. logger.LogError(ctx, fmt.Sprintf("bad response status code %d with empty error message, body: %s", resp.StatusCode, responseBodyPreview)) diff --git a/service/protected_fetch_client.go b/service/protected_fetch_client.go index 9d1d4cc87871..73ae79eee6d9 100644 --- a/service/protected_fetch_client.go +++ b/service/protected_fetch_client.go @@ -29,6 +29,7 @@ type ssrfProtectedRoundTripper struct { dialContext func(ctx context.Context, network, address string) (net.Conn, error) getProtection func() (*common.SSRFProtection, bool, error) proxy func(*http.Request) (*url.URL, error) + maxBodyBytes int64 mutex sync.Mutex transports map[string]*http.Transport @@ -55,10 +56,48 @@ func currentFetchProtection() (*common.SSRFProtection, bool, error) { return protection, true, nil } +func strictFetchProtection() (*common.SSRFProtection, bool, error) { + return &common.SSRFProtection{ + AllowPrivateIp: false, + DomainFilterMode: false, + IpFilterMode: false, + ApplyIPFilterForDomain: true, + }, true, nil +} + +func ValidateStrictSSRFProtectedFetchURL(urlStr string) error { + protection, _, _ := strictFetchProtection() + return protection.ValidateURL(urlStr) +} + +func checkStrictProtectedFetchRedirect(req *http.Request, via []*http.Request) error { + if err := ValidateStrictSSRFProtectedFetchURL(req.URL.String()); err != nil { + return fmt.Errorf("redirect blocked: %v", err) + } + if len(via) >= 10 { + return fmt.Errorf("stopped after 10 redirects") + } + return nil +} + func newProtectedFetchHTTPClient() *http.Client { return newProtectedFetchHTTPClientWithDialer(nil, nil, nil) } +const StrictSSRFProtectedResponseBodyLimitBytes int64 = 8 << 20 + +var strictSSRFProtectedHTTPClient = func() *http.Client { + noProxy := func(*http.Request) (*url.URL, error) { return nil, nil } + client := newProtectedFetchHTTPClientWithProxy(nil, nil, strictFetchProtection, noProxy) + client.Transport.(*ssrfProtectedRoundTripper).maxBodyBytes = StrictSSRFProtectedResponseBodyLimitBytes + client.CheckRedirect = checkStrictProtectedFetchRedirect + return client +}() + +func GetStrictSSRFProtectedHTTPClient() *http.Client { + return strictSSRFProtectedHTTPClient +} + func newProtectedFetchHTTPClientWithDialer(resolver ssrfResolver, dialContext func(ctx context.Context, network, address string) (net.Conn, error), getProtection func() (*common.SSRFProtection, bool, error)) *http.Client { return newProtectedFetchHTTPClientWithProxy(resolver, dialContext, getProtection, http.ProxyFromEnvironment) } @@ -101,15 +140,28 @@ func (t *ssrfProtectedRoundTripper) RoundTrip(req *http.Request) (*http.Response if req == nil || req.URL == nil { return nil, fmt.Errorf("invalid request") } - if err := ValidateSSRFProtectedFetchURL(req.URL.String()); err != nil { + protection, enabled, err := t.getProtection() + if err != nil { return nil, err } + if enabled { + if err := protection.ValidateURL(req.URL.String()); err != nil { + return nil, err + } + } proxyURL, err := t.proxy(req) if err != nil { return nil, err } - return t.transportFor(proxyURL).RoundTrip(req) + response, err := t.transportFor(proxyURL).RoundTrip(req) + if err != nil { + return response, err + } + if response != nil && response.Body != nil && t.maxBodyBytes > 0 { + response.Body = http.MaxBytesReader(nil, response.Body, t.maxBodyBytes) + } + return response, nil } func (t *ssrfProtectedRoundTripper) CloseIdleConnections() { diff --git a/service/protected_fetch_client_test.go b/service/protected_fetch_client_test.go index 3aa49c9ac79f..753e2561edf9 100644 --- a/service/protected_fetch_client_test.go +++ b/service/protected_fetch_client_test.go @@ -4,9 +4,12 @@ import ( "context" "errors" "fmt" + "io" "net" "net/http" + "net/http/httptest" "net/url" + "strings" "testing" "github.com/QuantumNous/new-api/common" @@ -315,3 +318,29 @@ func TestProtectedFetchRoundTripperReusesTransportPerProxy(t *testing.T) { require.True(t, direct.ForceAttemptHTTP2) require.False(t, direct.DisableKeepAlives) } + +func TestProtectedFetchRoundTripperLimitsResponseBody(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, strings.Repeat("x", 32)) + })) + t.Cleanup(server.Close) + + client := newProtectedFetchHTTPClientWithProxy( + nil, + nil, + func() (*common.SSRFProtection, bool, error) { return nil, false, nil }, + func(*http.Request) (*url.URL, error) { return nil, nil }, + ) + roundTripper, ok := client.Transport.(*ssrfProtectedRoundTripper) + require.True(t, ok) + roundTripper.maxBodyBytes = 8 + + response, err := client.Get(server.URL) + require.NoError(t, err) + t.Cleanup(func() { _ = response.Body.Close() }) + body, err := io.ReadAll(response.Body) + var maxBytesErr *http.MaxBytesError + require.ErrorAs(t, err, &maxBytesErr) + require.Equal(t, int64(8), maxBytesErr.Limit) + require.Equal(t, "xxxxxxxx", string(body)) +} diff --git a/setting/config/config.go b/setting/config/config.go index a82d934f6e5b..3fe04151bc10 100644 --- a/setting/config/config.go +++ b/setting/config/config.go @@ -38,6 +38,17 @@ func (cm *ConfigManager) Get(name string) interface{} { return cm.configs[name] } +// Snapshot copies a registered configuration while holding the manager read lock. +func (cm *ConfigManager) Snapshot(name string, destination any) error { + cm.mutex.RLock() + defer cm.mutex.RUnlock() + encoded, err := common.Marshal(cm.configs[name]) + if err != nil { + return err + } + return common.Unmarshal(encoded, destination) +} + // LoadFromDB 从数据库加载配置 func (cm *ConfigManager) LoadFromDB(options map[string]string) error { cm.mutex.Lock() diff --git a/setting/operation_setting/channel_contribution_setting.go b/setting/operation_setting/channel_contribution_setting.go new file mode 100644 index 000000000000..d3255dc3bd59 --- /dev/null +++ b/setting/operation_setting/channel_contribution_setting.go @@ -0,0 +1,277 @@ +package operation_setting + +import ( + "fmt" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/setting/config" +) + +const ChannelContributionSettingPrefix = "channel_contribution_setting." + +const DefaultChannelContributionAgreementContent = `# 渠道贡献协议 + +版本：2026-08-16 + +## 一、适用范围与接受 + +本协议适用于您通过本平台提交、测试、维护或撤回第三方 API 渠道（以下简称“贡献渠道”）的行为。您勾选“我已阅读并同意《渠道贡献协议》”并提交渠道，即表示您已阅读、理解并接受本协议全部内容。未勾选同意或所确认的协议版本已失效时，平台不会受理提交。 + +## 二、贡献资格与授权保证 + +1. 您确认对所提交的 API 端点、API Key 及相关账户拥有合法、完整且持续有效的使用权和授权，并有权允许平台按本协议约定进行测试、审核、接入和调用。 +2. 您不得提交盗用、泄露、共享越权、来源不明、通过欺骗取得，或受第三方条款限制而无权贡献的凭据。 +3. 因授权范围、凭据来源、账户归属或第三方权利产生争议时，您应及时配合核验；平台可暂停测试、拒绝审核、停用或删除相关渠道。 +4. 若您位于中国大陆地区，或贡献行为、上游账户、API 调用链路涉及中国大陆地区，您确认自己具备实施该贡献所需的民事行为能力、主体资格、网络与数据处理权限，并遵守适用的网络安全、数据安全、个人信息保护、跨境数据及生成式人工智能服务管理要求。 +5. 您不得通过贡献渠道提供依法需要但尚未取得许可、备案、批准或授权的服务，不得以个人贡献方式规避上游服务的地区限制、实名要求、网络接入规则或其他强制性要求。 +6. 平台可基于地区、主体身份、上游条款或合规要求，请求您补充资格、授权或用途说明；在核验完成前，平台可暂停受理或使用相关渠道。 + +## 三、API Key 的处理与使用 + +1. 为完成模型获取、连通性测试、审核、渠道创建、请求转发、健康巡检和故障排查，平台需要接收、存储并调用您提交的 API Key。 +2. 平台仅将 API Key 用于贡献渠道的管理与服务运行。平台的必要服务组件及经授权的管理员可能在审核、运维和故障排查过程中处理该凭据。 +3. 您应确保 API Key 具有适当的权限范围、额度和有效期，并及时处理上游账户中的余额不足、限额、封禁、过期或权限变化。 +4. 请勿提交与渠道调用无关的账户密码、支付凭据、个人身份资料或其他敏感信息。 + +## 四、模型获取、测试与持续巡检 + +1. 平台可使用您填写的 API 端点和 API Key 自动获取模型列表；自动获取结果仅供配置参考，您仍需核对模型名称、可用性和模型映射。 +2. 提交前，贡献渠道中的每个模型均须通过平台根据其端点能力要求的测试：聊天、Responses、Claude、Gemini 等支持流式传输的端点须同时通过非流式与流式测试；Embedding、Rerank 等不适用流式传输的端点仅测试适用的非流式模式。修改 API 端点、API Key、模型列表、模型映射、分组或其他影响调用的配置后，原测试结果失效，需重新测试。 +3. 平台可在审核后持续对渠道进行自动健康巡检，并根据检测结果将渠道标记为可用或不可用。测试通过仅代表测试时点满足条件，不构成持续可用承诺。 + +## 五、价格与提交条件 + +只有在贡献渠道的全部模型均已由管理员配置有效价格，且全部必需测试通过后，渠道才可提交审核。模型价格由管理员维护，您不得通过模型命名、映射或其他方式规避计费配置。 + +## 六、贡献性质与奖励 + +1. 渠道贡献为自愿行为。平台可按管理员当前配置的奖励比例，将通过贡献渠道成功完成并最终结算的实际计费额度乘以奖励比例，记入贡献者的独立“渠道贡献奖励余额”。奖励比例默认值为 0，平台可调整后仅对相应请求生效，不承诺固定比例、固定收益或最低调用量。 +2. 奖励以请求最终成功结算的实际额度为基数；失败请求、免费请求、测试请求、未产生正数结算额度的请求，以及贡献者本人使用自己贡献渠道产生的请求，不计入奖励。重复结算不会重复记账。 +3. 渠道贡献奖励余额与用户普通额度分开记录。贡献者可按平台提供的手动划转功能，将可用奖励余额划转为本人平台额度；完成划转后不可撤销。奖励余额及划转所得平台额度不属于现金、存款或可提现资产，不支持提现、转账给他人或兑换法定货币。 +4. 提交贡献不代表渠道必然通过审核，也不代表平台必须持续使用该渠道或维持特定调用量。渠道被停用、删除、撤回或变为不可用后，不再对后续请求产生奖励，但不影响此前已正确入账的奖励记录。 +5. 您保留 API Key 及相关账户中依法属于您的权利；本协议不转移 API Key 或上游账户的所有权。 + +## 七、审核与渠道配置 + +1. 管理员可根据连通性、模型能力、价格完整性、来源可信度、服务稳定性及平台运营需要，对贡献进行通过、拒绝、停用、恢复或删除处理。 +2. 渠道通过审核后，平台将按管理员配置写入渠道标签、可用分组、优先级和权重。上述配置可由管理员根据运行情况调整，不以贡献者提交时的展示值为准。 +3. 管理员可要求您补充说明或重新测试。未在合理期限内完成核验的，平台可拒绝或关闭该贡献。 + +## 八、不可用与自动删除 + +1. 当健康巡检失败或上游返回持续异常时，贡献渠道可被标记为“不可用”并停止参与请求分发；恢复检测通过后，平台可将其恢复为“已通过”。 +2. 渠道连续不可用达到管理员配置的时长后，平台可自动删除对应渠道，贡献历史将显示为“已删除”。当前展示的默认时长为 48 小时，实际以管理员实时配置为准。 +3. 自动删除后，如需再次贡献，应重新创建、测试并提交审核；原审核结果不会自动恢复。 + +## 九、撤回贡献 + +1. 您可在贡献记录处撤回任何状态的贡献。撤回确认后，平台立即取消待审核修订、停止该贡献渠道承接新流量，并删除对应正式渠道及其路由能力。 +2. 撤回时，平台将清除贡献修订中保存的 API Key；贡献历史、测试结果、审核记录、奖励流水及必要审计记录可继续保留，但不再包含可用于调用上游的有效凭据。 +3. 撤回不影响撤回生效前已经发生的调用、计费、奖励和审计记录，也不撤销已经完成的奖励余额划转。 + +## 十、上游服务与贡献者责任 + +1. 上游服务的接口、模型、价格、额度、地区限制、内容规则和服务条款可能随时变化。您应确保贡献渠道的持续授权和可用性，并在发现变化后及时更新或撤回。 +2. 因上游余额不足、限流、服务中断、账户封禁、接口变更、模型下线或第三方限制导致的不可用，由对应上游服务及账户状态决定；平台可据此停用或删除渠道。 +3. 您不得利用贡献功能干扰平台运行、绕过访问控制、提交恶意端点，或诱导平台访问与模型服务无关的系统和数据。 + +## 十一、记录与通知 + +平台可记录贡献配置、测试结果、审核结果、健康状态、协议版本和同意时间，用于渠道管理、故障排查、计费核对和审计。与贡献相关的状态变化及补充要求，可通过站内页面、系统通知或平台已提供的联系方式告知您。 + +## 十二、协议更新 + +1. 平台可根据功能、运营规则或上游要求更新本协议，并发布新的协议版本。 +2. 已提交的贡献保留其提交时确认的协议版本记录；新建、重新提交或发生需要重新确认的重大配置变更时，您须阅读并接受当时有效的最新版本。 +3. 若您不同意更新后的协议，请勿继续提交新的贡献，并可按本协议第九条申请撤回已有渠道。 + +## 十三、联系与解释 + +如对贡献渠道、审核结果、状态变化或本协议有疑问，请通过平台提供的支持渠道联系管理员。具体功能名称、状态展示和配置数值以平台实际页面及管理员配置为准。` + +type ChannelContributionSetting struct { + Tag string `json:"tag"` + AllowedGroups []string `json:"allowed_groups"` + AllowedChannelTypes []int `json:"allowed_channel_types"` + Priority int64 `json:"priority"` + Weight uint `json:"weight"` + UnavailableDeleteHours int `json:"unavailable_delete_hours"` + HealthCheckIntervalMinutes int `json:"health_check_interval_minutes"` + RewardBps int `json:"reward_bps"` + AgreementVersion string `json:"agreement_version"` + AgreementContent string `json:"agreement_content"` +} + +var supportedChannelContributionTypeList = []int{ + constant.ChannelTypeOpenAI, + constant.ChannelTypeOllama, + constant.ChannelTypeAnthropic, + constant.ChannelTypeAli, + constant.ChannelTypeOpenRouter, + constant.ChannelTypeTencent, + constant.ChannelTypeGemini, + constant.ChannelTypeMoonshot, + constant.ChannelTypeZhipu_v4, + constant.ChannelTypePerplexity, + constant.ChannelTypeLingYiWanWu, + constant.ChannelTypeCohere, + constant.ChannelTypeMiniMax, + constant.ChannelTypeSiliconFlow, + constant.ChannelTypeMistral, + constant.ChannelTypeDeepSeek, + constant.ChannelTypeXinference, + constant.ChannelTypeXai, + constant.ChannelTypeSub2API, + constant.ChannelTypeNewAPI, +} + +var channelContributionSetting = ChannelContributionSetting{ + Tag: "donate", + AllowedGroups: []string{"default"}, + AllowedChannelTypes: append([]int(nil), supportedChannelContributionTypeList...), + Priority: 100, + Weight: 0, + UnavailableDeleteHours: 48, + HealthCheckIntervalMinutes: 10, + RewardBps: 0, + AgreementVersion: "2026-08-16", + AgreementContent: DefaultChannelContributionAgreementContent, +} + +var supportedChannelContributionTypes = func() map[int]struct{} { + types := make(map[int]struct{}, len(supportedChannelContributionTypeList)) + for _, channelType := range supportedChannelContributionTypeList { + types[channelType] = struct{}{} + } + return types +}() + +func init() { + config.GlobalConfig.Register("channel_contribution_setting", &channelContributionSetting) +} + +func GetChannelContributionSetting() *ChannelContributionSetting { + setting := ChannelContributionSetting{} + if err := config.GlobalConfig.Snapshot("channel_contribution_setting", &setting); err != nil { + common.SysError("failed to snapshot channel contribution setting: " + err.Error()) + } + if strings.TrimSpace(setting.Tag) == "" { + setting.Tag = "donate" + } + if setting.UnavailableDeleteHours <= 0 { + setting.UnavailableDeleteHours = 48 + } + if setting.HealthCheckIntervalMinutes <= 0 { + setting.HealthCheckIntervalMinutes = 10 + } + if strings.TrimSpace(setting.AgreementVersion) == "" { + setting.AgreementVersion = "2026-08-16" + } + if strings.TrimSpace(setting.AgreementContent) == "" { + setting.AgreementContent = DefaultChannelContributionAgreementContent + } + return &setting +} + +func (setting *ChannelContributionSetting) IsGroupAllowed(group string) bool { + group = strings.TrimSpace(group) + for _, allowed := range setting.AllowedGroups { + if strings.TrimSpace(allowed) == group { + return true + } + } + return false +} + +func (setting *ChannelContributionSetting) IsChannelTypeAllowed(channelType int) bool { + if !IsChannelContributionTypeSupported(channelType) { + return false + } + for _, allowed := range setting.AllowedChannelTypes { + if allowed == channelType { + return true + } + } + return false +} + +func IsChannelContributionTypeSupported(channelType int) bool { + _, ok := supportedChannelContributionTypes[channelType] + return ok +} + +func GetSupportedChannelContributionTypes() []int { + return append([]int(nil), supportedChannelContributionTypeList...) +} + +func ValidateChannelContributionOption(key string, value string) error { + if !strings.HasPrefix(key, ChannelContributionSettingPrefix) { + return nil + } + + field := strings.TrimPrefix(key, ChannelContributionSettingPrefix) + switch field { + case "allowed_groups": + var groups []string + if err := common.UnmarshalJsonStr(value, &groups); err != nil { + return fmt.Errorf("allowed_groups must be a JSON string array: %w", err) + } + seen := make(map[string]struct{}, len(groups)) + for _, group := range groups { + group = strings.TrimSpace(group) + if group == "" || len(group) > 64 { + return fmt.Errorf("allowed_groups contains an invalid group") + } + if _, exists := seen[group]; exists { + return fmt.Errorf("allowed_groups contains duplicate group %q", group) + } + seen[group] = struct{}{} + } + case "allowed_channel_types": + var channelTypes []int + if err := common.UnmarshalJsonStr(value, &channelTypes); err != nil { + return fmt.Errorf("allowed_channel_types must be a JSON integer array: %w", err) + } + seen := make(map[int]struct{}, len(channelTypes)) + for _, channelType := range channelTypes { + if !IsChannelContributionTypeSupported(channelType) { + return fmt.Errorf("unsupported channel type %d", channelType) + } + if _, exists := seen[channelType]; exists { + return fmt.Errorf("allowed_channel_types contains duplicate type %d", channelType) + } + seen[channelType] = struct{}{} + } + case "tag": + if strings.TrimSpace(value) == "" || len(strings.TrimSpace(value)) > 64 { + return fmt.Errorf("tag must contain 1 to 64 characters") + } + case "agreement_version": + if strings.TrimSpace(value) == "" || len(strings.TrimSpace(value)) > 64 { + return fmt.Errorf("agreement_version must contain 1 to 64 characters") + } + case "agreement_content": + if strings.TrimSpace(value) == "" || len(value) > 100_000 { + return fmt.Errorf("agreement_content must contain 1 to 100000 characters") + } + case "unavailable_delete_hours", "health_check_interval_minutes": + parsed, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil || parsed <= 0 { + return fmt.Errorf("%s must be a positive integer", field) + } + case "reward_bps": + parsed, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil || parsed < 0 || parsed > 10_000 { + return fmt.Errorf("reward_bps must be an integer from 0 to 10000") + } + case "priority", "weight": + parsed, err := strconv.ParseInt(strings.TrimSpace(value), 10, 64) + if err != nil || parsed < 0 { + return fmt.Errorf("%s must be a non-negative integer", field) + } + } + return nil +} diff --git a/setting/operation_setting/channel_contribution_setting_test.go b/setting/operation_setting/channel_contribution_setting_test.go new file mode 100644 index 000000000000..a34823581a63 --- /dev/null +++ b/setting/operation_setting/channel_contribution_setting_test.go @@ -0,0 +1,83 @@ +package operation_setting + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/setting/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValidateChannelContributionOptionRestrictsChannelTypesToSupportedSubset(t *testing.T) { + assert.NoError(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"allowed_channel_types", + `[1,14,60]`, + )) + assert.Error(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"allowed_channel_types", + `[3]`, + )) + assert.Error(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"allowed_channel_types", + `[1,1]`, + )) + assert.Error(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"allowed_channel_types", + `["1"]`, + )) +} + +func TestValidateChannelContributionOptionRejectsInvalidGroupsAndDurations(t *testing.T) { + assert.NoError(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"allowed_groups", + `["default","vip"]`, + )) + assert.Error(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"allowed_groups", + `["default"," default "]`, + )) + assert.Error(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"allowed_groups", + `[""]`, + )) + assert.NoError(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"unavailable_delete_hours", + "48", + )) + assert.Error(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"unavailable_delete_hours", + "0", + )) + assert.Error(t, ValidateChannelContributionOption( + ChannelContributionSettingPrefix+"unavailable_delete_hours", + "48 hours", + )) +} + +func TestGetChannelContributionSettingReturnsDeepCopy(t *testing.T) { + previous := make(map[string]string) + require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error { + if strings.HasPrefix(key, ChannelContributionSettingPrefix) { + previous[key] = value + } + return nil + })) + t.Cleanup(func() { + require.NoError(t, config.GlobalConfig.LoadFromDB(previous)) + }) + require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{ + ChannelContributionSettingPrefix + "allowed_groups": `["fixture"]`, + ChannelContributionSettingPrefix + "allowed_channel_types": `[1]`, + })) + + first := GetChannelContributionSetting() + require.NotEmpty(t, first.AllowedGroups) + require.NotEmpty(t, first.AllowedChannelTypes) + first.AllowedGroups[0] = "mutated" + first.AllowedChannelTypes[0] = -1 + + second := GetChannelContributionSetting() + assert.NotEqual(t, "mutated", second.AllowedGroups[0]) + assert.NotEqual(t, -1, second.AllowedChannelTypes[0]) +} diff --git a/web/src/features/channel-contributions/__tests__/api-contract.test.ts b/web/src/features/channel-contributions/__tests__/api-contract.test.ts new file mode 100644 index 000000000000..e752a0a77412 --- /dev/null +++ b/web/src/features/channel-contributions/__tests__/api-contract.test.ts @@ -0,0 +1,90 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { afterEach, assert, describe, test } from 'vitest' + +import { api } from '@/lib/http-client' + +import { + getAdminChannelContributions, + getChannelContributionRewardTransfers, + getChannelContributionRewards, + getChannelContributions, +} from '../api' + +type ApiGet = ( + url: string, + config?: unknown +) => Promise<{ data: { success: boolean; data: unknown } }> + +const apiClient = api as unknown as { get: ApiGet } +const originalGet = apiClient.get + +afterEach(() => { + apiClient.get = originalGet +}) + +describe('channel contribution API contract', () => { + test('uses the backend p parameter for every paginated endpoint', async () => { + const calls: Array<{ url: string; config?: unknown }> = [] + apiClient.get = async (url, config) => { + calls.push({ url, config }) + return { data: { success: true, data: { items: [], total: 0 } } } + } + + await getChannelContributions({ page: 3, page_size: 20 }) + await getAdminChannelContributions({ + page: 4, + page_size: 25, + status: 'pending', + }) + await getChannelContributionRewards({ page: 5, page_size: 50 }) + await getChannelContributionRewardTransfers({ page: 6, page_size: 10 }) + + assert.deepEqual(calls, [ + { + url: '/api/channel-contributions', + config: { + params: { p: 3, page_size: 20, status: undefined }, + disableDuplicate: true, + }, + }, + { + url: '/api/channel-contributions/admin', + config: { + params: { p: 4, page_size: 25, status: 'pending' }, + disableDuplicate: true, + }, + }, + { + url: '/api/channel-contributions/rewards', + config: { + params: { p: 5, page_size: 50 }, + disableDuplicate: true, + }, + }, + { + url: '/api/channel-contributions/reward-transfers', + config: { + params: { p: 6, page_size: 10 }, + disableDuplicate: true, + }, + }, + ]) + }) +}) diff --git a/web/src/features/channel-contributions/admin.tsx b/web/src/features/channel-contributions/admin.tsx new file mode 100644 index 000000000000..8539d58426ca --- /dev/null +++ b/web/src/features/channel-contributions/admin.tsx @@ -0,0 +1,119 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { Link } from '@tanstack/react-router' +import { HeartHandshake, ListChecks, Settings, UserRound } from 'lucide-react' +import { useState } from 'react' +import { useTranslation } from 'react-i18next' + +import { SectionPageLayout } from '@/components/layout' +import { Button } from '@/components/ui/button' +import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs' + +import { AdminContributionDetail } from './components/admin-contribution-detail' +import { AdminContributionList } from './components/admin-contribution-list' +import { AdminContributionSettings } from './components/admin-contribution-settings' + +type AdminTab = 'pending' | 'all' | 'settings' + +export function ChannelContributionAdmin() { + const { t } = useTranslation() + const [tab, setTab] = useState('pending') + const [selectedId, setSelectedId] = useState(null) + + return ( + <> + + + {t('Channel Contribution Review')} + + + + + +
+ setTab(value as AdminTab)} + > +
+ + + + + + + + +
+ + + + + + + + + + + + + +
+ + + + {selectedId ? ( + { + if (!open) setSelectedId(null) + }} + /> + ) : null} + + ) +} diff --git a/web/src/features/channel-contributions/api.ts b/web/src/features/channel-contributions/api.ts new file mode 100644 index 000000000000..1e92c67d8559 --- /dev/null +++ b/web/src/features/channel-contributions/api.ts @@ -0,0 +1,261 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { api } from '@/lib/http-client' + +import type { + ApiResponse, + ChannelContribution, + ChannelContributionAdminSettings, + ChannelContributionConfig, + ChannelContributionFetchModelsResult, + ChannelContributionList, + ChannelContributionPayload, + ChannelContributionRewardSummary, + ChannelContributionRewardTransfer, + ChannelContributionRewardTransferList, + ChannelContributionSubmitPayload, + ChannelContributionTestRun, +} from './types' + +const basePath = '/api/channel-contributions' + +export async function getChannelContributionConfig(): Promise< + ApiResponse +> { + const response = await api.get>( + `${basePath}/config` + ) + return response.data +} + +export async function getChannelContributions(params?: { + page?: number + page_size?: number + status?: string +}): Promise> { + const query = params + ? { p: params.page, page_size: params.page_size, status: params.status } + : undefined + const response = await api.get< + ApiResponse + >(basePath, { params: query, disableDuplicate: true }) + return response.data +} + +export async function getChannelContribution( + id: number +): Promise> { + const response = await api.get>( + `${basePath}/${id}`, + { disableDuplicate: true } + ) + return response.data +} + +export async function createChannelContribution( + payload: ChannelContributionPayload +): Promise> { + const response = await api.post>( + basePath, + payload + ) + return response.data +} + +export async function updateChannelContribution( + id: number, + payload: ChannelContributionPayload +): Promise> { + const response = await api.put>( + `${basePath}/${id}`, + payload + ) + return response.data +} + +export async function fetchChannelContributionModels( + id: number +): Promise> { + const response = await api.post< + ApiResponse + >(`${basePath}/${id}/fetch-models`) + return response.data +} + +export async function createChannelContributionTestRun( + id: number +): Promise> { + const response = await api.post>( + `${basePath}/${id}/test-runs` + ) + return response.data +} + +export async function getChannelContributionTestRun( + id: number, + runId: number | string +): Promise> { + const response = await api.get>( + `${basePath}/${id}/test-runs/${encodeURIComponent(runId)}`, + { disableDuplicate: true } + ) + return response.data +} + +export async function submitChannelContribution( + id: number, + payload: ChannelContributionSubmitPayload, + turnstile?: string +): Promise> { + const response = await api.post>( + `${basePath}/${id}/submit`, + payload, + { params: turnstile ? { turnstile } : undefined } + ) + return response.data +} + +export async function withdrawChannelContribution( + id: number +): Promise> { + const response = await api.post>( + `${basePath}/${id}/withdraw` + ) + return response.data +} + +export async function getChannelContributionRewards(params?: { + page?: number + page_size?: number +}): Promise> { + const query = params + ? { p: params.page, page_size: params.page_size } + : undefined + const response = await api.get>( + `${basePath}/rewards`, + { params: query, disableDuplicate: true } + ) + return response.data +} + +export async function getChannelContributionRewardTransfers(params?: { + page?: number + page_size?: number +}): Promise> { + const query = params + ? { p: params.page, page_size: params.page_size } + : undefined + const response = await api.get< + ApiResponse + >(`${basePath}/reward-transfers`, { + params: query, + disableDuplicate: true, + }) + return response.data +} + +export async function createChannelContributionRewardTransfer( + amount: number +): Promise> { + const response = await api.post< + ApiResponse + >(`${basePath}/reward-transfers`, { amount }) + return response.data +} + +export async function getAdminChannelContributions(params?: { + page?: number + page_size?: number + status?: string +}): Promise> { + const query = params + ? { p: params.page, page_size: params.page_size, status: params.status } + : undefined + const response = await api.get< + ApiResponse + >(`${basePath}/admin`, { params: query, disableDuplicate: true }) + return response.data +} + +export async function getAdminChannelContribution( + id: number +): Promise> { + const response = await api.get>( + `${basePath}/admin/${id}`, + { disableDuplicate: true } + ) + return response.data +} + +export async function createAdminChannelContributionTestRun( + id: number +): Promise> { + const response = await api.post>( + `${basePath}/admin/${id}/test-runs` + ) + return response.data +} + +export async function approveChannelContribution( + id: number, + testRunId: number | string +): Promise> { + const response = await api.post>( + `${basePath}/admin/${id}/approve`, + { test_run_id: testRunId } + ) + return response.data +} + +export async function rejectChannelContribution( + id: number, + reason: string +): Promise> { + const response = await api.post>( + `${basePath}/admin/${id}/reject`, + { reason } + ) + return response.data +} + +export async function deleteAdminChannelContribution( + id: number +): Promise { + const response = await api.delete(`${basePath}/admin/${id}`) + return response.data +} + +export async function getChannelContributionAdminSettings(): Promise< + ApiResponse +> { + const response = await api.get>( + `${basePath}/admin/settings` + ) + return response.data +} + +export async function updateChannelContributionAdminSettings( + payload: ChannelContributionAdminSettings +): Promise> { + const response = await api.put>( + `${basePath}/admin/settings`, + payload + ) + return response.data +} diff --git a/web/src/features/channel-contributions/components/__tests__/turnstile-submission.test.tsx b/web/src/features/channel-contributions/components/__tests__/turnstile-submission.test.tsx new file mode 100644 index 000000000000..40ef367408b2 --- /dev/null +++ b/web/src/features/channel-contributions/components/__tests__/turnstile-submission.test.tsx @@ -0,0 +1,281 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { Window } from 'happy-dom' +import { afterAll, afterEach, assert, describe, test } from 'vitest' + +const domWindow = new Window() +Object.defineProperty(domWindow, 'PointerEvent', { + configurable: true, + value: domWindow.MouseEvent, +}) +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'KeyboardEvent', + 'PointerEvent', + 'MouseEvent', + 'FocusEvent', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +const { act, useState } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { api } = await import('@/lib/http-client') +const { submitChannelContribution } = await import('../../api') +const { executeTurnstileSubmission } = await import('../../lib') +const { ContributionSubmissionControls } = + await import('../submission-controls') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { en: { translation: {} } }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +type ApiPost = ( + url: string, + data?: unknown, + config?: unknown +) => Promise<{ data: unknown }> +const apiClient = api as unknown as { post: ApiPost } +const originalPost = apiClient.post +let renderedRoot: ReturnType | null = null +let renderedHost: HTMLDivElement | null = null + +function findSubmitButton(): HTMLButtonElement { + const button = [ + ...document.querySelectorAll('button'), + ].find((candidate) => candidate.textContent?.includes('Submit for review')) + assert.ok(button) + return button +} + +async function renderControls(props: { + token: string + onSubmit: () => void + onExpire?: () => void + onAgreementCheckedChange?: (checked: boolean) => void + onOpenAgreement?: () => void +}) { + const host = document.createElement('div') + document.body.append(host) + const root = createRoot(host) + renderedHost = host + renderedRoot = root + await act(async () => { + root.render( + + undefined) + } + onOpenAgreement={props.onOpenAgreement ?? (() => undefined)} + isTurnstileEnabled + turnstileSiteKey='test-site-key' + turnstileToken={props.token} + turnstileWidgetKey={0} + onTurnstileVerify={() => undefined} + onTurnstileExpire={props.onExpire ?? (() => undefined)} + onSubmit={props.onSubmit} + submitting={false} + /> + + ) + await Promise.resolve() + }) +} + +afterEach(async () => { + apiClient.post = originalPost + if (renderedRoot) { + await act(async () => renderedRoot?.unmount()) + } + renderedHost?.remove() + renderedRoot = null + renderedHost = null + document.body.replaceChildren() + delete (window as unknown as Window & { turnstile?: unknown }).turnstile +}) + +afterAll(() => { + domWindow.close() +}) + +describe('channel contribution Turnstile submission', () => { + test('disables submit and does not call the handler when verification is empty', async () => { + let calls = 0 + await renderControls({ token: '', onSubmit: () => calls++ }) + const button = findSubmitButton() + + assert.equal(button.disabled, true) + button.click() + assert.equal(calls, 0) + }) + + test('sends a verified token as the submit query parameter', async () => { + let captured: { url?: string; config?: unknown } = {} + apiClient.post = async (url, _data, config) => { + captured = { url, config } + return { data: { success: true, data: { id: 7 } } } + } + + await submitChannelContribution( + 7, + { + test_run_id: 91, + agreement_version: '2026-08-16', + agreement_accepted: true, + }, + 'verified-token' + ) + + assert.equal(captured.url, '/api/channel-contributions/7/submit') + assert.deepEqual(captured.config, { + params: { turnstile: 'verified-token' }, + }) + }) + + test('opens the agreement without changing the checkbox state', async () => { + let agreementChanges = 0 + let agreementOpens = 0 + await renderControls({ + token: 'verified-token', + onSubmit: () => undefined, + onAgreementCheckedChange: () => agreementChanges++, + onOpenAgreement: () => agreementOpens++, + }) + const agreementButton = [ + ...document.querySelectorAll('button'), + ].find((button) => + button.textContent?.includes('Channel Contribution Agreement') + ) + assert.ok(agreementButton) + + agreementButton.click() + + assert.equal(agreementOpens, 1) + assert.equal(agreementChanges, 0) + }) + + test('resets the token and widget after success and business failure', async () => { + for (const success of [true, false]) { + let token = 'verified-token' + let widgetKey = 4 + const execution = await executeTurnstileSubmission({ + enabled: true, + token, + submit: async (submittedToken) => ({ success, submittedToken }), + reset: () => { + token = '' + widgetKey++ + }, + }) + + assert.equal(execution.called, true) + assert.equal(execution.result?.submittedToken, 'verified-token') + assert.equal(token, '') + assert.equal(widgetKey, 5) + } + }) + + test('clears verification when the Turnstile widget expires', async () => { + let renderOptions: Record | null = null + window.turnstile = { + render: (_element, options) => { + renderOptions = options + }, + } + + function Harness() { + const [token, setToken] = useState('verified-token') + const [widgetKey, setWidgetKey] = useState(0) + return ( + undefined} + onOpenAgreement={() => undefined} + isTurnstileEnabled + turnstileSiteKey='test-site-key' + turnstileToken={token} + turnstileWidgetKey={widgetKey} + onTurnstileVerify={setToken} + onTurnstileExpire={() => { + setToken('') + setWidgetKey((current) => current + 1) + }} + onSubmit={() => undefined} + submitting={false} + /> + ) + } + + const host = document.createElement('div') + document.body.append(host) + const root = createRoot(host) + renderedHost = host + renderedRoot = root + await act(async () => { + root.render( + + + + ) + await Promise.resolve() + }) + assert.equal(findSubmitButton().disabled, false) + assert.ok(renderOptions) + const expire = renderOptions['expired-callback'] + assert.equal(typeof expire, 'function') + + await act(async () => { + ;(expire as () => void)() + await Promise.resolve() + }) + assert.equal(findSubmitButton().disabled, true) + }) +}) diff --git a/web/src/features/channel-contributions/components/admin-contribution-detail.tsx b/web/src/features/channel-contributions/components/admin-contribution-detail.tsx new file mode 100644 index 000000000000..ae318531ced8 --- /dev/null +++ b/web/src/features/channel-contributions/components/admin-contribution-detail.tsx @@ -0,0 +1,475 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query' +import { Check, FlaskConical, Loader2, Trash2, X } from 'lucide-react' +import { useState } from 'react' +import { useTranslation } from 'react-i18next' +import { toast } from 'sonner' + +import { ConfirmDialog } from '@/components/confirm-dialog' +import { Button } from '@/components/ui/button' +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from '@/components/ui/dialog' +import { Label } from '@/components/ui/label' +import { ScrollArea } from '@/components/ui/scroll-area' +import { Textarea } from '@/components/ui/textarea' + +import { + approveChannelContribution, + createAdminChannelContributionTestRun, + deleteAdminChannelContribution, + getAdminChannelContribution, + rejectChannelContribution, +} from '../api' +import { + formatContributionModelMapping, + formatContributionTimestamp, + getContributionName, + getContributionRevision, + getContributionTestRun, + getTestRunId, + hasPendingContributionRevision, + isTestRunActive, + parseContributionModels, + testRunPassed, +} from '../lib' +import type { ChannelContributionTestRun } from '../types' +import { + ContributionRevisionStatusBadge, + ContributionStatusBadge, +} from './contribution-status' +import { ContributionTestMatrix } from './test-matrix' + +function DetailValue(props: { label: string; children: React.ReactNode }) { + return ( +
+
{props.label}
+
{props.children}
+
+ ) +} + +export function AdminContributionDetail(props: { + id: number + open: boolean + onOpenChange: (open: boolean) => void +}) { + const { t } = useTranslation() + const queryClient = useQueryClient() + const [createdAdminRun, setCreatedAdminRun] = + useState(null) + const [rejectOpen, setRejectOpen] = useState(false) + const [rejectReason, setRejectReason] = useState('') + const [deleteOpen, setDeleteOpen] = useState(false) + + const detailQuery = useQuery({ + queryKey: ['channel-contributions', 'admin', 'detail', props.id], + queryFn: async () => { + const response = await getAdminChannelContribution(props.id) + if (!response.success || !response.data) { + throw new Error(response.message || t('Failed to load contribution')) + } + return response.data + }, + enabled: props.open, + refetchInterval: (query) => { + const embedded = query.state.data?.latest_test_run + return isTestRunActive(createdAdminRun) || isTestRunActive(embedded) + ? 1500 + : false + }, + }) + const contribution = detailQuery.data + const revision = getContributionRevision(contribution) + const pendingReview = hasPendingContributionRevision(contribution) + const embeddedRun = getContributionTestRun(contribution) + const createdRunId = getTestRunId(createdAdminRun) + const embeddedRunId = getTestRunId(embeddedRun) + let adminRun = createdAdminRun + if (embeddedRun?.actor_type === 'admin' && !createdAdminRun) { + adminRun = embeddedRun + } else if (createdRunId && createdRunId === embeddedRunId) { + adminRun = embeddedRun + } + + const testMutation = useMutation({ + mutationFn: createAdminChannelContributionTestRun, + }) + const approveMutation = useMutation({ + mutationFn: (testRunId: number | string) => + approveChannelContribution(props.id, testRunId), + }) + const rejectMutation = useMutation({ + mutationFn: (reason: string) => rejectChannelContribution(props.id, reason), + }) + const deleteMutation = useMutation({ + mutationFn: () => deleteAdminChannelContribution(props.id), + }) + + const refreshLists = async () => { + await queryClient.invalidateQueries({ + queryKey: ['channel-contributions', 'admin'], + }) + } + + const handleTest = async () => { + try { + const response = await testMutation.mutateAsync(props.id) + if (!response.success || !response.data) { + toast.error(response.message || t('Failed to start model tests')) + return + } + setCreatedAdminRun(response.data) + toast.success(t('Administrator verification started')) + } catch (error) { + toast.error( + error instanceof Error + ? error.message + : t('Failed to start model tests') + ) + } + } + + const handleApprove = async () => { + const runId = getTestRunId(adminRun) + if (!runId) return + try { + const response = await approveMutation.mutateAsync(runId) + if (!response.success) { + toast.error(response.message || t('Failed to approve contribution')) + return + } + await refreshLists() + props.onOpenChange(false) + toast.success(t('Contribution approved')) + } catch (error) { + toast.error( + error instanceof Error + ? error.message + : t('Failed to approve contribution') + ) + } + } + + const handleReject = async () => { + if (!rejectReason.trim()) return + try { + const response = await rejectMutation.mutateAsync(rejectReason.trim()) + if (!response.success) { + toast.error(response.message || t('Failed to reject contribution')) + return + } + setRejectOpen(false) + await refreshLists() + props.onOpenChange(false) + toast.success(t('Contribution rejected')) + } catch (error) { + toast.error( + error instanceof Error + ? error.message + : t('Failed to reject contribution') + ) + } + } + + const handleDelete = async () => { + try { + const response = await deleteMutation.mutateAsync() + if (!response.success) { + toast.error(response.message || t('Failed to delete contribution')) + return + } + setDeleteOpen(false) + await refreshLists() + props.onOpenChange(false) + toast.success(t('Contribution deleted')) + } catch (error) { + toast.error( + error instanceof Error + ? error.message + : t('Failed to delete contribution') + ) + } + } + + const busy = + testMutation.isPending || + approveMutation.isPending || + rejectMutation.isPending || + deleteMutation.isPending || + isTestRunActive(adminRun) + const approveReady = pendingReview && testRunPassed(adminRun) + + return ( + <> + + + +
+ + {contribution + ? getContributionName(contribution) + : t('Contribution review')} + + {contribution ? ( + <> + + {contribution.revision_status === 'pending' && + contribution.status !== 'pending' ? ( + + ) : null} + + ) : null} +
+ + {t( + 'Run an independent administrator test before approving this revision.' + )} + + + + +
+ {detailQuery.isLoading ? ( +
+
+ ) : null} + {!detailQuery.isLoading && + (detailQuery.error || !contribution || !revision) ? ( +

+ {detailQuery.error instanceof Error + ? detailQuery.error.message + : t('Contribution details are incomplete')} +

+ ) : null} + {!detailQuery.isLoading && + !detailQuery.error && + contribution && + revision ? ( + <> +
+

+ {t('Connection details')} +

+
+ + {contribution.username || '-'} · {t('ID')}{' '} + {contribution.user_id ?? '-'} + + + {revision.type} + + + {revision.group} + + + + {revision.base_url} + + + + {revision.revision_number ?? '-'} + + + {formatContributionTimestamp( + revision.submitted_at || contribution.submitted_at + )} + +
+
+ +
+
+

{t('Models')}

+ + {t('{{count}} models', { + count: parseContributionModels(revision.models) + .length, + })} + +
+
+ {parseContributionModels(revision.models).map((model) => ( + + {model} + + ))} +
+
+

+ {t('Model Mapping')} +

+ 
+                        {formatContributionModelMapping(
+                          revision.model_mapping
+                        ) || '{}'}
+
+
+
+ + {contribution.review_reason ? ( +
+

+ {t('Review note')} +

+

+ {contribution.review_reason} +

+
+ ) : null} + +
+
+
+

+ {t('Administrator verification')} +

+

+ {t( + 'Approval is bound to this administrator test run ID.' + )} +

+
+ +
+ +
+ + ) : null} +
+ + + + + + + + + + + + + + {t('Reject contribution')} + + {t( + 'Give the contributor a clear reason they can address before resubmitting.' + )} + + +
+ +