Merge branch 'upstream-main' into feature/improve-param-override

# Conflicts:
#	relay/channel/api_request_test.go
#	relay/common/override_test.go
#	web/src/components/table/channels/modals/EditChannelModal.jsx
This commit is contained in:
Seefs
2026-02-25 13:39:54 +08:00
124 changed files with 7729 additions and 1971 deletions

View File

@@ -145,6 +145,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
return
}
channelProxy := ""
if channelID > 0 {
ch, err := model.GetChannelById(channelID, false)
if err != nil {
@@ -159,6 +160,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel type is not Codex"})
return
}
channelProxy = ch.GetSetting().Proxy
}
session := sessions.Default(c)
@@ -176,7 +178,7 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second)
defer cancel()
tokenRes, err := service.ExchangeCodexAuthorizationCode(ctx, code, verifier)
tokenRes, err := service.ExchangeCodexAuthorizationCodeWithProxy(ctx, code, verifier, channelProxy)
if err != nil {
common.SysError("failed to exchange codex authorization code: " + err.Error())
c.JSON(http.StatusOK, gin.H{"success": false, "message": "授权码交换失败,请重试"})

View File

@@ -2,7 +2,6 @@ package controller
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strconv"
@@ -80,7 +79,7 @@ func GetCodexChannelUsage(c *gin.Context) {
refreshCtx, refreshCancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
defer refreshCancel()
res, refreshErr := service.RefreshCodexOAuthToken(refreshCtx, oauthKey.RefreshToken)
res, refreshErr := service.RefreshCodexOAuthTokenWithProxy(refreshCtx, oauthKey.RefreshToken, ch.GetSetting().Proxy)
if refreshErr == nil {
oauthKey.AccessToken = res.AccessToken
oauthKey.RefreshToken = res.RefreshToken
@@ -109,7 +108,7 @@ func GetCodexChannelUsage(c *gin.Context) {
}
var payload any
if json.Unmarshal(body, &payload) != nil {
if common.Unmarshal(body, &payload) != nil {
payload = string(body)
}

View File

@@ -1,8 +1,13 @@
package controller
import (
"context"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
@@ -16,6 +21,7 @@ type CustomOAuthProviderResponse struct {
Id int `json:"id"`
Name string `json:"name"`
Slug string `json:"slug"`
Icon string `json:"icon"`
Enabled bool `json:"enabled"`
ClientId string `json:"client_id"`
AuthorizationEndpoint string `json:"authorization_endpoint"`
@@ -28,6 +34,16 @@ type CustomOAuthProviderResponse struct {
EmailField string `json:"email_field"`
WellKnown string `json:"well_known"`
AuthStyle int `json:"auth_style"`
AccessPolicy string `json:"access_policy"`
AccessDeniedMessage string `json:"access_denied_message"`
}
type UserOAuthBindingResponse struct {
ProviderId int `json:"provider_id"`
ProviderName string `json:"provider_name"`
ProviderSlug string `json:"provider_slug"`
ProviderIcon string `json:"provider_icon"`
ProviderUserId string `json:"provider_user_id"`
}
func toCustomOAuthProviderResponse(p *model.CustomOAuthProvider) *CustomOAuthProviderResponse {
@@ -35,6 +51,7 @@ func toCustomOAuthProviderResponse(p *model.CustomOAuthProvider) *CustomOAuthPro
Id: p.Id,
Name: p.Name,
Slug: p.Slug,
Icon: p.Icon,
Enabled: p.Enabled,
ClientId: p.ClientId,
AuthorizationEndpoint: p.AuthorizationEndpoint,
@@ -47,6 +64,8 @@ func toCustomOAuthProviderResponse(p *model.CustomOAuthProvider) *CustomOAuthPro
EmailField: p.EmailField,
WellKnown: p.WellKnown,
AuthStyle: p.AuthStyle,
AccessPolicy: p.AccessPolicy,
AccessDeniedMessage: p.AccessDeniedMessage,
}
}
@@ -96,6 +115,7 @@ func GetCustomOAuthProvider(c *gin.Context) {
type CreateCustomOAuthProviderRequest struct {
Name string `json:"name" binding:"required"`
Slug string `json:"slug" binding:"required"`
Icon string `json:"icon"`
Enabled bool `json:"enabled"`
ClientId string `json:"client_id" binding:"required"`
ClientSecret string `json:"client_secret" binding:"required"`
@@ -109,6 +129,85 @@ type CreateCustomOAuthProviderRequest struct {
EmailField string `json:"email_field"`
WellKnown string `json:"well_known"`
AuthStyle int `json:"auth_style"`
AccessPolicy string `json:"access_policy"`
AccessDeniedMessage string `json:"access_denied_message"`
}
type FetchCustomOAuthDiscoveryRequest struct {
WellKnownURL string `json:"well_known_url"`
IssuerURL string `json:"issuer_url"`
}
// FetchCustomOAuthDiscovery fetches OIDC discovery document via backend (root-only route)
func FetchCustomOAuthDiscovery(c *gin.Context) {
var req FetchCustomOAuthDiscoveryRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ApiErrorMsg(c, "无效的请求参数: "+err.Error())
return
}
wellKnownURL := strings.TrimSpace(req.WellKnownURL)
issuerURL := strings.TrimSpace(req.IssuerURL)
if wellKnownURL == "" && issuerURL == "" {
common.ApiErrorMsg(c, "请先填写 Discovery URL 或 Issuer URL")
return
}
targetURL := wellKnownURL
if targetURL == "" {
targetURL = strings.TrimRight(issuerURL, "/") + "/.well-known/openid-configuration"
}
targetURL = strings.TrimSpace(targetURL)
parsedURL, err := url.Parse(targetURL)
if err != nil || parsedURL.Host == "" || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") {
common.ApiErrorMsg(c, "Discovery URL 无效,仅支持 http/https")
return
}
ctx, cancel := context.WithTimeout(c.Request.Context(), 20*time.Second)
defer cancel()
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, targetURL, nil)
if err != nil {
common.ApiErrorMsg(c, "创建 Discovery 请求失败: "+err.Error())
return
}
httpReq.Header.Set("Accept", "application/json")
client := &http.Client{Timeout: 20 * time.Second}
resp, err := client.Do(httpReq)
if err != nil {
common.ApiErrorMsg(c, "获取 Discovery 配置失败: "+err.Error())
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
message := strings.TrimSpace(string(body))
if message == "" {
message = resp.Status
}
common.ApiErrorMsg(c, "获取 Discovery 配置失败: "+message)
return
}
var discovery map[string]any
if err = common.DecodeJson(resp.Body, &discovery); err != nil {
common.ApiErrorMsg(c, "解析 Discovery 配置失败: "+err.Error())
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": gin.H{
"well_known_url": targetURL,
"discovery": discovery,
},
})
}
// CreateCustomOAuthProvider creates a new custom OAuth provider
@@ -134,6 +233,7 @@ func CreateCustomOAuthProvider(c *gin.Context) {
provider := &model.CustomOAuthProvider{
Name: req.Name,
Slug: req.Slug,
Icon: req.Icon,
Enabled: req.Enabled,
ClientId: req.ClientId,
ClientSecret: req.ClientSecret,
@@ -147,6 +247,8 @@ func CreateCustomOAuthProvider(c *gin.Context) {
EmailField: req.EmailField,
WellKnown: req.WellKnown,
AuthStyle: req.AuthStyle,
AccessPolicy: req.AccessPolicy,
AccessDeniedMessage: req.AccessDeniedMessage,
}
if err := model.CreateCustomOAuthProvider(provider); err != nil {
@@ -168,9 +270,10 @@ func CreateCustomOAuthProvider(c *gin.Context) {
type UpdateCustomOAuthProviderRequest struct {
Name string `json:"name"`
Slug string `json:"slug"`
Enabled *bool `json:"enabled"` // Optional: if nil, keep existing
Icon *string `json:"icon"` // Optional: if nil, keep existing
Enabled *bool `json:"enabled"` // Optional: if nil, keep existing
ClientId string `json:"client_id"`
ClientSecret string `json:"client_secret"` // Optional: if empty, keep existing
ClientSecret string `json:"client_secret"` // Optional: if empty, keep existing
AuthorizationEndpoint string `json:"authorization_endpoint"`
TokenEndpoint string `json:"token_endpoint"`
UserInfoEndpoint string `json:"user_info_endpoint"`
@@ -181,6 +284,8 @@ type UpdateCustomOAuthProviderRequest struct {
EmailField string `json:"email_field"`
WellKnown *string `json:"well_known"` // Optional: if nil, keep existing
AuthStyle *int `json:"auth_style"` // Optional: if nil, keep existing
AccessPolicy *string `json:"access_policy"` // Optional: if nil, keep existing
AccessDeniedMessage *string `json:"access_denied_message"` // Optional: if nil, keep existing
}
// UpdateCustomOAuthProvider updates an existing custom OAuth provider
@@ -227,6 +332,9 @@ func UpdateCustomOAuthProvider(c *gin.Context) {
if req.Slug != "" {
provider.Slug = req.Slug
}
if req.Icon != nil {
provider.Icon = *req.Icon
}
if req.Enabled != nil {
provider.Enabled = *req.Enabled
}
@@ -266,6 +374,12 @@ func UpdateCustomOAuthProvider(c *gin.Context) {
if req.AuthStyle != nil {
provider.AuthStyle = *req.AuthStyle
}
if req.AccessPolicy != nil {
provider.AccessPolicy = *req.AccessPolicy
}
if req.AccessDeniedMessage != nil {
provider.AccessDeniedMessage = *req.AccessDeniedMessage
}
if err := model.UpdateCustomOAuthProvider(provider); err != nil {
common.ApiError(c, err)
@@ -327,6 +441,30 @@ func DeleteCustomOAuthProvider(c *gin.Context) {
})
}
func buildUserOAuthBindingsResponse(userId int) ([]UserOAuthBindingResponse, error) {
bindings, err := model.GetUserOAuthBindingsByUserId(userId)
if err != nil {
return nil, err
}
response := make([]UserOAuthBindingResponse, 0, len(bindings))
for _, binding := range bindings {
provider, err := model.GetCustomOAuthProviderById(binding.ProviderId)
if err != nil {
continue
}
response = append(response, UserOAuthBindingResponse{
ProviderId: binding.ProviderId,
ProviderName: provider.Name,
ProviderSlug: provider.Slug,
ProviderIcon: provider.Icon,
ProviderUserId: binding.ProviderUserId,
})
}
return response, nil
}
// GetUserOAuthBindings returns all OAuth bindings for the current user
func GetUserOAuthBindings(c *gin.Context) {
userId := c.GetInt("id")
@@ -335,32 +473,43 @@ func GetUserOAuthBindings(c *gin.Context) {
return
}
bindings, err := model.GetUserOAuthBindingsByUserId(userId)
response, err := buildUserOAuthBindingsResponse(userId)
if err != nil {
common.ApiError(c, err)
return
}
// Build response with provider info
type BindingResponse struct {
ProviderId int `json:"provider_id"`
ProviderName string `json:"provider_name"`
ProviderSlug string `json:"provider_slug"`
ProviderUserId string `json:"provider_user_id"`
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": response,
})
}
func GetUserOAuthBindingsByAdmin(c *gin.Context) {
userIdStr := c.Param("id")
userId, err := strconv.Atoi(userIdStr)
if err != nil {
common.ApiErrorMsg(c, "invalid user id")
return
}
response := make([]BindingResponse, 0)
for _, binding := range bindings {
provider, err := model.GetCustomOAuthProviderById(binding.ProviderId)
if err != nil {
continue // Skip if provider not found
}
response = append(response, BindingResponse{
ProviderId: binding.ProviderId,
ProviderName: provider.Name,
ProviderSlug: provider.Slug,
ProviderUserId: binding.ProviderUserId,
})
targetUser, err := model.GetUserById(userId, false)
if err != nil {
common.ApiError(c, err)
return
}
myRole := c.GetInt("role")
if myRole <= targetUser.Role && myRole != common.RoleRootUser {
common.ApiErrorMsg(c, "no permission")
return
}
response, err := buildUserOAuthBindingsResponse(userId)
if err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{
@@ -395,3 +544,41 @@ func UnbindCustomOAuth(c *gin.Context) {
"message": "解绑成功",
})
}
func UnbindCustomOAuthByAdmin(c *gin.Context) {
userIdStr := c.Param("id")
userId, err := strconv.Atoi(userIdStr)
if err != nil {
common.ApiErrorMsg(c, "invalid user id")
return
}
targetUser, err := model.GetUserById(userId, false)
if err != nil {
common.ApiError(c, err)
return
}
myRole := c.GetInt("role")
if myRole <= targetUser.Role && myRole != common.RoleRootUser {
common.ApiErrorMsg(c, "no permission")
return
}
providerIdStr := c.Param("provider_id")
providerId, err := strconv.Atoi(providerIdStr)
if err != nil {
common.ApiErrorMsg(c, "invalid provider id")
return
}
if err := model.DeleteUserOAuthBinding(userId, providerId); err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "success",
})
}

View File

@@ -105,13 +105,13 @@ func UpdateMidjourneyTaskBulk() {
}
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Get Task parse body error: %v", err))
logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error: %v", err))
continue
}
var responseItems []dto.MidjourneyDto
err = json.Unmarshal(responseBody, &responseItems)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Get Task parse body error2: %v, body: %s", err, string(responseBody)))
logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error2: %v, body: %s", err, string(responseBody)))
continue
}
resp.Body.Close()
@@ -130,6 +130,7 @@ func UpdateMidjourneyTaskBulk() {
if !checkMjTaskNeedUpdate(task, responseItem) {
continue
}
preStatus := task.Status
task.Code = 1
task.Progress = responseItem.Progress
task.PromptEn = responseItem.PromptEn
@@ -172,18 +173,26 @@ func UpdateMidjourneyTaskBulk() {
shouldReturnQuota = true
}
}
err = task.Update()
won, err := task.UpdateWithStatus(preStatus)
if err != nil {
logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error())
} else {
if shouldReturnQuota {
err = model.IncreaseUserQuota(task.UserId, task.Quota, false)
if err != nil {
logger.LogError(ctx, "fail to increase user quota: "+err.Error())
}
logContent := fmt.Sprintf("构图失败 %s补偿 %s", task.MjId, logger.LogQuota(task.Quota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
} else if won && shouldReturnQuota {
err = model.IncreaseUserQuota(task.UserId, task.Quota, false)
if err != nil {
logger.LogError(ctx, "fail to increase user quota: "+err.Error())
}
model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
UserId: task.UserId,
LogType: model.LogTypeRefund,
Content: "",
ChannelId: task.ChannelId,
ModelName: service.CovertMjpActionToModelName(task.Action),
Quota: task.Quota,
Other: map[string]interface{}{
"task_id": task.MjId,
"reason": "构图失败",
},
})
}
}
}

View File

@@ -134,8 +134,10 @@ func GetStatus(c *gin.Context) {
customProviders := oauth.GetEnabledCustomProviders()
if len(customProviders) > 0 {
type CustomOAuthInfo struct {
Id int `json:"id"`
Name string `json:"name"`
Slug string `json:"slug"`
Icon string `json:"icon"`
ClientId string `json:"client_id"`
AuthorizationEndpoint string `json:"authorization_endpoint"`
Scopes string `json:"scopes"`
@@ -144,8 +146,10 @@ func GetStatus(c *gin.Context) {
for _, p := range customProviders {
config := p.GetConfig()
providersInfo = append(providersInfo, CustomOAuthInfo{
Id: config.Id,
Name: config.Name,
Slug: config.Slug,
Icon: config.Icon,
ClientId: config.ClientId,
AuthorizationEndpoint: config.AuthorizationEndpoint,
Scopes: config.Scopes,

View File

@@ -237,6 +237,16 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
// Set up new user
user.Username = provider.GetProviderPrefix() + strconv.Itoa(model.GetMaxUserId()+1)
if oauthUser.Username != "" {
if exists, err := model.CheckUserExistOrDeleted(oauthUser.Username, ""); err == nil && !exists {
// 防止索引退化
if len(oauthUser.Username) <= model.UserNameMaxLength {
user.Username = oauthUser.Username
}
}
}
if oauthUser.DisplayName != "" {
user.DisplayName = oauthUser.DisplayName
} else if oauthUser.Username != "" {
@@ -295,12 +305,12 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
// Set the provider user ID on the user model and update
provider.SetProviderUserID(user, oauthUser.ProviderUserID)
if err := tx.Model(user).Updates(map[string]interface{}{
"github_id": user.GitHubId,
"discord_id": user.DiscordId,
"oidc_id": user.OidcId,
"linux_do_id": user.LinuxDOId,
"wechat_id": user.WeChatId,
"telegram_id": user.TelegramId,
"github_id": user.GitHubId,
"discord_id": user.DiscordId,
"oidc_id": user.OidcId,
"linux_do_id": user.LinuxDOId,
"wechat_id": user.WeChatId,
"telegram_id": user.TelegramId,
}).Error; err != nil {
return err
}
@@ -340,6 +350,8 @@ func handleOAuthError(c *gin.Context, err error) {
} else {
common.ApiErrorI18n(c, e.MsgKey)
}
case *oauth.AccessDeniedError:
common.ApiErrorMsg(c, e.Message)
case *oauth.TrustLevelError:
common.ApiErrorI18n(c, i18n.MsgOAuthTrustLevelLow)
default:

View File

@@ -455,72 +455,147 @@ func RelayNotFound(c *gin.Context) {
})
}
func RelayTask(c *gin.Context) {
retryTimes := common.RetryTimes
channelId := c.GetInt("channel_id")
c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)})
func RelayTaskFetch(c *gin.Context) {
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, &dto.TaskError{
Code: "gen_relay_info_failed",
Message: err.Error(),
StatusCode: http.StatusInternalServerError,
})
return
}
taskErr := taskRelayHandler(c, relayInfo)
if taskErr == nil {
retryTimes = 0
if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
respondTaskError(c, taskErr)
}
}
func RelayTask(c *gin.Context) {
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, &dto.TaskError{
Code: "gen_relay_info_failed",
Message: err.Error(),
StatusCode: http.StatusInternalServerError,
})
return
}
if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
respondTaskError(c, taskErr)
return
}
var result *relay.TaskSubmitResult
var taskErr *dto.TaskError
defer func() {
if taskErr != nil && relayInfo.Billing != nil {
relayInfo.Billing.Refund(c)
}
}()
retryParam := &service.RetryParam{
Ctx: c,
TokenGroup: relayInfo.TokenGroup,
ModelName: relayInfo.OriginModelName,
Retry: common.GetPointer(0),
}
for ; shouldRetryTaskRelay(c, channelId, taskErr, retryTimes) && retryParam.GetRetry() < retryTimes; retryParam.IncreaseRetry() {
channel, newAPIError := getChannel(c, relayInfo, retryParam)
if newAPIError != nil {
logger.LogError(c, fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", newAPIError.Error()))
taskErr = service.TaskErrorWrapperLocal(newAPIError.Err, "get_channel_failed", http.StatusInternalServerError)
break
}
channelId = channel.Id
useChannel := c.GetStringSlice("use_channel")
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
c.Set("use_channel", useChannel)
logger.LogInfo(c, fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, retryParam.GetRetry()))
//middleware.SetupContextForSelectedChannel(c, channel, originalModel)
bodyStorage, err := common.GetBodyStorage(c)
if err != nil {
if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusRequestEntityTooLarge)
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
var channel *model.Channel
if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
channel = lockedCh
if retryParam.GetRetry() > 0 {
if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
break
}
}
} else {
var channelErr *types.NewAPIError
channel, channelErr = getChannel(c, relayInfo, retryParam)
if channelErr != nil {
logger.LogError(c, channelErr.Error())
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
break
}
}
addUsedChannel(c, channel.Id)
bodyStorage, bodyErr := common.GetBodyStorage(c)
if bodyErr != nil {
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
} else {
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest)
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest)
}
break
}
c.Request.Body = io.NopCloser(bodyStorage)
taskErr = taskRelayHandler(c, relayInfo)
result, taskErr = relay.RelayTaskSubmit(c, relayInfo)
if taskErr == nil {
break
}
if !taskErr.LocalError {
processChannelError(c,
*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
}
if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {
break
}
}
useChannel := c.GetStringSlice("use_channel")
if len(useChannel) > 1 {
retryLogStr := fmt.Sprintf("重试:%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
logger.LogInfo(c, retryLogStr)
}
if taskErr != nil {
if taskErr.StatusCode == http.StatusTooManyRequests {
taskErr.Message = "当前分组上游负载已饱和,请稍后再试"
// ── 成功:结算 + 日志 + 插入任务 ──
if taskErr == nil {
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
common.SysError("settle task billing error: " + settleErr.Error())
}
c.JSON(taskErr.StatusCode, taskErr)
service.LogTaskConsumption(c, relayInfo)
task := model.InitTask(result.Platform, relayInfo)
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
task.PrivateData.BillingSource = relayInfo.BillingSource
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
task.PrivateData.TokenId = relayInfo.TokenId
task.PrivateData.BillingContext = &model.TaskBillingContext{
ModelPrice: relayInfo.PriceData.ModelPrice,
GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
ModelRatio: relayInfo.PriceData.ModelRatio,
OtherRatios: relayInfo.PriceData.OtherRatios,
OriginModelName: relayInfo.OriginModelName,
PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName),
}
task.Quota = result.Quota
task.Data = result.TaskData
task.Action = relayInfo.Action
if insertErr := task.Insert(); insertErr != nil {
common.SysError("insert task error: " + insertErr.Error())
}
}
if taskErr != nil {
respondTaskError(c, taskErr)
}
}
func taskRelayHandler(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dto.TaskError {
var err *dto.TaskError
switch relayInfo.RelayMode {
case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID, relayconstant.RelayModeVideoFetchByID:
err = relay.RelayTaskFetch(c, relayInfo.RelayMode)
default:
err = relay.RelayTaskSubmit(c, relayInfo)
// respondTaskError 统一输出 Task 错误响应(含 429 限流提示改写)
func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
if taskErr.StatusCode == http.StatusTooManyRequests {
taskErr.Message = "当前分组上游负载已饱和,请稍后再试"
}
return err
c.JSON(taskErr.StatusCode, taskErr)
}
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
@@ -544,7 +619,7 @@ func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError,
}
if taskErr.StatusCode/100 == 5 {
// 超时不重试
if taskErr.StatusCode == 504 || taskErr.StatusCode == 524 {
if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
return false
}
return true

View File

@@ -1,231 +1,22 @@
package controller
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"sort"
"strconv"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
// UpdateTaskBulk 薄入口,实际轮询逻辑在 service 层
func UpdateTaskBulk() {
//revocer
//imageModel := "midjourney"
for {
time.Sleep(time.Duration(15) * time.Second)
common.SysLog("任务进度轮询开始")
ctx := context.TODO()
allTasks := model.GetAllUnFinishSyncTasks(constant.TaskQueryLimit)
platformTask := make(map[constant.TaskPlatform][]*model.Task)
for _, t := range allTasks {
platformTask[t.Platform] = append(platformTask[t.Platform], t)
}
for platform, tasks := range platformTask {
if len(tasks) == 0 {
continue
}
taskChannelM := make(map[int][]string)
taskM := make(map[string]*model.Task)
nullTaskIds := make([]int64, 0)
for _, task := range tasks {
if task.TaskID == "" {
// 统计失败的未完成任务
nullTaskIds = append(nullTaskIds, task.ID)
continue
}
taskM[task.TaskID] = task
taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], task.TaskID)
}
if len(nullTaskIds) > 0 {
err := model.TaskBulkUpdateByID(nullTaskIds, map[string]any{
"status": "FAILURE",
"progress": "100%",
})
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Fix null task_id task error: %v", err))
} else {
logger.LogInfo(ctx, fmt.Sprintf("Fix null task_id task success: %v", nullTaskIds))
}
}
if len(taskChannelM) == 0 {
continue
}
UpdateTaskByPlatform(platform, taskChannelM, taskM)
}
common.SysLog("任务进度轮询完成")
}
}
func UpdateTaskByPlatform(platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) {
switch platform {
case constant.TaskPlatformMidjourney:
//_ = UpdateMidjourneyTaskAll(context.Background(), tasks)
case constant.TaskPlatformSuno:
_ = UpdateSunoTaskAll(context.Background(), taskChannelM, taskM)
default:
if err := UpdateVideoTaskAll(context.Background(), platform, taskChannelM, taskM); err != nil {
common.SysLog(fmt.Sprintf("UpdateVideoTaskAll fail: %s", err))
}
}
}
func UpdateSunoTaskAll(ctx context.Context, taskChannelM map[int][]string, taskM map[string]*model.Task) error {
for channelId, taskIds := range taskChannelM {
err := updateSunoTaskAll(ctx, channelId, taskIds, taskM)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("渠道 #%d 更新异步任务失败: %s", channelId, err.Error()))
}
}
return nil
}
func updateSunoTaskAll(ctx context.Context, channelId int, taskIds []string, taskM map[string]*model.Task) error {
logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds)))
if len(taskIds) == 0 {
return nil
}
channel, err := model.CacheGetChannel(channelId)
if err != nil {
common.SysLog(fmt.Sprintf("CacheGetChannel: %v", err))
err = model.TaskBulkUpdate(taskIds, map[string]any{
"fail_reason": fmt.Sprintf("获取渠道信息失败请联系管理员渠道ID%d", channelId),
"status": "FAILURE",
"progress": "100%",
})
if err != nil {
common.SysLog(fmt.Sprintf("UpdateMidjourneyTask error2: %v", err))
}
return err
}
adaptor := relay.GetTaskAdaptor(constant.TaskPlatformSuno)
if adaptor == nil {
return errors.New("adaptor not found")
}
proxy := channel.GetSetting().Proxy
resp, err := adaptor.FetchTask(*channel.BaseURL, channel.Key, map[string]any{
"ids": taskIds,
}, proxy)
if err != nil {
common.SysLog(fmt.Sprintf("Get Task Do req error: %v", err))
return err
}
if resp.StatusCode != http.StatusOK {
logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode))
return errors.New(fmt.Sprintf("Get Task status code: %d", resp.StatusCode))
}
defer resp.Body.Close()
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
common.SysLog(fmt.Sprintf("Get Task parse body error: %v", err))
return err
}
var responseItems dto.TaskResponse[[]dto.SunoDataResponse]
err = json.Unmarshal(responseBody, &responseItems)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Get Task parse body error2: %v, body: %s", err, string(responseBody)))
return err
}
if !responseItems.IsSuccess() {
common.SysLog(fmt.Sprintf("渠道 #%d 未完成的任务有: %d, 成功获取到任务数: %s", channelId, len(taskIds), string(responseBody)))
return err
}
for _, responseItem := range responseItems.Data {
task := taskM[responseItem.TaskID]
if !checkTaskNeedUpdate(task, responseItem) {
continue
}
task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status)
task.FailReason = lo.If(responseItem.FailReason != "", responseItem.FailReason).Else(task.FailReason)
task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime)
task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime)
task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime)
if responseItem.FailReason != "" || task.Status == model.TaskStatusFailure {
logger.LogInfo(ctx, task.TaskID+" 构建失败,"+task.FailReason)
task.Progress = "100%"
//err = model.CacheUpdateUserQuota(task.UserId) ?
if err != nil {
logger.LogError(ctx, "error update user quota cache: "+err.Error())
} else {
quota := task.Quota
if quota != 0 {
err = model.IncreaseUserQuota(task.UserId, quota, false)
if err != nil {
logger.LogError(ctx, "fail to increase user quota: "+err.Error())
}
logContent := fmt.Sprintf("异步任务执行失败 %s补偿 %s", task.TaskID, logger.LogQuota(quota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
}
}
}
if responseItem.Status == model.TaskStatusSuccess {
task.Progress = "100%"
}
task.Data = responseItem.Data
err = task.Update()
if err != nil {
common.SysLog("UpdateMidjourneyTask task error: " + err.Error())
}
}
return nil
}
func checkTaskNeedUpdate(oldTask *model.Task, newTask dto.SunoDataResponse) bool {
if oldTask.SubmitTime != newTask.SubmitTime {
return true
}
if oldTask.StartTime != newTask.StartTime {
return true
}
if oldTask.FinishTime != newTask.FinishTime {
return true
}
if string(oldTask.Status) != newTask.Status {
return true
}
if oldTask.FailReason != newTask.FailReason {
return true
}
if oldTask.FinishTime != newTask.FinishTime {
return true
}
if (oldTask.Status == model.TaskStatusFailure || oldTask.Status == model.TaskStatusSuccess) && oldTask.Progress != "100%" {
return true
}
oldData, _ := json.Marshal(oldTask.Data)
newData, _ := json.Marshal(newTask.Data)
sort.Slice(oldData, func(i, j int) bool {
return oldData[i] < oldData[j]
})
sort.Slice(newData, func(i, j int) bool {
return newData[i] < newData[j]
})
if string(oldData) != string(newData) {
return true
}
return false
service.TaskPollingLoop()
}
func GetAllTask(c *gin.Context) {
@@ -247,7 +38,7 @@ func GetAllTask(c *gin.Context) {
items := model.TaskGetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
total := model.TaskCountAllTasks(queryParams)
pageInfo.SetTotal(int(total))
pageInfo.SetItems(items)
pageInfo.SetItems(tasksToDto(items, true))
common.ApiSuccess(c, pageInfo)
}
@@ -271,6 +62,33 @@ func GetUserTask(c *gin.Context) {
items := model.TaskGetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
total := model.TaskCountAllUserTask(userId, queryParams)
pageInfo.SetTotal(int(total))
pageInfo.SetItems(items)
pageInfo.SetItems(tasksToDto(items, false))
common.ApiSuccess(c, pageInfo)
}
func tasksToDto(tasks []*model.Task, fillUser bool) []*dto.TaskDto {
var userIdMap map[int]*model.UserBase
if fillUser {
userIdMap = make(map[int]*model.UserBase)
userIds := types.NewSet[int]()
for _, task := range tasks {
userIds.Add(task.UserId)
}
for _, userId := range userIds.Items() {
cacheUser, err := model.GetUserCache(userId)
if err == nil {
userIdMap[userId] = cacheUser
}
}
}
result := make([]*dto.TaskDto, len(tasks))
for i, task := range tasks {
if fillUser {
if user, ok := userIdMap[task.UserId]; ok {
task.Username = user.Username
}
}
result[i] = relay.TaskModel2Dto(task)
}
return result
}

View File

@@ -1,313 +0,0 @@
package controller
import (
"context"
"encoding/json"
"fmt"
"io"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/ratio_setting"
)
func UpdateVideoTaskAll(ctx context.Context, platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) error {
for channelId, taskIds := range taskChannelM {
if err := updateVideoTaskAll(ctx, platform, channelId, taskIds, taskM); err != nil {
logger.LogError(ctx, fmt.Sprintf("Channel #%d failed to update video async tasks: %s", channelId, err.Error()))
}
}
return nil
}
func updateVideoTaskAll(ctx context.Context, platform constant.TaskPlatform, channelId int, taskIds []string, taskM map[string]*model.Task) error {
logger.LogInfo(ctx, fmt.Sprintf("Channel #%d pending video tasks: %d", channelId, len(taskIds)))
if len(taskIds) == 0 {
return nil
}
cacheGetChannel, err := model.CacheGetChannel(channelId)
if err != nil {
errUpdate := model.TaskBulkUpdate(taskIds, map[string]any{
"fail_reason": fmt.Sprintf("Failed to get channel info, channel ID: %d", channelId),
"status": "FAILURE",
"progress": "100%",
})
if errUpdate != nil {
common.SysLog(fmt.Sprintf("UpdateVideoTask error: %v", errUpdate))
}
return fmt.Errorf("CacheGetChannel failed: %w", err)
}
adaptor := relay.GetTaskAdaptor(platform)
if adaptor == nil {
return fmt.Errorf("video adaptor not found")
}
info := &relaycommon.RelayInfo{}
info.ChannelMeta = &relaycommon.ChannelMeta{
ChannelBaseUrl: cacheGetChannel.GetBaseURL(),
}
info.ApiKey = cacheGetChannel.Key
adaptor.Init(info)
for _, taskId := range taskIds {
if err := updateVideoSingleTask(ctx, adaptor, cacheGetChannel, taskId, taskM); err != nil {
logger.LogError(ctx, fmt.Sprintf("Failed to update video task %s: %s", taskId, err.Error()))
}
}
return nil
}
func updateVideoSingleTask(ctx context.Context, adaptor channel.TaskAdaptor, channel *model.Channel, taskId string, taskM map[string]*model.Task) error {
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
proxy := channel.GetSetting().Proxy
task := taskM[taskId]
if task == nil {
logger.LogError(ctx, fmt.Sprintf("Task %s not found in taskM", taskId))
return fmt.Errorf("task %s not found", taskId)
}
key := channel.Key
privateData := task.PrivateData
if privateData.Key != "" {
key = privateData.Key
}
resp, err := adaptor.FetchTask(baseURL, key, map[string]any{
"task_id": taskId,
"action": task.Action,
}, proxy)
if err != nil {
return fmt.Errorf("fetchTask failed for task %s: %w", taskId, err)
}
//if resp.StatusCode != http.StatusOK {
//return fmt.Errorf("get Video Task status code: %d", resp.StatusCode)
//}
defer resp.Body.Close()
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return fmt.Errorf("readAll failed for task %s: %w", taskId, err)
}
logger.LogDebug(ctx, fmt.Sprintf("UpdateVideoSingleTask response: %s", string(responseBody)))
taskResult := &relaycommon.TaskInfo{}
// try parse as New API response format
var responseItems dto.TaskResponse[model.Task]
if err = common.Unmarshal(responseBody, &responseItems); err == nil && responseItems.IsSuccess() {
logger.LogDebug(ctx, fmt.Sprintf("UpdateVideoSingleTask parsed as new api response format: %+v", responseItems))
t := responseItems.Data
taskResult.TaskID = t.TaskID
taskResult.Status = string(t.Status)
taskResult.Url = t.FailReason
taskResult.Progress = t.Progress
taskResult.Reason = t.FailReason
task.Data = t.Data
} else if taskResult, err = adaptor.ParseTaskResult(responseBody); err != nil {
return fmt.Errorf("parseTaskResult failed for task %s: %w", taskId, err)
} else {
task.Data = redactVideoResponseBody(responseBody)
}
logger.LogDebug(ctx, fmt.Sprintf("UpdateVideoSingleTask taskResult: %+v", taskResult))
now := time.Now().Unix()
if taskResult.Status == "" {
//return fmt.Errorf("task %s status is empty", taskId)
taskResult = relaycommon.FailTaskInfo("upstream returned empty status")
}
// 记录原本的状态,防止重复退款
shouldRefund := false
quota := task.Quota
preStatus := task.Status
task.Status = model.TaskStatus(taskResult.Status)
switch taskResult.Status {
case model.TaskStatusSubmitted:
task.Progress = "10%"
case model.TaskStatusQueued:
task.Progress = "20%"
case model.TaskStatusInProgress:
task.Progress = "30%"
if task.StartTime == 0 {
task.StartTime = now
}
case model.TaskStatusSuccess:
task.Progress = "100%"
if task.FinishTime == 0 {
task.FinishTime = now
}
if !(len(taskResult.Url) > 5 && taskResult.Url[:5] == "data:") {
task.FailReason = taskResult.Url
}
// 如果返回了 total_tokens 并且配置了模型倍率(非固定价格),则重新计费
if taskResult.TotalTokens > 0 {
// 获取模型名称
var taskData map[string]interface{}
if err := json.Unmarshal(task.Data, &taskData); err == nil {
if modelName, ok := taskData["model"].(string); ok && modelName != "" {
// 获取模型价格和倍率
modelRatio, hasRatioSetting, _ := ratio_setting.GetModelRatio(modelName)
// 只有配置了倍率(非固定价格)时才按 token 重新计费
if hasRatioSetting && modelRatio > 0 {
// 获取用户和组的倍率信息
group := task.Group
if group == "" {
user, err := model.GetUserById(task.UserId, false)
if err == nil {
group = user.Group
}
}
if group != "" {
groupRatio := ratio_setting.GetGroupRatio(group)
userGroupRatio, hasUserGroupRatio := ratio_setting.GetGroupGroupRatio(group, group)
var finalGroupRatio float64
if hasUserGroupRatio {
finalGroupRatio = userGroupRatio
} else {
finalGroupRatio = groupRatio
}
// 计算实际应扣费额度: totalTokens * modelRatio * groupRatio
actualQuota := int(float64(taskResult.TotalTokens) * modelRatio * finalGroupRatio)
// 计算差额
preConsumedQuota := task.Quota
quotaDelta := actualQuota - preConsumedQuota
if quotaDelta > 0 {
// 需要补扣费
logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费后补扣费:%s实际消耗%s预扣费%stokens%d",
task.TaskID,
logger.LogQuota(quotaDelta),
logger.LogQuota(actualQuota),
logger.LogQuota(preConsumedQuota),
taskResult.TotalTokens,
))
if err := model.DecreaseUserQuota(task.UserId, quotaDelta); err != nil {
logger.LogError(ctx, fmt.Sprintf("补扣费失败: %s", err.Error()))
} else {
model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta)
model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta)
task.Quota = actualQuota // 更新任务记录的实际扣费额度
// 记录消费日志
logContent := fmt.Sprintf("视频任务成功补扣费,模型倍率 %.2f,分组倍率 %.2ftokens %d预扣费 %s实际扣费 %s补扣费 %s",
modelRatio, finalGroupRatio, taskResult.TotalTokens,
logger.LogQuota(preConsumedQuota), logger.LogQuota(actualQuota), logger.LogQuota(quotaDelta))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
}
} else if quotaDelta < 0 {
// 需要退还多扣的费用
refundQuota := -quotaDelta
logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费后返还:%s实际消耗%s预扣费%stokens%d",
task.TaskID,
logger.LogQuota(refundQuota),
logger.LogQuota(actualQuota),
logger.LogQuota(preConsumedQuota),
taskResult.TotalTokens,
))
if err := model.IncreaseUserQuota(task.UserId, refundQuota, false); err != nil {
logger.LogError(ctx, fmt.Sprintf("退还预扣费失败: %s", err.Error()))
} else {
task.Quota = actualQuota // 更新任务记录的实际扣费额度
// 记录退款日志
logContent := fmt.Sprintf("视频任务成功退还多扣费用,模型倍率 %.2f,分组倍率 %.2ftokens %d预扣费 %s实际扣费 %s退还 %s",
modelRatio, finalGroupRatio, taskResult.TotalTokens,
logger.LogQuota(preConsumedQuota), logger.LogQuota(actualQuota), logger.LogQuota(refundQuota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
}
} else {
// quotaDelta == 0, 预扣费刚好准确
logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费准确(%stokens%d",
task.TaskID, logger.LogQuota(actualQuota), taskResult.TotalTokens))
}
}
}
}
}
}
case model.TaskStatusFailure:
logger.LogJson(ctx, fmt.Sprintf("Task %s failed", taskId), task)
task.Status = model.TaskStatusFailure
task.Progress = "100%"
if task.FinishTime == 0 {
task.FinishTime = now
}
task.FailReason = taskResult.Reason
logger.LogInfo(ctx, fmt.Sprintf("Task %s failed: %s", task.TaskID, task.FailReason))
taskResult.Progress = "100%"
if quota != 0 {
if preStatus != model.TaskStatusFailure {
shouldRefund = true
} else {
logger.LogWarn(ctx, fmt.Sprintf("Task %s already in failure status, skip refund", task.TaskID))
}
}
default:
return fmt.Errorf("unknown task status %s for task %s", taskResult.Status, taskId)
}
if taskResult.Progress != "" {
task.Progress = taskResult.Progress
}
if err := task.Update(); err != nil {
common.SysLog("UpdateVideoTask task error: " + err.Error())
shouldRefund = false
}
if shouldRefund {
// 任务失败且之前状态不是失败才退还额度,防止重复退还
if err := model.IncreaseUserQuota(task.UserId, quota, false); err != nil {
logger.LogWarn(ctx, "Failed to increase user quota: "+err.Error())
}
logContent := fmt.Sprintf("Video async task failed %s, refund %s", task.TaskID, logger.LogQuota(quota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
}
return nil
}
func redactVideoResponseBody(body []byte) []byte {
var m map[string]any
if err := json.Unmarshal(body, &m); err != nil {
return body
}
resp, _ := m["response"].(map[string]any)
if resp != nil {
delete(resp, "bytesBase64Encoded")
if v, ok := resp["video"].(string); ok {
resp["video"] = truncateBase64(v)
}
if vs, ok := resp["videos"].([]any); ok {
for i := range vs {
if vm, ok := vs[i].(map[string]any); ok {
delete(vm, "bytesBase64Encoded")
}
}
}
}
b, err := json.Marshal(m)
if err != nil {
return body
}
return b
}
func truncateBase64(s string) string {
const maxKeep = 256
if len(s) <= maxKeep {
return s
}
return s[:maxKeep] + "..."
}

View File

@@ -582,6 +582,44 @@ func UpdateUser(c *gin.Context) {
return
}
func AdminClearUserBinding(c *gin.Context) {
id, err := strconv.Atoi(c.Param("id"))
if err != nil {
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
return
}
bindingType := strings.ToLower(strings.TrimSpace(c.Param("binding_type")))
if bindingType == "" {
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
return
}
user, err := model.GetUserById(id, false)
if err != nil {
common.ApiError(c, err)
return
}
myRole := c.GetInt("role")
if myRole <= user.Role && myRole != common.RoleRootUser {
common.ApiErrorI18n(c, i18n.MsgUserNoPermissionSameLevel)
return
}
if err := user.ClearBinding(bindingType); err != nil {
common.ApiError(c, err)
return
}
model.RecordLog(user.Id, model.LogTypeManage, fmt.Sprintf("admin cleared %s binding for user %s", bindingType, user.Username))
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "success",
})
}
func UpdateSelf(c *gin.Context) {
var requestData map[string]interface{}
err := json.NewDecoder(c.Request.Body).Decode(&requestData)

View File

@@ -16,59 +16,44 @@ import (
"github.com/gin-gonic/gin"
)
// videoProxyError returns a standardized OpenAI-style error response.
func videoProxyError(c *gin.Context, status int, errType, message string) {
c.JSON(status, gin.H{
"error": gin.H{
"message": message,
"type": errType,
},
})
}
func VideoProxy(c *gin.Context) {
taskID := c.Param("task_id")
if taskID == "" {
c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{
"message": "task_id is required",
"type": "invalid_request_error",
},
})
videoProxyError(c, http.StatusBadRequest, "invalid_request_error", "task_id is required")
return
}
task, exists, err := model.GetByOnlyTaskId(taskID)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to query task %s: %s", taskID, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to query task",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to query task")
return
}
if !exists || task == nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get task %s: %v", taskID, err))
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"message": "Task not found",
"type": "invalid_request_error",
},
})
videoProxyError(c, http.StatusNotFound, "invalid_request_error", "Task not found")
return
}
if task.Status != model.TaskStatusSuccess {
c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{
"message": fmt.Sprintf("Task is not completed yet, current status: %s", task.Status),
"type": "invalid_request_error",
},
})
videoProxyError(c, http.StatusBadRequest, "invalid_request_error",
fmt.Sprintf("Task is not completed yet, current status: %s", task.Status))
return
}
channel, err := model.CacheGetChannel(task.ChannelId)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get task %s: not found", taskID))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to retrieve channel information",
"type": "server_error",
},
})
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get channel for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to retrieve channel information")
return
}
baseURL := channel.GetBaseURL()
@@ -81,12 +66,7 @@ func VideoProxy(c *gin.Context) {
client, err := service.GetHttpClientWithProxy(proxy)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create proxy client for task %s: %s", taskID, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to create proxy client",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy client")
return
}
@@ -95,12 +75,7 @@ func VideoProxy(c *gin.Context) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create request: %s", err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to create proxy request",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request")
return
}
@@ -109,68 +84,43 @@ func VideoProxy(c *gin.Context) {
apiKey := task.PrivateData.Key
if apiKey == "" {
logger.LogError(c.Request.Context(), fmt.Sprintf("Missing stored API key for Gemini task %s", taskID))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "API key not stored for task",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "API key not stored for task")
return
}
videoURL, err = getGeminiVideoURL(channel, task, apiKey)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Gemini video URL for task %s: %s", taskID, err.Error()))
c.JSON(http.StatusBadGateway, gin.H{
"error": gin.H{
"message": "Failed to resolve Gemini video URL",
"type": "server_error",
},
})
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Gemini video URL")
return
}
req.Header.Set("x-goog-api-key", apiKey)
case constant.ChannelTypeOpenAI, constant.ChannelTypeSora:
videoURL = fmt.Sprintf("%s/v1/videos/%s/content", baseURL, task.TaskID)
videoURL = fmt.Sprintf("%s/v1/videos/%s/content", baseURL, task.GetUpstreamTaskID())
req.Header.Set("Authorization", "Bearer "+channel.Key)
default:
// Video URL is directly in task.FailReason
videoURL = task.FailReason
// Video URL is stored in PrivateData.ResultURL (fallback to FailReason for old data)
videoURL = task.GetResultURL()
}
req.URL, err = url.Parse(videoURL)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to parse URL %s: %s", videoURL, err.Error()))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Failed to create proxy request",
"type": "server_error",
},
})
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request")
return
}
resp, err := client.Do(req)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to fetch video from %s: %s", videoURL, err.Error()))
c.JSON(http.StatusBadGateway, gin.H{
"error": gin.H{
"message": "Failed to fetch video content",
"type": "server_error",
},
})
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content")
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
logger.LogError(c.Request.Context(), fmt.Sprintf("Upstream returned status %d for %s", resp.StatusCode, videoURL))
c.JSON(http.StatusBadGateway, gin.H{
"error": gin.H{
"message": fmt.Sprintf("Upstream service returned status %d", resp.StatusCode),
"type": "server_error",
},
})
videoProxyError(c, http.StatusBadGateway, "server_error",
fmt.Sprintf("Upstream service returned status %d", resp.StatusCode))
return
}
@@ -180,10 +130,9 @@ func VideoProxy(c *gin.Context) {
}
}
c.Writer.Header().Set("Cache-Control", "public, max-age=86400") // Cache for 24 hours
c.Writer.Header().Set("Cache-Control", "public, max-age=86400")
c.Writer.WriteHeader(resp.StatusCode)
_, err = io.Copy(c.Writer, resp.Body)
if err != nil {
if _, err = io.Copy(c.Writer, resp.Body); err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream video content: %s", err.Error()))
}
}

View File

@@ -1,12 +1,12 @@
package controller
import (
"encoding/json"
"fmt"
"io"
"strconv"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
@@ -37,7 +37,7 @@ func getGeminiVideoURL(channel *model.Channel, task *model.Task, apiKey string)
proxy := channel.GetSetting().Proxy
resp, err := adaptor.FetchTask(baseURL, apiKey, map[string]any{
"task_id": task.TaskID,
"task_id": task.GetUpstreamTaskID(),
"action": task.Action,
}, proxy)
if err != nil {
@@ -71,7 +71,7 @@ func extractGeminiVideoURLFromTaskData(task *model.Task) string {
return ""
}
var payload map[string]any
if err := json.Unmarshal(task.Data, &payload); err != nil {
if err := common.Unmarshal(task.Data, &payload); err != nil {
return ""
}
return extractGeminiVideoURLFromMap(payload)
@@ -79,7 +79,7 @@ func extractGeminiVideoURLFromTaskData(task *model.Task) string {
func extractGeminiVideoURLFromPayload(body []byte) string {
var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil {
if err := common.Unmarshal(body, &payload); err != nil {
return ""
}
return extractGeminiVideoURLFromMap(payload)