Merge remote-tracking branch 'origin/main' into feature/openai-live-gateway
This commit is contained in:
commit
5374ce2a0f
61 files changed
+1809
-199
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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: "",
|
||||
},
|
||||
{
|
||||
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",
|
||||
},
|
||||
}
|
||||
|
||||
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": ""}, 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)
|
||||
}
|
||||
@@ -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("", 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() {
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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{}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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 隔离,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(`{
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
@@ -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);
|
||||
}
|
||||
</style>
|
||||
|
||||
<style>
|
||||
/* Enhanced Markdown Styles */
|
||||
.markdown-body {
|
||||
@apply text-[15px] leading-[1.75];
|
||||
@apply text-gray-700 dark:text-gray-300;
|
||||
}
|
||||
|
||||
.markdown-body h1 {
|
||||
@apply mb-6 mt-8 border-b border-gray-200 pb-3 text-3xl font-bold text-gray-900 dark:border-dark-600 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body h2 {
|
||||
@apply mb-4 mt-7 border-b border-gray-100 pb-2 text-2xl font-bold text-gray-900 dark:border-dark-700 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body h3 {
|
||||
@apply mb-3 mt-6 text-xl font-semibold text-gray-900 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body h4 {
|
||||
@apply mb-2 mt-5 text-lg font-semibold text-gray-900 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body p {
|
||||
@apply mb-4 leading-relaxed;
|
||||
}
|
||||
|
||||
.markdown-body a {
|
||||
@apply font-medium text-blue-600 underline decoration-blue-600/30 decoration-2 underline-offset-2 transition-all hover:decoration-blue-600 dark:text-blue-400 dark:decoration-blue-400/30 dark:hover:decoration-blue-400;
|
||||
}
|
||||
|
||||
.markdown-body ul,
|
||||
.markdown-body ol {
|
||||
@apply mb-4 ml-6 space-y-2;
|
||||
}
|
||||
|
||||
.markdown-body ul {
|
||||
@apply list-disc;
|
||||
}
|
||||
|
||||
.markdown-body ol {
|
||||
@apply list-decimal;
|
||||
}
|
||||
|
||||
.markdown-body li {
|
||||
@apply leading-relaxed;
|
||||
@apply pl-2;
|
||||
}
|
||||
|
||||
.markdown-body li::marker {
|
||||
@apply text-blue-600 dark:text-blue-400;
|
||||
}
|
||||
|
||||
.markdown-body blockquote {
|
||||
@apply relative my-5 border-l-4 border-blue-500 bg-blue-50/50 py-3 pl-5 pr-4 italic text-gray-700 dark:border-blue-400 dark:bg-blue-900/10 dark:text-gray-300;
|
||||
}
|
||||
|
||||
.markdown-body blockquote::before {
|
||||
content: '"';
|
||||
@apply absolute -left-1 top-0 text-5xl font-serif text-blue-500/20 dark:text-blue-400/20;
|
||||
}
|
||||
|
||||
.markdown-body code {
|
||||
@apply rounded-lg bg-gray-100 px-2 py-1 text-[13px] font-mono text-pink-600 dark:bg-dark-700 dark:text-pink-400;
|
||||
}
|
||||
|
||||
.markdown-body pre {
|
||||
@apply my-5 overflow-x-auto rounded-xl border border-gray-200 bg-gray-50 p-5 dark:border-dark-600 dark:bg-dark-900/50;
|
||||
}
|
||||
|
||||
.markdown-body pre code {
|
||||
@apply bg-transparent p-0 text-[13px] text-gray-800 dark:text-gray-200;
|
||||
}
|
||||
|
||||
.markdown-body hr {
|
||||
@apply my-8 border-0 border-t-2 border-gray-200 dark:border-dark-700;
|
||||
}
|
||||
|
||||
.markdown-body table {
|
||||
@apply mb-5 w-full overflow-hidden rounded-lg border border-gray-200 dark:border-dark-600;
|
||||
}
|
||||
|
||||
.markdown-body th,
|
||||
.markdown-body td {
|
||||
@apply border-r border-b border-gray-200 px-4 py-3 text-left dark:border-dark-600;
|
||||
}
|
||||
|
||||
.markdown-body th:last-child,
|
||||
.markdown-body td:last-child {
|
||||
@apply border-r-0;
|
||||
}
|
||||
|
||||
.markdown-body tr:last-child td {
|
||||
@apply border-b-0;
|
||||
}
|
||||
|
||||
.markdown-body th {
|
||||
@apply bg-gradient-to-br from-blue-50 to-indigo-50 font-semibold text-gray-900 dark:from-blue-900/20 dark:to-indigo-900/10 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body tbody tr {
|
||||
@apply transition-colors hover:bg-gray-50 dark:hover:bg-dark-700/30;
|
||||
}
|
||||
|
||||
.markdown-body img {
|
||||
@apply my-5 max-w-full rounded-xl border border-gray-200 shadow-md dark:border-dark-600;
|
||||
}
|
||||
|
||||
.markdown-body strong {
|
||||
@apply font-semibold text-gray-900 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body em {
|
||||
@apply italic text-gray-600 dark:text-gray-400;
|
||||
}
|
||||
</style>
|
||||
@@ -2,7 +2,7 @@
|
||||
<Teleport to="body">
|
||||
<Transition name="popup-fade">
|
||||
<div
|
||||
v-if="announcementStore.currentPopup"
|
||||
v-if="displayedAnnouncement"
|
||||
class="fixed inset-0 z-[120] flex items-start justify-center overflow-y-auto bg-gradient-to-br from-black/70 via-black/60 to-black/70 p-4 pt-[8vh] backdrop-blur-md"
|
||||
>
|
||||
<div
|
||||
@@ -35,7 +35,7 @@
|
||||
|
||||
<!-- Title -->
|
||||
<h2 class="mb-2 text-2xl font-bold leading-tight text-gray-900 dark:text-white">
|
||||
{{ announcementStore.currentPopup.title }}
|
||||
{{ displayedAnnouncement.title }}
|
||||
</h2>
|
||||
|
||||
<!-- Time -->
|
||||
@@ -43,7 +43,7 @@
|
||||
<svg class="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
|
||||
<path stroke-linecap="round" stroke-linejoin="round" d="M12 8v4l3 3m6-3a9 9 0 11-18 0 9 9 0 0118 0z" />
|
||||
</svg>
|
||||
<time>{{ formatRelativeWithDateTime(announcementStore.currentPopup.created_at) }}</time>
|
||||
<time>{{ formatRelativeWithDateTime(displayedAnnouncement.created_at) }}</time>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@@ -66,13 +66,17 @@
|
||||
<div class="flex items-center justify-end">
|
||||
<button
|
||||
@click="handleDismiss"
|
||||
data-testid="announcement-popup-dismiss"
|
||||
class="rounded-xl bg-gradient-to-r from-amber-500 to-orange-600 px-6 py-2.5 text-sm font-medium text-white shadow-lg shadow-amber-500/30 transition-all hover:shadow-xl hover:scale-105"
|
||||
>
|
||||
<span class="flex items-center gap-2">
|
||||
<svg class="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
|
||||
<svg v-if="preview" class="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
|
||||
<path stroke-linecap="round" stroke-linejoin="round" d="M6 18L18 6M6 6l12 12" />
|
||||
</svg>
|
||||
<svg v-else class="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
|
||||
<path stroke-linecap="round" stroke-linejoin="round" d="M5 13l4 4L19 7" />
|
||||
</svg>
|
||||
{{ t('announcements.markRead') }}
|
||||
{{ preview ? t('common.close') : t('announcements.markRead') }}
|
||||
</span>
|
||||
</button>
|
||||
</div>
|
||||
@@ -84,15 +88,34 @@
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import { computed, watch } from 'vue'
|
||||
import { computed, onBeforeUnmount, watch } from 'vue'
|
||||
import { useI18n } from 'vue-i18n'
|
||||
import { marked } from 'marked'
|
||||
import DOMPurify from 'dompurify'
|
||||
import { useAnnouncementStore } from '@/stores/announcements'
|
||||
import { formatRelativeWithDateTime } from '@/utils/format'
|
||||
import type { Announcement, UserAnnouncement } from '@/types'
|
||||
import '@/styles/announcement-markdown.css'
|
||||
|
||||
type PreviewAnnouncement = Pick<Announcement | UserAnnouncement, 'title' | 'content' | 'created_at'>
|
||||
|
||||
const props = withDefaults(defineProps<{
|
||||
announcement?: PreviewAnnouncement | null
|
||||
preview?: boolean
|
||||
}>(), {
|
||||
announcement: null,
|
||||
preview: false,
|
||||
})
|
||||
|
||||
const emit = defineEmits<{
|
||||
close: []
|
||||
}>()
|
||||
|
||||
const { t } = useI18n()
|
||||
const announcementStore = useAnnouncementStore()
|
||||
const displayedAnnouncement = computed(() => (
|
||||
props.preview ? props.announcement : announcementStore.currentPopup
|
||||
))
|
||||
|
||||
marked.setOptions({
|
||||
breaks: true,
|
||||
@@ -100,25 +123,38 @@ marked.setOptions({
|
||||
})
|
||||
|
||||
const renderedContent = computed(() => {
|
||||
const content = announcementStore.currentPopup?.content
|
||||
const content = displayedAnnouncement.value?.content
|
||||
if (!content) return ''
|
||||
const html = marked.parse(content) as string
|
||||
return DOMPurify.sanitize(html)
|
||||
})
|
||||
|
||||
function handleDismiss() {
|
||||
if (props.preview) {
|
||||
emit('close')
|
||||
return
|
||||
}
|
||||
announcementStore.dismissPopup()
|
||||
}
|
||||
|
||||
// Manage body overflow — only set, never unset (bell component handles restore)
|
||||
watch(
|
||||
() => announcementStore.currentPopup,
|
||||
displayedAnnouncement,
|
||||
(popup) => {
|
||||
if (popup) {
|
||||
document.body.style.overflow = 'hidden'
|
||||
} else if (props.preview) {
|
||||
document.body.style.overflow = ''
|
||||
}
|
||||
}
|
||||
},
|
||||
{ immediate: true },
|
||||
)
|
||||
|
||||
onBeforeUnmount(() => {
|
||||
if (props.preview) {
|
||||
document.body.style.overflow = ''
|
||||
}
|
||||
})
|
||||
</script>
|
||||
|
||||
<style scoped>
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
import { mount } from '@vue/test-utils'
|
||||
import { createPinia, setActivePinia } from 'pinia'
|
||||
import { readFileSync } from 'node:fs'
|
||||
import { resolve } from 'node:path'
|
||||
|
||||
import AnnouncementPopup from '../AnnouncementPopup.vue'
|
||||
import { useAnnouncementStore } from '@/stores/announcements'
|
||||
|
||||
const announcementMarkdownStyles = readFileSync(
|
||||
resolve(process.cwd(), 'src/styles/announcement-markdown.css'),
|
||||
'utf8',
|
||||
)
|
||||
|
||||
vi.mock('vue-i18n', async () => {
|
||||
const actual = await vi.importActual<typeof import('vue-i18n')>('vue-i18n')
|
||||
return {
|
||||
...actual,
|
||||
useI18n: () => ({
|
||||
t: (key: string) => key,
|
||||
}),
|
||||
}
|
||||
})
|
||||
|
||||
const announcement = {
|
||||
id: 1,
|
||||
title: 'Preview announcement',
|
||||
content: '## Preview heading\n\n<div>HTML content</div><script>window.__xss = true</script>',
|
||||
status: 'draft' as const,
|
||||
notify_mode: 'popup' as const,
|
||||
targeting: { any_of: [] },
|
||||
created_at: '2026-07-24T07:30:00Z',
|
||||
updated_at: '2026-07-24T07:30:00Z',
|
||||
}
|
||||
|
||||
describe('AnnouncementPopup', () => {
|
||||
beforeEach(() => {
|
||||
setActivePinia(createPinia())
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
document.body.innerHTML = ''
|
||||
document.body.style.overflow = ''
|
||||
})
|
||||
|
||||
it('renders mixed Markdown and HTML inside the shared styled container', async () => {
|
||||
const store = useAnnouncementStore()
|
||||
store.currentPopup = {
|
||||
id: 1,
|
||||
title: 'Mixed content announcement',
|
||||
content: [
|
||||
'## Markdown heading',
|
||||
'',
|
||||
'<div><h3>HTML heading</h3><ul><li>HTML list item</li></ul></div>',
|
||||
'',
|
||||
'<table><thead><tr><th>Status</th></tr></thead><tbody><tr><td>OK</td></tr></tbody></table>',
|
||||
'<script>window.__announcementXss = true</script>',
|
||||
].join('\n'),
|
||||
notify_mode: 'popup',
|
||||
created_at: '2026-07-24T07:30:00Z',
|
||||
updated_at: '2026-07-24T07:30:00Z',
|
||||
}
|
||||
|
||||
const wrapper = mount(AnnouncementPopup)
|
||||
await wrapper.vm.$nextTick()
|
||||
|
||||
const content = document.body.querySelector('.markdown-body')
|
||||
expect(content?.querySelector('h2')?.textContent).toBe('Markdown heading')
|
||||
expect(content?.querySelector('h3')?.textContent).toBe('HTML heading')
|
||||
expect(content?.querySelector('li')?.textContent).toBe('HTML list item')
|
||||
expect(content?.querySelector('table td')?.textContent).toBe('OK')
|
||||
expect(content?.querySelector('script')).toBeNull()
|
||||
|
||||
wrapper.unmount()
|
||||
})
|
||||
|
||||
it.each(['h2', 'h3', 'ul', 'li', 'blockquote', 'table', 'th', 'td', 'code'])(
|
||||
'loads a shared style rule for mixed-content <%s> elements',
|
||||
(element) => {
|
||||
expect(announcementMarkdownStyles).toContain(`.markdown-body ${element}`)
|
||||
},
|
||||
)
|
||||
|
||||
it('previews an admin announcement without marking it as read', async () => {
|
||||
const store = useAnnouncementStore()
|
||||
const dismissPopup = vi.spyOn(store, 'dismissPopup')
|
||||
const wrapper = mount(AnnouncementPopup, {
|
||||
props: {
|
||||
announcement,
|
||||
preview: true,
|
||||
},
|
||||
})
|
||||
|
||||
expect(document.body.textContent).toContain('Preview announcement')
|
||||
expect(document.body.querySelector('.markdown-body h2')?.textContent).toBe('Preview heading')
|
||||
expect(document.body.querySelector('.markdown-body script')).toBeNull()
|
||||
expect(document.body.textContent).toContain('common.close')
|
||||
|
||||
const dismissButton = document.body.querySelector<HTMLButtonElement>(
|
||||
'[data-testid="announcement-popup-dismiss"]',
|
||||
)
|
||||
dismissButton?.click()
|
||||
await wrapper.vm.$nextTick()
|
||||
|
||||
expect(wrapper.emitted('close')).toHaveLength(1)
|
||||
expect(dismissPopup).not.toHaveBeenCalled()
|
||||
|
||||
await wrapper.setProps({ announcement: null })
|
||||
expect(document.body.style.overflow).toBe('')
|
||||
wrapper.unmount()
|
||||
})
|
||||
|
||||
it('keeps the existing user popup dismissal behavior', async () => {
|
||||
const store = useAnnouncementStore()
|
||||
store.currentPopup = announcement
|
||||
const dismissPopup = vi.spyOn(store, 'dismissPopup').mockResolvedValue()
|
||||
const wrapper = mount(AnnouncementPopup)
|
||||
|
||||
const dismissButton = document.body.querySelector<HTMLButtonElement>(
|
||||
'[data-testid="announcement-popup-dismiss"]',
|
||||
)
|
||||
dismissButton?.click()
|
||||
await wrapper.vm.$nextTick()
|
||||
|
||||
expect(dismissPopup).toHaveBeenCalledTimes(1)
|
||||
expect(wrapper.emitted('close')).toBeUndefined()
|
||||
wrapper.unmount()
|
||||
})
|
||||
})
|
||||
@@ -433,6 +433,7 @@ export default {
|
||||
timeImmediate: 'Immediate',
|
||||
timeNever: 'Never',
|
||||
readStatus: 'Read Status',
|
||||
preview: 'Preview',
|
||||
eligible: 'Eligible',
|
||||
readAt: 'Read at',
|
||||
unread: 'Unread',
|
||||
|
||||
@@ -430,6 +430,7 @@ export default {
|
||||
timeImmediate: '立即',
|
||||
timeNever: '永久',
|
||||
readStatus: '已读情况',
|
||||
preview: '预览',
|
||||
eligible: '符合条件',
|
||||
readAt: '已读时间',
|
||||
unread: '未读',
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
/* Shared rich-text styles for announcement popup and bell detail content. */
|
||||
.markdown-body {
|
||||
@apply text-[15px] leading-[1.75];
|
||||
@apply text-gray-700 dark:text-gray-300;
|
||||
}
|
||||
|
||||
.markdown-body h1 {
|
||||
@apply mb-6 mt-8 border-b border-gray-200 pb-3 text-3xl font-bold text-gray-900 dark:border-dark-600 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body h2 {
|
||||
@apply mb-4 mt-7 border-b border-gray-100 pb-2 text-2xl font-bold text-gray-900 dark:border-dark-700 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body h3 {
|
||||
@apply mb-3 mt-6 text-xl font-semibold text-gray-900 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body h4 {
|
||||
@apply mb-2 mt-5 text-lg font-semibold text-gray-900 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body p {
|
||||
@apply mb-4 leading-relaxed;
|
||||
}
|
||||
|
||||
.markdown-body a {
|
||||
@apply font-medium text-blue-600 underline decoration-blue-600/30 decoration-2 underline-offset-2 transition-all hover:decoration-blue-600 dark:text-blue-400 dark:decoration-blue-400/30 dark:hover:decoration-blue-400;
|
||||
}
|
||||
|
||||
.markdown-body ul,
|
||||
.markdown-body ol {
|
||||
@apply mb-4 ml-6 space-y-2;
|
||||
}
|
||||
|
||||
.markdown-body ul {
|
||||
@apply list-disc;
|
||||
}
|
||||
|
||||
.markdown-body ol {
|
||||
@apply list-decimal;
|
||||
}
|
||||
|
||||
.markdown-body li {
|
||||
@apply pl-2 leading-relaxed;
|
||||
}
|
||||
|
||||
.markdown-body li::marker {
|
||||
@apply text-blue-600 dark:text-blue-400;
|
||||
}
|
||||
|
||||
.markdown-body blockquote {
|
||||
@apply relative my-5 border-l-4 border-blue-500 bg-blue-50/50 py-3 pl-5 pr-4 italic text-gray-700 dark:border-blue-400 dark:bg-blue-900/10 dark:text-gray-300;
|
||||
}
|
||||
|
||||
.markdown-body blockquote::before {
|
||||
content: '"';
|
||||
@apply absolute -left-1 top-0 font-serif text-5xl text-blue-500/20 dark:text-blue-400/20;
|
||||
}
|
||||
|
||||
.markdown-body code {
|
||||
@apply rounded-lg bg-gray-100 px-2 py-1 font-mono text-[13px] text-pink-600 dark:bg-dark-700 dark:text-pink-400;
|
||||
}
|
||||
|
||||
.markdown-body pre {
|
||||
@apply my-5 overflow-x-auto rounded-xl border border-gray-200 bg-gray-50 p-5 dark:border-dark-600 dark:bg-dark-900/50;
|
||||
}
|
||||
|
||||
.markdown-body pre code {
|
||||
@apply bg-transparent p-0 text-[13px] text-gray-800 dark:text-gray-200;
|
||||
}
|
||||
|
||||
.markdown-body hr {
|
||||
@apply my-8 border-0 border-t-2 border-gray-200 dark:border-dark-700;
|
||||
}
|
||||
|
||||
.markdown-body table {
|
||||
@apply mb-5 w-full overflow-hidden rounded-lg border border-gray-200 dark:border-dark-600;
|
||||
}
|
||||
|
||||
.markdown-body th,
|
||||
.markdown-body td {
|
||||
@apply border-b border-r border-gray-200 px-4 py-3 text-left dark:border-dark-600;
|
||||
}
|
||||
|
||||
.markdown-body th:last-child,
|
||||
.markdown-body td:last-child {
|
||||
@apply border-r-0;
|
||||
}
|
||||
|
||||
.markdown-body tr:last-child td {
|
||||
@apply border-b-0;
|
||||
}
|
||||
|
||||
.markdown-body th {
|
||||
@apply bg-gradient-to-br from-blue-50 to-indigo-50 font-semibold text-gray-900 dark:from-blue-900/20 dark:to-indigo-900/10 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body tbody tr {
|
||||
@apply transition-colors hover:bg-gray-50 dark:hover:bg-dark-700/30;
|
||||
}
|
||||
|
||||
.markdown-body img {
|
||||
@apply my-5 max-w-full rounded-xl border border-gray-200 shadow-md dark:border-dark-600;
|
||||
}
|
||||
|
||||
.markdown-body strong {
|
||||
@apply font-semibold text-gray-900 dark:text-white;
|
||||
}
|
||||
|
||||
.markdown-body em {
|
||||
@apply italic text-gray-600 dark:text-gray-400;
|
||||
}
|
||||
@@ -114,12 +114,19 @@
|
||||
|
||||
<template #cell-actions="{ row }">
|
||||
<div class="flex items-center space-x-1">
|
||||
<button
|
||||
@click="openPreview(row)"
|
||||
class="flex flex-col items-center gap-0.5 rounded-lg p-1.5 text-gray-500 transition-colors hover:bg-blue-50 hover:text-blue-600 dark:hover:bg-blue-900/20 dark:hover:text-blue-400"
|
||||
:title="t('admin.announcements.preview')"
|
||||
>
|
||||
<Icon name="eye" size="sm" />
|
||||
</button>
|
||||
<button
|
||||
@click="openReadStatus(row)"
|
||||
class="flex flex-col items-center gap-0.5 rounded-lg p-1.5 text-gray-500 transition-colors hover:bg-blue-50 hover:text-blue-600 dark:hover:bg-blue-900/20 dark:hover:text-blue-400"
|
||||
:title="t('admin.announcements.readStatus')"
|
||||
>
|
||||
<Icon name="eye" size="sm" />
|
||||
<Icon name="chartBar" size="sm" />
|
||||
</button>
|
||||
<button
|
||||
@click="openEditDialog(row)"
|
||||
@@ -240,6 +247,12 @@
|
||||
:announcement-id="readStatusAnnouncementId"
|
||||
@close="showReadStatusDialog = false"
|
||||
/>
|
||||
|
||||
<AnnouncementPopup
|
||||
:announcement="previewAnnouncement"
|
||||
preview
|
||||
@close="previewAnnouncement = null"
|
||||
/>
|
||||
</AppLayout>
|
||||
</template>
|
||||
|
||||
@@ -265,6 +278,7 @@ import Icon from '@/components/icons/Icon.vue'
|
||||
|
||||
import AnnouncementTargetingEditor from '@/components/admin/announcements/AnnouncementTargetingEditor.vue'
|
||||
import AnnouncementReadStatusDialog from '@/components/admin/announcements/AnnouncementReadStatusDialog.vue'
|
||||
import AnnouncementPopup from '@/components/common/AnnouncementPopup.vue'
|
||||
|
||||
const { t } = useI18n()
|
||||
const appStore = useAppStore()
|
||||
@@ -588,6 +602,11 @@ async function confirmDelete() {
|
||||
// ===== Read status =====
|
||||
const showReadStatusDialog = ref(false)
|
||||
const readStatusAnnouncementId = ref<number | null>(null)
|
||||
const previewAnnouncement = ref<Announcement | null>(null)
|
||||
|
||||
function openPreview(row: Announcement) {
|
||||
previewAnnouncement.value = row
|
||||
}
|
||||
|
||||
function openReadStatus(row: Announcement) {
|
||||
readStatusAnnouncementId.value = row.id
|
||||
|
||||
@@ -51,9 +51,9 @@
|
||||
<div class="mt-5 grid gap-4 md:grid-cols-2">
|
||||
<div class="space-y-2">
|
||||
<p class="text-sm font-medium text-gray-700 dark:text-gray-300">{{ t('affiliate.yourCode') }}</p>
|
||||
<div class="flex items-center gap-2 rounded-xl border border-gray-200 bg-gray-50 px-3 py-2 dark:border-dark-700 dark:bg-dark-900">
|
||||
<code class="flex-1 truncate text-sm font-semibold text-gray-900 dark:text-white">{{ detail.aff_code }}</code>
|
||||
<button class="btn btn-secondary btn-sm" @click="copyCode">
|
||||
<div class="flex flex-col items-stretch gap-2 rounded-xl border border-gray-200 bg-gray-50 px-3 py-2 dark:border-dark-700 dark:bg-dark-900 sm:flex-row sm:items-center">
|
||||
<code class="min-w-0 break-all text-sm font-semibold text-gray-900 dark:text-white sm:flex-1 sm:truncate">{{ detail.aff_code }}</code>
|
||||
<button class="btn btn-secondary btn-sm w-full sm:w-auto sm:shrink-0" @click="copyCode">
|
||||
<Icon name="copy" size="sm" />
|
||||
<span>{{ t('affiliate.copyCode') }}</span>
|
||||
</button>
|
||||
@@ -62,9 +62,9 @@
|
||||
|
||||
<div class="space-y-2">
|
||||
<p class="text-sm font-medium text-gray-700 dark:text-gray-300">{{ t('affiliate.inviteLink') }}</p>
|
||||
<div class="flex items-center gap-2 rounded-xl border border-gray-200 bg-gray-50 px-3 py-2 dark:border-dark-700 dark:bg-dark-900">
|
||||
<code class="flex-1 truncate text-sm text-gray-700 dark:text-gray-300">{{ inviteLink }}</code>
|
||||
<button class="btn btn-secondary btn-sm" @click="copyInviteLink">
|
||||
<div class="flex flex-col items-stretch gap-2 rounded-xl border border-gray-200 bg-gray-50 px-3 py-2 dark:border-dark-700 dark:bg-dark-900 sm:flex-row sm:items-center">
|
||||
<code class="min-w-0 break-all text-sm text-gray-700 dark:text-gray-300 sm:flex-1 sm:truncate">{{ inviteLink }}</code>
|
||||
<button class="btn btn-secondary btn-sm w-full sm:w-auto sm:shrink-0" @click="copyInviteLink">
|
||||
<Icon name="copy" size="sm" />
|
||||
<span>{{ t('affiliate.copyLink') }}</span>
|
||||
</button>
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
import { flushPromises, mount } from '@vue/test-utils'
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import AffiliateView from '../AffiliateView.vue'
|
||||
|
||||
const { copyToClipboard, getAffiliateDetail } = vi.hoisted(() => ({
|
||||
copyToClipboard: vi.fn(),
|
||||
getAffiliateDetail: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/api/user', () => ({
|
||||
default: {
|
||||
getAffiliateDetail,
|
||||
transferAffiliateQuota: vi.fn(),
|
||||
},
|
||||
}))
|
||||
|
||||
vi.mock('@/stores/app', () => ({
|
||||
useAppStore: () => ({
|
||||
showError: vi.fn(),
|
||||
showSuccess: vi.fn(),
|
||||
}),
|
||||
}))
|
||||
|
||||
vi.mock('@/stores/auth', () => ({
|
||||
useAuthStore: () => ({
|
||||
refreshUser: vi.fn(),
|
||||
}),
|
||||
}))
|
||||
|
||||
vi.mock('@/composables/useClipboard', () => ({
|
||||
useClipboard: () => ({ copyToClipboard }),
|
||||
}))
|
||||
|
||||
vi.mock('vue-i18n', async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import('vue-i18n')>()
|
||||
return {
|
||||
...actual,
|
||||
useI18n: () => ({
|
||||
t: (key: string) => key,
|
||||
}),
|
||||
}
|
||||
})
|
||||
|
||||
describe('AffiliateView', () => {
|
||||
const affiliateCode = 'affiliate-code-that-is-long-enough-to-overflow-a-mobile-viewport'
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
copyToClipboard.mockResolvedValue(true)
|
||||
getAffiliateDetail.mockResolvedValue({
|
||||
user_id: 1,
|
||||
aff_code: affiliateCode,
|
||||
inviter_id: null,
|
||||
aff_count: 0,
|
||||
aff_quota: 0,
|
||||
aff_frozen_quota: 0,
|
||||
aff_history_quota: 0,
|
||||
effective_rebate_rate_percent: 10,
|
||||
invitees: [],
|
||||
})
|
||||
})
|
||||
|
||||
it('stacks long values and copy controls on mobile while retaining desktop rows', async () => {
|
||||
const wrapper = mount(AffiliateView, {
|
||||
global: {
|
||||
stubs: {
|
||||
AppLayout: { template: '<main><slot /></main>' },
|
||||
Icon: true,
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
await flushPromises()
|
||||
|
||||
const values = wrapper.findAll('code')
|
||||
expect(values).toHaveLength(2)
|
||||
for (const value of values) {
|
||||
expect(value.classes()).toEqual(expect.arrayContaining([
|
||||
'min-w-0',
|
||||
'break-all',
|
||||
'sm:flex-1',
|
||||
'sm:truncate',
|
||||
]))
|
||||
expect(Array.from(value.element.parentElement?.classList ?? [])).toEqual(expect.arrayContaining([
|
||||
'flex-col',
|
||||
'items-stretch',
|
||||
'sm:flex-row',
|
||||
'sm:items-center',
|
||||
]))
|
||||
}
|
||||
|
||||
const copyButtons = wrapper.findAll('button').filter((button) =>
|
||||
['affiliate.copyCode', 'affiliate.copyLink'].includes(button.text()),
|
||||
)
|
||||
expect(copyButtons).toHaveLength(2)
|
||||
for (const button of copyButtons) {
|
||||
expect(button.classes()).toEqual(expect.arrayContaining([
|
||||
'w-full',
|
||||
'sm:w-auto',
|
||||
'sm:shrink-0',
|
||||
]))
|
||||
}
|
||||
|
||||
await copyButtons[0].trigger('click')
|
||||
await copyButtons[1].trigger('click')
|
||||
await flushPromises()
|
||||
|
||||
expect(copyToClipboard).toHaveBeenNthCalledWith(1, affiliateCode, 'affiliate.codeCopied')
|
||||
expect(copyToClipboard).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
`${window.location.origin}/register?aff=${encodeURIComponent(affiliateCode)}`,
|
||||
'affiliate.linkCopied',
|
||||
)
|
||||
})
|
||||
})
|
||||
Reference in new issue
Block a user