From ccf457592b8ae5e0f5d75b43aacb75cc6cd5259f Mon Sep 17 00:00:00 2001 From: chengqian Date: Wed, 3 Jun 2026 17:03:09 +0800 Subject: [PATCH] =?UTF-8?q?fix(relay):=E4=BF=AE=E5=A4=8D=E5=A4=9A=E4=B8=AA?= =?UTF-8?q?=E5=B9=B6=E5=8F=91=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/cmd/cmd.go | 5 ++ internal/handler/relay/relay.go | 14 +++- internal/logic/relay/apikey.go | 18 ++++- internal/logic/relay/provider.go | 102 ++++++++++++------------ internal/logic/task/async_provider.go | 3 +- relay/channel/baidu_v2/adaptor.go | 12 ++- relay/channel/claude/adaptor.go | 2 +- relay/channel/claude/response.go | 65 +++------------ relay/channel/coze/adaptor.go | 47 ++--------- relay/channel/dify/adaptor.go | 54 ++++--------- relay/channel/gemini/response.go | 30 +++++-- relay/channel/ollama/adaptor.go | 7 +- relay/channel/openai/adaptor.go | 38 +-------- relay/channel/openai/claude_response.go | 4 +- relay/channel/openai/completion.go | 3 +- relay/channel/openai/converter.go | 2 +- relay/channel/openai/embedding.go | 12 +-- relay/channel/openai/gemini_response.go | 4 +- relay/channel/openai/rerank.go | 9 ++- relay/channel/openai/responses.go | 4 +- relay/channel/openai/stream.go | 4 +- relay/channel/tencent/adaptor.go | 9 +-- relay/channel/vertex/auth.go | 61 +++++++++----- relay/common/task_provider.go | 2 +- relay/handler/relay_handler.go | 26 ++++++ relay/handler/task_handler.go | 26 ++++-- relay/helper/model_name.go | 34 ++++++++ relay/helper/stream.go | 44 +++++++++- relay/override/header.go | 2 +- relay/override/operations.go | 2 +- relay/override/override.go | 78 ++---------------- 31 files changed, 357 insertions(+), 366 deletions(-) create mode 100644 relay/helper/model_name.go diff --git a/internal/cmd/cmd.go b/internal/cmd/cmd.go index bfba4a0..ee41977 100644 --- a/internal/cmd/cmd.go +++ b/internal/cmd/cmd.go @@ -31,6 +31,7 @@ import ( "github.com/qianfree/team-api/internal/handler/relay" setupHandler "github.com/qianfree/team-api/internal/handler/setup" "github.com/qianfree/team-api/internal/plugin" + "github.com/qianfree/team-api/relay/scheduler" "github.com/qianfree/team-api/web" ) @@ -164,6 +165,10 @@ var ( billing.CleanExpiredPreDeducts(ctx) return nil }) + cs.Register("affinity_cache_cleanup", "*/5 * * * *", func(ctx context.Context) error { + scheduler.GetGlobalAffinity().CleanExpired() + return nil + }) cs.StartBackground(ctx) // Initialize plugin system (must be after CronScheduler init) diff --git a/internal/handler/relay/relay.go b/internal/handler/relay/relay.go index f77c965..c3730f1 100644 --- a/internal/handler/relay/relay.go +++ b/internal/handler/relay/relay.go @@ -25,7 +25,7 @@ var ( // handleRelayChatCompletions 处理 /v1/chat/completions func HandleChatCompletions(r *ghttp.Request) { // DEBUG: 记录请求入口 - g.Log().Infof(r.Context(), "[HandleChatCompletions] Entry: method=%s, path=%s, uri=%s", + g.Log().Debugf(r.Context(), "[HandleChatCompletions] Entry: method=%s, path=%s, uri=%s", r.Method, r.URL.Path, r.RequestURI) body := r.GetBody() @@ -166,7 +166,7 @@ func HandleCompletions(r *ghttp.Request) { // HandleResponses 处理 /v1/responses（OpenAI Responses API） func HandleResponses(r *ghttp.Request) { // DEBUG: 记录请求入口 - g.Log().Infof(r.Context(), "[HandleResponses] Entry: method=%s, path=%s, uri=%s", + g.Log().Debugf(r.Context(), "[HandleResponses] Entry: method=%s, path=%s, uri=%s", r.Method, r.URL.Path, r.RequestURI) body := r.GetBody() @@ -235,7 +235,7 @@ func HandleMessages(r *ghttp.Request) { // HandleGeminiGenerateContent 处理 /v1beta/models/{model}:generateContent（Gemini 原生格式） func HandleGeminiGenerateContent(r *ghttp.Request) { - g.Log().Infof(r.Context(), "[HandleGeminiGenerateContent] Entry: method=%s, path=%s, uri=%s", + g.Log().Debugf(r.Context(), "[HandleGeminiGenerateContent] Entry: method=%s, path=%s, uri=%s", r.Method, r.URL.Path, r.RequestURI) body := r.GetBody() @@ -287,7 +287,13 @@ func HandleGeminiGenerateContent(r *ghttp.Request) { rawBody["model"] = modelJSON streamJSON, _ := json.Marshal(isStream) rawBody["stream"] = streamJSON - modifiedBody, _ := json.Marshal(rawBody) + modifiedBody, err := json.Marshal(rawBody) + if err != nil { + r.Response.WriteJsonExit(g.Map{ + "error": g.Map{"code": 400, "message": "failed to encode request body", "status": "INVALID_ARGUMENT"}, + }) + return + } // 构造路径（使 Path2RelayMode 能识别为 Gemini 模式） path := "/models/" + modelAction diff --git a/internal/logic/relay/apikey.go b/internal/logic/relay/apikey.go index e682e26..53d88f2 100644 --- a/internal/logic/relay/apikey.go +++ b/internal/logic/relay/apikey.go @@ -2,6 +2,8 @@ package relay import ( "context" + "crypto/sha256" + "encoding/hex" "encoding/json" "github.com/qianfree/team-api/internal/dao" "time" @@ -24,6 +26,7 @@ type ApiKeyInfo struct { Status string KeyType string // personal 或 project ProjectID int64 // 项目密钥关联的项目 ID，个人密钥为 0 + KeyHash string // SHA-256(rawKey)，缓存命中时用于全键校验 } // apiKeyCache API Key 缓存实例（TTL 60s） @@ -41,11 +44,17 @@ func DefaultChannelSettings() common.ChannelSettings { func GetEncryptionKey() []byte { key := g.Cfg().MustGet(context.Background(), "security.encryptionKey").String() if key == "" { - key = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + panic("security.encryptionKey is required but not configured") } return uc.MustGetEncryptionKey(key) } +// apiKeyHash 计算 API Key 的 SHA-256 哈希 +func apiKeyHash(rawKey string) string { + h := sha256.Sum256([]byte(rawKey)) + return hex.EncodeToString(h[:]) +} + // ValidateApiKey 验证 API Key 并返回认证信息 func ValidateApiKey(ctx context.Context, rawKey string) (*ApiKeyInfo, error) { if len(rawKey) < 12 { @@ -54,10 +63,11 @@ func ValidateApiKey(ctx context.Context, rawKey string) (*ApiKeyInfo, error) { } prefix := rawKey[:12] + keyHash := apiKeyHash(rawKey) - // 尝试从缓存获取 + // 尝试从缓存获取（需校验完整 key 哈希，防止前缀碰撞） var cached ApiKeyInfo - if apiKeyCache.GetJSON(ctx, prefix, &cached) { + if apiKeyCache.GetJSON(ctx, prefix, &cached) && cached.KeyHash == keyHash { if cached.Status == "active" { return &cached, nil } @@ -114,6 +124,7 @@ func ValidateApiKey(ctx context.Context, rawKey string) (*ApiKeyInfo, error) { Status: "expired", KeyType: k.KeyType, ProjectID: k.ProjectID, + KeyHash: keyHash, } apiKeyCache.Set(ctx, prefix, info) return nil, consts.ErrKeyExpired @@ -139,6 +150,7 @@ func ValidateApiKey(ctx context.Context, rawKey string) (*ApiKeyInfo, error) { Status: "active", KeyType: k.KeyType, ProjectID: k.ProjectID, + KeyHash: keyHash, } apiKeyCache.Set(ctx, prefix, info) diff --git a/internal/logic/relay/provider.go b/internal/logic/relay/provider.go index 53bd71b..8902868 100644 --- a/internal/logic/relay/provider.go +++ b/internal/logic/relay/provider.go @@ -368,19 +368,15 @@ func (p *DataProviderImpl) GetModelMapping(ctx context.Context, modelName string Status string `json:"status"` } - var model modelRow + var model *modelRow err := dao.MdlModels.Ctx(ctx). Where("model_id", modelName). Fields("model_id, category, status"). Scan(&model) if err != nil { - // sql.ErrNoRows 表示模型不存在，转换为业务错误 - if strings.Contains(err.Error(), "no rows in result set") { - return "", "", common.ErrModelNotFound - } return "", "", err } - if model.ModelId == "" { + if model == nil { return "", "", common.ErrModelNotFound } if model.Status == "offline" { @@ -532,37 +528,37 @@ func (p *DataProviderImpl) RecordAudit(ctx context.Context, record *common.Audit } } - insertData := g.Map{ - "tenant_id": record.TenantID, - "user_id": record.UserID, - "api_key_id": record.ApiKeyID, - "project_id": record.ProjectID, - "request_id": record.RequestID, - "method": record.Method, - "path": record.Path, - "query_params": record.QueryParams, - "status_code": record.StatusCode, - "client_ip": record.ClientIP, - "user_agent": record.UserAgent, - "request_body": sysReq, - "response_body": sysResp, - "tenant_request_body": tntReq, - "tenant_response_body": tntResp, - "latency_ms": record.LatencyMs, - "first_token_ms": record.FirstTokenMs, - "audit_level": globalLevel, - "tenant_audit_level": tenantLevel, - "request_headers": requestHeadersJSON, - "response_headers": responseHeadersJSON, - "forwarding_trace": forwardingTraceJSON, + insertData := do.AudRequestLogs{ + TenantId: record.TenantID, + UserId: record.UserID, + ApiKeyId: record.ApiKeyID, + ProjectId: record.ProjectID, + RequestId: record.RequestID, + Method: record.Method, + Path: record.Path, + QueryParams: record.QueryParams, + StatusCode: record.StatusCode, + ClientIp: record.ClientIP, + UserAgent: record.UserAgent, + RequestBody: sysReq, + ResponseBody: sysResp, + TenantRequestBody: tntReq, + TenantResponseBody: tntResp, + LatencyMs: record.LatencyMs, + FirstTokenMs: record.FirstTokenMs, + AuditLevel: globalLevel, + TenantAuditLevel: tenantLevel, + RequestHeaders: requestHeadersJSON, + ResponseHeaders: responseHeadersJSON, + ForwardingTrace: forwardingTraceJSON, } // 异步任务字段 if record.TaskID != "" { - insertData["task_id"] = record.TaskID + insertData.TaskId = record.TaskID } if record.TaskStatus != "" { - insertData["task_status"] = record.TaskStatus + insertData.TaskStatus = record.TaskStatus } _, insertErr := dao.AudRequestLogs.Ctx(bgCtx).Data(insertData).Insert() @@ -609,14 +605,13 @@ func (p *DataProviderImpl) UpdateTaskAudit(ctx context.Context, record *common.A upstreamHeadersJSON = string(b) } - updateData := g.Map{ - "task_status": record.TaskStatus, - "task_result": taskResult, - "task_upstream_headers": upstreamHeadersJSON, - "updated_at": time.Now(), + updateData := do.AudRequestLogs{ + TaskStatus: record.TaskStatus, + TaskResult: taskResult, + TaskUpstreamHeaders: upstreamHeadersJSON, } if record.TaskCompletedAt != nil { - updateData["task_completed_at"] = record.TaskCompletedAt + updateData.TaskCompletedAt = gtime.NewFromTime(*record.TaskCompletedAt) } _, err := dao.AudRequestLogs.Ctx(bgCtx). @@ -706,11 +701,14 @@ func (p *DataProviderImpl) IncrementConsecutiveFailure(ctx context.Context, chan // ResetConsecutiveFailure 实现 DataProvider.ResetConsecutiveFailure func (p *DataProviderImpl) ResetConsecutiveFailure(ctx context.Context, channelID int64) { - dao.ChnHealthScores.Ctx(ctx). + _, err := dao.ChnHealthScores.Ctx(ctx). Where("channel_id", channelID). Data(do.ChnHealthScores{ ConsecutiveFailures: 0, }).Update() + if err != nil { + g.Log().Warningf(ctx, "[DataProvider] ResetConsecutiveFailure failed: channelID=%d, err=%v", channelID, err) + } } // GetAvailableModels 实现 DataProvider.GetAvailableModels @@ -760,9 +758,11 @@ func (p *DataProviderImpl) GetAvailableModels(ctx context.Context, tenantID, api var explicitIDs []struct { ModelID int64 `json:"model_id"` } - dao.MdlTenantModels.Ctx(ctx). + if err := dao.MdlTenantModels.Ctx(ctx). Where("tenant_id", tenantID).Where("enabled", true). - Fields("model_id").Scan(&explicitIDs) + Fields("model_id").Scan(&explicitIDs); err != nil { + g.Log().Warningf(ctx, "[DataProvider] GetAvailableModels scan tenant models failed: tenantID=%d, err=%v", tenantID, err) + } for _, id := range explicitIDs { allModelIDs = append(allModelIDs, id.ModelID) } @@ -938,9 +938,12 @@ func parseCapabilitiesJSON(raw string) map[string]bool { // incrementConsecutiveFailure 递增连续失败计数 func incrementConsecutiveFailure(ctx context.Context, channelID int64) { - g.DB().Exec(ctx, - "UPDATE chn_health_scores SET consecutive_failures = consecutive_failures + 1, updated_at = ? WHERE channel_id = ? AND consecutive_failures < 10", - time.Now(), channelID) + _, err := g.DB().Exec(ctx, + "UPDATE chn_health_scores SET consecutive_failures = consecutive_failures + 1, updated_at = NOW() WHERE channel_id = ? AND consecutive_failures < 10", + channelID) + if err != nil { + g.Log().Warningf(ctx, "[DataProvider] incrementConsecutiveFailure failed: channelID=%d, err=%v", channelID, err) + } } // modelInfoCached 模型信息缓存 @@ -983,10 +986,12 @@ func getChannelKey(ctx context.Context, channelID int64) (string, error) { } // 更新最后使用时间（用于监控） - dao.ChnChannelKeys.Ctx(ctx). + if _, err := dao.ChnChannelKeys.Ctx(ctx). Where("id", key.ID). Data(do.ChnChannelKeys{LastUsedAt: gtime.Now()}). - Update() + Update(); err != nil { + g.Log().Debugf(ctx, "[DataProvider] update channel key LastUsedAt failed: keyID=%d, err=%v", key.ID, err) + } encKey := GetEncryptionKey() decrypted, err := uc.DecryptString(encKey, key.EncryptedKey) @@ -1089,18 +1094,15 @@ func (p *DataProviderImpl) GetModelDeprecationInfo(ctx context.Context, modelNam ReplacementModel string `json:"replacement_model"` } - var row depRow + var row *depRow err := dao.MdlModels.Ctx(ctx). Where("model_id", modelName). Fields("status, sunset_date, replacement_model"). Scan(&row) if err != nil { - if strings.Contains(err.Error(), "no rows in result set") { - return nil, common.ErrModelNotFound - } return nil, err } - if row.Status == "" { + if row == nil { return nil, common.ErrModelNotFound } if row.Status != "deprecated" { diff --git a/internal/logic/task/async_provider.go b/internal/logic/task/async_provider.go index f25ef86..87fb86e 100644 --- a/internal/logic/task/async_provider.go +++ b/internal/logic/task/async_provider.go @@ -169,7 +169,7 @@ func (p *AsyncProvider) GetTaskByPublicID(ctx context.Context, publicTaskID stri } // GetTaskByPublicIDAndUser 根据公开任务 ID + 用户 ID 查询 -func (p *AsyncProvider) GetTaskByPublicIDAndUser(ctx context.Context, publicTaskID string, userID int64) (*common.AsyncTask, error) { +func (p *AsyncProvider) GetTaskByPublicIDAndUser(ctx context.Context, publicTaskID string, userID int64, tenantID int64) (*common.AsyncTask, error) { var row *struct { ID int64 `json:"id"` PublicTaskID string `json:"public_task_id"` @@ -199,6 +199,7 @@ func (p *AsyncProvider) GetTaskByPublicIDAndUser(ctx context.Context, publicTask err := dao.TskModelTasks.Ctx(ctx). Where("public_task_id", publicTaskID). Where("user_id", userID). + Where("tenant_id", tenantID). Scan(&row) if err != nil { return nil, gerror.Wrapf(err, "query async task failed: public_id=%s user_id=%d", publicTaskID, userID) diff --git a/relay/channel/baidu_v2/adaptor.go b/relay/channel/baidu_v2/adaptor.go index 27ac935..f61f961 100644 --- a/relay/channel/baidu_v2/adaptor.go +++ b/relay/channel/baidu_v2/adaptor.go @@ -68,6 +68,15 @@ func (a *Adaptor) SetupRequestHeader(header http.Header, info *common.RelayInfo) // ConvertRequest 转换请求体。 // 如果模型名以 "-search" 结尾，去掉后缀并注入 web_search 配置。 func (a *Adaptor) ConvertRequest(ctx context.Context, info *common.RelayInfo, requestBody []byte) (io.Reader, error) { + // 非 OpenAI 格式先转换为 OpenAI + if info.InboundFormat != "" && info.InboundFormat != constant.RelayFormatOpenAI { + converted, err := openai.ConvertToOpenAI(requestBody, info) + if err != nil { + return nil, err + } + requestBody = converted + } + var rawMap map[string]json.RawMessage if err := json.Unmarshal(requestBody, &rawMap); err != nil { return bytes.NewReader(requestBody), nil @@ -95,7 +104,8 @@ func (a *Adaptor) ConvertRequest(ctx context.Context, info *common.RelayInfo, re } // 设置模型名 - rawMap["model"] = json.RawMessage(`"` + modelName + `"`) + modelJSON, _ := json.Marshal(modelName) + rawMap["model"] = json.RawMessage(modelJSON) converted, err := json.Marshal(rawMap) if err != nil { diff --git a/relay/channel/claude/adaptor.go b/relay/channel/claude/adaptor.go index a245367..f910bd4 100644 --- a/relay/channel/claude/adaptor.go +++ b/relay/channel/claude/adaptor.go @@ -195,7 +195,7 @@ func injectClaudeEffort(r io.Reader, info *common.RelayInfo) io.Reader { // 仅在客户端未显式设置 thinking 时注入 if _, exists := rawMap["thinking"]; !exists { - rawMap["thinking"] = json.RawMessage(fmt.Sprintf(`{"type":"adaptive"}`)) + rawMap["thinking"] = json.RawMessage(`{"type":"adaptive"}`) } result, err := json.Marshal(rawMap) diff --git a/relay/channel/claude/response.go b/relay/channel/claude/response.go index 42ef088..1a91c83 100644 --- a/relay/channel/claude/response.go +++ b/relay/channel/claude/response.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/qianfree/team-api/relay/common" @@ -63,7 +64,8 @@ func (a *Adaptor) handleStreamToOpenAI(ctx context.Context, resp *http.Response, } helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() scanner := bufio.NewScanner(resp.Body) buf := make([]byte, 0, 64*1024) @@ -158,8 +160,9 @@ func (a *Adaptor) handleStreamToOpenAI(ctx context.Context, resp *http.Response, // 脱敏思考，OpenAI 格式无等价物，忽略 case "tool_use": toolCall := dto.ToolCall{ - ID: event.ContentBlock.ID, - Type: "function", + Index: toolCallIdx, + ID: event.ContentBlock.ID, + Type: "function", Function: dto.FunctionCall{ Name: event.ContentBlock.Name, Arguments: "", @@ -348,7 +351,7 @@ func (a *Adaptor) handleClaudeNativeNonStream(ctx context.Context, resp *http.Re } if info.ChannelMeta.IsModelMapped { - body = replaceModelName(body, info.OriginModelName) + body = helper.ReplaceModelName(body, info.OriginModelName) } writer.Header().Set("Content-Type", "application/json") @@ -379,7 +382,8 @@ func (a *Adaptor) handleClaudeNativeStream(ctx context.Context, resp *http.Respo } helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() reader := bufio.NewReaderSize(resp.Body, 64*1024) var usage dto.ClaudeUsage @@ -437,7 +441,7 @@ func (a *Adaptor) handleClaudeNativeStream(ctx context.Context, resp *http.Respo } if info.ChannelMeta.IsModelMapped { - replaced := string(replaceModelName([]byte(data), info.OriginModelName)) + replaced := string(helper.ReplaceModelName([]byte(data), info.OriginModelName)) line = fmt.Sprintf("data: %s\n", replaced) } } @@ -562,55 +566,6 @@ func writeStreamChunk(w http.ResponseWriter, chunk *dto.ChatCompletionStreamResp _ = helper.WriteSSEData(w, string(data)) } -// replaceModelName 替换 JSON 中 "model" 字段的值 -func replaceModelName(body []byte, modelName string) []byte { - fieldPrefix := []byte(`"model":"`) - replacement := append([]byte(`"model":"`), modelName...) - replacement = append(replacement, '"') - - result := make([]byte, 0, len(body)) - i := 0 - for i < len(body) { - idx := indexOf(body[i:], fieldPrefix) - if idx == -1 { - result = append(result, body[i:]...) - break - } - result = append(result, body[i:i+idx+len(fieldPrefix)]...) - i += idx + len(fieldPrefix) - - endQuote := indexOf(body[i:], []byte{'"'}) - if endQuote == -1 { - result = append(result, body[i:]...) - break - } - i += endQuote + 1 - result = append(result, replacement...) - } - return result -} - -func indexOf(s, sep []byte) int { - for i := 0; i <= len(s)-len(sep); i++ { - if equal(s[i:i+len(sep)], sep) { - return i - } - } - return -1 -} - -func equal(a, b []byte) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if a[i] != b[i] { - return false - } - } - return true -} - func joinTextPartsResponse(parts []string) any { if len(parts) == 0 { return nil diff --git a/relay/channel/coze/adaptor.go b/relay/channel/coze/adaptor.go index 0efce54..06d41ca 100644 --- a/relay/channel/coze/adaptor.go +++ b/relay/channel/coze/adaptor.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/qianfree/team-api/relay/channel/openai" @@ -130,7 +131,8 @@ func (a *Adaptor) DoResponse(ctx context.Context, resp *http.Response, info *com // handleStreamResponse 流式模式：逐事件读取 Coze SSE，转换为 OpenAI SSE 格式输出 func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, info *common.RelayInfo, writer http.ResponseWriter) (*common.Usage, error) { helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() streamStatus := info.StreamStatus if streamStatus == nil { @@ -180,9 +182,9 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, if msg.Type != "answer" { continue } - completionTokens += estimateTokens(msg.Content) + completionTokens += helper.EstimateTokens(msg.Content) - chunk := buildOpenAIStreamChunk(chatID, createdAt, modelName, msg.Content, nil) + chunk := helper.BuildOpenAIStreamChunk(chatID, createdAt, modelName, msg.Content, nil) chunkJSON, err := json.Marshal(chunk) if err != nil { continue @@ -195,7 +197,7 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, case "conversation.message.completed": // 完成事件，发送带 finish_reason 的最终 chunk finishReason := "stop" - chunk := buildOpenAIStreamChunk(chatID, createdAt, modelName, "", &finishReason) + chunk := helper.BuildOpenAIStreamChunk(chatID, createdAt, modelName, "", &finishReason) chunkJSON, _ := json.Marshal(chunk) _ = helper.WriteSSEData(writer, string(chunkJSON)) @@ -286,7 +288,7 @@ func (a *Adaptor) handleNonStreamResponse(ctx context.Context, resp *http.Respon } content := fullContent.String() - completionTokens := estimateTokens(content) + completionTokens := helper.EstimateTokens(content) // 构建 OpenAI 非流式响应 chatResp := dto.ChatCompletionResponse{ @@ -325,41 +327,6 @@ func (a *Adaptor) handleNonStreamResponse(ctx context.Context, resp *http.Respon }, nil } -// buildOpenAIStreamChunk 构建 OpenAI 格式的流式响应块 -func buildOpenAIStreamChunk(id string, created int64, model string, content string, finishReason *string) dto.ChatCompletionStreamResponse { - delta := dto.Message{} - if content != "" { - delta.Content = content - } - if finishReason == nil { - delta.Role = "assistant" - } - - return dto.ChatCompletionStreamResponse{ - ID: id, - Object: "chat.completion.chunk", - Created: created, - Model: model, - Choices: []dto.StreamChoice{ - { - Index: 0, - Delta: delta, - FinishReason: finishReason, - }, - }, - } -} - -// estimateTokens 粗略估算 token 数（英文按空格分词，中文按字符计数）。 -// 这是一个简化实现，实际 token 计数应由上游返回或使用 tiktoken 等库计算。 -func estimateTokens(text string) int { - if text == "" { - return 0 - } - // 粗略估算：每 4 个字符约 1 个 token - return (len(text) + 3) / 4 -} - func (a *Adaptor) GetChannelName() string { return ChannelName } diff --git a/relay/channel/dify/adaptor.go b/relay/channel/dify/adaptor.go index 0178d76..9330c6c 100644 --- a/relay/channel/dify/adaptor.go +++ b/relay/channel/dify/adaptor.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/qianfree/team-api/relay/channel/openai" @@ -162,7 +163,7 @@ func (a *Adaptor) handleNonStreamResponse(ctx context.Context, resp *http.Respon // 如果 Dify 未返回用量信息，使用粗略估算 if chatResp.Usage.TotalTokens == 0 { - chatResp.Usage.CompletionTokens = estimateTokens(difyResp.Answer) + chatResp.Usage.CompletionTokens = helper.EstimateTokens(difyResp.Answer) chatResp.Usage.TotalTokens = chatResp.Usage.CompletionTokens } @@ -185,7 +186,8 @@ func (a *Adaptor) handleNonStreamResponse(ctx context.Context, resp *http.Respon // handleStreamResponse 处理 Dify streaming 模式 SSE 响应 func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, info *common.RelayInfo, writer http.ResponseWriter) (*common.Usage, error) { helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() streamStatus := info.StreamStatus if streamStatus == nil { @@ -233,7 +235,7 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, continue } - chunk := buildOpenAIStreamChunk(chatID, createdAt, modelName, event.Answer, nil) + chunk := helper.BuildOpenAIStreamChunk(chatID, createdAt, modelName, event.Answer, nil) chunkJSON, err := json.Marshal(chunk) if err != nil { continue @@ -251,7 +253,7 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, // 发送带 finish_reason 的最终 chunk finishReason := "stop" - chunk := buildOpenAIStreamChunk(chatID, createdAt, modelName, "", &finishReason) + chunk := helper.BuildOpenAIStreamChunk(chatID, createdAt, modelName, "", &finishReason) chunkJSON, _ := json.Marshal(chunk) _ = helper.WriteSSEData(writer, string(chunkJSON)) @@ -271,43 +273,19 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, return &usage, fmt.Errorf("read dify stream failed: %w", err) } - streamStatus.SetEndReason(common.StreamEndReasonDone, nil) + // 流异常结束（未收到 message_end），补充兜底的结束 chunk 和 [DONE] + if streamStatus.GetEndReason() == "" { + finishReason := "stop" + chunk := helper.BuildOpenAIStreamChunk(chatID, createdAt, modelName, "", &finishReason) + chunkJSON, _ := json.Marshal(chunk) + _ = helper.WriteSSEData(writer, string(chunkJSON)) + _ = helper.WriteSSEData(writer, "[DONE]") + streamStatus.SetEndReason(common.StreamEndReasonDone, nil) + } + return &usage, nil } -// buildOpenAIStreamChunk 构建 OpenAI 格式的流式响应块 -func buildOpenAIStreamChunk(id string, created int64, model string, content string, finishReason *string) dto.ChatCompletionStreamResponse { - delta := dto.Message{} - if content != "" { - delta.Content = content - } - if finishReason == nil { - delta.Role = "assistant" - } - - return dto.ChatCompletionStreamResponse{ - ID: id, - Object: "chat.completion.chunk", - Created: created, - Model: model, - Choices: []dto.StreamChoice{ - { - Index: 0, - Delta: delta, - FinishReason: finishReason, - }, - }, - } -} - -// estimateTokens 粗略估算 token 数（每 4 个字符约 1 个 token） -func estimateTokens(text string) int { - if text == "" { - return 0 - } - return (len(text) + 3) / 4 -} - func (a *Adaptor) GetChannelName() string { return ChannelName } diff --git a/relay/channel/gemini/response.go b/relay/channel/gemini/response.go index 2f6c46e..dba371f 100644 --- a/relay/channel/gemini/response.go +++ b/relay/channel/gemini/response.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/qianfree/team-api/relay/common" @@ -204,7 +205,8 @@ func (a *Adaptor) handleStreamToOpenAI(ctx context.Context, resp *http.Response, } helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() isCA := a.isCodeAssistActive() scanner := bufio.NewScanner(resp.Body) @@ -261,6 +263,20 @@ func (a *Adaptor) handleStreamToOpenAI(ctx context.Context, resp *http.Response, // 检查 promptFeedback.blockReason（流式安全过滤） if geminiResp.PromptFeedback != nil && geminiResp.PromptFeedback.BlockReason != "" { + // SSE 头已发送，需先发送结束 chunk 和 [DONE] 避免客户端挂起 + filterReason := "content_filter" + endChunk := dto.ChatCompletionStreamResponse{ + ID: fmt.Sprintf("chatcmpl-%s", info.RequestID), + Object: "chat.completion.chunk", + Model: info.OriginModelName, + Choices: []dto.StreamChoice{{ + Index: 0, + FinishReason: &filterReason, + }}, + } + writeStreamChunk(writer, &endChunk) + _ = helper.WriteSSEData(writer, "[DONE]") + info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil) return nil, constant.NewRequestError( fmt.Sprintf("request blocked by Gemini safety filter: %s", geminiResp.PromptFeedback.BlockReason), nil, ) @@ -464,16 +480,18 @@ func (a *Adaptor) handleStreamToOpenAI(ctx context.Context, resp *http.Response, TotalTokens: totalUsage.TotalTokenCount, } } - writeStreamChunk(writer, &endChunk) - - helper.WriteSSEData(writer, "[DONE]") - info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil) if err := scanner.Err(); err != nil && err != io.EOF && ctx.Err() == nil { + writeStreamChunk(writer, &endChunk) + _ = helper.WriteSSEData(writer, "[DONE]") info.StreamStatus.SetEndReason(common.StreamEndReasonError, err) - return &common.Usage{}, fmt.Errorf("stream scanner error: %w", err) + return geminiUsageToCommon(&totalUsage), fmt.Errorf("stream scanner error: %w", err) } + writeStreamChunk(writer, &endChunk) + _ = helper.WriteSSEData(writer, "[DONE]") + info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil) + return geminiUsageToCommon(&totalUsage), nil } diff --git a/relay/channel/ollama/adaptor.go b/relay/channel/ollama/adaptor.go index 70c1993..11b9026 100644 --- a/relay/channel/ollama/adaptor.go +++ b/relay/channel/ollama/adaptor.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/qianfree/team-api/relay/channel/openai" @@ -213,7 +214,8 @@ func (a *Adaptor) handleChatStreamResponse(ctx context.Context, resp *http.Respo } helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() scanner := bufio.NewScanner(resp.Body) buf := make([]byte, 0, 64*1024) @@ -360,7 +362,8 @@ func (a *Adaptor) handleGenerateStreamResponse(ctx context.Context, resp *http.R } helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() scanner := bufio.NewScanner(resp.Body) buf := make([]byte, 0, 64*1024) diff --git a/relay/channel/openai/adaptor.go b/relay/channel/openai/adaptor.go index e2e122c..2405e0b 100644 --- a/relay/channel/openai/adaptor.go +++ b/relay/channel/openai/adaptor.go @@ -15,6 +15,7 @@ import ( "github.com/qianfree/team-api/relay/common" "github.com/qianfree/team-api/relay/constant" "github.com/qianfree/team-api/relay/dto" + "github.com/qianfree/team-api/relay/helper" "github.com/qianfree/team-api/relay/override" ) @@ -229,7 +230,7 @@ func (a *Adaptor) DoRequest(ctx context.Context, info *common.RelayInfo, request // DoResponse 处理上游响应，按客户端格式和 RelayMode 分发 func (a *Adaptor) DoResponse(ctx context.Context, resp *http.Response, info *common.RelayInfo, writer http.ResponseWriter) (*common.Usage, error) { clientFormat := info.GetOriginalClientFormat() - g.Log().Infof(ctx, "[OpenAI.Adaptor.DoResponse] Entry: clientFormat=%s, relayMode=%d, isStream=%v, statusCode=%d", + g.Log().Debugf(ctx, "[OpenAI.Adaptor.DoResponse] Entry: clientFormat=%s, relayMode=%d, isStream=%v, statusCode=%d", clientFormat, info.RelayMode, info.IsStream, resp.StatusCode) // 根据客户端格式转换响应 @@ -312,7 +313,7 @@ func (a *Adaptor) handleChatNonStreamResponse(ctx context.Context, resp *http.Re } if info.ChannelMeta.IsModelMapped { - body = replaceModelName(body, info.OriginModelName) + body = helper.ReplaceModelName(body, info.OriginModelName) } writer.Header().Set("Content-Type", "application/json") @@ -363,39 +364,6 @@ func (a *Adaptor) handleCompletionStreamResponse(ctx context.Context, resp *http return StreamHandlerForCompletions(ctx, resp, info, writer) } -// replaceModelName 替换响应 JSON 中 "model":"xxx" 的值为指定模型名 -func replaceModelName(body []byte, modelName string) []byte { - target := []byte(`"model":"`) - replacement := append(target, modelName...) - replacement = append(replacement, '"') - return replaceModelNameInJSON(body, target, replacement) -} - -// replaceModelNameInJSON 替换 JSON 中 "model" 字段的值 -func replaceModelNameInJSON(body, fieldPrefix, replacement []byte) []byte { - result := make([]byte, 0, len(body)) - i := 0 - for i < len(body) { - idx := bytes.Index(body[i:], fieldPrefix) - if idx == -1 { - result = append(result, body[i:]...) - break - } - result = append(result, body[i:i+idx+len(fieldPrefix)]...) - i += idx + len(fieldPrefix) - - // 跳到引号结束 - endQuote := bytes.IndexByte(body[i:], '"') - if endQuote == -1 { - result = append(result, body[i:]...) - break - } - i += endQuote + 1 - result = append(result, replacement...) - } - return result -} - // GetChannelName 返回渠道名称 func (a *Adaptor) GetChannelName() string { return "OpenAI" diff --git a/relay/channel/openai/claude_response.go b/relay/channel/openai/claude_response.go index bfc53e4..6785cc5 100644 --- a/relay/channel/openai/claude_response.go +++ b/relay/channel/openai/claude_response.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/qianfree/team-api/relay/common" @@ -78,7 +79,8 @@ func handleClaudeInboundStream(ctx context.Context, resp *http.Response, info *c } helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() scanner := bufio.NewScanner(resp.Body) buf := make([]byte, 0, 64*1024) diff --git a/relay/channel/openai/completion.go b/relay/channel/openai/completion.go index 43cb734..ab05b77 100644 --- a/relay/channel/openai/completion.go +++ b/relay/channel/openai/completion.go @@ -9,6 +9,7 @@ import ( "github.com/qianfree/team-api/relay/common" "github.com/qianfree/team-api/relay/constant" "github.com/qianfree/team-api/relay/dto" + "github.com/qianfree/team-api/relay/helper" ) // handleCompletionNonStreamResponse 处理 Completions 非流式响应 @@ -29,7 +30,7 @@ func (a *Adaptor) handleCompletionNonStreamResponse(ctx context.Context, resp *h } if info.ChannelMeta.IsModelMapped { - body = replaceModelName(body, info.OriginModelName) + body = helper.ReplaceModelName(body, info.OriginModelName) } writer.Header().Set("Content-Type", "application/json") diff --git a/relay/channel/openai/converter.go b/relay/channel/openai/converter.go index 82ac033..e59ca01 100644 --- a/relay/channel/openai/converter.go +++ b/relay/channel/openai/converter.go @@ -1408,7 +1408,7 @@ func HandleResponsesNonStreamToChat(ctx context.Context, resp *http.Response, in if resp.StatusCode != http.StatusOK { writer.Header().Set("Content-Type", "application/json") writer.WriteHeader(resp.StatusCode) - writer.Write(body) + _, _ = writer.Write(body) return nil, constant.NewUpstreamError(resp.StatusCode, string(body), nil) } diff --git a/relay/channel/openai/embedding.go b/relay/channel/openai/embedding.go index 1f991c1..222c0c1 100644 --- a/relay/channel/openai/embedding.go +++ b/relay/channel/openai/embedding.go @@ -9,6 +9,7 @@ import ( "github.com/qianfree/team-api/relay/common" "github.com/qianfree/team-api/relay/constant" "github.com/qianfree/team-api/relay/dto" + "github.com/qianfree/team-api/relay/helper" ) // handleEmbeddingResponse 处理 Embeddings 非流式响应 @@ -25,9 +26,13 @@ func (a *Adaptor) handleEmbeddingResponse(ctx context.Context, resp *http.Respon } if info.ChannelMeta.IsModelMapped { - body = replaceModelName(body, info.OriginModelName) + body = helper.ReplaceModelName(body, info.OriginModelName) } + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusOK) + _, _ = writer.Write(body) + var embedResp dto.EmbeddingResponse if err := json.Unmarshal(body, &embedResp); err == nil { return &common.Usage{ @@ -39,10 +44,5 @@ func (a *Adaptor) handleEmbeddingResponse(ctx context.Context, resp *http.Respon }, nil } - // 解析失败也透传 body - writer.Header().Set("Content-Type", "application/json") - writer.WriteHeader(resp.StatusCode) - _, _ = writer.Write(body) - return &common.Usage{}, nil } diff --git a/relay/channel/openai/gemini_response.go b/relay/channel/openai/gemini_response.go index de94f5f..652c5d0 100644 --- a/relay/channel/openai/gemini_response.go +++ b/relay/channel/openai/gemini_response.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/qianfree/team-api/relay/common" @@ -63,7 +64,8 @@ func handleGeminiInboundStream(ctx context.Context, resp *http.Response, info *c } helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() scanner := bufio.NewScanner(resp.Body) buf := make([]byte, 0, 64*1024) diff --git a/relay/channel/openai/rerank.go b/relay/channel/openai/rerank.go index 18783f4..8beecdb 100644 --- a/relay/channel/openai/rerank.go +++ b/relay/channel/openai/rerank.go @@ -9,6 +9,7 @@ import ( "github.com/qianfree/team-api/relay/common" "github.com/qianfree/team-api/relay/constant" "github.com/qianfree/team-api/relay/dto" + "github.com/qianfree/team-api/relay/helper" "github.com/gogf/gf/v2/frame/g" ) @@ -31,7 +32,7 @@ func handleRerankResponse(ctx context.Context, resp *http.Response, info *common } if info.ChannelMeta.IsModelMapped { - body = replaceModelName(body, info.OriginModelName) + body = helper.ReplaceModelName(body, info.OriginModelName) } writer.Header().Set("Content-Type", "application/json") @@ -39,15 +40,15 @@ func handleRerankResponse(ctx context.Context, resp *http.Response, info *common _, _ = writer.Write(body) // 提取 usage - usage := extractRerankUsage(body) + usage := extractRerankUsage(ctx, body) return usage, nil } // extractRerankUsage 从重排响应中提取 token 使用量 -func extractRerankUsage(body []byte) *common.Usage { +func extractRerankUsage(ctx context.Context, body []byte) *common.Usage { var rerankResp dto.RerankResponse if err := json.Unmarshal(body, &rerankResp); err != nil { - g.Log().Warningf(nil, "[OpenAI.extractRerankUsage] unmarshal rerank response failed: %v", err) + g.Log().Warningf(ctx, "[OpenAI.extractRerankUsage] unmarshal rerank response failed: %v", err) return &common.Usage{} } diff --git a/relay/channel/openai/responses.go b/relay/channel/openai/responses.go index d0ef7eb..d5fa995 100644 --- a/relay/channel/openai/responses.go +++ b/relay/channel/openai/responses.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/qianfree/team-api/relay/common" @@ -170,7 +171,8 @@ func (a *Adaptor) handleResponsesInboundStream(ctx context.Context, resp *http.R } helper.SetEventStreamHeaders(writer) - defer helper.PingTicker(writer, 15*time.Second)() + var writeMu sync.Mutex + defer helper.PingTicker(writer, 15*time.Second, &writeMu)() scanner := bufio.NewScanner(resp.Body) buf := make([]byte, 0, 64*1024) diff --git a/relay/channel/openai/stream.go b/relay/channel/openai/stream.go index 9ac9e7a..f74d6d9 100644 --- a/relay/channel/openai/stream.go +++ b/relay/channel/openai/stream.go @@ -19,7 +19,7 @@ func StreamHandler(ctx context.Context, resp *http.Response, info *common.RelayI helper.StreamScannerHandler(ctx, resp, info, writer, func(data string, sr *helper.StreamResult) { if info.ChannelMeta.IsModelMapped { - data = string(replaceModelName([]byte(data), info.OriginModelName)) + data = string(helper.ReplaceModelName([]byte(data), info.OriginModelName)) } var streamResp dto.ChatCompletionStreamResponse @@ -91,7 +91,7 @@ func StreamHandlerForCompletions(ctx context.Context, resp *http.Response, info helper.StreamScannerHandler(ctx, resp, info, writer, func(data string, sr *helper.StreamResult) { if info.ChannelMeta.IsModelMapped { - data = string(replaceModelName([]byte(data), info.OriginModelName)) + data = string(helper.ReplaceModelName([]byte(data), info.OriginModelName)) } var streamResp dto.CompletionsStreamResponse diff --git a/relay/channel/tencent/adaptor.go b/relay/channel/tencent/adaptor.go index 697f511..9d2355a 100644 --- a/relay/channel/tencent/adaptor.go +++ b/relay/channel/tencent/adaptor.go @@ -85,14 +85,7 @@ func (a *Adaptor) SetupRequestHeader(header http.Header, info *common.RelayInfo) header.Set("X-TC-Action", action) header.Set("X-TC-Version", apiVersion) - timestamp := time.Now().Unix() - header.Set("X-TC-Timestamp", fmt.Sprintf("%d", timestamp)) - - // 如果有 secretKey 则计算签名；否则跳过（代理网关场景） - if a.secretKey != "" { - // 签名需要请求体，此处在 DoRequest 中重新计算 - // SetupRequestHeader 只设置非签名头，签名在 DoRequest 中补充 - } + // 时间戳和签名统一在 DoRequest 中设置（签名依赖请求体，SetupRequestHeader 阶段无法计算） return nil } diff --git a/relay/channel/vertex/auth.go b/relay/channel/vertex/auth.go index 402786a..f951ee9 100644 --- a/relay/channel/vertex/auth.go +++ b/relay/channel/vertex/auth.go @@ -16,6 +16,8 @@ import ( "strings" "sync" "time" + + "golang.org/x/sync/singleflight" ) // serviceAccountKey Google 服务账号 JSON 密钥结构 @@ -41,9 +43,12 @@ type tokenEntry struct { // tokenCache 全局令牌缓存（key 为 client_email） var tokenCache sync.Map +// sfGroup 合并并发令牌刷新请求 +var sfGroup singleflight.Group + // getVertexAccessToken 获取 Vertex AI 访问令牌。 // 通过服务账号 JSON 密钥签发 JWT，然后通过 Google OAuth2 端点交换访问令牌。 -// 令牌缓存在 sync.Map 中，未过期时复用。 +// 令牌缓存在 sync.Map 中，未过期时复用。并发刷新通过 singleflight 合并。 func getVertexAccessToken(serviceAccountJSON string) (string, error) { var key serviceAccountKey if err := json.Unmarshal([]byte(serviceAccountJSON), &key); err != nil { @@ -58,32 +63,44 @@ func getVertexAccessToken(serviceAccountJSON string) (string, error) { } } - // 解析 RSA 私钥 - rsaKey, err := parseRSAPrivateKey(key.PrivateKey) - if err != nil { - return "", fmt.Errorf("parse private key failed: %w", err) - } + // 使用 singleflight 合并并发刷新 + val, err, _ := sfGroup.Do(key.ClientEmail, func() (any, error) { + // 双重检查：进入 singleflight 后再次检查缓存 + if cached, ok := tokenCache.Load(key.ClientEmail); ok { + entry := cached.(*tokenEntry) + if time.Now().Before(entry.expiry.Add(-60 * time.Second)) { + return entry.accessToken, nil + } + } - // 创建 JWT - now := time.Now() - jwtToken, err := createSignedJWT(key.ClientEmail, rsaKey, now) - if err != nil { - return "", fmt.Errorf("create JWT failed: %w", err) - } + rsaKey, err := parseRSAPrivateKey(key.PrivateKey) + if err != nil { + return "", fmt.Errorf("parse private key failed: %w", err) + } - // 交换访问令牌 - accessToken, expiresIn, err := exchangeJWTForToken(jwtToken) - if err != nil { - return "", fmt.Errorf("exchange token failed: %w", err) - } + now := time.Now() + jwtToken, err := createSignedJWT(key.ClientEmail, rsaKey, now) + if err != nil { + return "", fmt.Errorf("create JWT failed: %w", err) + } - // 缓存令牌 - tokenCache.Store(key.ClientEmail, &tokenEntry{ - accessToken: accessToken, - expiry: now.Add(time.Duration(expiresIn) * time.Second), + accessToken, expiresIn, err := exchangeJWTForToken(jwtToken) + if err != nil { + return "", fmt.Errorf("exchange token failed: %w", err) + } + + tokenCache.Store(key.ClientEmail, &tokenEntry{ + accessToken: accessToken, + expiry: now.Add(time.Duration(expiresIn) * time.Second), + }) + + return accessToken, nil }) + if err != nil { + return "", err + } - return accessToken, nil + return val.(string), nil } // parseServiceAccountProjectID 从服务账号 JSON 中提取 project_id diff --git a/relay/common/task_provider.go b/relay/common/task_provider.go index a88186a..3e61def 100644 --- a/relay/common/task_provider.go +++ b/relay/common/task_provider.go @@ -21,7 +21,7 @@ type TaskDataProvider interface { GetTaskByPublicID(ctx context.Context, publicTaskID string) (*AsyncTask, error) // GetTaskByPublicIDAndUser 根据公开任务 ID + 用户 ID 查询（客户端查询） - GetTaskByPublicIDAndUser(ctx context.Context, publicTaskID string, userID int64) (*AsyncTask, error) + GetTaskByPublicIDAndUser(ctx context.Context, publicTaskID string, userID int64, tenantID int64) (*AsyncTask, error) // GetNonTerminalTasks 获取所有非终态任务（轮询用） GetNonTerminalTasks(ctx context.Context, limit int) ([]*AsyncTask, error) diff --git a/relay/handler/relay_handler.go b/relay/handler/relay_handler.go index 30d8668..a768270 100644 --- a/relay/handler/relay_handler.go +++ b/relay/handler/relay_handler.go @@ -330,6 +330,9 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He return nil, nil, err } + // 在写入任何响应体之前设置限流/弃用 header（流式响应中 WriteHeader 后无法追加） + setPreResponseHeaders(rc.Writer, v.billingResult) + // Phase 2: 资源准备（并发控制 + 监控注册 + 预扣） if billing != nil { if !billing.AcquireConcurrent(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, v.modelName) { @@ -436,6 +439,7 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He if constant.IsRetryable(err) && attempt < maxRetries { recordChannelError(rc, selection, v.modelName, attempt, false, err, info.LatencyMs()) appendHop(trace, hop, false, err.Error(), info.LatencyMs()) + settleCancel() continue } @@ -494,6 +498,7 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He } recordChannelError(rc, selection, v.modelName, attempt, false, err, info.LatencyMs()) appendHop(trace, hop, false, err.Error(), info.LatencyMs()) + settleCancel() continue } @@ -769,6 +774,27 @@ func tokenDetailField(details *common.TokenDetails, getter func(*common.TokenDet return getter(details) } +// setPreResponseHeaders 在写入响应体之前设置限流和弃用 header +func setPreResponseHeaders(w http.ResponseWriter, br *BillingResult) { + if br == nil { + return + } + if info := br.RateLimitInfo; info != nil { + w.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", info.Limit)) + w.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", info.Remaining)) + w.Header().Set("X-RateLimit-Reset", fmt.Sprintf("%d", info.ResetAt)) + } + if dep := br.Deprecation; dep != nil { + w.Header().Set("Deprecation", "true") + if dep.SunsetDate != "" { + w.Header().Set("Sunset", dep.SunsetDate) + } + if dep.ReplacementModel != "" { + w.Header().Set("Link", fmt.Sprintf("; rel=\"successor-version\"", dep.ReplacementModel)) + } + } +} + // requestType 根据 isStream 返回请求类型 func requestType(isStream bool) int { if isStream { diff --git a/relay/handler/task_handler.go b/relay/handler/task_handler.go index 937940b..494ff73 100644 --- a/relay/handler/task_handler.go +++ b/relay/handler/task_handler.go @@ -130,7 +130,9 @@ func HandleTaskSubmit( // 7. 构建并发送请求 requestBody, err := adaptor.BuildRequestBody(ctx, info, body) if err != nil { - billingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount) + if err := billingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount); err != nil { + g.Log().Errorf(ctx, "HandleTaskSubmit: SettleTaskFailed error, requestID=%s, amount=%.6f, err=%v", rc.RequestID, preDeductAmount, err) + } g.Log().Errorf(ctx, "HandleTaskSubmit: build request failed, model=%s, err=%v", modelName, err) writeTaskError(rc.Writer, http.StatusInternalServerError, "build request failed: "+err.Error(), "") return @@ -138,7 +140,9 @@ func HandleTaskSubmit( resp, err := adaptor.DoRequest(ctx, info, requestBody) if err != nil { - billingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount) + if err := billingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount); err != nil { + g.Log().Errorf(ctx, "HandleTaskSubmit: SettleTaskFailed error, requestID=%s, amount=%.6f, err=%v", rc.RequestID, preDeductAmount, err) + } g.Log().Errorf(ctx, "HandleTaskSubmit: upstream request failed, model=%s, err=%v", modelName, err) writeTaskError(rc.Writer, http.StatusBadGateway, "upstream request failed: "+err.Error(), "") return @@ -150,7 +154,9 @@ func HandleTaskSubmit( // 8. 解析响应 upstreamTaskID, taskData, taskErr := adaptor.DoResponse(ctx, resp, info) if taskErr != nil { - billingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount) + if err := billingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount); err != nil { + g.Log().Errorf(ctx, "HandleTaskSubmit: SettleTaskFailed error, requestID=%s, amount=%.6f, err=%v", rc.RequestID, preDeductAmount, err) + } g.Log().Warningf(ctx, "HandleTaskSubmit: upstream response error, model=%s, status=%d, message=%q, body=%s", modelName, taskErr.StatusCode, taskErr.Message, string(taskData)) writeTaskError(rc.Writer, taskErr.StatusCode, taskErr.Message, taskErr.ErrCode) return @@ -214,7 +220,9 @@ func HandleTaskSubmit( } if err := dataProvider.CreateTask(ctx, task); err != nil { - billingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount) + if err := billingProvider.SettleTaskFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount); err != nil { + g.Log().Errorf(ctx, "HandleTaskSubmit: SettleTaskFailed error, requestID=%s, amount=%.6f, err=%v", rc.RequestID, preDeductAmount, err) + } g.Log().Errorf(ctx, "HandleTaskSubmit: create task record failed, publicTaskID=%s, model=%s, err=%v", publicTaskID, modelName, err) writeTaskError(rc.Writer, http.StatusInternalServerError, "create task record failed: "+err.Error(), "") return @@ -240,7 +248,7 @@ func HandleTaskFetch( rc *TaskRelayContext, dataProvider common.TaskDataProvider, ) { - task, err := dataProvider.GetTaskByPublicIDAndUser(ctx, publicTaskID, rc.UserID) + task, err := dataProvider.GetTaskByPublicIDAndUser(ctx, publicTaskID, rc.UserID, rc.TenantID) if err != nil { g.Log().Errorf(ctx, "HandleTaskFetch: query task failed, publicTaskID=%s, err=%v", publicTaskID, err) writeTaskError(rc.Writer, http.StatusInternalServerError, "query task failed", "") @@ -300,8 +308,10 @@ func generatePublicTaskID() string { } func randomHex(n int) string { - b := make([]byte, n) - rand.Read(b) + b := make([]byte, (n+1)/2) + if _, err := rand.Read(b); err != nil { + panic("crypto/rand failed: " + err.Error()) + } return hex.EncodeToString(b)[:n] } @@ -313,7 +323,7 @@ func HandleMjImageProxy( dataProvider common.TaskDataProvider, writer http.ResponseWriter, ) { - task, err := dataProvider.GetTaskByPublicIDAndUser(ctx, publicTaskID, rc.UserID) + task, err := dataProvider.GetTaskByPublicIDAndUser(ctx, publicTaskID, rc.UserID, rc.TenantID) if err != nil { g.Log().Errorf(ctx, "HandleMjImageProxy: query task failed, publicTaskID=%s, err=%v", publicTaskID, err) writeTaskError(writer, http.StatusInternalServerError, "query task failed", "") diff --git a/relay/helper/model_name.go b/relay/helper/model_name.go new file mode 100644 index 0000000..b80d936 --- /dev/null +++ b/relay/helper/model_name.go @@ -0,0 +1,34 @@ +package helper + +import "bytes" + +// ReplaceModelName 替换 JSON 响应中 "model":"xxx" 的值为指定模型名。 +// 使用字节级操作避免 JSON 反序列化再序列化的开销。 +func ReplaceModelName(body []byte, modelName string) []byte { + fieldPrefix := []byte(`"model":"`) + replacement := make([]byte, 0, len(fieldPrefix)+len(modelName)+1) + replacement = append(replacement, fieldPrefix...) + replacement = append(replacement, modelName...) + replacement = append(replacement, '"') + + result := make([]byte, 0, len(body)) + i := 0 + for i < len(body) { + idx := bytes.Index(body[i:], fieldPrefix) + if idx == -1 { + result = append(result, body[i:]...) + break + } + result = append(result, body[i:i+idx+len(fieldPrefix)]...) + i += idx + len(fieldPrefix) + + endQuote := bytes.IndexByte(body[i:], '"') + if endQuote == -1 { + result = append(result, body[i:]...) + break + } + i += endQuote + 1 + result = append(result, replacement...) + } + return result +} diff --git a/relay/helper/stream.go b/relay/helper/stream.go index 290d870..20c0c8a 100644 --- a/relay/helper/stream.go +++ b/relay/helper/stream.go @@ -4,7 +4,10 @@ import ( "fmt" "net/http" "strings" + "sync" "time" + + "github.com/qianfree/team-api/relay/dto" ) // SetEventStreamHeaders 设置 SSE 必要的响应头 @@ -80,9 +83,43 @@ func ExtractSSEData(line string) (string, bool) { return strings.TrimSpace(data), true } +// BuildOpenAIStreamChunk 构建 OpenAI 格式的流式响应块 +func BuildOpenAIStreamChunk(id string, created int64, model string, content string, finishReason *string) dto.ChatCompletionStreamResponse { + delta := dto.Message{} + if content != "" { + delta.Content = content + } + if finishReason == nil { + delta.Role = "assistant" + } + + return dto.ChatCompletionStreamResponse{ + ID: id, + Object: "chat.completion.chunk", + Created: created, + Model: model, + Choices: []dto.StreamChoice{ + { + Index: 0, + Delta: delta, + FinishReason: finishReason, + }, + }, + } +} + +// EstimateTokens 粗略估算 token 数（每 4 个字符约 1 个 token） +func EstimateTokens(text string) int { + if text == "" { + return 0 + } + return (len(text) + 3) / 4 +} + // PingTicker 在后台定期发送 SSE 保活注释 +// mu 用于与主循环的写操作串行化，避免并发写 ResponseWriter // 返回一个 stop 函数用于停止 goroutine -func PingTicker(w http.ResponseWriter, interval time.Duration) (stop func()) { +func PingTicker(w http.ResponseWriter, interval time.Duration, mu *sync.Mutex) (stop func()) { ticker := time.NewTicker(interval) done := make(chan struct{}) @@ -91,7 +128,10 @@ func PingTicker(w http.ResponseWriter, interval time.Duration) (stop func()) { for { select { case <-ticker.C: - if err := WriteSSEPing(w); err != nil { + mu.Lock() + err := WriteSSEPing(w) + mu.Unlock() + if err != nil { return } case <-done: diff --git a/relay/override/header.go b/relay/override/header.go index 7208794..80c5e12 100644 --- a/relay/override/header.go +++ b/relay/override/header.go @@ -82,7 +82,7 @@ func resolveHeaderPlaceholder(template string, info *common.RelayInfo) string { result := template // {api_key} → 渠道 API Key - result = replaceAllStr(result, "{api_key}", info.ChannelMeta.ApiKey) + result = strings.ReplaceAll(result, "{api_key}", info.ChannelMeta.ApiKey) // {client_header:xxx} → 客户端原始请求 header for { diff --git a/relay/override/operations.go b/relay/override/operations.go index 5a9af69..b2ebba2 100644 --- a/relay/override/operations.go +++ b/relay/override/operations.go @@ -140,7 +140,7 @@ func opReplace(data []byte, path string, value any) ([]byte, error) { return data, nil } - replaced := replaceAllStr(target.Str, oldStr, newStr) + replaced := strings.ReplaceAll(target.Str, oldStr, newStr) return sjson.SetBytes(data, path, replaced) } diff --git a/relay/override/override.go b/relay/override/override.go index 03f9028..eb40fed 100644 --- a/relay/override/override.go +++ b/relay/override/override.go @@ -4,6 +4,8 @@ import ( "encoding/json" "errors" "fmt" + "strconv" + "strings" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -345,7 +347,7 @@ func evaluateCondition(target any, mode string, expected any) bool { } return targetStr[len(targetStr)-len(expectedStr):] == expectedStr case "contains": - return containsStr(targetStr, expectedStr) + return strings.Contains(targetStr, expectedStr) case "gt": return compareNumeric(targetStr, expectedStr) > 0 case "gte": @@ -372,9 +374,9 @@ func resolveContextPath(path string, ctx map[string]any) string { result := path for k, v := range ctx { placeholder := "{" + k + "}" - if containsStr(result, placeholder) { + if strings.Contains(result, placeholder) { val := fmt.Sprintf("%v", v) - result = replaceAllStr(result, placeholder, val) + result = strings.ReplaceAll(result, placeholder, val) } } return result @@ -461,48 +463,10 @@ func valueToJSON(v any) json.RawMessage { return b } -func containsStr(s, substr string) bool { - return len(s) >= len(substr) && (len(substr) == 0 || containsStrImpl(s, substr)) -} - -func containsStrImpl(s, substr string) bool { - for i := 0; i <= len(s)-len(substr); i++ { - if s[i:i+len(substr)] == substr { - return true - } - } - return false -} - -func replaceAllStr(s, old, new string) string { - result := make([]byte, 0, len(s)) - for { - idx := indexOfStr(s, old) - if idx == -1 { - result = append(result, s...) - break - } - result = append(result, s[:idx]...) - result = append(result, new...) - s = s[idx+len(old):] - } - return string(result) -} - -func indexOfStr(s, substr string) int { - for i := 0; i <= len(s)-len(substr); i++ { - if s[i:i+len(substr)] == substr { - return i - } - } - return -1 -} - func compareNumeric(a, b string) int { - af, aok := parseFloat(a) - bf, bok := parseFloat(b) - if !aok || !bok { - // 退化为字符串比较 + af, aErr := strconv.ParseFloat(a, 64) + bf, bErr := strconv.ParseFloat(b, 64) + if aErr != nil || bErr != nil { if a < b { return -1 } @@ -520,32 +484,6 @@ func compareNumeric(a, b string) int { return 0 } -func parseFloat(s string) (float64, bool) { - var f float64 - n := 0 - for _, c := range s { - if c >= '0' && c <= '9' { - f = f*10 + float64(c-'0') - n++ - } else if c == '.' && n > 0 { - dec := 1.0 - for i := indexOfStr(s, ".") + 1; i < len(s); i++ { - cc := s[i] - if cc >= '0' && cc <= '9' { - f += float64(cc-'0') / dec - dec *= 10 - } else { - return 0, false - } - } - return f, true - } else { - return 0, false - } - } - return f, true -} - func opReturnError(value any) error { m, ok := value.(map[string]any) if !ok {