303 lines
12 KiB
Go
303 lines
12 KiB
Go
package relay
|
||
|
||
import (
|
||
crand "crypto/rand"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"net"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/net/ghttp"
|
||
|
||
lcommon "github.com/qianfree/team-api/internal/logic/common"
|
||
"github.com/qianfree/team-api/internal/logic/task"
|
||
relay_common "github.com/qianfree/team-api/relay/common"
|
||
relay_constant "github.com/qianfree/team-api/relay/constant"
|
||
relay_handler "github.com/qianfree/team-api/relay/handler"
|
||
)
|
||
|
||
// HandleSyncImageSubmit 同步图片厂商的「异步化」提交管线。
|
||
//
|
||
// 由 HandleAliImageSubmit 在识别到同步厂商（ProviderTypeToTaskPlatform 未命中）时调用：
|
||
// 执行 gate → 强制非流式 → 估价 → 预扣 → 建 QUEUED 任务行 → 入 worker 队列，
|
||
// 客户端毫秒级拿到 task_id 后即释放连接，真正的上游长连接由 task 包的 worker 池后台持有。
|
||
//
|
||
// 与 relay_handler.HandleTaskSubmit 的区别：不在提交阶段调用上游（同步厂商无上游任务 ID），
|
||
// 渠道也不在提交时绑定（channel_id=0，worker dispatch 时按最新健康度/容量重选）。
|
||
func HandleSyncImageSubmit(r *ghttp.Request, body []byte, rc *relay_handler.TaskRelayContext, channelMeta *relay_common.ChannelMeta) {
|
||
ctx := r.Context()
|
||
w := rc.Writer
|
||
|
||
modelName := extractModelName(body)
|
||
if modelName == "" {
|
||
writeSyncImageError(w, http.StatusBadRequest, "model is required")
|
||
return
|
||
}
|
||
|
||
// 1. gate（复刻 HandleTaskSubmit 的 scope / IP 白名单 / 限流 / key 并发）
|
||
if rc.Scope == "read_only" {
|
||
writeSyncImageError(w, http.StatusForbidden, "API key scope denied")
|
||
return
|
||
}
|
||
if !checkSyncImageIPWhitelist(rc.KeyIpWhitelist, rc.ClientIP) {
|
||
writeSyncImageError(w, http.StatusForbidden, "IP address is not allowed")
|
||
return
|
||
}
|
||
if allowed, level, _, _, _ := taskBillingProvider.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps); !allowed {
|
||
writeSyncImageError(w, http.StatusTooManyRequests, fmt.Sprintf("rate limit exceeded at %s level", level))
|
||
return
|
||
}
|
||
if !taskBillingProvider.AcquireApiKeyConcurrent(ctx, rc.ApiKeyID, rc.KeyConcurrency) {
|
||
writeSyncImageError(w, http.StatusTooManyRequests, "API key concurrent request limit exceeded")
|
||
return
|
||
}
|
||
// 提交很快即返回，并发槽在提交结束即释放（上游长连接由 worker 持有，不占此槽）。
|
||
defer taskBillingProvider.ReleaseApiKeyConcurrent(ctx, rc.ApiKeyID)
|
||
|
||
// Fast-fail：对象存储未配置、且本请求**必然**需要存储保存结果（b64 结果 / 强制 re-host）时，
|
||
// 提交阶段即返回友好提示——避免白打一次上游生成（生成成功却因无处保存而失败退款，且用户
|
||
// 要空等 10–60s 才知道是配置问题）。url 透传类请求（未开启 re-host 且非 b64）不需要存储，
|
||
// 不在此拦截，仍可正常完成。
|
||
if !lcommon.IsStorageConfigured(ctx) {
|
||
rehostOn := lcommon.Config().GetBool(ctx, "sync_image_rehost_url")
|
||
if rehostOn || requestForcesB64(body, modelName) {
|
||
writeSyncImageError(w, http.StatusServiceUnavailable,
|
||
"平台尚未配置对象存储（OSS/S3/COS），无法保存生成的图片，请联系管理员在系统设置中配置对象存储后重试")
|
||
return
|
||
}
|
||
}
|
||
|
||
// 2. 强制非流式：剥离 stream 字段，防 DoResponse 走 SSE 分支
|
||
cleanBody := stripStreamField(body)
|
||
|
||
// 3. 估价（图片走 per_request 价；nil body = 参数倍率不参与，同步图片链路 ratios 不随任务持久化）
|
||
estimatedCost, err := taskBillingProvider.EstimateTaskCost(ctx, rc.TenantID, modelName, nil, nil)
|
||
if err != nil {
|
||
writeSyncImageError(w, http.StatusInternalServerError, "estimate cost failed: "+err.Error())
|
||
return
|
||
}
|
||
|
||
// 4. 配额 + 预扣
|
||
if err := taskBillingProvider.CheckApiKeyQuota(ctx, rc.ApiKeyID, estimatedCost); err != nil {
|
||
writeSyncImageError(w, http.StatusPaymentRequired, "API key quota exceeded")
|
||
return
|
||
}
|
||
preDeduct, err := taskBillingProvider.PreDeductTask(ctx, rc.TenantID, rc.RequestID, estimatedCost, modelName)
|
||
if err != nil {
|
||
writeSyncImageError(w, http.StatusPaymentRequired, "insufficient balance")
|
||
return
|
||
}
|
||
|
||
// 5. 建 QUEUED 任务行（channel_id=0，dispatch 时绑定）
|
||
// 生成任务 ID 依赖 crypto/rand，失败时此处尚未建任务行、无法交给超时兜底网结算，
|
||
// 必须就地退掉上一步的预扣，避免冻结额被永久占用。
|
||
publicTaskID, err := generateSyncImagePublicID()
|
||
if err != nil {
|
||
_ = taskBillingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeduct)
|
||
g.Log().Errorf(ctx, "sync_image: generate task ID failed, refunded pre-deduct: %v", err)
|
||
writeSyncImageError(w, http.StatusInternalServerError, "generate task ID failed")
|
||
return
|
||
}
|
||
now := time.Now()
|
||
|
||
// 从渠道类型推导实际的供应商平台（ali / gemini / volcengine / openai 等），用于任务日志展示。
|
||
providerType := relay_constant.ProviderType(channelMeta.ChannelType)
|
||
actualPlatform, ok := relay_constant.ProviderTypeToTaskPlatform(providerType)
|
||
if !ok {
|
||
// 未命中映射（如 openai、claude 等）：用供应商类型的可读名作平台名。
|
||
// 注意不能写 TaskPlatform(providerType)——那是 int→string 的 rune 转换，会得到
|
||
// 不可见控制字符（如 providerType=1 → "\x01"）而非 "openai"。取 String() 并转小写，
|
||
// 与已映射平台（"ali"/"gemini"）的小写风格保持一致。
|
||
actualPlatform = relay_constant.TaskPlatform(strings.ToLower(providerType.String()))
|
||
}
|
||
|
||
privateData, _ := json.Marshal(map[string]any{
|
||
"task_type": string(relay_constant.TaskPlatformSyncImage),
|
||
"billing_context": map[string]any{
|
||
"ratios": nil,
|
||
"model_name": modelName,
|
||
// pre_deduct 快照落 JSONB，读取端为 float64；decimal 默认带引号 MarshalJSON
|
||
// 会导致读取端反序列化失败，故显式转 float64。
|
||
"pre_deduct": preDeduct.InexactFloat64(),
|
||
},
|
||
})
|
||
asyncTask := &relay_common.AsyncTask{
|
||
PublicTaskID: publicTaskID,
|
||
RequestID: rc.RequestID,
|
||
Platform: string(actualPlatform),
|
||
Action: "generate",
|
||
Status: "QUEUED",
|
||
Progress: "0%",
|
||
TenantID: rc.TenantID,
|
||
UserID: rc.UserID,
|
||
ApiKeyID: rc.ApiKeyID,
|
||
ChannelID: 0,
|
||
ModelName: modelName,
|
||
PreDeductAmount: preDeduct,
|
||
PrivateData: privateData,
|
||
SubmitTime: &now,
|
||
}
|
||
if err := taskDataProvider.CreateTask(ctx, asyncTask); err != nil {
|
||
_ = taskBillingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeduct)
|
||
writeSyncImageError(w, http.StatusInternalServerError, "create task record failed: "+err.Error())
|
||
return
|
||
}
|
||
|
||
// CreateTask 不回填自增 ID（Postgres 无 LastInsertId），回查取任务行主键供 worker CAS 使用。
|
||
created, err := taskDataProvider.GetTaskByPublicID(ctx, publicTaskID)
|
||
if err != nil || created == nil {
|
||
// 此时任务行已 CreateTask 成功落库（QUEUED、billing_settled=false、active 计数已 +1），
|
||
// 但缺主键无法在这里安全地 CAS 收尾。**不能**在此直接退款：UnfreezePreDeduct 非幂等，
|
||
// 与超时兜底网 handleTimedOutTasks 会形成二次退款（侵蚀他人冻结额、可用余额虚高）。
|
||
// 故此处只返回错误，把「CAS→FAILURE + 退款 + DecrActiveTask」交给超时兜底网做恰好一次结算
|
||
// （handleTimedOutTasks 无平台过滤，会捕获这类 QUEUED 孤儿行）。
|
||
g.Log().Errorf(ctx, "sync_image: load created task %s failed, deferring settlement to timeout sweeper: %v", publicTaskID, err)
|
||
writeSyncImageError(w, http.StatusInternalServerError, "create task record failed")
|
||
return
|
||
}
|
||
rc.TaskID = publicTaskID
|
||
|
||
// 6. 入队（非阻塞）；队列满 → 退款 + CAS→FAILURE + DecrActive + 429
|
||
job := &task.SyncImageJob{
|
||
TaskID: created.ID,
|
||
PublicTaskID: publicTaskID,
|
||
RequestID: rc.RequestID,
|
||
TenantID: rc.TenantID,
|
||
UserID: rc.UserID,
|
||
ApiKeyID: rc.ApiKeyID,
|
||
ProjectID: rc.ProjectID,
|
||
Model: modelName,
|
||
RequestBody: cleanBody,
|
||
PreDeductAmount: preDeduct,
|
||
Ratios: nil,
|
||
SubmitTime: now,
|
||
}
|
||
if !task.EnqueueSyncImageJob(job) {
|
||
_ = taskBillingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeduct)
|
||
_ = taskDataProvider.UpdateTaskCAS(ctx, &relay_common.AsyncTask{
|
||
ID: created.ID,
|
||
Status: "FAILURE",
|
||
FailReason: "sync image worker queue full",
|
||
BillingSettled: true,
|
||
FinishTime: &now,
|
||
}, "QUEUED")
|
||
task.DecrActiveTask()
|
||
rc.TaskID = ""
|
||
writeSyncImageError(w, http.StatusTooManyRequests, "server busy, please retry later")
|
||
return
|
||
}
|
||
|
||
// 7. 返回 QUEUED（客户端连接在此释放）
|
||
writeSyncImageJSON(w, http.StatusOK, map[string]any{
|
||
"id": publicTaskID,
|
||
"status": "QUEUED",
|
||
"model": modelName,
|
||
"created_at": now.Unix(),
|
||
})
|
||
}
|
||
|
||
// requestForcesB64 判断该图片请求是否**必然**产生需要 re-host 的 base64 结果，即使
|
||
// 未开启 re-host 开关也一定用到对象存储：
|
||
// - response_format=b64_json：显式要求 base64；
|
||
// - gpt-image 系列：上游只返回 b64_json（不支持 url）。
|
||
//
|
||
// 其余情况（默认 / response_format=url）在未开启 re-host 时可 url 透传，不强制需要存储，
|
||
// 因此不在提交阶段 fast-fail，避免误伤纯 url 透传的模型。
|
||
func requestForcesB64(body []byte, model string) bool {
|
||
var req struct {
|
||
ResponseFormat string `json:"response_format"`
|
||
}
|
||
_ = json.Unmarshal(body, &req)
|
||
if strings.EqualFold(strings.TrimSpace(req.ResponseFormat), "b64_json") {
|
||
return true
|
||
}
|
||
return strings.HasPrefix(strings.ToLower(model), "gpt-image")
|
||
}
|
||
|
||
// stripStreamField 移除请求体中的 stream 字段（强制非流式）。
|
||
func stripStreamField(body []byte) []byte {
|
||
var m map[string]json.RawMessage
|
||
if err := json.Unmarshal(body, &m); err != nil {
|
||
return body
|
||
}
|
||
if _, ok := m["stream"]; !ok {
|
||
return body
|
||
}
|
||
delete(m, "stream")
|
||
cleaned, err := json.Marshal(m)
|
||
if err != nil {
|
||
return body
|
||
}
|
||
return cleaned
|
||
}
|
||
|
||
// checkSyncImageIPWhitelist 复刻 HandleTaskSubmit 的 IP 白名单校验（支持精确 IP 与 CIDR）。
|
||
func checkSyncImageIPWhitelist(whitelist, clientIP string) bool {
|
||
if whitelist == "" {
|
||
return true
|
||
}
|
||
host, _, err := net.SplitHostPort(clientIP)
|
||
if err != nil {
|
||
host = clientIP
|
||
}
|
||
host = strings.TrimPrefix(strings.TrimSuffix(host, "]"), "[")
|
||
parsedIP := net.ParseIP(host)
|
||
if parsedIP == nil {
|
||
return false
|
||
}
|
||
for _, item := range strings.Split(whitelist, ",") {
|
||
item = strings.TrimSpace(item)
|
||
if item == "" {
|
||
continue
|
||
}
|
||
if item == host {
|
||
return true
|
||
}
|
||
if strings.Contains(item, "/") {
|
||
if _, cidr, err := net.ParseCIDR(item); err == nil && cidr.Contains(parsedIP) {
|
||
return true
|
||
}
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func generateSyncImagePublicID() (string, error) {
|
||
b := make([]byte, 16)
|
||
if _, err := crand.Read(b); err != nil {
|
||
return "", err
|
||
}
|
||
return "task_" + hex.EncodeToString(b), nil
|
||
}
|
||
|
||
func writeSyncImageError(w http.ResponseWriter, statusCode int, message string) {
|
||
writeSyncImageJSON(w, statusCode, map[string]any{
|
||
"error": map[string]any{
|
||
"type": "invalid_request_error",
|
||
"message": message,
|
||
},
|
||
})
|
||
}
|
||
|
||
// writeSyncImageErrorWithCode 与 writeSyncImageError 相同，但在错误体中附带机器可读的 code，
|
||
// 供前端做条件分支（如在线体验「异步被禁用 → 优雅降级到同步端点」）。
|
||
func writeSyncImageErrorWithCode(w http.ResponseWriter, statusCode int, code, message string) {
|
||
writeSyncImageJSON(w, statusCode, map[string]any{
|
||
"error": map[string]any{
|
||
"type": "invalid_request_error",
|
||
"code": code,
|
||
"message": message,
|
||
},
|
||
})
|
||
}
|
||
|
||
func writeSyncImageJSON(w http.ResponseWriter, statusCode int, data any) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
w.WriteHeader(statusCode)
|
||
_ = json.NewEncoder(w).Encode(data)
|
||
}
|