diff --git a/backend/internal/handler/batch_image_handler.go b/backend/internal/handler/batch_image_handler.go index b93c5a65ae..62990d251a 100644 --- a/backend/internal/handler/batch_image_handler.go +++ b/backend/internal/handler/batch_image_handler.go @@ -42,6 +42,9 @@ func (h *BatchImageHandler) Submit(c *gin.Context) { if !h.checkSecurityAuditBeforeSubmit(c, &req) { return } + if sessionID := service.ExtractClientSessionID(c); sessionID != "" { + req.SessionID = &sessionID + } got, err := h.service.Submit(c.Request.Context(), owner, req, c.GetHeader("Idempotency-Key")) if err != nil { batchImageError(c, err) diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 567d12c13a..04c811cedf 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -670,6 +670,7 @@ func usageLogFromServiceUser(l *service.UsageLog) UsageLog { MediaType: l.MediaType, UserAgent: l.UserAgent, IPAddress: l.IPAddress, + SessionID: l.SessionID, CacheTTLOverridden: l.CacheTTLOverridden, BillingMode: l.BillingMode, CreatedAt: l.CreatedAt, diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index d52f81afe3..d5ba4e9db6 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -523,6 +523,9 @@ type UsageLog struct { UserAgent *string `json:"user_agent"` // IPAddress is visible to the owner of the usage record. IPAddress *string `json:"ip_address,omitempty"` + // SessionID is the explicit client-provided request correlation identifier + // (e.g. the session_id / X-Session-Id headers). Omitted when absent. + SessionID *string `json:"session_id,omitempty"` // Cache TTL Override 标记 CacheTTLOverridden bool `json:"cache_ttl_overridden"` diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 1ea81d1fb5..3d014f9f45 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -535,6 +535,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, @@ -547,6 +548,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { UpstreamEndpoint: upstreamEndpoint, UserAgent: userAgent, IPAddress: clientIP, + SessionID: sessionID, RequestPayloadHash: requestPayloadHash, ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, @@ -970,6 +972,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey) + sessionID := service.ExtractClientSessionID(c) h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, @@ -982,6 +985,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { UpstreamEndpoint: upstreamEndpoint, UserAgent: userAgent, IPAddress: clientIP, + SessionID: sessionID, RequestPayloadHash: requestPayloadHash, ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index 3cfb9026e2..7061325f8a 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -298,6 +298,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, @@ -312,6 +313,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { reqLog.Error("gateway.cc.record_usage_failed", diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index bd545d2cf6..8f5d8531bb 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -276,6 +276,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: result, @@ -290,6 +291,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { reqLog.Error("gateway.responses.record_usage_failed", diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 7b349ae1db..40470ac418 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -532,6 +532,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { // ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。 forceCacheBilling := fs.ForceCacheBilling quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsageWithLongContext(ctx, &service.RecordUsageLongContextInput{ Result: result, @@ -549,6 +550,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { LongContextMultiplier: 2.0, // 超出部分双倍计费 ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { logger.L().With( diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 0522d689b1..fd0697884e 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -465,6 +465,7 @@ func recordGrokMediaUsage( ) { userAgent := c.GetHeader("User-Agent") clientIP := ip.GetClientIP(c) + sessionID := service.ExtractClientSessionID(c) payloadForHash := body if len(payloadForHash) == 0 && strings.TrimSpace(requestID) != "" { payloadForHash = []byte(requestID) @@ -493,6 +494,7 @@ func recordGrokMediaUsage( RequestPayloadHash: service.HashUsageRequestPayload(payloadForHash), APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, + SessionID: sessionID, ChannelUsageFields: channelUsageFields, }); err != nil { logger.L().With( diff --git a/backend/internal/handler/openai_alpha_search.go b/backend/internal/handler/openai_alpha_search.go index eca8931e87..c76b7d9fa5 100644 --- a/backend/internal/handler/openai_alpha_search.go +++ b/backend/internal/handler/openai_alpha_search.go @@ -233,6 +233,7 @@ func (h *OpenAIGatewayHandler) recordAlphaSearchUsage( ) { userAgent := c.GetHeader("User-Agent") clientIP := ip.GetClientIP(c) + sessionID := service.ExtractClientSessionID(c) requestPayloadHash := service.HashUsageRequestPayload(body) inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) @@ -252,6 +253,7 @@ func (h *OpenAIGatewayHandler) recordAlphaSearchUsage( RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, + SessionID: sessionID, ChannelUsageFields: channelMapping.ToUsageFields(requestedModel, result.UpstreamModel), }); err != nil { logger.L().With( diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 7477057113..39c12253b1 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -340,6 +340,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) cyberBlocked := service.GetOpsCyberPolicy(c) != nil h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { @@ -355,6 +356,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { IPAddress: clientIP, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index b7bf64ce0e..3ed44b12b8 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -243,6 +243,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ @@ -257,6 +258,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { IPAddress: clientIP, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { logger.L().With( diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 1e5949f36f..761f291e90 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -627,6 +627,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 cyberBlocked := service.GetOpsCyberPolicy(c) != nil @@ -644,6 +645,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { @@ -1138,6 +1140,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) cyberBlocked := service.GetOpsCyberPolicy(c) != nil h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { @@ -1154,6 +1157,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMappingMsg, reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { @@ -1862,6 +1866,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { inboundEndpoint := GetInboundEndpoint(c) upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) cyberBlocked := service.GetOpsCyberPolicy(c) != nil h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) { if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{ @@ -1877,6 +1882,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMappingWS, reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { @@ -2884,6 +2890,8 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey userAgent = c.GetHeader("User-Agent") clientIPStr = strings.TrimSpace(ip.GetClientIP(c)) } + // 提前拍成标量，避免在下方 goroutine 内访问 gin.Context。 + sessionID := service.ExtractClientSessionID(c) apiKeyPrefix := "" if apiKey != nil { apiKeyPrefix = keyPrefix(apiKey.Key, 8) @@ -2940,6 +2948,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey UpstreamEndpoint: upstreamEndpoint, UserAgent: userAgent, IPAddress: clientIPStr, + SessionID: sessionID, RequestPayloadHash: requestPayloadHash, APIKeyService: apiKeySvc, ChannelUsageFields: channelFields, diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 7ae070f6c6..72e2f462ad 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -91,6 +91,8 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { zap.Bool("stream", parsed.Stream), zap.Bool("multipart", parsed.Multipart), zap.String("capability", string(parsed.RequiredCapability)), + zap.String("img_quality", parsed.Quality), + zap.String("img_size", parsed.Size), ) if !service.GroupAllowsImageGeneration(apiKey.Group) { @@ -374,6 +376,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { if result != nil { upstreamModel = result.UpstreamModel } + sessionID := service.ExtractClientSessionID(c) h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ Result: result, @@ -388,6 +391,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, + SessionID: sessionID, ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, requestModel, upstreamModel), }); err != nil { logger.L().With( diff --git a/backend/internal/handler/openai_images_failover_test.go b/backend/internal/handler/openai_images_failover_test.go index e5101d0148..8b07cb9a52 100644 --- a/backend/internal/handler/openai_images_failover_test.go +++ b/backend/internal/handler/openai_images_failover_test.go @@ -12,11 +12,14 @@ import ( "testing" "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" ) type openAIImagesFailoverAccountRepo struct { @@ -154,8 +157,10 @@ func TestOpenAIGatewayHandlerImages_ServerErrorFailsOverAndReturnsClearErrorWhen ) handler.maxAccountSwitches = 10 - body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) - req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","quality":"high","size":"1536x1024"}`) + core, observedLogs := observer.New(zap.DebugLevel) + requestCtx := logger.IntoContext(context.Background(), zap.New(core)) + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)).WithContext(requestCtx) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) @@ -173,6 +178,16 @@ func TestOpenAIGatewayHandlerImages_ServerErrorFailsOverAndReturnsClearErrorWhen handler.Images(c) + accountSelectingLogs := observedLogs.FilterMessage("openai.images.account_selecting").All() + require.NotEmpty(t, accountSelectingLogs) + loggedFields := make(map[string]string) + for _, field := range accountSelectingLogs[0].Context { + loggedFields[field.Key] = field.String + } + require.Equal(t, "high", loggedFields["img_quality"]) + require.Equal(t, "1536x1024", loggedFields["img_size"]) + require.NotContains(t, loggedFields, "prompt") + require.Equal(t, []int64{1, 2}, upstream.calls()) require.Equal(t, http.StatusBadGateway, rec.Code) require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String()) diff --git a/backend/internal/repository/batch_image_repo.go b/backend/internal/repository/batch_image_repo.go index 6531f75fd5..7c8991c8d4 100644 --- a/backend/internal/repository/batch_image_repo.go +++ b/backend/internal/repository/batch_image_repo.go @@ -750,7 +750,7 @@ INSERT INTO batch_image_jobs ( batch_discount_multiplier, hold_multiplier, billable_unit_price, hold_unit_price, pricing_snapshot_version, currency, hold_id, - idempotency_key, request_hash, manifest_hash, retry_count, output_expires_at + idempotency_key, request_hash, manifest_hash, retry_count, session_id, output_expires_at ) VALUES ( $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, @@ -760,7 +760,7 @@ INSERT INTO batch_image_jobs ( $25, $26, $27, $28, $29, $30, $31, - $32, $33, $34, $35, $36 + $32, $33, $34, $35, $36, $37 ) RETURNING `+batchImageJobColumns, params.BatchID, params.UserID, params.APIKeyID, params.AccountID, params.Provider, params.Model, params.TaskName, params.ParentBatchID, params.Status, @@ -771,7 +771,7 @@ RETURNING `+batchImageJobColumns, params.BatchDiscountMultiplier, params.HoldMultiplier, params.BillableUnitPrice, params.HoldUnitPrice, params.PricingSnapshotVersion, params.Currency, params.HoldID, - params.IdempotencyKey, params.RequestHash, params.ManifestHash, params.RetryCount, params.OutputExpiresAt, + params.IdempotencyKey, params.RequestHash, params.ManifestHash, params.RetryCount, params.SessionID, params.OutputExpiresAt, )) } @@ -825,7 +825,7 @@ batch_discount_multiplier, hold_multiplier, billable_unit_price, hold_unit_price pricing_snapshot_version, currency, hold_id, idempotency_key, request_hash, manifest_hash, -retry_count, version, output_expires_at, input_deleted_at, output_deleted_at, downloaded_at, user_deleted_at, +retry_count, version, session_id, output_expires_at, input_deleted_at, output_deleted_at, downloaded_at, user_deleted_at, last_error_code, last_error_message, created_at, updated_at, submitted_at, started_at, finished_at, settled_at` @@ -838,6 +838,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) { var parentBatchID sql.NullString var holdAmount, actualCost sql.NullFloat64 var holdID, idempotencyKey, requestHash, manifestHash sql.NullString + var sessionID sql.NullString var outputExpiresAt, inputDeletedAt, outputDeletedAt, downloadedAt, userDeletedAt sql.NullTime var lastErrorCode, lastErrorMessage sql.NullString var submittedAt, startedAt, finishedAt, settledAt sql.NullTime @@ -852,7 +853,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) { &job.PricingSnapshotVersion, &job.Currency, &holdID, &idempotencyKey, &requestHash, &manifestHash, - &job.RetryCount, &job.Version, &outputExpiresAt, &inputDeletedAt, &outputDeletedAt, &downloadedAt, &userDeletedAt, + &job.RetryCount, &job.Version, &sessionID, &outputExpiresAt, &inputDeletedAt, &outputDeletedAt, &downloadedAt, &userDeletedAt, &lastErrorCode, &lastErrorMessage, &job.CreatedAt, &job.UpdatedAt, &submittedAt, &startedAt, &finishedAt, &settledAt, ) @@ -874,6 +875,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) { job.IdempotencyKey = batchImageNullStringPtr(idempotencyKey) job.RequestHash = batchImageNullStringPtr(requestHash) job.ManifestHash = batchImageNullStringPtr(manifestHash) + job.SessionID = batchImageNullStringPtr(sessionID) job.OutputExpiresAt = batchImageNullTimePtr(outputExpiresAt) job.InputDeletedAt = batchImageNullTimePtr(inputDeletedAt) job.OutputDeletedAt = batchImageNullTimePtr(outputDeletedAt) diff --git a/backend/internal/repository/usage_log_repo_insert.go b/backend/internal/repository/usage_log_repo_insert.go index 54d24b5957..809b0652b5 100644 --- a/backend/internal/repository/usage_log_repo_insert.go +++ b/backend/internal/repository/usage_log_repo_insert.go @@ -79,6 +79,7 @@ var usageLogInsertArgTypes = [...]string{ "text", // billing_tier "text", // billing_mode "numeric", // account_stats_cost + "text", // session_id "timestamptz", // created_at } @@ -274,6 +275,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, billing_tier, billing_mode, account_stats_cost, + session_id, created_at ) VALUES ( $1, $2, $3, $4, $5, $6, $7, @@ -281,7 +283,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57 ) ON CONFLICT (request_id, api_key_id) DO NOTHING RETURNING id, created_at @@ -728,10 +730,13 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage billing_tier, billing_mode, account_stats_cost, + session_id, created_at ) AS (VALUES `) - args := make([]any, 0, len(keys)*56) + // Each batch row prepends the synthetic input_index before the 57 + // usage-log column values. + args := make([]any, 0, len(keys)*58) argPos := 1 for idx, key := range keys { if idx > 0 { @@ -815,6 +820,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage billing_tier, billing_mode, account_stats_cost, + session_id, created_at ) SELECT @@ -873,6 +879,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage billing_tier, billing_mode, account_stats_cost, + session_id, created_at FROM input ON CONFLICT (request_id, api_key_id) DO NOTHING @@ -971,10 +978,11 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( billing_tier, billing_mode, account_stats_cost, + session_id, created_at ) AS (VALUES `) - args := make([]any, 0, len(preparedList)*56) + args := make([]any, 0, len(preparedList)*57) argPos := 1 for idx, prepared := range preparedList { if idx > 0 { @@ -1055,6 +1063,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( billing_tier, billing_mode, account_stats_cost, + session_id, created_at ) SELECT @@ -1113,6 +1122,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( billing_tier, billing_mode, account_stats_cost, + session_id, created_at FROM input ON CONFLICT (request_id, api_key_id) DO NOTHING @@ -1179,6 +1189,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared billing_tier, billing_mode, account_stats_cost, + session_id, created_at ) VALUES ( $1, $2, $3, $4, $5, $6, $7, @@ -1186,7 +1197,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57 ) ON CONFLICT (request_id, api_key_id) DO NOTHING `, prepared.args...) @@ -1227,6 +1238,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared { modelMappingChain := nullString(log.ModelMappingChain) billingTier := nullString(log.BillingTier) billingMode := nullString(log.BillingMode) + sessionID := nullString(log.SessionID) requestedModel := strings.TrimSpace(log.RequestedModel) if requestedModel == "" { requestedModel = strings.TrimSpace(log.Model) @@ -1299,6 +1311,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared { billingTier, billingMode, log.AccountStatsCost, // account_stats_cost + sessionID, // session_id createdAt, }, } diff --git a/backend/internal/repository/usage_log_repo_query.go b/backend/internal/repository/usage_log_repo_query.go index fb0e03ed60..88b7eeaf35 100644 --- a/backend/internal/repository/usage_log_repo_query.go +++ b/backend/internal/repository/usage_log_repo_query.go @@ -19,7 +19,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/service" ) -const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, image_input_tokens, image_input_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" +const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, image_input_tokens, image_input_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, session_id, created_at" func (r *usageLogRepository) GetByID(ctx context.Context, id int64) (log *service.UsageLog, err error) { query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE id = $1" @@ -481,6 +481,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e billingTier sql.NullString billingMode sql.NullString accountStatsCost sql.NullFloat64 + sessionID sql.NullString createdAt time.Time ) @@ -541,6 +542,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e &billingTier, &billingMode, &accountStatsCost, + &sessionID, &createdAt, ); err != nil { return nil, err @@ -661,6 +663,9 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e if accountStatsCost.Valid { log.AccountStatsCost = &accountStatsCost.Float64 } + if sessionID.Valid { + log.SessionID = &sessionID.String + } return log, nil } diff --git a/backend/internal/repository/usage_log_repo_request_type_test.go b/backend/internal/repository/usage_log_repo_request_type_test.go index f039029608..aff472b0cb 100644 --- a/backend/internal/repository/usage_log_repo_request_type_test.go +++ b/backend/internal/repository/usage_log_repo_request_type_test.go @@ -96,6 +96,7 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) { sqlmock.AnyArg(), // billing_tier sqlmock.AnyArg(), // billing_mode sqlmock.AnyArg(), // account_stats_cost + sqlmock.AnyArg(), // session_id createdAt, ). WillReturnRows(sqlmock.NewRows([]string{"id", "created_at"}).AddRow(int64(99), createdAt)) @@ -185,6 +186,7 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) { sqlmock.AnyArg(), // billing_tier sqlmock.AnyArg(), // billing_mode sqlmock.AnyArg(), // account_stats_cost + sqlmock.AnyArg(), // session_id createdAt, ). WillReturnRows(sqlmock.NewRows([]string{"id", "created_at"}).AddRow(int64(100), createdAt)) @@ -826,6 +828,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, sql.NullString{}, sql.NullFloat64{}, + sql.NullString{}, now, }}) require.NoError(t, err) @@ -900,6 +903,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, // billing_tier sql.NullString{}, // billing_mode sql.NullFloat64{}, // account_stats_cost + sql.NullString{}, // session_id now, }}) require.NoError(t, err) @@ -957,6 +961,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, // billing_tier sql.NullString{}, // billing_mode sql.NullFloat64{}, // account_stats_cost + sql.NullString{}, // session_id now, }}) require.NoError(t, err) @@ -1014,6 +1019,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullString{}, // billing_tier sql.NullString{}, // billing_mode sql.NullFloat64{}, // account_stats_cost + sql.NullString{}, // session_id now, }}) require.NoError(t, err) diff --git a/backend/internal/repository/usage_log_session_id_integration_test.go b/backend/internal/repository/usage_log_session_id_integration_test.go new file mode 100644 index 0000000000..76778b4162 --- /dev/null +++ b/backend/internal/repository/usage_log_session_id_integration_test.go @@ -0,0 +1,71 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +// TestUsageLog_SessionIDPersistence proves session_id round-trips from insert to +// read and is omitted (NULL) when absent. +func TestUsageLog_SessionIDPersistence(t *testing.T) { + ctx := context.Background() + client := testEntClient(t) + repo := newUsageLogRepositoryWithSQL(client, integrationDB) + + user := mustCreateUser(t, client, &service.User{Email: "session-id-" + uuid.NewString() + "@example.com"}) + apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-session-" + uuid.NewString(), Name: "k"}) + account := mustCreateAccount(t, client, &service.Account{Name: "acc-session-" + uuid.NewString()}) + + sessionID := "sess-" + uuid.NewString() + + withSession := &service.UsageLog{ + UserID: user.ID, + APIKeyID: apiKey.ID, + AccountID: account.ID, + RequestID: uuid.NewString(), + Model: "claude-3", + InputTokens: 10, + OutputTokens: 5, + TotalCost: 1.0, + ActualCost: 1.0, + SessionID: &sessionID, + CreatedAt: time.Now().UTC(), + } + _, err := repo.Create(ctx, withSession) + require.NoError(t, err) + require.NotZero(t, withSession.ID) + + withoutSession := &service.UsageLog{ + UserID: user.ID, + APIKeyID: apiKey.ID, + AccountID: account.ID, + RequestID: uuid.NewString(), + Model: "claude-3", + InputTokens: 7, + OutputTokens: 3, + TotalCost: 0.5, + ActualCost: 0.5, + CreatedAt: time.Now().UTC(), + } + _, err = repo.Create(ctx, withoutSession) + require.NoError(t, err) + + // Round-trip: session id survives insert → read. + got, err := repo.GetByID(ctx, withSession.ID) + require.NoError(t, err) + require.NotNil(t, got.SessionID) + require.Equal(t, sessionID, *got.SessionID) + + // Omission: absent session id reads back as nil (NULL), not empty string. + gotNone, err := repo.GetByID(ctx, withoutSession.ID) + require.NoError(t, err) + require.Nil(t, gotNone.SessionID) +} diff --git a/backend/internal/repository/usage_log_session_id_unit_test.go b/backend/internal/repository/usage_log_session_id_unit_test.go new file mode 100644 index 0000000000..7804dce306 --- /dev/null +++ b/backend/internal/repository/usage_log_session_id_unit_test.go @@ -0,0 +1,91 @@ +//go:build unit + +package repository + +import ( + "database/sql" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func newSessionIDUsageLog(sessionID *string) *service.UsageLog { + return &service.UsageLog{ + UserID: 1, + APIKeyID: 2, + AccountID: 3, + RequestID: "req-session-id", + Model: "claude-3", + InputTokens: 10, + OutputTokens: 5, + TotalCost: 1.0, + ActualCost: 1.0, + SessionID: sessionID, + CreatedAt: time.Now().UTC(), + } +} + +// TestPrepareUsageLogInsert_SessionIDArgWiring pins the session_id column to the +// arg slice / arg-type table so the five INSERT column lists stay in sync. session_id +// is the penultimate arg (created_at is always last). +func TestPrepareUsageLogInsert_SessionIDArgWiring(t *testing.T) { + require.Len(t, usageLogInsertArgTypes, 57, "arg-type table must include session_id") + + sessionID := "sess-persisted-123" + prepared := prepareUsageLogInsert(newSessionIDUsageLog(&sessionID)) + + require.Len(t, prepared.args, len(usageLogInsertArgTypes), + "prepared args must match the arg-type table length") + + // created_at is last; session_id is the arg immediately before it. + sessionArg := prepared.args[len(prepared.args)-2] + ns, ok := sessionArg.(sql.NullString) + require.True(t, ok, "session_id arg should be a sql.NullString, got %T", sessionArg) + require.True(t, ns.Valid) + require.Equal(t, sessionID, ns.String) + + require.Equal(t, "text", usageLogInsertArgTypes[len(usageLogInsertArgTypes)-2], + "session_id arg type must be text") +} + +// TestPrepareUsageLogInsert_SessionIDNullWhenAbsent proves an absent session id is +// persisted as SQL NULL rather than an empty string. +func TestPrepareUsageLogInsert_SessionIDNullWhenAbsent(t *testing.T) { + prepared := prepareUsageLogInsert(newSessionIDUsageLog(nil)) + sessionArg := prepared.args[len(prepared.args)-2] + ns, ok := sessionArg.(sql.NullString) + require.True(t, ok, "session_id arg should be a sql.NullString, got %T", sessionArg) + require.False(t, ns.Valid, "absent session id must be NULL, not empty string") + + empty := "" + preparedEmpty := prepareUsageLogInsert(newSessionIDUsageLog(&empty)) + nsEmpty := preparedEmpty.args[len(preparedEmpty.args)-2].(sql.NullString) + require.False(t, nsEmpty.Valid, "empty session id must also be NULL") +} + +// TestUsageLogInsertQueries_IncludeSessionID guards that every generated INSERT path +// and the SELECT column list reference session_id. +func TestUsageLogInsertQueries_IncludeSessionID(t *testing.T) { + require.Contains(t, usageLogSelectColumns, "session_id", + "SELECT column list must include session_id") + + sessionID := "sess-in-query" + log := newSessionIDUsageLog(&sessionID) + prepared := prepareUsageLogInsert(log) + key := usageLogBatchKey(log.RequestID, log.APIKeyID) + + batchQuery, batchArgs := buildUsageLogBatchInsertQuery([]string{key}, + map[string]usageLogInsertPrepared{key: prepared}) + require.Contains(t, batchQuery, "session_id") + // Two column references (INSERT column list + SELECT ... FROM input) plus the CTE def. + require.GreaterOrEqual(t, strings.Count(batchQuery, "session_id"), 3) + require.Len(t, batchArgs, len(prepared.args)+1, + "batch args include the synthetic input_index before usage-log values") + + bestEffortQuery, bestEffortArgs := buildUsageLogBestEffortInsertQuery([]usageLogInsertPrepared{prepared}) + require.Contains(t, bestEffortQuery, "session_id") + require.Len(t, bestEffortArgs, len(prepared.args)) +} diff --git a/backend/internal/service/batch_image.go b/backend/internal/service/batch_image.go index af0206b8b0..130428f18e 100644 --- a/backend/internal/service/batch_image.go +++ b/backend/internal/service/batch_image.go @@ -137,6 +137,7 @@ type BatchImageJob struct { IdempotencyKey *string RequestHash *string ManifestHash *string + SessionID *string RetryCount int Version int @@ -196,6 +197,7 @@ type CreateBatchImageJobParams struct { IdempotencyKey *string RequestHash *string ManifestHash *string + SessionID *string RetryCount int diff --git a/backend/internal/service/batch_image_processor_test.go b/backend/internal/service/batch_image_processor_test.go index f54a1a99ec..ca3a583c99 100644 --- a/backend/internal/service/batch_image_processor_test.go +++ b/backend/internal/service/batch_image_processor_test.go @@ -425,6 +425,7 @@ func (r *fakeBatchImageRepository) CreateBatchImageJob(_ context.Context, params Currency: params.Currency, IdempotencyKey: params.IdempotencyKey, RequestHash: params.RequestHash, + SessionID: params.SessionID, CreatedAt: time.Now(), } r.jobs[job.BatchID] = job diff --git a/backend/internal/service/batch_image_public.go b/backend/internal/service/batch_image_public.go index 2960d4cd9f..580d1f4934 100644 --- a/backend/internal/service/batch_image_public.go +++ b/backend/internal/service/batch_image_public.go @@ -57,6 +57,7 @@ type BatchImageSubmitRequest struct { AspectRatio string `json:"aspect_ratio"` ImageSize string `json:"image_size"` Metadata map[string]string `json:"metadata"` + SessionID *string `json:"-"` } type BatchImageSubmitItem struct { @@ -281,6 +282,7 @@ func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOw HoldID: &holdID, IdempotencyKey: batchImageOptionalStringPtr(idempotencyKey), RequestHash: batchImageStringPtr(requestHash), + SessionID: normalized.SessionID, }) if err != nil { return nil, err diff --git a/backend/internal/service/batch_image_public_test.go b/backend/internal/service/batch_image_public_test.go index 2e5b930a7b..a69468f846 100644 --- a/backend/internal/service/batch_image_public_test.go +++ b/backend/internal/service/batch_image_public_test.go @@ -27,8 +27,10 @@ func TestBatchImagePublicService_Submit(t *testing.T) { t.Run("accepts valid request stores refs and enqueues once", func(t *testing.T) { svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) + req := validBatchImageSubmitRequest() + req.SessionID = batchImageStringPtr("batch-session-123") - got, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "") + got, err := svc.Submit(ctx, testBatchImageOwner(), req, "") require.NoError(t, err) require.Equal(t, "image.batch", got.Object) require.Equal(t, "queued", got.Status) @@ -61,6 +63,7 @@ func TestBatchImagePublicService_Submit(t *testing.T) { require.InDelta(t, 0.6, job.HoldMultiplier, 1e-12) require.InDelta(t, 0.125, job.BillableUnitPrice, 1e-12) require.InDelta(t, 0.15, job.HoldUnitPrice, 1e-12) + require.Equal(t, "batch-session-123", batchImageDerefString(job.SessionID)) }) t.Run("combines user group image rate account rate discount and hold margin", func(t *testing.T) { @@ -389,15 +392,18 @@ func TestBatchImagePublicService_Submit(t *testing.T) { }) t.Run("idempotency returns same batch without provider resubmit", func(t *testing.T) { - svc, _, queue, gemini, _ := newTestBatchImagePublicService(true) + svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true) req := validBatchImageSubmitRequest() + req.SessionID = batchImageStringPtr("original-session") first, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key") require.NoError(t, err) + req.SessionID = batchImageStringPtr("retry-session") second, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key") require.NoError(t, err) require.Equal(t, first.ID, second.ID) + require.Equal(t, "original-session", batchImageDerefString(repo.jobs[first.ID].SessionID)) require.Len(t, gemini.submits, 1) require.Equal(t, []string{first.ID}, queue.enqueued) }) diff --git a/backend/internal/service/batch_image_settlement.go b/backend/internal/service/batch_image_settlement.go index 26d2abb0fd..e9b8a3685e 100644 --- a/backend/internal/service/batch_image_settlement.go +++ b/backend/internal/service/batch_image_settlement.go @@ -277,6 +277,7 @@ func (s *BatchImageSettlementService) recordUsageLog(ctx context.Context, job *B RequestType: RequestTypeSync, BillingMode: &billingMode, ImageSize: &imageSize, + SessionID: job.SessionID, CreatedAt: createdAt, } writeUsageLogBestEffort(ctx, s.UsageLogRepo, usageLog, "service.batch_image_settlement") diff --git a/backend/internal/service/batch_image_settlement_test.go b/backend/internal/service/batch_image_settlement_test.go index a1fe3e899d..f1fc3a15eb 100644 --- a/backend/internal/service/batch_image_settlement_test.go +++ b/backend/internal/service/batch_image_settlement_test.go @@ -18,9 +18,14 @@ func TestBatchImageSettlementService_SettlesAndChargesSuccessfulImagesOnly(t *te job.SuccessCount = 3 job.FailCount = 2 job.ItemCount = 5 + job.SessionID = batchImageStringPtr("batch-settlement-session") repo.jobs[job.BatchID] = job billing := &fakeBatchImageBillingRepo{} - svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}} + usageLogs := &openAIRecordUsageLogRepoStub{} + svc := &BatchImageSettlementService{ + Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}, + UsageLogRepo: usageLogs, + } result, err := svc.Settle(context.Background(), job.BatchID) require.NoError(t, err) @@ -32,6 +37,7 @@ func TestBatchImageSettlementService_SettlesAndChargesSuccessfulImagesOnly(t *te require.Equal(t, 0.75, *repo.jobs[job.BatchID].ActualCost) require.NotEmpty(t, batchImageDerefString(repo.jobs[job.BatchID].ManifestHash)) require.NotNil(t, repo.jobs[job.BatchID].SettledAt) + require.Equal(t, "batch-settlement-session", batchImageDerefString(usageLogs.lastLog.SessionID)) require.Len(t, billing.captures, 1) require.Equal(t, int64(321), billing.captures[0].APIKeyID) require.Equal(t, job.UserID, billing.captures[0].UserID) diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 1606f0f3c1..f445814d01 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -46,6 +46,7 @@ type RecordUsageInput struct { UpstreamEndpoint string // 上游端点（标准化后的上游路径） UserAgent string // 请求的 User-Agent IPAddress string // 请求的客户端 IP 地址 + SessionID string // 客户端显式会话标识（session_id / X-Session-Id 等请求头），仅用于用量行会话关联 RequestPayloadHash string // 请求体语义哈希，用于降低 request_id 误复用时的静默误去重风险 ForceCacheBilling bool // 强制缓存计费：将 input_tokens 转为 cache_read 计费（用于粘性会话切换） APIKeyService APIKeyQuotaUpdater // 可选：用于更新API Key配额 @@ -572,6 +573,7 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu UpstreamEndpoint: input.UpstreamEndpoint, UserAgent: input.UserAgent, IPAddress: input.IPAddress, + SessionID: input.SessionID, RequestPayloadHash: input.RequestPayloadHash, ForceCacheBilling: input.ForceCacheBilling, APIKeyService: input.APIKeyService, @@ -591,6 +593,7 @@ type RecordUsageLongContextInput struct { UpstreamEndpoint string // 上游端点（标准化后的上游路径） UserAgent string // 请求的 User-Agent IPAddress string // 请求的客户端 IP 地址 + SessionID string // 客户端显式会话标识（session_id / X-Session-Id 等请求头），仅用于用量行会话关联 RequestPayloadHash string // 请求体语义哈希，用于降低 request_id 误复用时的静默误去重风险 LongContextThreshold int // 长上下文阈值（如 200000） LongContextMultiplier float64 // 超出阈值部分的倍率（如 2.0） @@ -613,6 +616,7 @@ func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input * UpstreamEndpoint: input.UpstreamEndpoint, UserAgent: input.UserAgent, IPAddress: input.IPAddress, + SessionID: input.SessionID, RequestPayloadHash: input.RequestPayloadHash, ForceCacheBilling: input.ForceCacheBilling, APIKeyService: input.APIKeyService, @@ -635,6 +639,7 @@ type recordUsageCoreInput struct { UpstreamEndpoint string UserAgent string IPAddress string + SessionID string RequestPayloadHash string ForceCacheBilling bool APIKeyService APIKeyQuotaUpdater @@ -1010,6 +1015,7 @@ func (s *GatewayService) buildRecordUsageLog( ModelMappingChain: optionalTrimmedStringPtr(input.ModelMappingChain), UserAgent: optionalTrimmedStringPtr(input.UserAgent), IPAddress: optionalTrimmedStringPtr(input.IPAddress), + SessionID: optionalTrimmedStringPtr(input.SessionID), GroupID: apiKey.GroupID, SubscriptionID: optionalSubscriptionID(subscription), CreatedAt: time.Now(), diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index de996649b8..5f868a5d66 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -473,7 +473,7 @@ func geminiResponseToChatCompletions( rawData []byte, usageOverride *ClaudeUsage, ) (*apicompat.ChatCompletionsResponse, *ClaudeUsage, error) { - claudeRespMap, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, rawData) + claudeRespMap, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, rawData, true) if usageOverride != nil && (usageOverride.InputTokens > 0 || usageOverride.OutputTokens > 0 || usageOverride.CacheReadInputTokens > 0) { usage = usageOverride if usageMap, ok := claudeRespMap["usage"].(map[string]any); ok { diff --git a/backend/internal/service/gemini_chat_completions_compat_service_test.go b/backend/internal/service/gemini_chat_completions_compat_service_test.go new file mode 100644 index 0000000000..e4ebf7a838 --- /dev/null +++ b/backend/internal/service/gemini_chat_completions_compat_service_test.go @@ -0,0 +1,163 @@ +package service + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGeminiResponseToChatCompletionsPreservesInlineData(t *testing.T) { + tests := []struct { + name string + parts []any + want string + }{ + { + name: "image only", + parts: []any{ + map[string]any{"inlineData": map[string]any{"mimeType": "image/png", "data": "aW1hZ2U="}}, + }, + want: "![image](data:image/png;base64,aW1hZ2U=)", + }, + { + name: "text and image", + parts: []any{ + map[string]any{"text": "rendered image:\n"}, + map[string]any{"inlineData": map[string]any{"mimeType": "image/webp", "data": "d2VicA=="}}, + }, + want: "rendered image:\n![image](data:image/webp;base64,d2VicA==)", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + geminiResp := map[string]any{ + "candidates": []any{map[string]any{ + "content": map[string]any{"parts": tt.parts}, + "finishReason": "STOP", + }}, + } + rawData, err := json.Marshal(geminiResp) + require.NoError(t, err) + + got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil) + require.NoError(t, err) + require.Len(t, got.Choices, 1) + + var content string + require.NoError(t, json.Unmarshal(got.Choices[0].Message.Content, &content)) + require.Equal(t, tt.want, content) + require.Equal(t, "stop", got.Choices[0].FinishReason) + }) + } +} + +func TestGeminiResponseToChatCompletionsOmitsInvalidInlineData(t *testing.T) { + tests := []struct { + name string + inlineData map[string]any + }{ + { + name: "unsupported MIME type", + inlineData: map[string]any{"mimeType": "image/svg+xml", "data": "PHN2Zz48L3N2Zz4="}, + }, + { + name: "malformed MIME type", + inlineData: map[string]any{"mimeType": "image/png; charset=utf-8", "data": "aW1hZ2U="}, + }, + { + name: "malformed base64", + inlineData: map[string]any{"mimeType": "image/png", "data": "not-valid-base64!!!"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + geminiResp := map[string]any{ + "candidates": []any{map[string]any{ + "content": map[string]any{"parts": []any{map[string]any{"text": "before"}, map[string]any{"inlineData": tt.inlineData}, map[string]any{"text": "after"}}}, + "finishReason": "STOP", + }}, + } + rawData, err := json.Marshal(geminiResp) + require.NoError(t, err) + + got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil) + require.NoError(t, err) + + var content string + require.NoError(t, json.Unmarshal(got.Choices[0].Message.Content, &content)) + require.Equal(t, "beforeafter", content) + }) + } +} + +func TestConvertGeminiToClaudeMessageOmitsInlineDataForAnthropicMessages(t *testing.T) { + geminiResp := map[string]any{ + "candidates": []any{map[string]any{ + "content": map[string]any{"parts": []any{ + map[string]any{"text": "before"}, + map[string]any{"inlineData": map[string]any{"mimeType": "image/png", "data": "aW1hZ2U="}}, + map[string]any{"functionCall": map[string]any{"name": "get_weather", "args": map[string]any{"city": "Paris"}}}, + map[string]any{"text": "after"}, + }}, + "finishReason": "STOP", + }}, + } + rawData, err := json.Marshal(geminiResp) + require.NoError(t, err) + + withInlineData, _ := convertGeminiToClaudeMessage(geminiResp, "gemini-test", rawData, true) + contentWithInlineData, ok := withInlineData["content"].([]any) + require.True(t, ok) + require.Len(t, contentWithInlineData, 4) + require.Equal(t, map[string]any{"type": "text", "text": "before"}, contentWithInlineData[0]) + require.Equal(t, map[string]any{"type": "text", "text": "![image](data:image/png;base64,aW1hZ2U=)"}, contentWithInlineData[1]) + toolUse, ok := contentWithInlineData[2].(map[string]any) + require.True(t, ok) + require.Equal(t, "tool_use", toolUse["type"]) + require.Equal(t, "get_weather", toolUse["name"]) + require.Equal(t, map[string]any{"type": "text", "text": "after"}, contentWithInlineData[3]) + + withoutInlineData, _ := convertGeminiToClaudeMessage(geminiResp, "gemini-test", rawData, false) + contentWithoutInlineData, ok := withoutInlineData["content"].([]any) + require.True(t, ok) + require.Len(t, contentWithoutInlineData, 3) + require.Equal(t, map[string]any{"type": "text", "text": "before"}, contentWithoutInlineData[0]) + toolUseWithoutInlineData, ok := contentWithoutInlineData[1].(map[string]any) + require.True(t, ok) + require.Equal(t, "tool_use", toolUseWithoutInlineData["type"]) + require.Equal(t, "get_weather", toolUseWithoutInlineData["name"]) + require.Equal(t, map[string]any{"type": "text", "text": "after"}, contentWithoutInlineData[2]) +} + +func TestGeminiResponseToChatCompletionsRetainsTextAndToolBehavior(t *testing.T) { + geminiResp := map[string]any{ + "candidates": []any{map[string]any{ + "content": map[string]any{"parts": []any{ + map[string]any{"text": "checking"}, + map[string]any{"functionCall": map[string]any{ + "name": "get_weather", + "args": map[string]any{"city": "Paris"}, + }}, + }}, + "finishReason": "STOP", + }}, + } + rawData, err := json.Marshal(geminiResp) + require.NoError(t, err) + + got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil) + require.NoError(t, err) + require.Len(t, got.Choices, 1) + + choice := got.Choices[0] + var content string + require.NoError(t, json.Unmarshal(choice.Message.Content, &content)) + require.Equal(t, "checking", content) + require.Equal(t, "tool_calls", choice.FinishReason) + require.Len(t, choice.Message.ToolCalls, 1) + require.Equal(t, "get_weather", choice.Message.ToolCalls[0].Function.Name) + require.JSONEq(t, `{"city":"Paris"}`, choice.Message.ToolCalls[0].Function.Arguments) +} diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 2aaba767e9..1a3821847d 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -5,6 +5,7 @@ import ( "bytes" "context" "crypto/rand" + "encoding/base64" "encoding/hex" "encoding/json" "errors" @@ -1068,7 +1069,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to read upstream stream") } collectedBytes, _ := json.Marshal(collected) - claudeResp, usageObj2 := convertGeminiToClaudeMessage(collected, originalModel, collectedBytes) + claudeResp, usageObj2 := convertGeminiToClaudeMessage(collected, originalModel, collectedBytes, false) c.JSON(http.StatusOK, claudeResp) usage = usageObj2 if usageObj != nil && (usageObj.InputTokens > 0 || usageObj.OutputTokens > 0) { @@ -1965,7 +1966,7 @@ func (s *GeminiMessagesCompatService) handleNonStreamingResponse(c *gin.Context, return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response") } - claudeResp, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, unwrappedBody) + claudeResp, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, unwrappedBody, false) c.JSON(http.StatusOK, claudeResp) return usage, nil @@ -2717,7 +2718,7 @@ func unwrapGeminiResponse(raw []byte) ([]byte, error) { return raw, nil } -func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel string, rawData []byte) (map[string]any, *ClaudeUsage) { +func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel string, rawData []byte, includeInlineData bool) (map[string]any, *ClaudeUsage) { usage := extractGeminiUsage(rawData) if usage == nil { usage = &ClaudeUsage{} @@ -2740,6 +2741,16 @@ func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel strin "text": text, }) } + if inlineData, ok := pm["inlineData"].(map[string]any); includeInlineData && ok { + mimeType, _ := inlineData["mimeType"].(string) + data, _ := inlineData["data"].(string) + if isGeminiInlineImageMIMEType(mimeType) && isValidBase64(data) { + contentBlocks = append(contentBlocks, map[string]any{ + "type": "text", + "text": fmt.Sprintf("![image](data:%s;base64,%s)", mimeType, data), + }) + } + } if fc, ok := pm["functionCall"].(map[string]any); ok { name, _ := fc["name"].(string) if strings.TrimSpace(name) == "" { @@ -2782,6 +2793,23 @@ func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel strin return resp, usage } +func isGeminiInlineImageMIMEType(mimeType string) bool { + switch mimeType { + case "image/gif", "image/jpeg", "image/png", "image/webp": + return true + default: + return false + } +} + +func isValidBase64(data string) bool { + if data == "" { + return false + } + _, err := base64.StdEncoding.DecodeString(data) + return err == nil +} + func extractGeminiUsage(data []byte) *ClaudeUsage { usage := gjson.GetBytes(data, "usageMetadata") if !usage.Exists() { diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 991a1e5170..c61f317859 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -91,7 +91,12 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont if shouldDisable && !modelTempMatched { s.BlockAccountScheduling(account, time.Time{}, "upstream_disable") } - if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey && shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) { + // Pool-mode retryable upstream errors are already bounded by the request-local + // same-account retry budget. Recording the generic account+model transient + // cooldown here would block the next approved retry before that budget is used. + poolModeRetryable := account.IsPoolMode() && account.IsPoolModeRetryableStatus(statusCode) + if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey && + shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) && !poolModeRetryable { model := "" if len(canonicalModel) > 0 { model = canonicalModel[0] diff --git a/backend/internal/service/openai_account_runtime_block_fastpath_test.go b/backend/internal/service/openai_account_runtime_block_fastpath_test.go index ae4bfa91dc..09d4662468 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath_test.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath_test.go @@ -142,6 +142,89 @@ func TestOpenAIPoolModeTempRule_StopsSameAccountRetryAndIsolatesBlockToModel(t * require.False(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5")) } +func TestOpenAIPoolModeRetryable5xx_DoesNotCreateModelTransientBlock(t *testing.T) { + repo := &errorPolicyRepoStub{} + rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + gateway := &OpenAIGatewayService{rateLimitService: rateLimitService} + account := &Account{ + ID: 47, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "pool_mode": true, + "pool_mode_retry_status_codes": []any{float64(524)}, + }, + } + + for i := 0; i < 2; i++ { + shouldDisable := gateway.handleOpenAIAccountUpstreamError( + context.Background(), + account, + 524, + http.Header{}, + []byte(`{"error":{"message":"upstream timeout"}}`), + "gpt-5.4", + ) + require.False(t, shouldDisable) + } + + require.False(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4")) +} + +func TestOpenAIPoolModeNonRetryable5xx_StillCreatesModelTransientBlock(t *testing.T) { + repo := &errorPolicyRepoStub{} + rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + gateway := &OpenAIGatewayService{rateLimitService: rateLimitService} + account := &Account{ + ID: 48, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "pool_mode": true, + "pool_mode_retry_status_codes": []any{float64(http.StatusGatewayTimeout)}, + }, + } + + for i := 0; i < 2; i++ { + shouldDisable := gateway.handleOpenAIAccountUpstreamError( + context.Background(), + account, + http.StatusServiceUnavailable, + http.Header{}, + []byte(`{"error":{"message":"upstream unavailable"}}`), + "gpt-5.4", + ) + require.False(t, shouldDisable) + } + + require.True(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4")) +} + +func TestOpenAINonPoolAPIKey5xx_StillCreatesModelTransientBlock(t *testing.T) { + repo := &errorPolicyRepoStub{} + rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + gateway := &OpenAIGatewayService{rateLimitService: rateLimitService} + account := &Account{ + ID: 49, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + } + + for i := 0; i < 2; i++ { + shouldDisable := gateway.handleOpenAIAccountUpstreamError( + context.Background(), + account, + http.StatusGatewayTimeout, + http.Header{}, + []byte(`{"error":{"message":"upstream timeout"}}`), + "gpt-5.4", + ) + require.False(t, shouldDisable) + } + + require.True(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4")) +} + func TestOpenAIModelNotFound_DoesNotRuntimeBlockWholeAccount(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &OpenAIGatewayService{ diff --git a/backend/internal/service/openai_codex_transform.go b/backend/internal/service/openai_codex_transform.go index 2d6766af4f..c89fd038d0 100644 --- a/backend/internal/service/openai_codex_transform.go +++ b/backend/internal/service/openai_codex_transform.go @@ -1494,25 +1494,9 @@ func filterCodexInputWithOptions(input []any, opts codexInputFilterOptions) []an if !opts.PreserveReferences { ensureCopy() delete(newItem, "id") - } else if isCodexToolCallInputType(typ) { - // 续链模式下保留 id 以维持上下文引用，但 function_call 等 - // call-input 类 item 的 id 必须以 "fc" 开头（上游校验 - // "Expected an ID that begins with 'fc'"）。item_* 形式的 id - // 来自客户端回放，需要删除。 - // 注意：function_call_output 等 output 类的 id 无此约束，不动。 - if id, ok := m["id"].(string); ok && id != "" && !strings.HasPrefix(id, "fc") { - ensureCopy() - delete(newItem, "id") - } - } else if typ == "message" { - // 同理，message 类 item 的 id 必须以 "msg" 开头（上游校验 - // "Expected an ID that begins with 'msg'"）。item_* 形式的 id - // 来自客户端回放，需要删除。 - // 注意：不改写成 msg_*，改写出的 id 未必对应真实的上游对象。 - if id, ok := m["id"].(string); ok && id != "" && !strings.HasPrefix(id, "msg") { - ensureCopy() - delete(newItem, "id") - } + } else if id, ok := m["id"].(string); ok && shouldStripOpenAIResponsesInputItemID(typ, id) { + ensureCopy() + delete(newItem, "id") } filtered = append(filtered, newItem) diff --git a/backend/internal/service/openai_gateway_apikey_item_id_test.go b/backend/internal/service/openai_gateway_apikey_item_id_test.go new file mode 100644 index 0000000000..b7af70ec75 --- /dev/null +++ b/backend/internal/service/openai_gateway_apikey_item_id_test.go @@ -0,0 +1,90 @@ +//go:build unit + +package service + +import ( + "context" + "fmt" + "io" + "net/http" + "runtime" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestOpenAIGatewayService_APIKeyPassthrough_StripsInvalidInputItemIDs(t *testing.T) { + gin.SetMode(gin.TestMode) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"id":"resp_test","model":"gpt-5.6-sol","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`, + )), + }} + svc := newOpenAIImageGenerationControlTestService(upstream) + c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.144.1") + account := newOpenAIImageGenerationControlTestAccount() + account.Extra = map[string]any{"openai_passthrough": true} + + body := []byte(`{ + "model":"gpt-5.6-sol", + "stream":false, + "input":[ + {"type":"message","id":"item_bad_message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}, + {"type":"function_call","id":"item_bad_call","call_id":"call_123","name":"exec_command","arguments":"{}"}, + {"type":"message","id":"msg_valid","role":"user","content":[{"type":"input_text","text":"continue"}]}, + {"type":"function_call","id":"fc_valid","call_id":"call_456","name":"apply_patch","arguments":"{}"}, + {"type":"function_call_output","id":"item_output","call_id":"call_123","output":"done"}, + {"type":"web_search_call","id":"item_unconstrained"} + ] + }`) + + result, err := svc.Forward(context.Background(), c, account, body) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, upstream.lastReq) + + forwarded := upstream.lastBody + require.False(t, gjson.GetBytes(forwarded, "input.0.id").Exists()) + require.Equal(t, "hello", gjson.GetBytes(forwarded, "input.0.content.0.text").String()) + require.False(t, gjson.GetBytes(forwarded, "input.1.id").Exists()) + require.Equal(t, "call_123", gjson.GetBytes(forwarded, "input.1.call_id").String()) + require.Equal(t, "exec_command", gjson.GetBytes(forwarded, "input.1.name").String()) + require.Equal(t, "{}", gjson.GetBytes(forwarded, "input.1.arguments").String()) + require.Equal(t, "msg_valid", gjson.GetBytes(forwarded, "input.2.id").String()) + require.Equal(t, "fc_valid", gjson.GetBytes(forwarded, "input.3.id").String()) + require.Equal(t, "item_output", gjson.GetBytes(forwarded, "input.4.id").String()) + require.Equal(t, "call_123", gjson.GetBytes(forwarded, "input.4.call_id").String()) + require.Equal(t, "item_unconstrained", gjson.GetBytes(forwarded, "input.5.id").String()) +} + +func TestSanitizeOpenAIResponsesInputItemIDs_AllocationGrowthIsLinear(t *testing.T) { + makeBody := func(itemCount int) []byte { + items := make([]string, itemCount) + for i := range items { + items[i] = fmt.Sprintf(`{"type":"message","id":"item_%d","role":"user","content":[{"type":"input_text","text":"hello"}]}`, i) + } + return []byte(`{"model":"gpt-5.6-sol","input":[` + strings.Join(items, ",") + `]}`) + } + allocatedBytes := func(body []byte) uint64 { + runtime.GC() + var before, after runtime.MemStats + runtime.ReadMemStats(&before) + sanitized, changed, err := sanitizeOpenAIResponsesInputItemIDs(body) + runtime.ReadMemStats(&after) + require.NoError(t, err) + require.True(t, changed) + require.NotEmpty(t, sanitized) + return after.TotalAlloc - before.TotalAlloc + } + + smallAllocated := allocatedBytes(makeBody(20)) + largeAllocated := allocatedBytes(makeBody(200)) + require.Less(t, largeAllocated, smallAllocated*30, + "10x more input items must not cause quadratic whole-body allocation growth") +} diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 7a42a445ef..b3963535b1 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -72,6 +72,16 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco return nil, err } } + if shouldStripOpenAIResponsesInputNamespaces(account, wsDecision.Transport, passthroughEnabled) { + body, err = stripOpenAIResponsesInputNamespaces(body) + if err != nil { + setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{ + "type": "invalid_request_error", "message": err.Error(), "param": "input", + }}) + return nil, err + } + } originalBody := body requestView := newOpenAIRequestView(body) @@ -85,6 +95,19 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco if account.Type == AccountTypeAPIKey && !openai_compat.ShouldUseResponsesAPI(account.Extra) { return s.forwardResponsesViaRawChatCompletions(ctx, c, account, body) } + if account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey { + sanitizedBody, changed, sanitizeErr := sanitizeOpenAIResponsesInputItemIDs(body) + if sanitizeErr != nil { + return nil, fmt.Errorf("sanitize OpenAI Responses input item IDs: %w", sanitizeErr) + } + if changed { + body = sanitizedBody + originalBody = sanitizedBody + requestView = newOpenAIRequestView(sanitizedBody) + reqModel, reqStream, promptCacheKey = requestView.Model, requestView.Stream, requestView.PromptCacheKey + originalModel = reqModel + } + } compatMessagesBridge := isOpenAICompatMessagesBridgeBody(body) setOpenAICompatMessagesBridgeContext(c, compatMessagesBridge) diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index d887ff73fc..fdfd0bf321 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -663,7 +663,13 @@ var grokResponsesSupportedToolTypes = map[string]struct{}{ func sanitizeGrokResponsesTools(body []byte) ([]byte, error) { tools := gjson.GetBytes(body, "tools") - if !tools.Exists() || !tools.IsArray() { + if !tools.Exists() { + if gjson.GetBytes(body, "tool_choice").Exists() { + return sjson.DeleteBytes(body, "tool_choice") + } + return body, nil + } + if !tools.IsArray() { return body, nil } @@ -1363,7 +1369,7 @@ func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Contex case http.StatusTooManyRequests: // updateGrokUsageSnapshot installs both runtime and durable rate-limit state. default: - if statusCode >= 500 { + if statusCode >= 500 && !account.IsPoolMode() { s.tempUnscheduleGrok(ctx, account, 2*time.Minute, "grok upstream temporary error") } } diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index ca798d8141..170528ce76 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -125,6 +125,11 @@ func isGrokRequestContext(c *gin.Context) bool { if c == nil { return false } + if c.Request != nil { + if platform, ok := ResolvedTargetPlatformFromContext(c.Request.Context()); ok { + return platform == PlatformGrok + } + } v, exists := c.Get("api_key") if !exists { return false diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 3a8e347332..d13cfa19fc 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -209,6 +209,60 @@ func TestPatchGrokResponsesBodyDropsToolChoiceWhenNoSupportedToolsRemain(t *test require.False(t, gjson.GetBytes(patched, "tool_choice").Exists()) } +func TestSanitizeGrokResponsesToolsKeepsToolChoiceOnlyWithSupportedTools(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + body string + wantTools bool + wantToolChoice bool + }{ + { + name: "missing tools with string tool choice", + body: `{"input":"hello","tool_choice":"auto"}`, + }, + { + name: "missing tools with object tool choice", + body: `{"input":"hello","tool_choice":{"type":"function","name":"lookup"}}`, + }, + { + name: "empty tools", + body: `{"input":"hello","tools":[],"tool_choice":"auto"}`, + wantTools: true, + }, + { + name: "all tools unsupported", + body: `{"input":"hello","tools":[{"type":"namespace","name":"client_tools"}],"tool_choice":"auto"}`, + }, + { + name: "supported tool", + body: `{"input":"hello","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`, + wantTools: true, + wantToolChoice: true, + }, + { + name: "malformed non-array tools remain untouched", + body: `{"input":"hello","tools":{"type":"function","name":"lookup"},"tool_choice":"auto"}`, + wantTools: true, + wantToolChoice: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + patched, err := sanitizeGrokResponsesTools([]byte(tt.body)) + require.NoError(t, err) + require.True(t, json.Valid(patched)) + require.Equal(t, tt.wantTools, gjson.GetBytes(patched, "tools").Exists()) + require.Equal(t, tt.wantToolChoice, gjson.GetBytes(patched, "tool_choice").Exists()) + if tt.wantToolChoice { + require.Equal(t, "auto", gjson.GetBytes(patched, "tool_choice").String()) + } + }) + } +} + func TestPatchGrokResponsesBodyPromotesCodexAdditionalTools(t *testing.T) { t.Parallel() @@ -2475,6 +2529,43 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *test } } +func TestHandleGrokAccountUpstreamError5xxRespectsPoolMode(t *testing.T) { + t.Run("pool mode keeps scheduling state", func(t *testing.T) { + account := &Account{ + ID: 611, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "pool_mode": true, + }, + } + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadGateway, nil, nil) + + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Zero(t, repo.tempUnschedCalls) + require.Nil(t, account.TempUnschedulableUntil) + require.Empty(t, account.TempUnschedulableReason) + }) + + t.Run("non-pool mode keeps two minute cooldown", func(t *testing.T) { + account := &Account{ID: 612, Platform: PlatformGrok, Type: AccountTypeAPIKey} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + before := time.Now() + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadGateway, nil, nil) + + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Equal(t, 1, repo.tempUnschedCalls) + require.Equal(t, account.ID, repo.lastTempUnschedID) + require.Equal(t, "grok upstream temporary error", repo.lastTempUnschedReason) + require.WithinDuration(t, before.Add(2*time.Minute), repo.lastTempUnschedUntil, time.Second) + }) +} + func TestHandleGrokAccountUpstreamError429SetsRateLimitedFromRetryAfter(t *testing.T) { account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth} repo := &grokQuotaAccountRepo{} diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index a895b2ebf5..01316cb58a 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -26,6 +26,15 @@ const ( codeBuddyConversationHeader = "X-Conversation-ID" ) +var explicitOpenAIHeaderSessionNames = []string{ + "session_id", + "conversation_id", + openCodeSessionAffinityHeader, + openCodeSessionIDHeader, + openCodeNativeSessionHeader, + codeBuddyConversationHeader, +} + // explicitOpenAIHeaderSessionID resolves stable conversation identifiers sent // by OpenAI-compatible clients. Keep this list limited to session-scoped // fields: request/message IDs rotate every turn and would defeat sticky routing @@ -35,14 +44,7 @@ func explicitOpenAIHeaderSessionID(c *gin.Context) string { return "" } - for _, header := range []string{ - "session_id", - "conversation_id", - openCodeSessionAffinityHeader, - openCodeSessionIDHeader, - openCodeNativeSessionHeader, - codeBuddyConversationHeader, - } { + for _, header := range explicitOpenAIHeaderSessionNames { if sessionID := strings.TrimSpace(c.GetHeader(header)); sessionID != "" { return sessionID } diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 773b31d625..16429f7402 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -29,6 +29,7 @@ type OpenAIRecordUsageInput struct { UpstreamEndpoint string UserAgent string // 请求的 User-Agent IPAddress string // 请求的客户端 IP 地址 + SessionID string // 客户端显式会话标识（session_id / X-Session-Id 等请求头），仅用于用量行会话关联 RequestPayloadHash string APIKeyService APIKeyQuotaUpdater QuotaPlatform string // user×platform quota platform resolved by the handler before async billing. @@ -55,6 +56,7 @@ type CyberPolicyUsageInput struct { UpstreamEndpoint string UserAgent string IPAddress string + SessionID string RequestPayloadHash string APIKeyService APIKeyQuotaUpdater ChannelUsageFields @@ -89,6 +91,7 @@ func (s *OpenAIGatewayService) RecordCyberPolicyUsageLog(ctx context.Context, in UpstreamEndpoint: in.UpstreamEndpoint, UserAgent: in.UserAgent, IPAddress: in.IPAddress, + SessionID: in.SessionID, RequestPayloadHash: in.RequestPayloadHash, APIKeyService: in.APIKeyService, ChannelUsageFields: in.ChannelUsageFields, @@ -333,6 +336,9 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec usageLog.IPAddress = &input.IPAddress } + // 添加 SessionID（客户端显式会话标识；缺失/无效时保持 nil） + usageLog.SessionID = optionalTrimmedStringPtr(input.SessionID) + if apiKey.GroupID != nil { usageLog.GroupID = apiKey.GroupID } diff --git a/backend/internal/service/openai_oauth_passthrough_test.go b/backend/internal/service/openai_oauth_passthrough_test.go index d4f6e22203..43f8109f88 100644 --- a/backend/internal/service/openai_oauth_passthrough_test.go +++ b/backend/internal/service/openai_oauth_passthrough_test.go @@ -450,7 +450,10 @@ func TestOpenAIGatewayService_OAuthPassthrough_NamespaceRequestAndStreamResponse {"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"spawn_agent","description":"spawn","parameters":{"type":"object"}}]} ], "tool_choice":{"type":"function","name":"spawn_agent","namespace":"collaboration"}, - "input":[{"type":"function_call","call_id":"call_old","name":"spawn_agent","namespace":"collaboration","arguments":"{}"}] + "input":[ + {"type":"function_call","call_id":"call_old","name":"spawn_agent","namespace":"collaboration","arguments":"{}"}, + {"type":"message","role":"user","namespace":"residual","content":[{"type":"input_text","text":"keep","namespace":"nested"}]} + ] }`) upstreamSSE := strings.Join([]string{ @@ -488,6 +491,9 @@ func TestOpenAIGatewayService_OAuthPassthrough_NamespaceRequestAndStreamResponse require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice.namespace").Exists()) require.Equal(t, "collaboration__spawn_agent", gjson.GetBytes(upstream.lastBody, "input.0.name").String()) require.False(t, gjson.GetBytes(upstream.lastBody, "input.0.namespace").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "input.1.namespace").Exists()) + require.Equal(t, "nested", gjson.GetBytes(upstream.lastBody, "input.1.content.0.namespace").String()) + require.Len(t, upstream.bodies, 1) downstream := rec.Body.String() require.NotContains(t, downstream, "collaboration__spawn_agent") diff --git a/backend/internal/service/openai_responses_item_id.go b/backend/internal/service/openai_responses_item_id.go new file mode 100644 index 0000000000..27b7dea0f5 --- /dev/null +++ b/backend/internal/service/openai_responses_item_id.go @@ -0,0 +1,78 @@ +package service + +import ( + "fmt" + "strings" + + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +// Invalid replayed IDs are removed rather than rewritten because a fabricated +// msg/fc ID may point at a different upstream object. +func shouldStripOpenAIResponsesInputItemID(itemType, id string) bool { + if id == "" { + return false + } + if itemType == "message" { + return !strings.HasPrefix(id, "msg") + } + if isCodexToolCallInputType(itemType) { + return !strings.HasPrefix(id, "fc") + } + return false +} + +func sanitizeOpenAIResponsesInputItemIDs(body []byte) ([]byte, bool, error) { + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return body, false, nil + } + + items := make([][]byte, 0) + changed := false + var sanitizeErr error + index := 0 + input.ForEach(func(_, item gjson.Result) bool { + currentIndex := index + index++ + itemBody := []byte(item.Raw) + if item.IsObject() { + itemType := item.Get("type") + id := item.Get("id") + if itemType.Type == gjson.String && id.Type == gjson.String && + shouldStripOpenAIResponsesInputItemID(itemType.String(), id.String()) { + itemBody, sanitizeErr = sjson.DeleteBytes(itemBody, "id") + if sanitizeErr != nil { + sanitizeErr = fmt.Errorf("delete input.%d.id: %w", currentIndex, sanitizeErr) + return false + } + changed = true + } + } + items = append(items, itemBody) + return true + }) + if sanitizeErr != nil { + return nil, false, sanitizeErr + } + if !changed { + return body, false, nil + } + + rebuiltInput := make([]byte, 0, len(input.Raw)) + rebuiltInput = append(rebuiltInput, '[') + for i, item := range items { + if i > 0 { + rebuiltInput = append(rebuiltInput, ',') + } + rebuiltInput = append(rebuiltInput, item...) + } + rebuiltInput = append(rebuiltInput, ']') + + sanitized, err := sjson.SetRawBytes(body, "input", rebuiltInput) + if err != nil { + return nil, false, fmt.Errorf("replace sanitized input: %w", err) + } + return sanitized, true, nil +} diff --git a/backend/internal/service/openai_responses_namespace.go b/backend/internal/service/openai_responses_namespace.go index 4c3f433bc4..bd548b69ec 100644 --- a/backend/internal/service/openai_responses_namespace.go +++ b/backend/internal/service/openai_responses_namespace.go @@ -7,6 +7,8 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" ) const openAIResponsesNamespaceNamesContextKey = "openai_responses_namespace_names" @@ -26,6 +28,19 @@ func shouldFlattenOpenAIResponsesNamespaces(account *Account, transport OpenAIUp return true } +// shouldStripOpenAIResponsesInputNamespaces removes residual input item +// namespaces for OpenAI OAuth and API Key HTTP forwarding. Native WSv2 keeps +// namespaces because that protocol supports them and does not restore payloads. +func shouldStripOpenAIResponsesInputNamespaces(account *Account, transport OpenAIUpstreamTransport, passthroughEnabled bool) bool { + if account == nil || (!account.IsOpenAIOAuth() && !account.IsOpenAIApiKey()) { + return false + } + if transport == OpenAIUpstreamTransportResponsesWebsocketV2 && !passthroughEnabled { + return false + } + return true +} + func flattenOpenAIResponsesNamespaces(c *gin.Context, body []byte) ([]byte, error) { if !bytes.Contains(body, []byte(`"namespace"`)) { return body, nil @@ -49,6 +64,55 @@ func flattenOpenAIResponsesNamespaces(c *gin.Context, body []byte) ([]byte, erro return rebuilt, nil } +// stripOpenAIResponsesInputNamespaces removes namespace only from direct input +// array items. Namespace declarations and nested namespace fields are left +// untouched. Rebuilding the input array once keeps this linear for long +// histories and avoids decoding JSON numbers through float64. +func stripOpenAIResponsesInputNamespaces(body []byte) ([]byte, error) { + if !bytes.Contains(body, []byte(`"namespace"`)) { + return body, nil + } + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return body, nil + } + + var rebuilt bytes.Buffer + rebuilt.Grow(len(input.Raw)) + _ = rebuilt.WriteByte('[') + changed := false + first := true + var stripErr error + input.ForEach(func(_, item gjson.Result) bool { + if !first { + _ = rebuilt.WriteByte(',') + } + first = false + itemBody := []byte(item.Raw) + if item.IsObject() && item.Get("namespace").Exists() { + itemBody, stripErr = sjson.DeleteBytes(itemBody, "namespace") + if stripErr != nil { + return false + } + changed = true + } + _, _ = rebuilt.Write(itemBody) + return true + }) + if stripErr != nil { + return body, fmt.Errorf("delete OpenAI input namespace: %w", stripErr) + } + if !changed { + return body, nil + } + _ = rebuilt.WriteByte(']') + stripped, err := sjson.SetRawBytes(body, "input", rebuilt.Bytes()) + if err != nil { + return body, fmt.Errorf("replace OpenAI input after namespace deletion: %w", err) + } + return stripped, nil +} + func setOpenAIResponsesNamespaceNames(c *gin.Context, names map[string]apicompat.ResponsesNamespaceName) { if c != nil && len(names) > 0 { c.Set(openAIResponsesNamespaceNamesContextKey, names) diff --git a/backend/internal/service/openai_responses_namespace_test.go b/backend/internal/service/openai_responses_namespace_test.go index 00b8772d13..0b12355a3d 100644 --- a/backend/internal/service/openai_responses_namespace_test.go +++ b/backend/internal/service/openai_responses_namespace_test.go @@ -1,9 +1,11 @@ package service import ( + "strconv" "testing" "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" ) func TestShouldFlattenOpenAIResponsesNamespaces(t *testing.T) { @@ -34,3 +36,78 @@ func TestShouldFlattenOpenAIResponsesNamespaces(t *testing.T) { }) } } + +func TestShouldStripOpenAIResponsesInputNamespaces(t *testing.T) { + oauth := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth} + apiKey := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + setupToken := &Account{Platform: PlatformOpenAI, Type: AccountTypeSetupToken} + grokOAuth := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth} + + tests := []struct { + name string + account *Account + transport OpenAIUpstreamTransport + passthroughEnabled bool + want bool + }{ + {name: "oauth_http", account: oauth, transport: OpenAIUpstreamTransportHTTPSSE, want: true}, + {name: "apikey_http", account: apiKey, transport: OpenAIUpstreamTransportHTTPSSE, want: true}, + {name: "oauth_wsv2", account: oauth, transport: OpenAIUpstreamTransportResponsesWebsocketV2, want: false}, + {name: "apikey_wsv2", account: apiKey, transport: OpenAIUpstreamTransportResponsesWebsocketV2, want: false}, + {name: "oauth_wsv2_passthrough", account: oauth, transport: OpenAIUpstreamTransportResponsesWebsocketV2, passthroughEnabled: true, want: true}, + {name: "apikey_wsv2_passthrough", account: apiKey, transport: OpenAIUpstreamTransportResponsesWebsocketV2, passthroughEnabled: true, want: true}, + {name: "setup_token_http", account: setupToken, transport: OpenAIUpstreamTransportHTTPSSE, want: false}, + {name: "grok_oauth_http", account: grokOAuth, transport: OpenAIUpstreamTransportHTTPSSE, want: false}, + {name: "nil_account", account: nil, transport: OpenAIUpstreamTransportHTTPSSE, want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, shouldStripOpenAIResponsesInputNamespaces(tt.account, tt.transport, tt.passthroughEnabled)) + }) + } +} + +func TestStripOpenAIResponsesInputNamespaces(t *testing.T) { + body := []byte(`{ + "meta":9007199254740993, + "scientific":1.25e+42, + "escaped":"line\\n\\u003ctag\\u003e", + "tools":[{"type":"function","name":"keep","namespace":"tool-namespace"}], + "input":[ + {"type":"function_call","namespace":"n0","name":"one","content":{"namespace":"nested"},"large":9007199254740993}, + {"type":"message","namespace":"n1","content":[{"type":"input_text","text":"hello","namespace":"nested-content"}]}, + {"type":"custom_tool_call","namespace":"n2","input":"{}"}, + {"type":"function_call_output","namespace":"n3","output":"ok"}, + {"type":"item","namespace":"n4"}, + {"type":"item","namespace":"n5"}, + {"type":"item","namespace":"n6"}, + {"type":"item","namespace":"n7"} + ] + }`) + + stripped, err := stripOpenAIResponsesInputNamespaces(body) + require.NoError(t, err) + for index := 0; index < 8; index++ { + require.False(t, gjson.GetBytes(stripped, "input."+strconv.Itoa(index)+".namespace").Exists()) + } + require.Equal(t, "nested", gjson.GetBytes(stripped, "input.0.content.namespace").String()) + require.Equal(t, "nested-content", gjson.GetBytes(stripped, "input.1.content.0.namespace").String()) + require.Equal(t, "tool-namespace", gjson.GetBytes(stripped, "tools.0.namespace").String()) + require.Equal(t, gjson.GetBytes(body, "meta").Raw, gjson.GetBytes(stripped, "meta").Raw) + require.Equal(t, gjson.GetBytes(body, "scientific").Raw, gjson.GetBytes(stripped, "scientific").Raw) + require.Equal(t, gjson.GetBytes(body, "escaped").Raw, gjson.GetBytes(stripped, "escaped").Raw) + require.Equal(t, gjson.GetBytes(body, "input.0.large").Raw, gjson.GetBytes(stripped, "input.0.large").Raw) +} + +func TestStripOpenAIResponsesInputNamespacesLeavesOtherShapesByteExact(t *testing.T) { + tests := [][]byte{ + []byte(`{"input":"text","namespace":"top-level"}`), + []byte(`{"input":{"namespace":"single-object"}}`), + []byte(`{"input":[{"content":{"namespace":"nested-only"}}],"tools":[{"namespace":"keep"}]}`), + } + for _, body := range tests { + stripped, err := stripOpenAIResponsesInputNamespaces(body) + require.NoError(t, err) + require.Equal(t, body, stripped) + } +} diff --git a/backend/internal/service/openai_responses_rejected_field_retry_test.go b/backend/internal/service/openai_responses_rejected_field_retry_test.go index 96c20de508..724fd90c37 100644 --- a/backend/internal/service/openai_responses_rejected_field_retry_test.go +++ b/backend/internal/service/openai_responses_rejected_field_retry_test.go @@ -114,10 +114,9 @@ func TestNormalizeOpenAIResponsesRejectedFieldRetryBodyBindsMaxOutputTokensToRej require.False(t, gjson.GetBytes(retryBody, "max_output_tokens").Exists()) } -func TestOpenAIGatewayService_RetriesRejectedIndexedNamespaceField(t *testing.T) { - body := []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"function_call","name":"first","namespace":"keep","arguments":"{}"},{"type":"custom_tool_call","name":"second","namespace":"remove","input":"{}"}]}`) +func TestOpenAIGatewayService_APIKeyStripsAllIndexedNamespacesBeforeFirstForward(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"function_call","name":"first","namespace":"remove-first","arguments":"{}"},{"type":"custom_tool_call","name":"second","namespace":"remove-second","input":"{}"}]}`) upstream := &httpUpstreamRecorder{responses: []*http.Response{ - newOpenAIRejectedFieldTestResponse(http.StatusBadRequest, `{"error":{"code":"unknown_parameter","message":"Unknown parameter: 'input[1].namespace'.","param":"input[1].namespace","type":"invalid_request_error"}}`), newOpenAIRejectedFieldTestResponse(http.StatusOK, `{"output":[],"usage":{"input_tokens":1,"output_tokens":1,"input_tokens_details":{"cached_tokens":0}}}`), }} @@ -130,9 +129,44 @@ func TestOpenAIGatewayService_RetriesRejectedIndexedNamespaceField(t *testing.T) require.NoError(t, err) require.NotNil(t, result) - require.Len(t, upstream.bodies, 2) - require.Equal(t, "keep", gjson.GetBytes(upstream.bodies[1], "input.0.namespace").String()) - require.False(t, gjson.GetBytes(upstream.bodies[1], "input.1.namespace").Exists()) + require.Len(t, upstream.bodies, 1) + require.False(t, gjson.GetBytes(upstream.bodies[0], "input.0.namespace").Exists()) + require.False(t, gjson.GetBytes(upstream.bodies[0], "input.1.namespace").Exists()) +} + +func TestOpenAIGatewayService_OpenAIHTTPStripsInputNamespacesBeforeFirstForward(t *testing.T) { + accounts := []struct { + name string + account *Account + }{ + {name: "oauth", account: newOpenAIOAuthNamespaceTestAccount()}, + {name: "apikey", account: newOpenAIRejectedFieldTestAccount()}, + } + for _, tt := range accounts { + for _, path := range []string{"/v1/responses", "/v1/responses/compact"} { + t.Run(tt.name+path, func(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","stream":false,"instructions":"test","input":[{"type":"message","role":"user","namespace":"remove","content":[{"type":"input_text","text":"hello","namespace":"nested-keep"}]}]}`) + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + newOpenAIRejectedFieldTestResponse(http.StatusOK, `{"id":"resp_namespace_ok","output":[],"usage":{"input_tokens":1,"output_tokens":1,"input_tokens_details":{"cached_tokens":0}}}`), + }} + c := newOpenAIRejectedFieldTestContext(body) + c.Request.URL.Path = path + + result, err := newOpenAIRejectedFieldTestService(upstream).Forward( + context.Background(), + c, + tt.account, + body, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.bodies, 1, "namespace must be removed before the first upstream request") + require.False(t, gjson.GetBytes(upstream.bodies[0], "input.0.namespace").Exists()) + require.Equal(t, "nested-keep", gjson.GetBytes(upstream.bodies[0], "input.0.content.0.namespace").String()) + }) + } + } } func TestOpenAIGatewayService_RetriesExplicitMaxOutputTokensRejection(t *testing.T) { @@ -157,10 +191,9 @@ func TestOpenAIGatewayService_RetriesExplicitMaxOutputTokensRejection(t *testing require.Equal(t, "keep", gjson.GetBytes(upstream.bodies[1], "input.0.content.max_output_tokens").String()) } -func TestOpenAIGatewayService_ComposesDistinctRejectedFieldRetries(t *testing.T) { - body := []byte(`{"model":"gpt-5.5","stream":false,"max_output_tokens":2048,"input":[{"type":"function_call","name":"first","namespace":"keep","arguments":"{}"},{"type":"custom_tool_call","name":"second","namespace":"remove","input":"{}"}]}`) +func TestOpenAIGatewayService_ComposesProactiveNamespaceStripWithRejectedFieldRetry(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","stream":false,"max_output_tokens":2048,"input":[{"type":"function_call","name":"first","namespace":"remove-first","arguments":"{}"},{"type":"custom_tool_call","name":"second","namespace":"remove-second","input":"{}"}]}`) upstream := &httpUpstreamRecorder{responses: []*http.Response{ - newOpenAIRejectedFieldTestResponse(http.StatusBadRequest, `{"error":{"code":"unknown_parameter","message":"Unknown parameter: 'input[1].namespace'.","param":"input[1].namespace"}}`), newOpenAIRejectedFieldTestResponse(http.StatusBadRequest, `{"error":{"code":"unsupported_parameter","message":"Unsupported parameter: max_output_tokens","param":"max_output_tokens"}}`), newOpenAIRejectedFieldTestResponse(http.StatusOK, `{"output":[],"usage":{"input_tokens":1,"output_tokens":1,"input_tokens_details":{"cached_tokens":0}}}`), }} @@ -174,12 +207,13 @@ func TestOpenAIGatewayService_ComposesDistinctRejectedFieldRetries(t *testing.T) require.NoError(t, err) require.NotNil(t, result) - require.Len(t, upstream.bodies, 3) - require.True(t, gjson.GetBytes(upstream.bodies[0], "input.1.namespace").Exists()) - require.False(t, gjson.GetBytes(upstream.bodies[1], "input.1.namespace").Exists()) - require.Equal(t, int64(2048), gjson.GetBytes(upstream.bodies[1], "max_output_tokens").Int()) - require.False(t, gjson.GetBytes(upstream.bodies[2], "input.1.namespace").Exists()) - require.False(t, gjson.GetBytes(upstream.bodies[2], "max_output_tokens").Exists()) + require.Len(t, upstream.bodies, 2) + for _, forwardedBody := range upstream.bodies { + require.False(t, gjson.GetBytes(forwardedBody, "input.0.namespace").Exists()) + require.False(t, gjson.GetBytes(forwardedBody, "input.1.namespace").Exists()) + } + require.Equal(t, int64(2048), gjson.GetBytes(upstream.bodies[0], "max_output_tokens").Int()) + require.False(t, gjson.GetBytes(upstream.bodies[1], "max_output_tokens").Exists()) } func newOpenAIRejectedFieldTestService(upstream *httpUpstreamRecorder) *OpenAIGatewayService { @@ -221,6 +255,22 @@ func newOpenAIRejectedFieldTestAccount() *Account { } } +func newOpenAIOAuthNamespaceTestAccount() *Account { + return &Account{ + ID: 5108, + Name: "openai-oauth-namespace", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "oauth-token", + "chatgpt_account_id": "chatgpt-account", + }, + Status: StatusActive, + Schedulable: true, + } +} + func newOpenAIRejectedFieldTestResponse(status int, body string) *http.Response { return &http.Response{ StatusCode: status, diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index 5872f125bf..c7cbbc0b2f 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -733,7 +733,7 @@ func TestOpenAIGatewayService_Forward_WSv2_OAuthStoreFalseByDefault(t *testing.T }, } - body := []byte(`{"model":"gpt-5.1","stream":false,"store":true,"input":[{"type":"input_text","text":"hello"}]}`) + body := []byte(`{"model":"gpt-5.1","stream":false,"store":true,"input":[{"type":"input_text","text":"hello","namespace":"native-wsv2"}]}`) result, err := svc.Forward(context.Background(), c, account, body) require.NoError(t, err) require.NotNil(t, result) @@ -745,6 +745,7 @@ func TestOpenAIGatewayService_Forward_WSv2_OAuthStoreFalseByDefault(t *testing.T require.False(t, gjson.Get(requestJSON, "store").Bool(), "默认策略应将 OAuth store 置为 false") require.True(t, gjson.Get(requestJSON, "stream").Exists(), "WSv2 payload 应保留 stream 字段") require.True(t, gjson.Get(requestJSON, "stream").Bool(), "OAuth Codex 规范化后应强制 stream=true") + require.Equal(t, "native-wsv2", gjson.Get(requestJSON, "input.0.namespace").String(), "OAuth WSv2 应保留原生 namespace") require.Equal(t, openAIWSBetaV2Value, captureDialer.lastHeaders.Get("OpenAI-Beta")) require.Equal(t, "remote_compaction_v2", captureDialer.lastHeaders.Get("x-codex-beta-features")) // OAuth 账号的 session_id/conversation_id 应被 isolateOpenAISessionID 隔离， diff --git a/backend/internal/service/pricing_service.go b/backend/internal/service/pricing_service.go index 8bd916f4fc..636bd690cb 100644 --- a/backend/internal/service/pricing_service.go +++ b/backend/internal/service/pricing_service.go @@ -218,6 +218,11 @@ func (s *PricingService) Stop() { // startUpdateScheduler 启动定时更新调度器 func (s *PricingService) startUpdateScheduler() { + if s == nil || s.cfg == nil || strings.TrimSpace(s.cfg.Pricing.RemoteURL) == "" { + logger.LegacyPrintf("service.pricing", "%s", "[Pricing] Remote sync disabled: pricing remote URL is empty") + return + } + // 定期检查哈希更新 hashInterval := time.Duration(s.cfg.Pricing.HashCheckIntervalMinutes) * time.Minute if hashInterval < time.Minute { diff --git a/backend/internal/service/pricing_service_test.go b/backend/internal/service/pricing_service_test.go index b68d034638..88553f0fbf 100644 --- a/backend/internal/service/pricing_service_test.go +++ b/backend/internal/service/pricing_service_test.go @@ -5,11 +5,42 @@ import ( "os" "path/filepath" "testing" + "time" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/stretchr/testify/require" ) +func TestPricingSchedulerBlankRemoteURLDoesNotStart(t *testing.T) { + svc := NewPricingService(&config.Config{Pricing: config.PricingConfig{RemoteURL: " \t "}}, nil) + defer svc.Stop() + + svc.startUpdateScheduler() + done := make(chan struct{}) + go func() { + svc.wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(100 * time.Millisecond): + t.Fatal("blank remote URL must not start scheduler") + } +} + +func TestPricingNonEmptyInvalidRemoteURLStillReturnsValidationError(t *testing.T) { + svc := NewPricingService(&config.Config{Pricing: config.PricingConfig{ + RemoteURL: "://invalid", + DataDir: t.TempDir(), + }}, nil) + + err := svc.ForceUpdate() + + require.Error(t, err) + require.Contains(t, err.Error(), "invalid pricing url") +} + func TestParsePricingData_ParsesPriorityAndServiceTierFields(t *testing.T) { svc := &PricingService{} body := []byte(`{ diff --git a/backend/internal/service/session_id.go b/backend/internal/service/session_id.go new file mode 100644 index 0000000000..8cf39f056b --- /dev/null +++ b/backend/internal/service/session_id.go @@ -0,0 +1,75 @@ +package service + +import ( + "strings" + "unicode/utf8" + + "github.com/gin-gonic/gin" +) + +// maxPersistedSessionIDLength bounds the persisted client session identifier to the +// usage_logs.session_id column width (VARCHAR(255)). Longer values are rejected so +// distinct identifiers can never alias through truncation. +const maxPersistedSessionIDLength = 255 + +// clientSessionIDHeaders extends the OpenAI-compatible sticky-session signals with +// native protocol identifiers that are safe to persist but must not alter OpenAI +// scheduling behavior. +var clientSessionIDHeaders = append( + append([]string(nil), explicitOpenAIHeaderSessionNames...), + claudeCodeSessionHeader, +) + +// ExtractClientSessionID resolves the explicit client-provided session identifier from +// request headers for usage-log correlation and returns it sanitized. It is +// protocol-agnostic and shared by every gateway handler so all supported protocols +// record session_id through one seam. Returns "" when no valid identifier is present. +// +// This value feeds only usage_logs.session_id persistence. It does NOT affect sticky +// routing, account selection, request_id semantics, or upstream prompt caching, which +// keep their own (intentionally broader) session-signal resolution. +func ExtractClientSessionID(c *gin.Context) string { + if c == nil || c.Request == nil { + return "" + } + for _, header := range clientSessionIDHeaders { + if sessionID := sanitizeSessionID(c.GetHeader(header)); sessionID != "" { + return sessionID + } + } + if isGrokRequestContext(c) { + if sessionID := sanitizeSessionID(c.GetHeader(grokConversationIDHeader)); sessionID != "" { + return sessionID + } + } + return "" +} + +// sanitizeSessionID normalizes a raw client-supplied session identifier for safe +// persistence: it trims surrounding whitespace, rejects the value outright if it +// contains any control character (CR/LF/tab/NUL/…) so a log- or header-injection style +// payload cannot slip into stored correlation data, and rejects values longer than +// the DB column bound. Absent or invalid input yields "". +func sanitizeSessionID(raw string) string { + if !utf8.ValidString(raw) { + return "" + } + trimmed := strings.TrimSpace(raw) + if trimmed == "" { + return "" + } + count := 0 + for _, r := range trimmed { + if r < 0x20 || r == 0x7f { + // An explicit correlation id never legitimately contains control + // characters; drop the whole value rather than persist a mangled or + // partially-injected identifier. + return "" + } + count++ + if count > maxPersistedSessionIDLength { + return "" + } + } + return trimmed +} diff --git a/backend/internal/service/session_id_test.go b/backend/internal/service/session_id_test.go new file mode 100644 index 0000000000..1eb746d6ef --- /dev/null +++ b/backend/internal/service/session_id_test.go @@ -0,0 +1,155 @@ +//go:build unit + +package service + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func newSessionHeaderContext(t *testing.T, headers map[string]string) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + req := httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + for k, v := range headers { + req.Header.Set(k, v) + } + c.Request = req + return c +} + +func TestSanitizeSessionID(t *testing.T) { + longRunes := strings.Repeat("a", maxPersistedSessionIDLength+50) + multibyte := strings.Repeat("好", maxPersistedSessionIDLength+10) + + tests := []struct { + name string + in string + want string + }{ + {"empty", "", ""}, + {"whitespace only", " \t ", ""}, + {"trims surrounding whitespace", " sess-123 ", "sess-123"}, + {"plain value", "conv_abc-123.XYZ", "conv_abc-123.XYZ"}, + {"uuid", "550e8400-e29b-41d4-a716-446655440000", "550e8400-e29b-41d4-a716-446655440000"}, + {"reject CR", "sess\r123", ""}, + {"reject LF", "sess\n123", ""}, + {"reject CRLF injection", "sess-1\r\nSet-Cookie: x=y", ""}, + {"reject tab inside", "sess\t123", ""}, + {"reject NUL", "sess\x00123", ""}, + {"reject DEL", "sess\x7f123", ""}, + {"reject invalid UTF-8", string([]byte{'s', 'e', 's', 's', '-', 0xff}), ""}, + {"accepts value at column bound", strings.Repeat("b", maxPersistedSessionIDLength), strings.Repeat("b", maxPersistedSessionIDLength)}, + {"rejects overlong value", longRunes, ""}, + {"rejects overlong multibyte value", multibyte, ""}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := sanitizeSessionID(tc.in) + require.Equal(t, tc.want, got, "sanitizeSessionID(%q)", tc.in) + // Sanitized output must never exceed the DB column bound (rune-counted). + require.LessOrEqual(t, len([]rune(got)), maxPersistedSessionIDLength) + }) + } +} + +func TestExtractClientSessionID_NilContext(t *testing.T) { + require.Equal(t, "", ExtractClientSessionID(nil)) +} + +func TestExtractClientSessionID_NilRequest(t *testing.T) { + require.Equal(t, "", ExtractClientSessionID(&gin.Context{})) +} + +func TestExtractClientSessionID_AbsentReturnsEmpty(t *testing.T) { + c := newSessionHeaderContext(t, nil) + require.Equal(t, "", ExtractClientSessionID(c)) +} + +func TestExtractClientSessionID_SupportedHeaders(t *testing.T) { + tests := []struct { + name string + header string + value string + }{ + {"session_id", "session_id", "sess-A"}, + {"conversation_id", "conversation_id", "conv-B"}, + {"X-Session-Affinity", openCodeSessionAffinityHeader, "aff-C"}, + {"X-Session-Id", openCodeSessionIDHeader, "sid-D"}, + {"X-OpenCode-Session", openCodeNativeSessionHeader, "oc-E"}, + {"X-Conversation-ID", codeBuddyConversationHeader, "cb-F"}, + {"X-Claude-Code-Session-Id", claudeCodeSessionHeader, "cc-G"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + c := newSessionHeaderContext(t, map[string]string{tc.header: tc.value}) + require.Equal(t, tc.value, ExtractClientSessionID(c)) + }) + } +} + +func TestExtractClientSessionID_HeaderPrecedence(t *testing.T) { + // session_id ranks ahead of conversation_id and the X-* variants. + c := newSessionHeaderContext(t, map[string]string{ + "session_id": "primary", + "conversation_id": "secondary", + openCodeSessionIDHeader: "tertiary", + codeBuddyConversationHeader: "quaternary", + }) + require.Equal(t, "primary", ExtractClientSessionID(c)) +} + +func TestExtractClientSessionID_Sanitizes(t *testing.T) { + c := newSessionHeaderContext(t, map[string]string{openCodeSessionIDHeader: " clean-123 "}) + require.Equal(t, "clean-123", ExtractClientSessionID(c)) +} + +func TestExtractClientSessionID_IgnoresNonSessionHeaders(t *testing.T) { + // prompt_cache_key, request/message ids, and a Grok conversation header on a + // non-Grok request are NOT persisted as session_id. + c := newSessionHeaderContext(t, map[string]string{ + "prompt_cache_key": "cache-key-should-not-persist", + "X-Request-Id": "req-should-not-persist", + "x-grok-conv-id": "grok-conv-should-not-persist", + }) + require.Equal(t, "", ExtractClientSessionID(c)) +} + +func TestExtractClientSessionID_GrokConversationHeader(t *testing.T) { + c := newSessionHeaderContext(t, map[string]string{ + grokConversationIDHeader: "grok-native-session", + }) + c.Set("api_key", &APIKey{ + ID: 42, + Group: &Group{Platform: PlatformGrok}, + }) + + require.Equal(t, "grok-native-session", ExtractClientSessionID(c)) +} + +func TestExtractClientSessionID_GrokConversationHeaderForCompositeRoute(t *testing.T) { + c := newSessionHeaderContext(t, map[string]string{ + grokConversationIDHeader: "grok-composite-session", + }) + c.Set("api_key", &APIKey{ + ID: 43, + Group: &Group{Platform: PlatformComposite}, + }) + c.Request = c.Request.WithContext(WithResolvedTargetPlatform(context.Background(), PlatformGrok)) + + require.Equal(t, "grok-composite-session", ExtractClientSessionID(c)) +} + +func TestExtractClientSessionID_InjectionHeaderDropped(t *testing.T) { + // A supported header carrying a CRLF payload is rejected, not persisted mangled. + c := newSessionHeaderContext(t, map[string]string{"session_id": "abc"}) + c.Request.Header.Set("session_id", "abc\r\nX-Injected: 1") + require.Equal(t, "", ExtractClientSessionID(c)) +} diff --git a/backend/internal/service/usage_log.go b/backend/internal/service/usage_log.go index f0f586b977..b35d149ea8 100644 --- a/backend/internal/service/usage_log.go +++ b/backend/internal/service/usage_log.go @@ -170,6 +170,10 @@ type UsageLog struct { FirstTokenMs *int UserAgent *string IPAddress *string + // SessionID is the explicit client-provided request correlation identifier + // (e.g. the session_id / X-Session-Id headers). Nil when the client sent no + // valid session header. It is never derived from prompt_cache_key or content. + SessionID *string // Cache TTL Override 标记（管理员强制替换了缓存 TTL 计费） CacheTTLOverridden bool diff --git a/backend/migrations/187_add_usage_log_session_id.sql b/backend/migrations/187_add_usage_log_session_id.sql new file mode 100644 index 0000000000..7d6a4f6e21 --- /dev/null +++ b/backend/migrations/187_add_usage_log_session_id.sql @@ -0,0 +1,14 @@ +-- Persist the explicit client-provided request session/conversation identifier +-- (e.g. the session_id / X-Session-Id / X-Conversation-ID request headers) so +-- usage rows can be correlated by session for querying. +-- +-- Nullable with no default: on PostgreSQL 11+ this is a metadata-only change, so +-- it does NOT rewrite the (potentially large) usage_logs table. Absent/invalid +-- session identifiers stay NULL. Only explicit client correlation identity is +-- stored here; prompt_cache_key and content-derived sticky hashes are never +-- persisted as session_id. +ALTER TABLE usage_logs ADD COLUMN IF NOT EXISTS session_id VARCHAR(255); + +-- Batch image usage is recorded asynchronously after the originating HTTP +-- request has finished, so retain the same sanitized identifier on the job. +ALTER TABLE batch_image_jobs ADD COLUMN IF NOT EXISTS session_id VARCHAR(255); diff --git a/frontend/src/components/common/AnnouncementBell.vue b/frontend/src/components/common/AnnouncementBell.vue index cfbd62a0b8..e80e449f28 100644 --- a/frontend/src/components/common/AnnouncementBell.vue +++ b/frontend/src/components/common/AnnouncementBell.vue @@ -322,6 +322,7 @@ import { useAnnouncementStore } from '@/stores/announcements' import { formatRelativeTime, formatRelativeWithDateTime } from '@/utils/format' import type { UserAnnouncement } from '@/types' import Icon from '@/components/icons/Icon.vue' +import '@/styles/announcement-markdown.css' const { t } = useI18n() const appStore = useAppStore() @@ -471,120 +472,3 @@ watch( background: linear-gradient(to bottom, #6b7280, #4b5563); } - - diff --git a/frontend/src/components/common/AnnouncementPopup.vue b/frontend/src/components/common/AnnouncementPopup.vue index e26d9ee30b..99e10ed461 100644 --- a/frontend/src/components/common/AnnouncementPopup.vue +++ b/frontend/src/components/common/AnnouncementPopup.vue @@ -2,7 +2,7 @@

- {{ announcementStore.currentPopup.title }} + {{ displayedAnnouncement.title }}

@@ -43,7 +43,7 @@ - +
@@ -66,13 +66,17 @@
@@ -84,15 +88,34 @@