mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Compare commits
108
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
017bbcb901 | ||
|
|
9ecb340822 | ||
|
|
a0f41f95a0 | ||
|
|
a60a29549f | ||
|
|
2f3fed2fdb | ||
|
|
fe9dadfc7d | ||
|
|
41dcfec489 | ||
|
|
052e08e764 | ||
|
|
1f955af353 | ||
|
|
ccb027b515 | ||
|
|
f50b99b66c | ||
|
|
883a53fea0 | ||
|
|
b35a611538 | ||
|
|
1dca86f287 | ||
|
|
dd6cb410e6 | ||
|
|
27a9421e6a | ||
|
|
860782e6fb | ||
|
|
35232c3949 | ||
|
|
421da3f0d4 | ||
|
|
48970f61b4 | ||
|
|
f8984b8a93 | ||
|
|
6a8f921650 | ||
|
|
2cc459e24f | ||
|
|
8490a81867 | ||
|
|
6073704fd7 | ||
|
|
6a69dd051a | ||
|
|
2840cdeced | ||
|
|
aa3a3ea787 | ||
|
|
bd48b2b2d2 | ||
|
|
9a62841fd1 | ||
|
|
4c00df2e01 | ||
|
|
b5298fdde8 | ||
|
|
5a350108c4 | ||
|
|
1df5c39311 | ||
|
|
b3c7038cf4 | ||
|
|
8616d4e036 | ||
|
|
93c3da03eb | ||
|
|
68c2f5ae6d | ||
|
|
ff6b85e30b | ||
|
|
79c18ec836 | ||
|
|
790723e040 | ||
|
|
1ff2bd6d95 | ||
|
|
f592b336a3 | ||
|
|
37f35b12ea | ||
|
|
e4b34446af | ||
|
|
90ae81e9f8 | ||
|
|
9375e288ba | ||
|
|
06f70a6232 | ||
|
|
6906062187 | ||
|
|
dc4874a2f1 | ||
|
|
11608c51f0 | ||
|
|
8aa7fb8c2d | ||
|
|
00ae7c38e2 | ||
|
|
51a10e42f7 | ||
|
|
6655f4ec0e | ||
|
|
be64b553f9 | ||
|
|
66f07efd26 | ||
|
|
89e9976950 | ||
|
|
99d86e41ac | ||
|
|
ec334fda2c | ||
|
|
35d2ef0cbd | ||
|
|
512aa8bc93 | ||
|
|
f45126c8ec | ||
|
|
82d2552360 | ||
|
|
f652d0ed38 | ||
|
|
6b40449c3f | ||
|
|
02088be2d1 | ||
|
|
14483c9254 | ||
|
|
fd7af1d462 | ||
|
|
8b95dc94f8 | ||
|
|
7626110d0b | ||
|
|
3c5ea297ff | ||
|
|
994048e6a4 | ||
|
|
c802c45e37 | ||
|
|
f4f8ff04db | ||
|
|
510b680fca | ||
|
|
afbe51abf8 | ||
|
|
52c2f4cb10 | ||
|
|
630639f747 | ||
|
|
11ffee5543 | ||
|
|
85810365a6 | ||
|
|
f0ebc183e1 | ||
|
|
ca2746f3bb | ||
|
|
b31b435091 | ||
|
|
1d126c833b | ||
|
|
9461196864 | ||
|
|
a2007db707 | ||
|
|
97bdde3135 | ||
|
|
090b2cc000 | ||
|
|
82b092ffc2 | ||
|
|
5979521309 | ||
|
|
1ca7b2a519 | ||
|
|
00f6f4e4a8 | ||
|
|
e2c35cf394 | ||
|
|
a3eb7ef302 | ||
|
|
2fc91663d2 | ||
|
|
453b796df0 | ||
|
|
e878ab710f | ||
|
|
13fcfc0f3b | ||
|
|
0416c9ede4 | ||
|
|
ecba49f711 | ||
|
|
17ce21c038 | ||
|
|
b5b5f1f94b | ||
|
|
1af3269adb | ||
|
|
f867e3e2ec | ||
|
|
ece1ecccca | ||
|
|
b13200d7a1 | ||
|
|
ef51690b5a |
@@ -1 +1 @@
|
||||
0.2.7
|
||||
0.2.10
|
||||
|
||||
@@ -288,14 +288,15 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
auditLogHandler := admin.NewAuditLogHandler(auditLogService, totpService)
|
||||
upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db)
|
||||
openCodeGoUsageService := service.ProvideOpenCodeGoUsageService(accountRepository, httpUpstream, settingService, leaderLockCache, db)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService)
|
||||
claudeResetCreditService := service.NewClaudeResetCreditService(accountRepository, claudeTokenProvider, proxyRepository, settingService)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, claudeResetCreditService)
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig)
|
||||
legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService)
|
||||
coordinator := securityaudit.NewCoordinator(legacyEngine, promptService)
|
||||
gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator, compositeRouteResolver)
|
||||
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
|
||||
totpHandler := handler.NewTotpHandler(totpService)
|
||||
passkeyRepository := repository.NewPasskeyRepository(db)
|
||||
|
||||
@@ -89,6 +89,7 @@ type Config struct {
|
||||
Pricing PricingConfig `mapstructure:"pricing"`
|
||||
Gateway GatewayConfig `mapstructure:"gateway"`
|
||||
APIKeyAuth APIKeyAuthCacheConfig `mapstructure:"api_key_auth_cache"`
|
||||
APIKeyCreate APIKeyCreateConfig `mapstructure:"api_key_create"`
|
||||
SubscriptionCache SubscriptionCacheConfig `mapstructure:"subscription_cache"`
|
||||
SubscriptionMaintenance SubscriptionMaintenanceConfig `mapstructure:"subscription_maintenance"`
|
||||
Dashboard DashboardCacheConfig `mapstructure:"dashboard_cache"`
|
||||
@@ -1696,6 +1697,14 @@ type APIKeyAuthCacheConfig struct {
|
||||
InvalidAbuse InvalidAuthAbuseConfig `mapstructure:"invalid_abuse"`
|
||||
}
|
||||
|
||||
// APIKeyCreateConfig 用户创建 API Key 的防滥用限制(0 表示不限制)
|
||||
type APIKeyCreateConfig struct {
|
||||
// MaxActivePerUser 单个用户同时存在(未删除)的 API Key 上限
|
||||
MaxActivePerUser int `mapstructure:"max_active_per_user"`
|
||||
// MaxPerUserPerHour 单个用户每小时可创建的 API Key 次数(删除不返还次数)
|
||||
MaxPerUserPerHour int `mapstructure:"max_per_user_per_hour"`
|
||||
}
|
||||
|
||||
type InvalidAuthAbuseConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
Threshold int `mapstructure:"threshold"`
|
||||
@@ -2332,6 +2341,8 @@ func setDefaults() {
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.window_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.block_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.capacity", 16384)
|
||||
viper.SetDefault("api_key_create.max_active_per_user", 200)
|
||||
viper.SetDefault("api_key_create.max_per_user_per_hour", 60)
|
||||
|
||||
// Subscription auth L1 cache
|
||||
viper.SetDefault("subscription_cache.l1_size", 16384)
|
||||
@@ -2711,6 +2722,12 @@ func (c *Config) Validate() error {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_stream must be positive")
|
||||
}
|
||||
}
|
||||
if c.APIKeyCreate.MaxActivePerUser < 0 {
|
||||
return fmt.Errorf("api_key_create.max_active_per_user must be non-negative")
|
||||
}
|
||||
if c.APIKeyCreate.MaxPerUserPerHour < 0 {
|
||||
return fmt.Errorf("api_key_create.max_per_user_per_hour must be non-negative")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.Enabled {
|
||||
if c.APIKeyAuth.InvalidAbuse.Threshold < 10 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.threshold must be at least 10")
|
||||
|
||||
@@ -190,6 +190,8 @@ var DefaultBedrockModelMapping = map[string]string{
|
||||
"claude-opus-4-1": "us.anthropic.claude-opus-4-1-20250805-v1:0",
|
||||
"claude-opus-4-20250514": "us.anthropic.claude-opus-4-20250514-v1:0",
|
||||
// Claude Sonnet
|
||||
// Sonnet 5.5 is available on bedrock-runtime through Global inference only.
|
||||
"claude-sonnet-5-5": "global.anthropic.claude-sonnet-5-5",
|
||||
"claude-sonnet-5": "us.anthropic.claude-sonnet-5-v1",
|
||||
"claude-sonnet-4-6-thinking": "us.anthropic.claude-sonnet-4-6",
|
||||
"claude-sonnet-4-6": "us.anthropic.claude-sonnet-4-6",
|
||||
|
||||
@@ -109,9 +109,10 @@ func TestDefaultBedrockModelMapping_ContainsNewClaudeModels(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := map[string]string{
|
||||
"claude-fable-5-1": "anthropic.claude-fable-5-1",
|
||||
"claude-fable-5": "anthropic.claude-fable-5",
|
||||
"claude-opus-4-8": "us.anthropic.claude-opus-4-8-v1",
|
||||
"claude-fable-5-1": "anthropic.claude-fable-5-1",
|
||||
"claude-fable-5": "anthropic.claude-fable-5",
|
||||
"claude-opus-4-8": "us.anthropic.claude-opus-4-8-v1",
|
||||
"claude-sonnet-5-5": "global.anthropic.claude-sonnet-5-5",
|
||||
}
|
||||
for from, want := range cases {
|
||||
got, ok := DefaultBedrockModelMapping[from]
|
||||
|
||||
@@ -48,6 +48,7 @@ func NewOAuthHandler(oauthService *service.OAuthService) *OAuthHandler {
|
||||
|
||||
// AccountHandler handles admin account management
|
||||
type AccountHandler struct {
|
||||
claudeResetCredits claudeResetReader
|
||||
adminService service.AdminService
|
||||
oauthService *service.OAuthService
|
||||
openaiOAuthService *service.OpenAIOAuthService
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
type claudeResetReader interface {
|
||||
Query(context.Context, int64) (*service.ClaudeResetCredits, error)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) SetClaudeResetCreditService(s *service.ClaudeResetCreditService) {
|
||||
h.claudeResetCredits = s
|
||||
}
|
||||
func (h *AccountHandler) ClaudeResetCredits(c *gin.Context) {
|
||||
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h.claudeResetCredits == nil {
|
||||
response.Error(c, 503, "Claude reset service unavailable")
|
||||
return
|
||||
}
|
||||
status, err := h.claudeResetCredits.Query(c.Request.Context(), id)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, status)
|
||||
}
|
||||
@@ -500,7 +500,11 @@ func (h *DashboardHandler) GetUserUsageTrend(c *gin.Context) {
|
||||
limit = 12
|
||||
}
|
||||
|
||||
trend, hit, err := h.getUserUsageTrendCached(c.Request.Context(), startTime, endTime, granularity, limit)
|
||||
metric := c.DefaultQuery("metric", "tokens")
|
||||
if metric != "tokens" && metric != "actual_cost" {
|
||||
metric = "tokens"
|
||||
}
|
||||
trend, hit, err := h.getUserUsageTrendCached(c.Request.Context(), startTime, endTime, granularity, limit, metric)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get user usage trend")
|
||||
return
|
||||
|
||||
@@ -45,6 +45,7 @@ func (r *dashboardUsageRepoCacheProbe) GetUserUsageTrend(
|
||||
startTime, endTime time.Time,
|
||||
granularity string,
|
||||
limit int,
|
||||
metric string,
|
||||
) ([]usagestats.UserUsageTrendPoint, error) {
|
||||
r.usersTrendCalls.Add(1)
|
||||
return []usagestats.UserUsageTrendPoint{{
|
||||
@@ -115,4 +116,20 @@ func TestDashboardHandler_GetUserUsageTrend_UsesCache(t *testing.T) {
|
||||
require.Equal(t, http.StatusOK, rec2.Code)
|
||||
require.Equal(t, "hit", rec2.Header().Get("X-Snapshot-Cache"))
|
||||
require.Equal(t, int32(1), repo.usersTrendCalls.Load())
|
||||
|
||||
for _, tc := range []struct {
|
||||
metric, cache string
|
||||
calls int32
|
||||
}{
|
||||
{"actual_cost", "miss", 2},
|
||||
{"actual_cost", "hit", 2},
|
||||
{"tokens", "hit", 2},
|
||||
} {
|
||||
req := httptest.NewRequest(http.MethodGet, "/admin/dashboard/users-trend?start_date=2026-03-01&end_date=2026-03-07&granularity=day&limit=8&metric="+tc.metric, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, tc.cache, rec.Header().Get("X-Snapshot-Cache"))
|
||||
require.Equal(t, tc.calls, repo.usersTrendCalls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -53,6 +53,7 @@ type dashboardEntityTrendCacheKey struct {
|
||||
EndTime string `json:"end_time"`
|
||||
Granularity string `json:"granularity"`
|
||||
Limit int `json:"limit"`
|
||||
Metric string `json:"metric,omitempty"`
|
||||
}
|
||||
|
||||
func cacheStatusValue(hit bool) string {
|
||||
@@ -213,15 +214,16 @@ func (h *DashboardHandler) getAPIKeyUsageTrendCached(ctx context.Context, startT
|
||||
return trend, hit, err
|
||||
}
|
||||
|
||||
func (h *DashboardHandler) getUserUsageTrendCached(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.UserUsageTrendPoint, bool, error) {
|
||||
func (h *DashboardHandler) getUserUsageTrendCached(ctx context.Context, startTime, endTime time.Time, granularity string, limit int, metric string) ([]usagestats.UserUsageTrendPoint, bool, error) {
|
||||
key := mustMarshalDashboardCacheKey(dashboardEntityTrendCacheKey{
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
Granularity: granularity,
|
||||
Limit: limit,
|
||||
Metric: metric,
|
||||
})
|
||||
entry, hit, err := dashboardUsersTrendCache.GetOrLoad(key, func() (any, error) {
|
||||
return h.dashboardService.GetUserUsageTrend(ctx, startTime, endTime, granularity, limit)
|
||||
return h.dashboardService.GetUserUsageTrend(ctx, startTime, endTime, granularity, limit, metric)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, hit, err
|
||||
|
||||
@@ -243,7 +243,7 @@ func (h *DashboardHandler) buildSnapshotV2Response(
|
||||
}
|
||||
|
||||
if includeUsersTrend {
|
||||
usersTrend, _, err := h.getUserUsageTrendCached(ctx, startTime, endTime, granularity, usersTrendLimit)
|
||||
usersTrend, _, err := h.getUserUsageTrendCached(ctx, startTime, endTime, granularity, usersTrendLimit, "tokens")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to get user usage trend")
|
||||
}
|
||||
|
||||
@@ -265,6 +265,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
DefaultBalance: settings.DefaultBalance,
|
||||
RiskControlEnabled: settings.RiskControlEnabled,
|
||||
CyberSessionBlockEnabled: settings.CyberSessionBlockEnabled,
|
||||
CyberPolicyUserAllowlist: settings.CyberPolicyUserAllowlist,
|
||||
CyberSessionBlockTTLSeconds: settings.CyberSessionBlockTTLSeconds,
|
||||
AffiliateRebateRate: settings.AffiliateRebateRate,
|
||||
AffiliateRebateFreezeHours: settings.AffiliateRebateFreezeHours,
|
||||
|
||||
@@ -603,6 +603,9 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.RiskControlEnabled != after.RiskControlEnabled {
|
||||
changed = append(changed, "risk_control_enabled")
|
||||
}
|
||||
if before.CyberPolicyUserAllowlist != after.CyberPolicyUserAllowlist {
|
||||
changed = append(changed, "cyber_policy_user_allowlist")
|
||||
}
|
||||
if before.CyberSessionBlockEnabled != after.CyberSessionBlockEnabled {
|
||||
changed = append(changed, "cyber_session_block_enabled")
|
||||
}
|
||||
|
||||
@@ -364,8 +364,9 @@ type UpdateSettingsRequest struct {
|
||||
RiskControlEnabled *bool `json:"risk_control_enabled"`
|
||||
|
||||
// cyber 会话屏蔽开关 + TTL
|
||||
CyberSessionBlockEnabled *bool `json:"cyber_session_block_enabled"`
|
||||
CyberSessionBlockTTLSeconds *int `json:"cyber_session_block_ttl_seconds"`
|
||||
CyberSessionBlockEnabled *bool `json:"cyber_session_block_enabled"`
|
||||
CyberPolicyUserAllowlist *string `json:"cyber_policy_user_allowlist"`
|
||||
CyberSessionBlockTTLSeconds *int `json:"cyber_session_block_ttl_seconds"`
|
||||
|
||||
// OpenAI fast/flex policy (optional, only updated when provided)
|
||||
OpenAIFastPolicySettings *dto.OpenAIFastPolicySettings `json:"openai_fast_policy_settings,omitempty"`
|
||||
@@ -1503,6 +1504,13 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
if req.CyberPolicyUserAllowlist != nil {
|
||||
if _, err := service.ParseCyberPolicyUserAllowlist(*req.CyberPolicyUserAllowlist); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// cyber 会话屏蔽 TTL 校验:提供时必须 > 0
|
||||
if req.CyberSessionBlockTTLSeconds != nil && *req.CyberSessionBlockTTLSeconds <= 0 {
|
||||
response.BadRequest(c, "cyber_session_block_ttl_seconds must be > 0")
|
||||
@@ -2008,6 +2016,12 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
return previousSettings.RiskControlEnabled
|
||||
}(),
|
||||
CyberPolicyUserAllowlist: func() string {
|
||||
if req.CyberPolicyUserAllowlist != nil {
|
||||
return *req.CyberPolicyUserAllowlist
|
||||
}
|
||||
return previousSettings.CyberPolicyUserAllowlist
|
||||
}(),
|
||||
CyberSessionBlockEnabled: func() bool {
|
||||
if req.CyberSessionBlockEnabled != nil {
|
||||
return *req.CyberSessionBlockEnabled
|
||||
@@ -2429,6 +2443,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
|
||||
RiskControlEnabled: updatedSettings.RiskControlEnabled,
|
||||
CyberSessionBlockEnabled: updatedSettings.CyberSessionBlockEnabled,
|
||||
CyberPolicyUserAllowlist: updatedSettings.CyberPolicyUserAllowlist,
|
||||
CyberSessionBlockTTLSeconds: updatedSettings.CyberSessionBlockTTLSeconds,
|
||||
AccountSchedulingThresholds: updatedSettings.AccountSchedulingThresholds,
|
||||
AllowUserViewErrorRequests: updatedSettings.AllowUserViewErrorRequests,
|
||||
|
||||
@@ -2,10 +2,13 @@ package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -73,3 +76,65 @@ func TestConcurrencyErrorResponse(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoogleConcurrencyError(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
err error
|
||||
wantStatus int
|
||||
wantGStatus string
|
||||
wantMessage string
|
||||
}{
|
||||
{
|
||||
name: "client cancellation is 499",
|
||||
err: context.Canceled,
|
||||
wantStatus: statusClientClosedRequest,
|
||||
wantGStatus: "CANCELLED",
|
||||
wantMessage: "context canceled",
|
||||
},
|
||||
{
|
||||
name: "full wait queue is 429",
|
||||
err: &WaitQueueFullError{SlotType: "user"},
|
||||
wantStatus: http.StatusTooManyRequests,
|
||||
wantGStatus: "RESOURCE_EXHAUSTED",
|
||||
wantMessage: "Too many pending requests, please retry later",
|
||||
},
|
||||
{
|
||||
name: "slot wait timeout is 429",
|
||||
err: &ConcurrencyError{SlotType: "user", IsTimeout: true},
|
||||
wantStatus: http.StatusTooManyRequests,
|
||||
wantGStatus: "RESOURCE_EXHAUSTED",
|
||||
wantMessage: "Concurrency limit exceeded for user, please retry later",
|
||||
},
|
||||
{
|
||||
name: "acquire backend error is 503",
|
||||
err: errors.New("redis unavailable"),
|
||||
wantStatus: http.StatusServiceUnavailable,
|
||||
wantGStatus: "INTERNAL",
|
||||
wantMessage: "Service temporarily unavailable, please retry later",
|
||||
},
|
||||
}
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
googleConcurrencyError(c, tt.err, "user")
|
||||
|
||||
require.Equal(t, tt.wantStatus, rec.Code)
|
||||
var body struct {
|
||||
Error struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Status string `json:"status"`
|
||||
} `json:"error"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
|
||||
require.Equal(t, tt.wantStatus, body.Error.Code)
|
||||
require.Equal(t, tt.wantGStatus, body.Error.Status)
|
||||
require.Equal(t, tt.wantMessage, body.Error.Message)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -335,8 +335,9 @@ type SystemSettings struct {
|
||||
RiskControlEnabled bool `json:"risk_control_enabled"`
|
||||
|
||||
// cyber 会话屏蔽开关 + TTL
|
||||
CyberSessionBlockEnabled bool `json:"cyber_session_block_enabled"`
|
||||
CyberSessionBlockTTLSeconds int `json:"cyber_session_block_ttl_seconds"`
|
||||
CyberSessionBlockEnabled bool `json:"cyber_session_block_enabled"`
|
||||
CyberPolicyUserAllowlist string `json:"cyber_policy_user_allowlist"`
|
||||
CyberSessionBlockTTLSeconds int `json:"cyber_session_block_ttl_seconds"`
|
||||
|
||||
// Affiliate (邀请返利) feature switch
|
||||
AffiliateEnabled bool `json:"affiliate_enabled"`
|
||||
|
||||
@@ -330,6 +330,10 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gateway.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, service.PlatformGemini)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -660,6 +664,10 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
)
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, parsedReq.MetadataUserID, subject.UserID)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gateway.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, currentAPIKey, reqModel, reqModel, platform)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -1971,10 +1979,15 @@ func (h *GatewayHandler) handleStreamingAwareErrorWithCode(c *gin.Context, statu
|
||||
// Writer 已被写过时(ping 已 flush)走 streamStarted 分支,
|
||||
// 让 handleStreamingAwareError 通过 SSE 发协议合规的终止事件,
|
||||
// 否则下游收到的就是 silent EOF。
|
||||
// 客户端已断开时不补写,响应未提交则标记 499(见 failoverClientGone)。
|
||||
func (h *GatewayHandler) ensureForwardErrorResponse(c *gin.Context, streamStarted bool) bool {
|
||||
if c == nil || c.Writer == nil {
|
||||
return false
|
||||
}
|
||||
if c.Request != nil && errors.Is(c.Request.Context().Err(), context.Canceled) {
|
||||
failoverClientGone(c)
|
||||
return false
|
||||
}
|
||||
if service.IsResponseCommitted(c) {
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -93,6 +93,8 @@ func TestGatewayHandlerPreCancelledCompatibleRequestsDoNotSelectAccount(t *testi
|
||||
require.Zero(t, schedulerCache.snapshotCalls.Load(), "a cancelled request must stop before the account selector")
|
||||
_, selected := c.Get(opsAccountIDKey)
|
||||
require.False(t, selected)
|
||||
require.Equal(t, statusClientClosedRequest, c.Writer.Status(), "an uncommitted cancelled request is marked 499")
|
||||
require.Zero(t, recorder.Body.Len())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -101,8 +101,11 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/chat/completions is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
@@ -168,11 +171,15 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
}
|
||||
|
||||
for {
|
||||
if c.Request.Context().Err() != nil {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, selectionSessionHash, reqModel, fs.FailedAccountIDs, "", int64(0))
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gateway.cc.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, groupPlatform)
|
||||
cls = classifySelectionFailureError(err, cls)
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "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"
|
||||
)
|
||||
|
||||
// groupScopedSchedulerCache 按桶的分组只返回该分组的成员账号,并记录被查询过的分组。
|
||||
type groupScopedSchedulerCache struct {
|
||||
*fakeSchedulerCache
|
||||
mu sync.Mutex
|
||||
groupIDs []int64
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) GetSnapshot(_ context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
|
||||
c.mu.Lock()
|
||||
c.groupIDs = append(c.groupIDs, bucket.GroupID)
|
||||
c.mu.Unlock()
|
||||
var members []*service.Account
|
||||
for _, account := range c.accounts {
|
||||
for _, ag := range account.AccountGroups {
|
||||
if ag.GroupID == bucket.GroupID {
|
||||
members = append(members, account)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return members, true, nil
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) queriedGroupIDs() []int64 {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return append([]int64(nil), c.groupIDs...)
|
||||
}
|
||||
|
||||
type groupMapRepo struct {
|
||||
*fakeGroupRepo
|
||||
groups map[int64]*service.Group
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByID(_ context.Context, id int64) (*service.Group, error) {
|
||||
if group, ok := r.groups[id]; ok {
|
||||
return group, nil
|
||||
}
|
||||
return nil, service.ErrGroupNotFound
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByIDLite(ctx context.Context, id int64) (*service.Group, error) {
|
||||
return r.GetByID(ctx, id)
|
||||
}
|
||||
|
||||
func TestGatewayOpenAICompatibleHandlersClaudeCodeOnlyFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
const (
|
||||
primaryGroupID = int64(9300)
|
||||
fallbackGroupID = int64(9301)
|
||||
primaryAccountID = int64(9310)
|
||||
fallbackAccountID = int64(9311)
|
||||
)
|
||||
|
||||
// 账号不带 api_key:转发在取令牌时失败,请求不会触达上游,断言只看选号结果。
|
||||
newAccount := func(id, groupID int64) *service.Account {
|
||||
return &service.Account{
|
||||
ID: id, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}},
|
||||
}
|
||||
}
|
||||
|
||||
endpoints := []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
call func(*GatewayHandler, *gin.Context)
|
||||
}{
|
||||
{
|
||||
name: "responses", path: "/v1/responses",
|
||||
body: `{"model":"claude-sonnet-4-5","input":"hello","stream":false}`,
|
||||
call: (*GatewayHandler).Responses,
|
||||
},
|
||||
{
|
||||
name: "chat completions", path: "/v1/chat/completions",
|
||||
body: `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}],"stream":false}`,
|
||||
call: (*GatewayHandler).ChatCompletions,
|
||||
},
|
||||
}
|
||||
|
||||
for _, ep := range endpoints {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
hasFallback bool
|
||||
}{
|
||||
{name: "with fallback group", hasFallback: true},
|
||||
{name: "without fallback group", hasFallback: false},
|
||||
} {
|
||||
t.Run(ep.name+"/"+tc.name, func(t *testing.T) {
|
||||
primary := &service.Group{
|
||||
ID: primaryGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive, ClaudeCodeOnly: true,
|
||||
}
|
||||
if tc.hasFallback {
|
||||
fallbackID := fallbackGroupID
|
||||
primary.FallbackGroupID = &fallbackID
|
||||
}
|
||||
fallback := &service.Group{
|
||||
ID: fallbackGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
schedulerCache := &groupScopedSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{
|
||||
newAccount(primaryAccountID, primaryGroupID),
|
||||
newAccount(fallbackAccountID, fallbackGroupID),
|
||||
}}}
|
||||
gatewayService := service.NewGatewayService(
|
||||
nil, &groupMapRepo{fakeGroupRepo: &fakeGroupRepo{}, groups: map[int64]*service.Group{
|
||||
primaryGroupID: primary,
|
||||
fallbackGroupID: fallback,
|
||||
}}, nil, nil, nil, nil, nil, nil, nil,
|
||||
service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil),
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCacheService.Stop)
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
billingCacheService: billingCacheService,
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0),
|
||||
maxAccountSwitches: 1,
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
primaryGroupIDRef := primaryGroupID
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9320, UserID: 9330, GroupID: &primaryGroupIDRef, Group: primary, Status: service.StatusActive,
|
||||
User: &service.User{ID: 9330, Concurrency: 10, Balance: 100},
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
ctx := context.WithValue(context.Background(), ctxkey.Group, primary)
|
||||
req := httptest.NewRequest(http.MethodPost, ep.path, bytes.NewBufferString(ep.body)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
|
||||
ep.call(h, c)
|
||||
|
||||
selected, reachedSelection := c.Get(opsAccountIDKey)
|
||||
if !tc.hasFallback {
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "This group is restricted to Claude Code clients")
|
||||
require.False(t, reachedSelection, "a Claude Code only group without fallback must be rejected before account selection")
|
||||
require.Empty(t, schedulerCache.queriedGroupIDs())
|
||||
return
|
||||
}
|
||||
require.NotContains(t, recorder.Body.String(), "restricted to Claude Code clients")
|
||||
require.True(t, reachedSelection, "a Claude Code only group with fallback must reach account selection")
|
||||
require.Equal(t, fallbackAccountID, selected)
|
||||
queried := schedulerCache.queriedGroupIDs()
|
||||
require.Contains(t, queried, fallbackGroupID)
|
||||
require.NotContains(t, queried, primaryGroupID)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
@@ -70,6 +71,43 @@ func TestGatewayEnsureForwardErrorResponse_SkipsCommittedSSEError(t *testing.T)
|
||||
require.Equal(t, 1, strings.Count(w.Body.String(), "event: error"))
|
||||
}
|
||||
|
||||
// 客户端已断开且响应未提交:不补写 502,标记 499 供访问日志与 ops 归类。
|
||||
func TestGatewayEnsureForwardErrorResponse_SkipsCanceledClient(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, EndpointMessages, nil).WithContext(ctx)
|
||||
cancel()
|
||||
|
||||
h := &GatewayHandler{}
|
||||
wrote := h.ensureForwardErrorResponse(c, false)
|
||||
|
||||
require.False(t, wrote)
|
||||
require.Equal(t, statusClientClosedRequest, c.Writer.Status())
|
||||
require.Empty(t, w.Body.String())
|
||||
}
|
||||
|
||||
// 客户端已断开但流已开始:状态码已固化为 200,不再向断开的连接追加错误帧。
|
||||
func TestGatewayEnsureForwardErrorResponse_CanceledClientAfterStreamStartedAppendsNothing(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, EndpointMessages, nil).WithContext(ctx)
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
_, _ = c.Writer.WriteString(":\n\n")
|
||||
cancel()
|
||||
|
||||
h := &GatewayHandler{}
|
||||
wrote := h.ensureForwardErrorResponse(c, true)
|
||||
|
||||
require.False(t, wrote)
|
||||
require.Equal(t, http.StatusOK, c.Writer.Status())
|
||||
require.Equal(t, ":\n\n", w.Body.String())
|
||||
require.Empty(t, service.GetOpsStreamErrors(c))
|
||||
}
|
||||
|
||||
// case B 回归:Anthropic-backed /responses,Writer 已被写过时
|
||||
// ensureForwardErrorResponse 仍要发 response.failed。
|
||||
func TestGatewayEnsureForwardErrorResponse_ResponsesRouteAfterWrittenEmitsResponseFailed(t *testing.T) {
|
||||
|
||||
@@ -105,13 +105,11 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(requestCtx, apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction:
|
||||
// /v1/responses is never a Claude Code endpoint.
|
||||
// When claude_code_only is enabled, this endpoint is rejected.
|
||||
// The existing service-layer checkClaudeCodeRestriction handles degradation
|
||||
// to fallback groups when the Forward path calls SelectAccountForModelWithExclusions.
|
||||
// Here we just reject at handler level since /v1/responses clients can't be Claude Code.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/responses is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.responsesErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
@@ -170,11 +168,15 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
fs := NewFailoverState(h.maxAccountSwitches, false)
|
||||
|
||||
for {
|
||||
if requestCtx.Err() != nil {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(requestCtx, apiKey.GroupID, sessionHash, reqModel, fs.FailedAccountIDs, "", int64(0))
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gateway.responses.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, effectiveAPIKeyPlatform(c, apiKey))
|
||||
cls = classifySelectionFailureError(err, cls)
|
||||
|
||||
@@ -287,6 +287,18 @@ func TestGatewayModels_UnmappedOpenAIAccountsSupplementMappedModels(t *testing.T
|
||||
accounts: accounts[1:],
|
||||
want: []string{sparkModel, alias},
|
||||
},
|
||||
{
|
||||
// A passthrough account with a stale mapping behaves like an unmapped
|
||||
// one: it adds the defaults but never its own mapping keys, and it no
|
||||
// longer hides the aliases declared on ordinary accounts.
|
||||
name: "passthrough account contributes defaults without hiding mapped aliases",
|
||||
accounts: append([]service.Account{{
|
||||
ID: 5, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"stale-model": "stale-model"}},
|
||||
Extra: map[string]any{"openai_passthrough": true},
|
||||
}}, accounts[1:]...),
|
||||
want: append(openai.DefaultModelIDs(), alias),
|
||||
},
|
||||
{
|
||||
name: "unmapped accounts from another platform do not add defaults",
|
||||
accounts: append([]service.Account{{ID: 4, Platform: service.PlatformAnthropic}}, accounts[1:]...),
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "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"
|
||||
)
|
||||
|
||||
// clientCancelUpstream 模拟上游尚未响应时客户端断开:Do 期间取消请求 context,
|
||||
// 并按 net/http 的形态返回包裹 context.Canceled 的 *url.Error。
|
||||
type clientCancelUpstream struct {
|
||||
service.HTTPUpstream
|
||||
cancel context.CancelFunc
|
||||
calls atomic.Int32
|
||||
}
|
||||
|
||||
func (u *clientCancelUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
u.calls.Add(1)
|
||||
u.cancel()
|
||||
return nil, &url.Error{Op: "Post", URL: req.URL.String(), Err: context.Canceled}
|
||||
}
|
||||
|
||||
type geminiClientCancelFixture struct {
|
||||
handler *GatewayHandler
|
||||
group *service.Group
|
||||
apiKey *service.APIKey
|
||||
upstream *clientCancelUpstream
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func newGeminiClientCancelFixture(t *testing.T) *geminiClientCancelFixture {
|
||||
t.Helper()
|
||||
groupID := int64(9200)
|
||||
accountID := int64(9201)
|
||||
group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformGemini, Status: service.StatusActive}
|
||||
account := &service.Account{
|
||||
ID: accountID,
|
||||
Platform: service.PlatformGemini,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "test-key"},
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Priority: 1,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: accountID, GroupID: groupID}},
|
||||
}
|
||||
h, cleanup := newTestGatewayHandler(t, group, []*service.Account{account})
|
||||
t.Cleanup(cleanup)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
t.Cleanup(cancel)
|
||||
upstream := &clientCancelUpstream{cancel: cancel}
|
||||
h.geminiCompatService = service.NewGeminiMessagesCompatService(nil, nil, nil, nil, nil, nil, upstream, nil, &config.Config{})
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9202, UserID: 9203, GroupID: &groupID, Group: group, Status: service.StatusActive,
|
||||
User: &service.User{ID: 9203, Concurrency: 10, Balance: 100},
|
||||
}
|
||||
return &geminiClientCancelFixture{handler: h, group: group, apiKey: apiKey, upstream: upstream, ctx: ctx}
|
||||
}
|
||||
|
||||
// serve 经 ops 错误日志中间件执行请求,返回响应与 ops 队列中的条目数。
|
||||
func (f *geminiClientCancelFixture) serve(t *testing.T, route, path, body string, call func(*gin.Context)) (*httptest.ResponseRecorder, int64) {
|
||||
t.Helper()
|
||||
setupOpsErrorLogTestQueue(t, 4)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.POST(route, func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyAPIKey), f.apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: f.apiKey.UserID, Concurrency: 10})
|
||||
call(c)
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, path, bytes.NewBufferString(body))
|
||||
req = req.WithContext(context.WithValue(f.ctx, ctxkey.Group, f.group))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
return rec, OpsErrorLogQueueLength()
|
||||
}
|
||||
|
||||
// 原生入口:上游响应前客户端断开,响应未提交,标记 499;纯客户端取消不落 ops 错误日志。
|
||||
func TestGeminiV1BetaModels_ClientCancelBeforeUpstreamResponseMarks499(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
f := newGeminiClientCancelFixture(t)
|
||||
|
||||
rec, opsQueued := f.serve(t,
|
||||
"/v1beta/models/*modelAction",
|
||||
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
|
||||
`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
|
||||
f.handler.GeminiV1BetaModels,
|
||||
)
|
||||
|
||||
require.Equal(t, int32(1), f.upstream.calls.Load())
|
||||
require.Equal(t, statusClientClosedRequest, rec.Code)
|
||||
require.Zero(t, rec.Body.Len())
|
||||
require.Zero(t, opsQueued)
|
||||
}
|
||||
|
||||
// Chat Completions 兼容入口(Gemini 分组):同一场景标记 499,不再补写 502。
|
||||
func TestGatewayChatCompletions_GeminiClientCancelBeforeUpstreamResponseMarks499(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
f := newGeminiClientCancelFixture(t)
|
||||
|
||||
rec, opsQueued := f.serve(t,
|
||||
"/v1/chat/completions",
|
||||
"/v1/chat/completions",
|
||||
`{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"hi"}],"stream":true}`,
|
||||
f.handler.ChatCompletions,
|
||||
)
|
||||
|
||||
require.Equal(t, int32(1), f.upstream.calls.Load())
|
||||
require.Equal(t, statusClientClosedRequest, rec.Code)
|
||||
require.Zero(t, rec.Body.Len())
|
||||
require.Zero(t, opsQueued)
|
||||
}
|
||||
@@ -370,7 +370,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
userReleaseFunc, err := geminiConcurrency.AcquireUserSlotWithWait(c, authSubject.UserID, authSubject.Concurrency, stream, &streamStarted)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.user_slot_acquire_failed", zap.Error(err))
|
||||
googleError(c, http.StatusTooManyRequests, err.Error())
|
||||
googleConcurrencyError(c, err, "user")
|
||||
return
|
||||
}
|
||||
// 确保请求取消时也会释放槽位,避免长连接被动中断造成泄漏
|
||||
@@ -506,6 +506,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gemini.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, service.PlatformGemini)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -598,7 +602,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
googleError(c, http.StatusTooManyRequests, err.Error())
|
||||
googleConcurrencyError(c, err, "account")
|
||||
return
|
||||
}
|
||||
if accountWaitCounted {
|
||||
@@ -674,7 +678,8 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// ForwardNative already wrote the response
|
||||
// 转发层已写出错误响应;客户端断开时转发层不写,响应未提交则标记 499。
|
||||
failoverClientGone(c)
|
||||
reqLog.Error("gemini.forward_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
return
|
||||
}
|
||||
@@ -836,6 +841,13 @@ func googleError(c *gin.Context, status int, message string) {
|
||||
})
|
||||
}
|
||||
|
||||
// googleConcurrencyError 以 Google 错误格式回写并发槽获取失败,状态码与文案
|
||||
// 沿用 concurrencyErrorResponse 的统一映射(客户端断开为 499)。
|
||||
func googleConcurrencyError(c *gin.Context, err error, slotType string) {
|
||||
status, _, _, message := concurrencyErrorResponse(err, slotType)
|
||||
googleError(c, status, message)
|
||||
}
|
||||
|
||||
func writeUpstreamResponse(c *gin.Context, res *service.UpstreamHTTPResult) {
|
||||
if res == nil {
|
||||
googleError(c, http.StatusBadGateway, "Empty upstream response")
|
||||
|
||||
@@ -92,6 +92,8 @@ type modelPlazaGroup struct {
|
||||
// 不取分组/用户专属倍率。
|
||||
ImageRateIndependent bool `json:"image_rate_independent"`
|
||||
ImageRateMultiplier float64 `json:"image_rate_multiplier"`
|
||||
VideoRateIndependent bool `json:"video_rate_independent"`
|
||||
VideoRateMultiplier float64 `json:"video_rate_multiplier"`
|
||||
// 分组是否启用长上下文阶梯计费;关闭时模型实付列只展示最低档/基础价。
|
||||
LongContextPricingEnabled bool `json:"long_context_pricing_enabled"`
|
||||
Models []modelPlazaModel `json:"models"`
|
||||
@@ -211,6 +213,8 @@ func toModelPlazaGroupDTO(g *service.PlazaGroup, userRates map[int64]float64) mo
|
||||
IsExclusive: g.IsExclusive,
|
||||
ImageRateIndependent: g.ImageRateIndependent,
|
||||
ImageRateMultiplier: g.ImageRateMultiplier,
|
||||
VideoRateIndependent: g.VideoRateIndependent,
|
||||
VideoRateMultiplier: g.VideoRateMultiplier,
|
||||
LongContextPricingEnabled: g.LongContextPricingEnabled,
|
||||
Models: models,
|
||||
}
|
||||
|
||||
@@ -89,6 +89,7 @@ func TestToModelPlazaGroupDTO_UserRateAndFieldWhitelist(t *testing.T) {
|
||||
g := service.PlazaGroup{
|
||||
ID: 2, Name: "vip", Description: "d", Platform: "anthropic",
|
||||
SubscriptionType: "standard", RateMultiplier: 1, IsExclusive: true,
|
||||
VideoRateIndependent: true, VideoRateMultiplier: 0.7,
|
||||
Models: []service.PlazaModel{{
|
||||
Name: "claude-sonnet",
|
||||
Platform: "anthropic",
|
||||
@@ -115,11 +116,14 @@ func TestToModelPlazaGroupDTO_UserRateAndFieldWhitelist(t *testing.T) {
|
||||
"rate_multiplier", "user_rate_multiplier", "is_exclusive", "models",
|
||||
"peak_rate_enabled", "peak_start", "peak_end", "peak_rate_multiplier",
|
||||
"image_rate_independent", "image_rate_multiplier", "long_context_pricing_enabled",
|
||||
"video_rate_independent", "video_rate_multiplier",
|
||||
} {
|
||||
_, exists := decoded[key]
|
||||
require.Truef(t, exists, "plaza group DTO must expose %q", key)
|
||||
}
|
||||
require.InDelta(t, 0.5, decoded["user_rate_multiplier"].(float64), 1e-9)
|
||||
require.Equal(t, true, decoded["video_rate_independent"])
|
||||
require.Equal(t, 0.7, decoded["video_rate_multiplier"])
|
||||
|
||||
// 模型条目:pricing + official_pricing 并存;official 缺失字段输出 null 而非省略
|
||||
models := decoded["models"].([]any)
|
||||
|
||||
@@ -197,6 +197,10 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if !errors.As(err, &failoverErr) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account, openAIAccountScheduleModel(c, account, requestedModel, false, result), false, nil, err)
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_alpha_search.forward_aborted_client_disconnected", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
return
|
||||
}
|
||||
if c.Writer.Size() == writerSizeBeforeForward {
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func (h *OpenAIGatewayHandler) cyberPolicyLogOnly(c *gin.Context, apiKey *service.APIKey) bool {
|
||||
return h != nil && c != nil && c.Request != nil && h.gatewayService.CyberPolicyLogOnly(c.Request.Context(), apiKey)
|
||||
}
|
||||
|
||||
// Both HTTP and WebSocket admission skip existing blocks for trusted users.
|
||||
func (h *OpenAIGatewayHandler) findBlockedCyberSessionForAPIKey(c *gin.Context, apiKey *service.APIKey, body []byte) string {
|
||||
if apiKey == nil || h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return ""
|
||||
}
|
||||
return findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, body)
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestCyberAllowlistedUserBypassesExistingBlocksAndContinuesWebSocket(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := coderws.Accept(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
|
||||
defer cancel()
|
||||
for _, response := range []string{
|
||||
`{"type":"response.failed","response":{"id":"resp_cyber","model":"gpt-5.1","error":{"code":"cyber_policy","message":"blocked by upstream"}}}`,
|
||||
`{"type":"response.completed","response":{"id":"resp_ok","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||||
} {
|
||||
if _, _, err := conn.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
if err := conn.Write(ctx, coderws.MessageText, []byte(response)); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
// Let the downstream close so the completed frame can drain normally.
|
||||
_, _, _ = conn.Read(ctx)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
harness := newOpenAIWSPassthroughHandlerHarness(t, upstream.URL, map[string]string{
|
||||
service.SettingKeyCyberPolicyUserAllowlist: "1751",
|
||||
})
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5.1","prompt_cache_key":"trusted-session","input":"test"}`)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(string(payload)))
|
||||
explicitKey := service.CyberSessionExplicitBlockKey(harness.apiKey.ID, c, payload)
|
||||
store, ok := harness.gatewayCache.(service.CyberSessionBlockStore)
|
||||
require.True(t, ok)
|
||||
require.NoError(t, store.SetCyberSessionBlocked(context.Background(), "", []string{explicitKey}, time.Minute))
|
||||
for _, format := range []cyberSessionBlockFormat{cyberBlockFormatResponses, cyberBlockFormatChat, cyberBlockFormatAnthropic} {
|
||||
require.False(t, harness.handler.rejectIfCyberSessionBlocked(c, harness.apiKey, payload, "gpt-5.1", format))
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
require.NoError(t, harness.clientConn.Write(ctx, coderws.MessageText, payload))
|
||||
_, event, err := harness.clientConn.Read(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cyber_policy", gjson.GetBytes(event, "response.error.code").String())
|
||||
require.Eventually(t, func() bool {
|
||||
logs := harness.moderationRepo.logSnapshot()
|
||||
return len(logs) == 1 && logs[0].Mode == service.ContentModerationModeCyberLogOnly
|
||||
}, 3*time.Second, 10*time.Millisecond)
|
||||
matched, err := store.FindCyberSessionBlocked(ctx, service.CyberSessionTranscriptBlockKeys(harness.apiKey.ID, payload))
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, matched, "the cyber response must not write new transcript blocks")
|
||||
require.NoError(t, harness.clientConn.Write(ctx, coderws.MessageText, payload))
|
||||
_, event, err = harness.clientConn.Read(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
|
||||
_ = harness.clientConn.CloseNow()
|
||||
select {
|
||||
case <-harness.handlerDone:
|
||||
case <-ctx.Done():
|
||||
t.Fatal("handler did not exit")
|
||||
}
|
||||
}
|
||||
@@ -34,6 +34,7 @@ import (
|
||||
|
||||
// OpenAIGatewayHandler handles OpenAI API gateway requests
|
||||
type OpenAIGatewayHandler struct {
|
||||
compositeResolver *service.CompositeRouteResolver
|
||||
gatewayService *service.OpenAIGatewayService
|
||||
billingCacheService *service.BillingCacheService
|
||||
apiKeyService *service.APIKeyService
|
||||
@@ -2391,7 +2392,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
// 分组级模型白名单:首帧校验客户端模型,不通过则关闭连接并标记运维原因。
|
||||
// 必须在 ensureCompositeTargetPlatform(合成路由改写)之前执行。
|
||||
// 必须在合成路由解析和上游模型映射之前执行。
|
||||
// 与 HTTP 准入一致:帧内重复 model 键/大小写变体可能被上游按末值绑定,
|
||||
// 全部候选值逐一校验,任一未命中即拒绝。
|
||||
if blocked := blockedModelAllowlistCandidate(apiKey.Group, requestmodel.FromBodyCandidates("", "application/json", firstMessage)); blocked != "" {
|
||||
@@ -2400,7 +2401,21 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, fmt.Sprintf("Model %q is not available for this group", blocked))
|
||||
return
|
||||
}
|
||||
ensureCompositeTargetPlatform(c, apiKey, reqModel)
|
||||
// Keep the client model in the frame for admission, response identity and usage.
|
||||
// Apply the resolved upstream model through MapRequestModel on every turn.
|
||||
wsRouteModel := reqModel
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
decision, resolveErr := h.compositeResolver.Resolve(c.Request.Context(), apiKey.Group.ID, reqModel, service.CompositeRouteEndpointResponses)
|
||||
if resolveErr != nil {
|
||||
reqLog.Error("openai.websocket_composite_route_failed", zap.Error(resolveErr))
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "Failed to resolve composite model route")
|
||||
return
|
||||
}
|
||||
if decision.Matched {
|
||||
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
|
||||
wsRouteModel = decision.UpstreamModel
|
||||
}
|
||||
}
|
||||
ctx = c.Request.Context()
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
platform, ok := service.ResolvedTargetPlatformFromContext(ctx)
|
||||
@@ -2440,7 +2455,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
|
||||
// The first response.create frame is available here, so explicit IDs are
|
||||
// checked directly and body-derived sessions use the coarse scope gate.
|
||||
if cyberBlockKey := findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, firstMessage); cyberBlockKey != "" {
|
||||
if cyberBlockKey := h.findBlockedCyberSessionForAPIKey(c, apiKey, firstMessage); cyberBlockKey != "" {
|
||||
writeCyberSessionBlockedWSError(c.Request.Context(), wsConn)
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "session blocked by cyber-security policy")
|
||||
h.enqueueCyberSessionBlockedOpsEntry(c, apiKey, reqModel, cyberBlockKey)
|
||||
@@ -2464,8 +2479,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 解析渠道级模型映射
|
||||
channelMappingWS, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, reqModel)
|
||||
wsForwardModel := openAIChannelForwardModel(channelMappingWS, reqModel)
|
||||
channelMappingWS, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, wsRouteModel)
|
||||
wsForwardModel := openAIChannelForwardModel(channelMappingWS, wsRouteModel)
|
||||
|
||||
var currentUserRelease func()
|
||||
var currentAccountRelease func()
|
||||
@@ -2817,7 +2832,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// 连接级 cyber session gate 也在 BeforeRequest 先执行,使 native 与
|
||||
// passthrough ingress 都能在 BeforeTurn 及上游写入前无副作用地拒绝。
|
||||
// BeforeTurn 中保留同一检查作为防御式兜底。
|
||||
if cyberBlockedThisConn {
|
||||
if cyberBlockedThisConn && !h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil)
|
||||
}
|
||||
if turn == 1 {
|
||||
@@ -2856,7 +2871,16 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
model = reqModel
|
||||
}
|
||||
setOpsRequestContext(c, model, true)
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, model)
|
||||
routeModel := model
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
// The account, target platform and route context are connection-scoped.
|
||||
// A different public model needs a new connection and fresh resolution.
|
||||
if model != reqModel {
|
||||
return "", newOpenAIWSUnsupportedModelSwitchError(model)
|
||||
}
|
||||
routeModel = wsRouteModel
|
||||
}
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, routeModel)
|
||||
mappedModelUnchanged := false
|
||||
if previous := turnChannelMapping.Load(); previous != nil && previous.turn < turn {
|
||||
mappedModelUnchanged = strings.TrimSpace(previous.mapping.MappedModel) == strings.TrimSpace(mapping.MappedModel)
|
||||
@@ -2869,7 +2893,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
},
|
||||
BeforeTurn: func(turn int) error {
|
||||
// turn==1 的会话屏蔽已由握手层检查覆盖;连接内 flag 只拦截后续 turn。
|
||||
if cyberBlockedThisConn {
|
||||
if cyberBlockedThisConn && !h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil)
|
||||
}
|
||||
// 长连接跨峰谷/倍率刷新防护:每个 turn 按当前时刻重装门并复核
|
||||
@@ -2949,7 +2973,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
cyberBlockedThisConn, cyberBlockPendingAfterFailover = advanceOpenAIWSCyberBlockState(
|
||||
cyberBlockedThisConn,
|
||||
cyberBlockPendingAfterFailover,
|
||||
cyberMarked,
|
||||
cyberMarked && !h.cyberPolicyLogOnly(c, apiKey),
|
||||
turnErr,
|
||||
)
|
||||
if turnErr != nil {
|
||||
@@ -4045,7 +4069,7 @@ func (h *OpenAIGatewayHandler) rejectIfCyberSessionBlocked(c *gin.Context, apiKe
|
||||
if enabled, _ := h.gatewayService.CyberSessionBlockRuntime(c.Request.Context()); !enabled {
|
||||
return false
|
||||
}
|
||||
key := findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, body)
|
||||
key := h.findBlockedCyberSessionForAPIKey(c, apiKey, body)
|
||||
if key == "" {
|
||||
return false
|
||||
}
|
||||
@@ -4242,7 +4266,8 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
|
||||
ClientIP: clientIPStr,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if gwSvc != nil && apiKey != nil {
|
||||
cyberLogOnly := h.cyberPolicyLogOnly(c, apiKey)
|
||||
if gwSvc != nil && apiKey != nil && !cyberLogOnly {
|
||||
plan := buildCyberSessionBlockWritePlan(apiKey.ID, c, cyberBlockBody)
|
||||
if len(plan.keys) > 0 {
|
||||
blockCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
|
||||
@@ -4255,6 +4280,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
|
||||
defer cancel()
|
||||
if cmSvc != nil {
|
||||
cmSvc.RecordCyberPolicyEvent(ctx, service.CyberPolicyRecordInput{
|
||||
LogOnly: cyberLogOnly,
|
||||
RequestID: requestID,
|
||||
UserID: userID,
|
||||
UserEmail: userEmail,
|
||||
|
||||
@@ -1931,7 +1931,10 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject
|
||||
|
||||
type openAIResponsesWSUsageLogCase struct {
|
||||
simpleModeRejectAtRead int64
|
||||
compositeResolver *service.CompositeRouteResolver
|
||||
accountPlatform string
|
||||
closeReason string
|
||||
closeStatus coderws.StatusCode
|
||||
firstPayload string
|
||||
// midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
|
||||
// 回一个 response.completed,客户端按普通事件读取。
|
||||
@@ -2878,6 +2881,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
upstreamErrCh := make(chan error, 1)
|
||||
var channelSvc *service.ChannelService
|
||||
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if tc.accountPlatform == service.PlatformGrok {
|
||||
payload, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
upstreamErrCh <- err
|
||||
return
|
||||
}
|
||||
upstreamPayloadCh <- payload
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = fmt.Fprintf(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_grok_test\",\"model\":%q,\"usage\":{\"input_tokens\":2,\"output_tokens\":1}}}\n\n", gjson.GetBytes(payload, "model").String())
|
||||
return
|
||||
}
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
|
||||
CompressionMode: coderws.CompressionContextTakeover,
|
||||
})
|
||||
@@ -2945,6 +2959,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
}
|
||||
if tc.accountPlatform != "" {
|
||||
account.Platform = tc.accountPlatform
|
||||
}
|
||||
if strings.TrimSpace(tc.ingressMode) != "" {
|
||||
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
|
||||
}
|
||||
@@ -3000,7 +3017,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
service.NewBillingService(cfg, nil),
|
||||
nil,
|
||||
billingCacheSvc,
|
||||
nil,
|
||||
&compositeWSHTTPUpstream{},
|
||||
&service.DeferredService{},
|
||||
nil,
|
||||
nil,
|
||||
@@ -3021,6 +3038,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}
|
||||
h := &OpenAIGatewayHandler{
|
||||
cfg: cfg,
|
||||
compositeResolver: tc.compositeResolver,
|
||||
gatewayService: gatewaySvc,
|
||||
billingCacheService: billingCacheSvc,
|
||||
apiKeyService: &service.APIKeyService{},
|
||||
@@ -3069,6 +3087,12 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
if tc.closeReason == "" {
|
||||
tc.closeReason = "not available for this group"
|
||||
}
|
||||
if tc.closeStatus == 0 {
|
||||
tc.closeStatus = coderws.StatusPolicyViolation
|
||||
}
|
||||
if tc.firstFrameCloseExpected {
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, _, readErr := clientConn.Read(readCtx)
|
||||
@@ -3076,12 +3100,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
require.Error(t, readErr, "first frame should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
status := tc.closeStatus
|
||||
if status == 0 {
|
||||
status = coderws.StatusPolicyViolation
|
||||
}
|
||||
require.Equal(t, status, closeErr.Code)
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
require.Empty(t, upstreamPayloadCh, "rejected first frame must not reach upstream")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3115,12 +3144,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
require.Error(t, readErr, "second turn should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
status := tc.closeStatus
|
||||
if status == 0 {
|
||||
status = coderws.StatusPolicyViolation
|
||||
}
|
||||
require.Equal(t, status, closeErr.Code)
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
require.Len(t, upstreamPayloadCh, turnCount-1, "rejected turn must not reach upstream")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3149,6 +3183,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}
|
||||
}
|
||||
|
||||
if tc.accountPlatform == service.PlatformGrok {
|
||||
upstreamErrCh <- nil
|
||||
}
|
||||
select {
|
||||
case upstreamErr := <-upstreamErrCh:
|
||||
require.NoError(t, upstreamErr)
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type compositeWSRouteRepo struct {
|
||||
service.CompositeModelRouteRepository
|
||||
routes []service.CompositeModelRoute
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *compositeWSRouteRepo) ListByGroup(context.Context, int64, bool) ([]service.CompositeModelRoute, error) {
|
||||
return r.routes, r.err
|
||||
}
|
||||
|
||||
type compositeWSHTTPUpstream struct{ service.HTTPUpstream }
|
||||
|
||||
func (*compositeWSHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
return http.DefaultClient.Do(req)
|
||||
}
|
||||
|
||||
func compositeWSResolver(platform, endpoint, upstream string) *service.CompositeRouteResolver {
|
||||
return service.NewCompositeRouteResolver(&compositeWSRouteRepo{routes: []service.CompositeModelRoute{{
|
||||
GroupID: 4201, PublicModel: "public-alias", MatchType: service.CompositeRouteMatchExact,
|
||||
TargetPlatform: platform, Endpoint: endpoint, UpstreamModel: upstream, Enabled: true,
|
||||
}}})
|
||||
}
|
||||
func compositeWSGroup(models ...string) *service.Group {
|
||||
g := wsAllowlistGroup(len(models) > 0, models...)
|
||||
g.Platform = service.PlatformComposite
|
||||
return g
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeAlias(t *testing.T) {
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformGrok} {
|
||||
for _, endpoint := range []string{service.CompositeRouteEndpointResponses, service.CompositeRouteEndpointAny} {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
t.Run(platform+"/"+endpoint+"/"+mode, func(t *testing.T) {
|
||||
upstream := "gpt-5.4"
|
||||
if platform == service.PlatformGrok {
|
||||
upstream = "grok-4.3"
|
||||
}
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias","input":"hi"}`,
|
||||
secondPayload: `{"type":"response.create","input":"again"}`,
|
||||
group: compositeWSGroup("public-alias"), accountPlatform: platform, ingressMode: mode,
|
||||
compositeResolver: compositeWSResolver(platform, endpoint, upstream),
|
||||
})
|
||||
require.Len(t, got.upstreamPayloads, 2)
|
||||
for i, payload := range got.upstreamPayloads {
|
||||
require.Equal(t, upstream, gjson.GetBytes(payload, "model").String())
|
||||
require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[i], "response.model").String())
|
||||
require.Equal(t, "public-alias", got.logs[i].RequestedModel)
|
||||
require.NotNil(t, got.logs[i].UpstreamModel)
|
||||
require.Equal(t, upstream, *got.logs[i].UpstreamModel)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeRouteRejections(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, platform, endpoint, model, reason string
|
||||
repoErr error
|
||||
status coderws.StatusCode
|
||||
}{
|
||||
{name: "disallowed platform", platform: service.PlatformAnthropic, endpoint: "responses", model: "public-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "wrong endpoint", platform: service.PlatformOpenAI, endpoint: "messages", model: "public-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "unknown alias", platform: service.PlatformOpenAI, endpoint: "responses", model: "unknown-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "resolver error", model: "gpt-5.4", repoErr: errors.New("database unavailable"), reason: "Failed to resolve composite model route", status: coderws.StatusInternalError},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
resolver := compositeWSResolver(tc.platform, tc.endpoint, "gpt-5.4")
|
||||
if tc.repoErr != nil {
|
||||
resolver = service.NewCompositeRouteResolver(&compositeWSRouteRepo{err: tc.repoErr})
|
||||
}
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"` + tc.model + `"}`, group: compositeWSGroup(), compositeResolver: resolver,
|
||||
firstFrameCloseExpected: true, closeReason: tc.reason, closeStatus: tc.status,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeAdmissionUsesPublicModel(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`, group: compositeWSGroup("gpt-5.4"),
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"), firstFrameCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeModelSwitchRequiresReconnect(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
for _, model := range []string{"gpt-5.4", "grok-4.3", "unknown-alias"} {
|
||||
t.Run(mode+"/"+model, func(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`,
|
||||
secondPayload: `{"type":"response.create","model":"` + model + `"}`, group: compositeWSGroup(), ingressMode: mode,
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"),
|
||||
secondTurnCloseExpected: true, closeReason: "model switch requires reconnect",
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeChannelBilling(t *testing.T) {
|
||||
for _, source := range []string{service.BillingModelSourceRequested, service.BillingModelSourceChannelMapped, service.BillingModelSourceUpstream} {
|
||||
t.Run(source, func(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.6-sol"}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.6-sol"}`,
|
||||
group: compositeWSGroup("gpt-5.6-sol"),
|
||||
compositeResolver: service.NewCompositeRouteResolver(&compositeWSRouteRepo{routes: []service.CompositeModelRoute{{PublicModel: "gpt-5.6-sol", MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformOpenAI, Endpoint: service.CompositeRouteEndpointResponses, UpstreamModel: "route-target"}}}),
|
||||
channelMapping: map[string]string{"route-target": "gpt-5.4"}, billingModelSource: source,
|
||||
accountModelMapping: map[string]any{"gpt-5.4": "gpt-5.4"},
|
||||
})
|
||||
for i, log := range got.logs {
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamPayloads[i], "model").String())
|
||||
require.Equal(t, "gpt-5.6-sol", log.RequestedModel)
|
||||
require.Equal(t, "gpt-5.6-sol", log.Model)
|
||||
if source == service.BillingModelSourceRequested {
|
||||
require.InDelta(t, 40e-6, log.TotalCost, 1e-12)
|
||||
} else {
|
||||
require.InDelta(t, 20e-6, log.TotalCost, 1e-12)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeDetectorFallback(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4"}`, group: compositeWSGroup(),
|
||||
compositeResolver: service.NewCompositeRouteResolver(&compositeWSRouteRepo{}),
|
||||
})
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamFirstPayload, "model").String())
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeSessionModelSwitchRequiresReconnect(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`,
|
||||
midPayload: `{"type":"session.update","session":{"model":"grok-4.3"}}`,
|
||||
secondPayload: `{"type":"response.create"}`, group: compositeWSGroup(),
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"),
|
||||
secondTurnCloseExpected: true, closeReason: "model switch requires reconnect",
|
||||
})
|
||||
}
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
)
|
||||
|
||||
type openAIWSPassthroughHandlerHarness struct {
|
||||
handler *OpenAIGatewayHandler
|
||||
clientConn *coderws.Conn
|
||||
handlerDone <-chan struct{}
|
||||
moderationRepo *contentModerationHandlerTestRepo
|
||||
@@ -26,7 +27,7 @@ type openAIWSPassthroughHandlerHarness struct {
|
||||
apiKey *service.APIKey
|
||||
}
|
||||
|
||||
func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *openAIWSPassthroughHandlerHarness {
|
||||
func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string, settings ...map[string]string) *openAIWSPassthroughHandlerHarness {
|
||||
t.Helper()
|
||||
gatewayCache := testutil.NewRedisGatewayCache(t)
|
||||
|
||||
@@ -35,6 +36,11 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
service.SettingKeyCyberSessionBlockEnabled: "true",
|
||||
service.SettingKeyCyberSessionBlockTTLSeconds: "60",
|
||||
}}
|
||||
for _, overrides := range settings {
|
||||
for key, value := range overrides {
|
||||
settingRepo.values[key] = value
|
||||
}
|
||||
}
|
||||
moderationRepo := &contentModerationHandlerTestRepo{}
|
||||
moderationSvc := service.NewContentModerationService(settingRepo, moderationRepo, nil, nil, nil, nil, nil, nil)
|
||||
settingSvc := service.NewSettingService(settingRepo, nil)
|
||||
@@ -90,6 +96,7 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 1851,
|
||||
UserID: 1751,
|
||||
Name: "ws-cyber-key",
|
||||
Key: "sk-handler-cyber-test",
|
||||
GroupID: &groupID,
|
||||
@@ -116,6 +123,7 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
t.Cleanup(func() { _ = clientConn.CloseNow() })
|
||||
|
||||
return &openAIWSPassthroughHandlerHarness{
|
||||
handler: h,
|
||||
clientConn: clientConn,
|
||||
handlerDone: handlerDone,
|
||||
moderationRepo: moderationRepo,
|
||||
|
||||
@@ -1146,6 +1146,9 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
if shouldSkipOpsErrorLog(c.Request.Context(), ops, parsed.Message, string(body), c.Request.URL.Path) {
|
||||
return
|
||||
}
|
||||
if shouldSkipOpsClientClosed(c, ops, status) {
|
||||
return
|
||||
}
|
||||
|
||||
apiKey := getOpsAPIKey(c)
|
||||
|
||||
@@ -2491,6 +2494,21 @@ func shouldSkipOpsErrorLog(ctx context.Context, ops *service.OpsService, message
|
||||
return false
|
||||
}
|
||||
|
||||
// shouldSkipOpsClientClosed 按 IgnoreContextCanceled 过滤纯客户端取消的 499。
|
||||
// 499 表示客户端在响应提交前断开(见 failoverClientGone),通常不带
|
||||
// "context canceled" 文案,shouldSkipOpsErrorLog 的文本过滤命中不了。
|
||||
// 本次请求未观察到上游错误时是纯客户端取消;已有上游错误的 499 表示上游失败后
|
||||
// 客户端没等到换号结果就离开,仍按上游失败落库。
|
||||
func shouldSkipOpsClientClosed(c *gin.Context, ops *service.OpsService, status int) bool {
|
||||
if status != statusClientClosedRequest || ops == nil {
|
||||
return false
|
||||
}
|
||||
if !ops.OpsAdvancedSettingsSnapshot().IgnoreContextCanceled {
|
||||
return false
|
||||
}
|
||||
return !hasOpsUpstreamErrorContext(c)
|
||||
}
|
||||
|
||||
// shouldSkipOpsErrorLogForCyber:cyber_policy 命中的请求由 recordCyberPolicyIfMarked
|
||||
// 统一落一条 status=403 的错误请求,故中间件跳过自身落库,避免双写。
|
||||
func shouldSkipOpsErrorLogForCyber(c *gin.Context) bool {
|
||||
|
||||
@@ -2168,3 +2168,83 @@ func TestNormalizeOpsErrorType_KeepsGeminiInBandSignalTypes(t *testing.T) {
|
||||
require.Equal(t, errType, normalizeOpsErrorType(errType, "PROHIBITED_CONTENT"), errType)
|
||||
}
|
||||
}
|
||||
|
||||
type opsAdvancedSettingsRepoStub struct {
|
||||
service.SettingRepository
|
||||
advanced string
|
||||
}
|
||||
|
||||
func (r *opsAdvancedSettingsRepoStub) GetValue(context.Context, string) (string, error) {
|
||||
return "", service.ErrSettingNotFound
|
||||
}
|
||||
|
||||
func (r *opsAdvancedSettingsRepoStub) GetMultiple(context.Context, []string) (map[string]string, error) {
|
||||
return map[string]string{service.SettingKeyOpsAdvancedSettings: r.advanced}, nil
|
||||
}
|
||||
|
||||
func (r *opsAdvancedSettingsRepoStub) Set(context.Context, string, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// serveClientClosedRequest 以已取消的请求 context 走 failoverClientGone,复现网关标记 499 的收尾路径。
|
||||
func serveClientClosedRequest(t *testing.T, ops *service.OpsService, prepare func(c *gin.Context)) {
|
||||
t.Helper()
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.POST("/v1/messages", func(c *gin.Context) {
|
||||
if prepare != nil {
|
||||
prepare(c)
|
||||
}
|
||||
failoverClientGone(c)
|
||||
})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/v1/messages", nil).WithContext(ctx))
|
||||
require.Equal(t, statusClientClosedRequest, recorder.Code)
|
||||
require.Zero(t, recorder.Body.Len())
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_SkipsPureClientClosed(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 2)
|
||||
gin.SetMode(gin.TestMode)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
require.True(t, ops.OpsAdvancedSettingsSnapshot().IgnoreContextCanceled)
|
||||
|
||||
serveClientClosedRequest(t, ops, nil)
|
||||
|
||||
require.Zero(t, OpsErrorLogQueueLength(), "纯客户端取消的 499 按 IgnoreContextCanceled 跳过")
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_RecordsClientClosedAfterUpstreamFailure(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 2)
|
||||
gin.SetMode(gin.TestMode)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
serveClientClosedRequest(t, ops, func(c *gin.Context) {
|
||||
c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{{
|
||||
AccountID: 7, UpstreamStatusCode: 524, Kind: "failover", Message: "upstream timeout",
|
||||
}})
|
||||
})
|
||||
|
||||
require.Equal(t, int64(1), OpsErrorLogQueueLength(), "上游失败后客户端离开仍按上游失败落库")
|
||||
job := <-opsErrorLogQueue
|
||||
require.Equal(t, statusClientClosedRequest, job.entry.StatusCode)
|
||||
require.Equal(t, "upstream", job.entry.ErrorPhase)
|
||||
require.NotNil(t, job.entry.UpstreamStatusCode)
|
||||
require.Equal(t, 524, *job.entry.UpstreamStatusCode)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_RecordsClientClosedWhenIgnoreContextCanceledDisabled(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 2)
|
||||
gin.SetMode(gin.TestMode)
|
||||
settings := &opsAdvancedSettingsRepoStub{advanced: `{"ignore_context_canceled":false}`}
|
||||
ops := service.NewOpsService(nil, settings, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
require.False(t, ops.OpsAdvancedSettingsSnapshot().IgnoreContextCanceled)
|
||||
|
||||
serveClientClosedRequest(t, ops, nil)
|
||||
|
||||
require.Equal(t, int64(1), OpsErrorLogQueueLength())
|
||||
job := <-opsErrorLogQueue
|
||||
require.Equal(t, statusClientClosedRequest, job.entry.StatusCode)
|
||||
}
|
||||
|
||||
@@ -50,10 +50,12 @@ func ProvideAdminHandlers(
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService,
|
||||
ollamaCloudUsage *service.OllamaCloudUsageService,
|
||||
opencodeGoUsage *service.OpenCodeGoUsageService,
|
||||
claudeResetCredits *service.ClaudeResetCreditService,
|
||||
) *AdminHandlers {
|
||||
accountHandler.SetUpstreamBillingProbeService(upstreamBillingProbe)
|
||||
accountHandler.SetOllamaCloudUsageService(ollamaCloudUsage)
|
||||
accountHandler.SetOpenCodeGoUsageService(opencodeGoUsage)
|
||||
accountHandler.SetClaudeResetCreditService(claudeResetCredits)
|
||||
return &AdminHandlers{
|
||||
Dashboard: dashboardHandler,
|
||||
User: userHandler,
|
||||
@@ -132,10 +134,12 @@ func ProvideOpenAIGatewayHandler(
|
||||
grokQuotaService *service.GrokQuotaService,
|
||||
cfg *config.Config,
|
||||
coordinator *securityaudit.Coordinator,
|
||||
compositeResolver *service.CompositeRouteResolver,
|
||||
) *OpenAIGatewayHandler {
|
||||
gatewayService.SetPluginManager(pluginManager)
|
||||
h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService,
|
||||
usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg)
|
||||
h.compositeResolver = compositeResolver
|
||||
h.securityAuditCoordinator = coordinator
|
||||
h.grokMediaEligibilityProber = grokQuotaService
|
||||
return h
|
||||
|
||||
@@ -532,8 +532,8 @@ func buildParts(content json.RawMessage, toolIDToName map[string]string, allowDu
|
||||
}
|
||||
parts = append(parts, part)
|
||||
|
||||
case "image":
|
||||
if block.Source != nil && block.Source.Type == "base64" {
|
||||
case "image", "document":
|
||||
if block.Source != nil && block.Source.Type == "base64" && strings.TrimSpace(block.Source.Data) != "" {
|
||||
parts = append(parts, GeminiPart{
|
||||
InlineData: &GeminiInlineData{
|
||||
MimeType: block.Source.MediaType,
|
||||
|
||||
@@ -646,6 +646,20 @@ func TestGeminiToolConfig_DropsBuiltinsWhenClientFunctionsPresent(t *testing.T)
|
||||
})
|
||||
}
|
||||
|
||||
func TestBuildParts_DocumentBecomesInlineData(t *testing.T) {
|
||||
content := `[
|
||||
{"type":"text","text":"read this"},
|
||||
{"type":"document","source":{"type":"base64","media_type":"application/pdf","data":"JVBERi0="}}
|
||||
]`
|
||||
parts, stripped, err := buildParts(json.RawMessage(content), map[string]string{}, true)
|
||||
require.NoError(t, err)
|
||||
require.False(t, stripped)
|
||||
require.Len(t, parts, 2)
|
||||
require.NotNil(t, parts[1].InlineData)
|
||||
require.Equal(t, "application/pdf", parts[1].InlineData.MimeType)
|
||||
require.Equal(t, "JVBERi0=", parts[1].InlineData.Data)
|
||||
}
|
||||
|
||||
// TestToolConfigAlwaysPresent ensures toolConfig is always emitted, including for
|
||||
// reasoning models without any tools: upstream rejects requests that omit it.
|
||||
func TestToolConfigAlwaysPresent(t *testing.T) {
|
||||
|
||||
@@ -258,6 +258,18 @@ func cleanJSONSchemaRecursive(value any) any {
|
||||
if constVal, exists := schemaMap["const"]; exists {
|
||||
if _, hasEnum := schemaMap["enum"]; !hasEnum {
|
||||
schemaMap["enum"] = []any{constVal}
|
||||
} else if constant, ok := constVal.(string); ok {
|
||||
if existing, ok := schemaMap["enum"].([]any); ok {
|
||||
// Both constraints apply; preserve their intersection before dropping const.
|
||||
values := []any{}
|
||||
for _, value := range existing {
|
||||
if text, ok := value.(string); ok && text == constant {
|
||||
values = append(values, constant)
|
||||
break
|
||||
}
|
||||
}
|
||||
schemaMap["enum"] = values
|
||||
}
|
||||
}
|
||||
if _, hasType := schemaMap["type"]; !hasType {
|
||||
switch constVal.(type) {
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
package antigravity
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestBuildToolsPreservesStringConst(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
schema string
|
||||
want string
|
||||
}{
|
||||
{"typed", `{"type":"string","const":"browser"}`, `{"type":"string","enum":["browser"]}`},
|
||||
{"inferred", `{"const":"browser"}`, `{"type":"string","enum":["browser"]}`},
|
||||
{"empty", `{"type":"string","const":""}`, `{"type":"string","enum":[""]}`},
|
||||
{"existing enum", `{"type":"string","const":"browser","enum":["browser","shell"]}`, `{"type":"string","enum":["browser"]}`},
|
||||
{"conflicting enum", `{"type":"string","const":"browser","enum":["shell"]}`, `{"type":"string","enum":[]}`},
|
||||
{"ordinary enum", `{"type":"string","enum":["browser","shell"]}`, `{"type":"string","enum":["browser","shell"]}`},
|
||||
{"array items", `{"type":"array","items":{"type":"string","const":"browser"}}`, `{"type":"array","items":{"type":"string","enum":["browser"]}}`},
|
||||
{"nested property named const", `{"type":"object","properties":{"const":{"const":"browser"}}}`, `{"type":"object","properties":{"const":{"type":"string","enum":["browser"]}}}`},
|
||||
{"schema metadata", `{"type":"string","const":"browser","$schema":"https://json-schema.org/draft/2020-12/schema","description":"Tool kind"}`, `{"type":"string","enum":["browser"],"description":"Tool kind"}`},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var property map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(tc.schema), &property))
|
||||
tools := buildTools([]ClaudeTool{{
|
||||
Name: "dispatch",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{"action": property},
|
||||
},
|
||||
}})
|
||||
require.Len(t, tools, 1)
|
||||
require.Len(t, tools[0].FunctionDeclarations, 1)
|
||||
properties, ok := tools[0].FunctionDeclarations[0].Parameters["properties"].(map[string]any)
|
||||
require.True(t, ok, "tool parameters must contain a properties object")
|
||||
got, err := json.Marshal(properties["action"])
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, tc.want, string(got))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -40,6 +40,7 @@ type StreamingProcessor struct {
|
||||
outputTokens int
|
||||
cacheReadTokens int
|
||||
imageOutputTokens int
|
||||
hasContent bool
|
||||
}
|
||||
|
||||
// NewStreamingProcessor 创建流式响应处理器
|
||||
@@ -175,6 +176,11 @@ func (p *StreamingProcessor) MessageStartSent() bool {
|
||||
return p.messageStartSent
|
||||
}
|
||||
|
||||
// HasContent reports whether any substantive text, thinking, or tool calls were emitted.
|
||||
func (p *StreamingProcessor) HasContent() bool {
|
||||
return p.hasContent
|
||||
}
|
||||
|
||||
// emitMessageStart 发送 message_start 事件
|
||||
func (p *StreamingProcessor) emitMessageStart(v1Resp *V1InternalResponse) []byte {
|
||||
if p.messageStartSent {
|
||||
@@ -296,6 +302,7 @@ func (p *StreamingProcessor) processThinking(text, signature string) []byte {
|
||||
}
|
||||
|
||||
if text != "" {
|
||||
p.hasContent = true
|
||||
_, _ = result.Write(p.emitDelta("thinking_delta", map[string]any{
|
||||
"thinking": text,
|
||||
}))
|
||||
@@ -350,6 +357,7 @@ func (p *StreamingProcessor) processText(text, signature string) []byte {
|
||||
}))
|
||||
}
|
||||
|
||||
p.hasContent = true
|
||||
_, _ = result.Write(p.emitDelta("text_delta", map[string]any{
|
||||
"text": text,
|
||||
}))
|
||||
@@ -362,6 +370,7 @@ func (p *StreamingProcessor) processFunctionCall(fc *GeminiFunctionCall, signatu
|
||||
var result bytes.Buffer
|
||||
|
||||
p.usedTool = true
|
||||
p.hasContent = true
|
||||
|
||||
toolID := fc.ID
|
||||
if toolID == "" {
|
||||
|
||||
@@ -1136,9 +1136,24 @@ func TestAnthropicToResponses_ThinkingDisabled(t *testing.T) {
|
||||
|
||||
resp, err := AnthropicToResponses(req)
|
||||
require.NoError(t, err)
|
||||
// Default effort applies (medium) even when thinking is disabled.
|
||||
require.NotNil(t, resp.Reasoning)
|
||||
assert.Equal(t, "medium", resp.Reasoning.Effort)
|
||||
assert.Equal(t, "none", resp.Reasoning.Effort)
|
||||
assert.Empty(t, resp.Reasoning.Summary)
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_ThinkingDisabledOverridesOutputEffort(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "gpt-5.6-sol",
|
||||
MaxTokens: 1024,
|
||||
Messages: []AnthropicMessage{{Role: "user", Content: json.RawMessage(`"Hello"`)}},
|
||||
Thinking: &AnthropicThinking{Type: "disabled"},
|
||||
OutputConfig: &AnthropicOutputConfig{Effort: "max"},
|
||||
}
|
||||
|
||||
resp, err := AnthropicToResponses(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp.Reasoning)
|
||||
assert.Equal(t, "none", resp.Reasoning.Effort)
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_NoThinking(t *testing.T) {
|
||||
@@ -1867,6 +1882,64 @@ func TestOpus55SignedThinkingResponsesRoundTrip(t *testing.T) {
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSonnet55ResponsesThinkingAndSampling(t *testing.T) {
|
||||
for _, effort := range []string{"", "low", "medium", "high", "xhigh", "max", "none"} {
|
||||
req := &ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: effort}}
|
||||
out, err := ResponsesToAnthropicRequest(req)
|
||||
require.NoError(t, err, effort)
|
||||
require.Zero(t, out.Thinking.BudgetTokens)
|
||||
if effort == "none" {
|
||||
require.Equal(t, "between_tools", out.Thinking.Type)
|
||||
require.Equal(t, "low", out.OutputConfig.Effort)
|
||||
} else {
|
||||
require.Equal(t, "adaptive", out.Thinking.Type)
|
||||
if effort == "" {
|
||||
effort = "high"
|
||||
}
|
||||
require.Equal(t, effort, out.OutputConfig.Effort)
|
||||
}
|
||||
}
|
||||
for _, choice := range []string{`"required"`, `{"type":"function","name":"lookup"}`} {
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), ToolChoice: json.RawMessage(choice)})
|
||||
require.ErrorContains(t, err, "forced tool_choice")
|
||||
}
|
||||
temperature := 0.7
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature})
|
||||
require.ErrorContains(t, err, "temperature")
|
||||
topP := 0.5
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), TopP: &topP})
|
||||
require.ErrorContains(t, err, "top_p")
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: "minimal"}})
|
||||
require.ErrorContains(t, err, "reasoning effort")
|
||||
|
||||
temperature, topP = 1, 0.99
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature, TopP: &topP})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestSonnet55SignedThinkingResponsesRoundTrip(t *testing.T) {
|
||||
block := AnthropicContentBlock{Type: "thinking", Thinking: "", Signature: "signed-sonnet-block"}
|
||||
response := AnthropicToResponsesResponse(&AnthropicResponse{Model: "claude-sonnet-5-5", Content: []AnthropicContentBlock{block, {Type: "text", Text: "progress"}, {Type: "tool_use", ID: "toolu_1", Name: "lookup", Input: json.RawMessage(`{}`)}}})
|
||||
require.Len(t, response.Output, 3)
|
||||
require.NotEmpty(t, response.Output[0].EncryptedContent)
|
||||
require.Equal(t, "message", response.Output[1].Type)
|
||||
raw, err := json.Marshal(response.Output)
|
||||
require.NoError(t, err)
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(raw, &items))
|
||||
items = append(items, ResponsesInputItem{Type: "function_call_output", CallID: response.Output[2].CallID, Output: "ok"})
|
||||
raw, err = json.Marshal(items)
|
||||
require.NoError(t, err)
|
||||
converted, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: raw})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, converted.Messages, 2)
|
||||
var blocks []AnthropicContentBlock
|
||||
require.NoError(t, json.Unmarshal(converted.Messages[0].Content, &blocks))
|
||||
require.Equal(t, block, blocks[0])
|
||||
require.Equal(t, "text", blocks[1].Type)
|
||||
require.Equal(t, "tool_use", blocks[2].Type)
|
||||
}
|
||||
|
||||
func TestGPT6ChatSamplingAndCacheFields(t *testing.T) {
|
||||
temperature := 0.7
|
||||
for _, model := range []string{"gpt-6-sol", "gpt-6-luna"} {
|
||||
|
||||
@@ -57,17 +57,13 @@ func AnthropicToResponses(req *AnthropicRequest) (*ResponsesRequest, error) {
|
||||
out.Tools = convertAnthropicToolsToResponses(req.Tools)
|
||||
}
|
||||
|
||||
// Determine reasoning effort: only output_config.effort controls the
|
||||
// level; thinking.type is ignored. Default follows Codex CLI / airgate's
|
||||
// Anthropic bridge shape, which uses medium when unset.
|
||||
// Anthropic levels map 1:1 to OpenAI: low→low, medium→medium, high→high, max→xhigh.
|
||||
effort := "medium"
|
||||
if req.OutputConfig != nil && req.OutputConfig.Effort != "" {
|
||||
effort = req.OutputConfig.Effort
|
||||
}
|
||||
// An explicit thinking disable takes precedence over output_config.effort.
|
||||
effort := anthropicReasoningEffort(req)
|
||||
out.Reasoning = &ResponsesReasoning{
|
||||
Effort: mapAnthropicEffortToResponses(effort),
|
||||
Summary: "auto",
|
||||
Effort: effort,
|
||||
}
|
||||
if effort != "none" {
|
||||
out.Reasoning.Summary = "auto"
|
||||
}
|
||||
|
||||
// Convert tool_choice
|
||||
@@ -423,17 +419,22 @@ func extractAnthropicTextFromBlocks(blocks []AnthropicContentBlock) string {
|
||||
return strings.Join(parts, "\n\n")
|
||||
}
|
||||
|
||||
// mapAnthropicEffortToResponses converts Anthropic reasoning effort levels to
|
||||
// OpenAI Responses API effort levels.
|
||||
//
|
||||
// Both APIs default to "high". The mapping is 1:1 for shared levels;
|
||||
// only Anthropic's "max" (Opus 4.6 exclusive) maps to OpenAI's "xhigh"
|
||||
// (GPT-5.2+ exclusive) as both represent the highest reasoning tier.
|
||||
//
|
||||
// low → low
|
||||
// medium → medium
|
||||
// high → high
|
||||
// max → xhigh
|
||||
// anthropicReasoningEffort resolves the Anthropic request preference for both
|
||||
// OpenAI bridges. Explicitly disabled thinking overrides output_config.effort;
|
||||
// otherwise the bridge keeps its medium default.
|
||||
func anthropicReasoningEffort(req *AnthropicRequest) string {
|
||||
if req.Thinking != nil && req.Thinking.Type == "disabled" {
|
||||
return "none"
|
||||
}
|
||||
effort := "medium"
|
||||
if req.OutputConfig != nil && req.OutputConfig.Effort != "" {
|
||||
effort = req.OutputConfig.Effort
|
||||
}
|
||||
return mapAnthropicEffortToResponses(effort)
|
||||
}
|
||||
|
||||
// mapAnthropicEffortToResponses maps shared effort levels directly and maps
|
||||
// Anthropic's max to OpenAI's xhigh.
|
||||
func mapAnthropicEffortToResponses(effort string) string {
|
||||
if effort == "max" {
|
||||
return "xhigh"
|
||||
@@ -469,10 +470,38 @@ func boolPtr(v bool) *bool {
|
||||
|
||||
// isReasoningModel reports whether model is a reasoning model that does not
|
||||
// support sampling parameters (temperature, top_p) via the Responses API.
|
||||
// All gpt-5.x models are reasoning-only; the Responses API returns
|
||||
// "Unsupported parameter: temperature" if these fields are present.
|
||||
// GPT-5 and every later generation are reasoning-only; the Responses API
|
||||
// returns "Unsupported parameter: temperature" if these fields are present.
|
||||
//
|
||||
// Keyed on the generation number instead of a "gpt-5" prefix: pinning the
|
||||
// prefix meant each new family (gpt-6-astra and whatever follows) silently
|
||||
// fell through to the sampling branch and failed upstream on every compat
|
||||
// request until someone edited this line.
|
||||
func isReasoningModel(model string) bool {
|
||||
return strings.HasPrefix(model, "gpt-5") || openai.IsGPT6SolOrLunaModelSpelling(model)
|
||||
major, ok := openAIModelGeneration(model)
|
||||
return (ok && major >= 5) || openai.IsGPT6SolOrLunaModelSpelling(model)
|
||||
}
|
||||
|
||||
// openAIModelGeneration extracts N from a "gpt-N[.M][-suffix]" model id.
|
||||
// ok is false for non-GPT ids and for GPT families that carry no numeric
|
||||
// generation (gpt-image-1, gpt-audio, ...).
|
||||
func openAIModelGeneration(model string) (int, bool) {
|
||||
rest, ok := strings.CutPrefix(strings.ToLower(strings.TrimSpace(model)), "gpt-")
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
major, digits := 0, 0
|
||||
for _, r := range rest {
|
||||
if r < '0' || r > '9' {
|
||||
break
|
||||
}
|
||||
major = major*10 + int(r-'0')
|
||||
digits++
|
||||
}
|
||||
if digits == 0 {
|
||||
return 0, false
|
||||
}
|
||||
return major, true
|
||||
}
|
||||
|
||||
// normalizeToolParameters ensures the tool parameter schema is valid for
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
@@ -52,7 +53,7 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
for _, block := range resp.Content {
|
||||
switch block.Type {
|
||||
case "thinking", "redacted_thinking":
|
||||
if claude.IsOpus55(resp.Model) && (block.Signature != "" || block.Data != "") {
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(resp.Model)) && (block.Signature != "" || block.Data != "") {
|
||||
item := ResponsesOutput{Type: "reasoning", ID: generateItemID(), EncryptedContent: encodeAnthropicThinking(block)}
|
||||
if block.Thinking != "" {
|
||||
item.Summary = []ResponsesSummary{{Type: "summary_text", Text: block.Thinking}}
|
||||
@@ -71,7 +72,7 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
})
|
||||
}
|
||||
case "text":
|
||||
if claude.IsOpus55(resp.Model) && block.Text != "" {
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(resp.Model)) && block.Text != "" {
|
||||
outputs = append(outputs, ResponsesOutput{Type: "message", ID: generateItemID(), Role: "assistant", Status: "completed", Content: []ResponsesContentPart{{Type: "output_text", Text: block.Text}}})
|
||||
continue
|
||||
}
|
||||
@@ -198,6 +199,12 @@ type AnthropicEventToResponsesState struct {
|
||||
CurrentThinking AnthropicContentBlock
|
||||
PreserveThinkingSignatures bool
|
||||
|
||||
// PendingToolInput holds tool arguments that arrived complete on
|
||||
// content_block_start instead of as input_json_delta. It is only consumed at
|
||||
// content_block_stop, and only when no delta ever arrived, so a canonical
|
||||
// Anthropic stream keeps its exact event sequence.
|
||||
PendingToolInput string
|
||||
|
||||
// Outputs accumulates every closed output item so that response.completed
|
||||
// can carry the full output list. The OpenAI SDK's get_final_response()
|
||||
// parses the terminal event's response directly; without this, clients see
|
||||
@@ -278,7 +285,7 @@ func ResponsesEventToSSE(evt ResponsesStreamEvent) (string, error) {
|
||||
func anthToResHandleMessageStart(evt *AnthropicStreamEvent, state *AnthropicEventToResponsesState) []ResponsesStreamEvent {
|
||||
if evt.Message != nil {
|
||||
state.ResponseID = evt.Message.ID
|
||||
state.PreserveThinkingSignatures = state.PreserveThinkingSignatures || claude.IsOpus55(evt.Message.Model)
|
||||
state.PreserveThinkingSignatures = state.PreserveThinkingSignatures || claude.IsOpus55(evt.Message.Model) || claude.IsSonnet55(evt.Message.Model)
|
||||
if state.Model == "" {
|
||||
state.Model = evt.Message.Model
|
||||
}
|
||||
@@ -382,6 +389,12 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi
|
||||
state.CurrentItemType = "function_call"
|
||||
state.CurrentCallID = toResponsesCallID(evt.ContentBlock.ID)
|
||||
state.CurrentName = evt.ContentBlock.Name
|
||||
// The canonical Anthropic stream leaves input empty here and streams the
|
||||
// arguments as input_json_delta, but Anthropic-compatible relays may put
|
||||
// the complete arguments on this event and never send a delta. Keep them
|
||||
// as a seed rather than emitting now: a delta, if one follows, is
|
||||
// authoritative and must not be concatenated onto this JSON.
|
||||
state.PendingToolInput = seedToolArguments(evt.ContentBlock.Input)
|
||||
|
||||
events = append(events, makeResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
@@ -432,6 +445,9 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
if evt.Delta.PartialJSON == "" {
|
||||
return nil
|
||||
}
|
||||
// A real delta supersedes whatever content_block_start carried; keeping
|
||||
// both would splice two complete JSON documents together.
|
||||
state.PendingToolInput = ""
|
||||
state.CurrentArgs += evt.Delta.PartialJSON
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
@@ -467,21 +483,35 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic
|
||||
return events
|
||||
|
||||
case "function_call":
|
||||
var events []ResponsesStreamEvent
|
||||
// No delta ever arrived, so the arguments the upstream put on
|
||||
// content_block_start are all there is. Emit them as one delta here so
|
||||
// the done event below still repeats exactly what the deltas streamed.
|
||||
if state.CurrentArgs == "" && state.PendingToolInput != "" {
|
||||
state.CurrentArgs = state.PendingToolInput
|
||||
events = append(events, makeResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
Delta: state.PendingToolInput,
|
||||
ItemID: state.CurrentItemID,
|
||||
CallID: state.CurrentCallID,
|
||||
Name: state.CurrentName,
|
||||
}))
|
||||
}
|
||||
state.PendingToolInput = ""
|
||||
|
||||
// Emit function_call_arguments.done + output item done.
|
||||
// arguments must repeat exactly what the deltas already streamed for this
|
||||
// item: clients reconcile the done event against the accumulated
|
||||
// function_call_arguments.delta payloads and reject the call as
|
||||
// inconsistent_tool_call when the two disagree. Omitting the field left it
|
||||
// empty while the deltas carried the whole JSON.
|
||||
events := []ResponsesStreamEvent{
|
||||
makeResponsesEvent(state, "response.function_call_arguments.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
CallID: state.CurrentCallID,
|
||||
Name: state.CurrentName,
|
||||
Arguments: state.CurrentArgs,
|
||||
}),
|
||||
}
|
||||
events = append(events, makeResponsesEvent(state, "response.function_call_arguments.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
CallID: state.CurrentCallID,
|
||||
Name: state.CurrentName,
|
||||
Arguments: state.CurrentArgs,
|
||||
}))
|
||||
events = append(events, closeCurrentResponsesItem(state)...)
|
||||
return events
|
||||
|
||||
@@ -562,6 +592,19 @@ func anthropicResponsesStreamTerminalState(stopReason string) (string, *Response
|
||||
return "completed", nil
|
||||
}
|
||||
|
||||
// seedToolArguments normalizes a tool_use content block's inline input into a
|
||||
// seed for the streaming converter. Empty, absent and no-argument payloads
|
||||
// return "" so the existing "{}" fallback still applies and no empty delta is
|
||||
// synthesized.
|
||||
func seedToolArguments(input json.RawMessage) string {
|
||||
trimmed := strings.TrimSpace(string(input))
|
||||
switch trimmed {
|
||||
case "", "{}", "null":
|
||||
return ""
|
||||
}
|
||||
return trimmed
|
||||
}
|
||||
|
||||
func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []ResponsesStreamEvent {
|
||||
if state.CurrentItemType == "" {
|
||||
return nil
|
||||
@@ -605,6 +648,7 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response
|
||||
state.CurrentName = ""
|
||||
state.CurrentContent = nil
|
||||
state.CurrentArgs = ""
|
||||
state.PendingToolInput = ""
|
||||
state.CurrentSummary = ""
|
||||
state.CurrentThinking = AnthropicContentBlock{}
|
||||
state.TextAccum = ""
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// collectToolCallStreamEvents drives a tool_use block through the Anthropic →
|
||||
// Responses stream converter and returns every emitted event in order.
|
||||
func collectToolCallStreamEvents(t *testing.T, blockInput json.RawMessage, partials []string) []ResponsesStreamEvent {
|
||||
t.Helper()
|
||||
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
var all []ResponsesStreamEvent
|
||||
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
|
||||
Type: "message_start",
|
||||
Message: &AnthropicResponse{
|
||||
ID: "msg_tool_input",
|
||||
Model: "gemini-3.7-flash",
|
||||
Usage: AnthropicUsage{InputTokens: 12},
|
||||
},
|
||||
}, state)...)
|
||||
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
|
||||
Type: "content_block_start",
|
||||
ContentBlock: &AnthropicContentBlock{
|
||||
Type: "tool_use",
|
||||
ID: "toolu_01abc",
|
||||
Name: "eval",
|
||||
Input: blockInput,
|
||||
},
|
||||
}, state)...)
|
||||
|
||||
for _, partial := range partials {
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
|
||||
Type: "content_block_delta",
|
||||
Delta: &AnthropicDelta{Type: "input_json_delta", PartialJSON: partial},
|
||||
}, state)...)
|
||||
}
|
||||
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{Type: "content_block_stop"}, state)...)
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
|
||||
Type: "message_delta",
|
||||
Usage: &AnthropicUsage{OutputTokens: 9},
|
||||
}, state)...)
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{Type: "message_stop"}, state)...)
|
||||
|
||||
return all
|
||||
}
|
||||
|
||||
func concatArgumentDeltas(events []ResponsesStreamEvent) string {
|
||||
out := ""
|
||||
for _, e := range events {
|
||||
if e.Type == "response.function_call_arguments.delta" {
|
||||
out += e.Delta
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func findFunctionCallOutput(events []ResponsesStreamEvent) *ResponsesOutput {
|
||||
for i := range events {
|
||||
if events[i].Type != "response.output_item.done" || events[i].Item == nil {
|
||||
continue
|
||||
}
|
||||
if events[i].Item.Type == "function_call" {
|
||||
return events[i].Item
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Upstreams that are not the canonical Anthropic API (here: an Anthropic-compatible
|
||||
// relay fronting Gemini) put the complete tool arguments on content_block_start and
|
||||
// never emit an input_json_delta. The converter must not drop them: this repository
|
||||
// already reads ContentBlock.Input on the non-streaming path
|
||||
// (anthropicResponseToResponsesOutputs) and on the ChatCompletions bridge, so the
|
||||
// streaming path has to agree. Dropping it makes every client-side tool call arrive
|
||||
// with `{}` and the agent loop cannot proceed.
|
||||
func TestAnthropicEventToResponses_ToolInputOnContentBlockStart(t *testing.T) {
|
||||
const args = `{"language":"py","code":"print(1)"}`
|
||||
|
||||
events := collectToolCallStreamEvents(t, json.RawMessage(args), nil)
|
||||
|
||||
item := findFunctionCallOutput(events)
|
||||
require.NotNil(t, item, "a function_call output item must be emitted")
|
||||
assert.Equal(t, args, item.Arguments,
|
||||
"arguments carried on content_block_start must survive to output_item.done")
|
||||
assert.Equal(t, "eval", item.Name)
|
||||
|
||||
// Clients that accumulate from the event stream (rather than reading the
|
||||
// terminal item) must see the arguments too.
|
||||
assert.Equal(t, args, concatArgumentDeltas(events),
|
||||
"argument deltas must reconstruct the full arguments JSON")
|
||||
|
||||
done := findEvent(events, "response.function_call_arguments.done")
|
||||
require.NotNil(t, done, "function_call_arguments.done must be emitted")
|
||||
assert.Equal(t, args, done.Arguments,
|
||||
"function_call_arguments.done must carry the complete arguments")
|
||||
|
||||
completed := findEvent(events, "response.completed")
|
||||
require.NotNil(t, completed)
|
||||
require.NotNil(t, completed.Response)
|
||||
require.Len(t, completed.Response.Output, 1)
|
||||
assert.Equal(t, args, completed.Response.Output[0].Arguments,
|
||||
"response.completed must carry the same arguments")
|
||||
}
|
||||
|
||||
// The canonical Anthropic shape (empty input on content_block_start, arguments
|
||||
// streamed as input_json_delta) must keep its existing event sequence exactly.
|
||||
func TestAnthropicEventToResponses_ToolInputFromDeltasUnchanged(t *testing.T) {
|
||||
const args = `{"language":"py","code":"print(1)"}`
|
||||
|
||||
events := collectToolCallStreamEvents(t, json.RawMessage(`{}`),
|
||||
[]string{`{"language":"py",`, `"code":"print(1)"}`})
|
||||
|
||||
item := findFunctionCallOutput(events)
|
||||
require.NotNil(t, item)
|
||||
assert.Equal(t, args, item.Arguments)
|
||||
|
||||
// Exactly the two upstream deltas, not duplicated by the seed.
|
||||
var deltas []string
|
||||
for _, e := range events {
|
||||
if e.Type == "response.function_call_arguments.delta" {
|
||||
deltas = append(deltas, e.Delta)
|
||||
}
|
||||
}
|
||||
assert.Equal(t, []string{`{"language":"py",`, `"code":"print(1)"}`}, deltas,
|
||||
"canonical delta streaming must not gain or lose events")
|
||||
}
|
||||
|
||||
// An upstream that sends both a populated content_block_start and real deltas
|
||||
// must not have the two concatenated into malformed JSON.
|
||||
func TestAnthropicEventToResponses_ToolInputSeedNotDuplicatedByDeltas(t *testing.T) {
|
||||
const args = `{"language":"py","code":"print(1)"}`
|
||||
|
||||
events := collectToolCallStreamEvents(t, json.RawMessage(args),
|
||||
[]string{`{"language":"py",`, `"code":"print(1)"}`})
|
||||
|
||||
item := findFunctionCallOutput(events)
|
||||
require.NotNil(t, item)
|
||||
assert.Equal(t, args, item.Arguments, "deltas win over the content_block_start seed")
|
||||
assert.True(t, json.Valid([]byte(item.Arguments)), "arguments must stay valid JSON")
|
||||
assert.Equal(t, args, concatArgumentDeltas(events), "seed must not be replayed alongside deltas")
|
||||
}
|
||||
|
||||
// An empty or absent input on content_block_start with no deltas keeps the
|
||||
// existing "{}" fallback rather than emitting a spurious delta.
|
||||
func TestAnthropicEventToResponses_ToolInputEmptyStaysEmptyObject(t *testing.T) {
|
||||
for name, input := range map[string]json.RawMessage{
|
||||
"absent": nil,
|
||||
"empty": json.RawMessage(`{}`),
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
events := collectToolCallStreamEvents(t, input, nil)
|
||||
|
||||
item := findFunctionCallOutput(events)
|
||||
require.NotNil(t, item)
|
||||
assert.Equal(t, "{}", item.Arguments)
|
||||
assert.Empty(t, concatArgumentDeltas(events),
|
||||
"no argument delta should be synthesized when there are no arguments")
|
||||
})
|
||||
}
|
||||
}
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// driveToolUseThroughRestorer runs a tool_use block through the full outbound
|
||||
// chain the gateway actually uses: the Anthropic→Responses stream converter
|
||||
// feeding ResponsesClientToolStreamRestorer. Returns the events a client sees.
|
||||
func driveToolUseThroughRestorer(
|
||||
blockInput json.RawMessage,
|
||||
partials []string,
|
||||
mapping ResponsesClientToolMapping,
|
||||
) []ResponsesStreamEvent {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
restorer := NewResponsesClientToolStreamRestorer(mapping)
|
||||
|
||||
var clientEvents []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
for _, converted := range AnthropicEventToResponsesEvents(evt, state) {
|
||||
clientEvents = append(clientEvents, restorer.Restore(converted)...)
|
||||
}
|
||||
}
|
||||
|
||||
feed(&AnthropicStreamEvent{
|
||||
Type: "message_start",
|
||||
Message: &AnthropicResponse{
|
||||
ID: "msg_chain",
|
||||
Model: "gemini-3.7-flash",
|
||||
Usage: AnthropicUsage{InputTokens: 11},
|
||||
},
|
||||
})
|
||||
feed(&AnthropicStreamEvent{
|
||||
Type: "content_block_start",
|
||||
ContentBlock: &AnthropicContentBlock{
|
||||
Type: "tool_use",
|
||||
ID: "toolu_chain",
|
||||
Name: "eval",
|
||||
Input: blockInput,
|
||||
},
|
||||
})
|
||||
for _, partial := range partials {
|
||||
feed(&AnthropicStreamEvent{
|
||||
Type: "content_block_delta",
|
||||
Delta: &AnthropicDelta{Type: "input_json_delta", PartialJSON: partial},
|
||||
})
|
||||
}
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop"})
|
||||
feed(&AnthropicStreamEvent{Type: "message_delta", Usage: &AnthropicUsage{OutputTokens: 6}})
|
||||
feed(&AnthropicStreamEvent{Type: "message_stop"})
|
||||
|
||||
return clientEvents
|
||||
}
|
||||
|
||||
func firstEventOfType(events []ResponsesStreamEvent, typ string) *ResponsesStreamEvent {
|
||||
for i := range events {
|
||||
if events[i].Type == typ {
|
||||
return &events[i]
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// A client-side custom tool whose arguments arrived only on content_block_start
|
||||
// must still reach the client as a populated custom_tool_call. The restorer
|
||||
// suppresses function_call_arguments.delta for adapted tools and rebuilds the
|
||||
// input from the accumulated arguments, so a converter that drops the inline
|
||||
// input yields an empty tool call even though nothing errors.
|
||||
func TestToolInputOnContentBlockStart_SurvivesClientToolRestore(t *testing.T) {
|
||||
mapping := ResponsesClientToolMapping{CustomTools: map[string]bool{"eval": true}}
|
||||
|
||||
events := driveToolUseThroughRestorer(
|
||||
json.RawMessage(`{"input":"print(1)"}`), nil, mapping)
|
||||
|
||||
inputDone := firstEventOfType(events, "response.custom_tool_call_input.done")
|
||||
require.NotNil(t, inputDone, "custom_tool_call_input.done must be emitted")
|
||||
assert.Equal(t, "print(1)", inputDone.Input,
|
||||
"inline tool input must survive conversion and client-tool restoration")
|
||||
|
||||
var itemDone *ResponsesStreamEvent
|
||||
for i := range events {
|
||||
if events[i].Type == "response.output_item.done" && events[i].Item != nil &&
|
||||
events[i].Item.Type == "custom_tool_call" {
|
||||
itemDone = &events[i]
|
||||
}
|
||||
}
|
||||
require.NotNil(t, itemDone, "a custom_tool_call item must be closed")
|
||||
assert.Equal(t, "print(1)", itemDone.Item.Input)
|
||||
}
|
||||
|
||||
// The same block delivered as a plain (non-adapted) function tool must carry
|
||||
// its arguments through untouched.
|
||||
func TestToolInputOnContentBlockStart_PlainFunctionToolChain(t *testing.T) {
|
||||
const args = `{"language":"py","code":"print(1)"}`
|
||||
|
||||
events := driveToolUseThroughRestorer(
|
||||
json.RawMessage(args), nil, ResponsesClientToolMapping{})
|
||||
|
||||
done := firstEventOfType(events, "response.function_call_arguments.done")
|
||||
require.NotNil(t, done)
|
||||
assert.Equal(t, args, done.Arguments)
|
||||
|
||||
var itemDone *ResponsesStreamEvent
|
||||
for i := range events {
|
||||
if events[i].Type == "response.output_item.done" && events[i].Item != nil &&
|
||||
events[i].Item.Type == "function_call" {
|
||||
itemDone = &events[i]
|
||||
}
|
||||
}
|
||||
require.NotNil(t, itemDone)
|
||||
assert.Equal(t, args, itemDone.Item.Arguments)
|
||||
}
|
||||
|
||||
// Canonical delta streaming through the same chain must keep working, so the
|
||||
// seed path cannot be mistaken for the only source of arguments.
|
||||
func TestToolInputFromDeltas_SurvivesClientToolRestore(t *testing.T) {
|
||||
mapping := ResponsesClientToolMapping{CustomTools: map[string]bool{"eval": true}}
|
||||
|
||||
events := driveToolUseThroughRestorer(
|
||||
json.RawMessage(`{}`),
|
||||
[]string{`{"input":"pri`, `nt(1)"}`},
|
||||
mapping)
|
||||
|
||||
inputDone := firstEventOfType(events, "response.custom_tool_call_input.done")
|
||||
require.NotNil(t, inputDone)
|
||||
assert.Equal(t, "print(1)", inputDone.Input)
|
||||
}
|
||||
@@ -99,13 +99,8 @@ func AnthropicToChatCompletionsRequest(req *AnthropicRequest) (*ChatCompletionsR
|
||||
}
|
||||
}
|
||||
|
||||
// Reasoning effort: output_config.effort maps 1:1 (max→xhigh). thinking.type
|
||||
// itself is ignored (the Responses bridge behaves identically).
|
||||
effort := "medium"
|
||||
if req.OutputConfig != nil && req.OutputConfig.Effort != "" {
|
||||
effort = req.OutputConfig.Effort
|
||||
}
|
||||
out.ReasoningEffort = mapAnthropicEffortToResponses(effort)
|
||||
// Match the Responses bridge, including an explicit thinking disable.
|
||||
out.ReasoningEffort = anthropicReasoningEffort(req)
|
||||
|
||||
parallelToolCalls := true
|
||||
out.ParallelToolCalls = ¶llelToolCalls
|
||||
|
||||
@@ -150,6 +150,20 @@ func TestAnthropicToChatCompletionsRequest_ThinkingDropped(t *testing.T) {
|
||||
require.Empty(t, out.Messages[0].ReasoningContent)
|
||||
}
|
||||
|
||||
func TestAnthropicToChatCompletionsRequest_ThinkingDisabledOverridesOutputEffort(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "gpt-5.6-luna",
|
||||
MaxTokens: 1024,
|
||||
Messages: []AnthropicMessage{{Role: "user", Content: json.RawMessage(`"Hello"`)}},
|
||||
Thinking: &AnthropicThinking{Type: "disabled"},
|
||||
OutputConfig: &AnthropicOutputConfig{Effort: "max"},
|
||||
}
|
||||
|
||||
out, err := AnthropicToChatCompletionsRequest(req)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "none", out.ReasoningEffort)
|
||||
}
|
||||
|
||||
func TestAnthropicToChatCompletionsRequest_ToolChoiceAuto(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "claude-sonnet-4-20250514",
|
||||
|
||||
@@ -92,6 +92,56 @@ func TestChatCompletionsToResponses_SystemMessage(t *testing.T) {
|
||||
assert.Equal(t, "user", items[1].Role)
|
||||
}
|
||||
|
||||
func TestChatCompletionsToResponses_MessageTypesWithReasoningAndTools(t *testing.T) {
|
||||
req := &ChatCompletionsRequest{
|
||||
Model: "step-5-preview",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "system", Content: json.RawMessage(`"You are helpful."`)},
|
||||
{Role: "user", Content: json.RawMessage(`"Check the directory"`)},
|
||||
{
|
||||
Role: "assistant",
|
||||
Content: json.RawMessage(`"I will check."`),
|
||||
ReasoningContent: "Need to inspect the directory.",
|
||||
ToolCalls: []ChatToolCall{{
|
||||
ID: "call_1", Type: "function",
|
||||
Function: ChatFunctionCall{Name: "bash", Arguments: `{"cmd":"pwd"}`},
|
||||
}},
|
||||
},
|
||||
{Role: "tool", ToolCallID: "call_1", Content: json.RawMessage(`"/tmp"`)},
|
||||
{Role: "assistant", ReasoningContent: "The tool returned /tmp."},
|
||||
{Role: "user", Content: json.RawMessage(`"Continue"`)},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := ChatCompletionsToResponses(req)
|
||||
require.NoError(t, err)
|
||||
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(resp.Input, &items))
|
||||
require.Len(t, items, 7)
|
||||
for i, want := range []struct{ typ, role string }{
|
||||
{"message", "system"},
|
||||
{"message", "user"},
|
||||
{"message", "assistant"},
|
||||
{"function_call", ""},
|
||||
{"function_call_output", ""},
|
||||
{"message", "assistant"},
|
||||
{"message", "user"},
|
||||
} {
|
||||
assert.Equal(t, want.typ, items[i].Type, "input item %d type", i)
|
||||
assert.Equal(t, want.role, items[i].Role, "input item %d role", i)
|
||||
}
|
||||
var assistantContent, reasoningOnlyContent []ResponsesContentPart
|
||||
require.NoError(t, json.Unmarshal(items[2].Content, &assistantContent))
|
||||
require.NoError(t, json.Unmarshal(items[5].Content, &reasoningOnlyContent))
|
||||
require.Len(t, assistantContent, 1)
|
||||
require.Len(t, reasoningOnlyContent, 1)
|
||||
assert.Equal(t, "<thinking>Need to inspect the directory.</thinking>\nI will check.", assistantContent[0].Text)
|
||||
assert.Equal(t, "<thinking>The tool returned /tmp.</thinking>", reasoningOnlyContent[0].Text)
|
||||
assert.Equal(t, "call_1", items[3].CallID)
|
||||
assert.Equal(t, "call_1", items[4].CallID)
|
||||
}
|
||||
|
||||
func TestChatCompletionsToResponses_ToolCalls(t *testing.T) {
|
||||
req := &ChatCompletionsRequest{
|
||||
Model: "gpt-4o",
|
||||
|
||||
@@ -143,7 +143,7 @@ func chatSystemToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []ResponsesInputItem{{Role: "system", Content: content}}, nil
|
||||
return []ResponsesInputItem{{Type: "message", Role: "system", Content: content}}, nil
|
||||
}
|
||||
|
||||
// chatUserToResponses converts a user message, handling both plain strings and
|
||||
@@ -157,7 +157,7 @@ func chatUserToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []ResponsesInputItem{{Role: "user", Content: content}}, nil
|
||||
return []ResponsesInputItem{{Type: "message", Role: "user", Content: content}}, nil
|
||||
}
|
||||
|
||||
// chatAssistantToResponses converts an assistant message. If there is both
|
||||
@@ -192,7 +192,7 @@ func chatAssistantToResponses(m ChatMessage) ([]ResponsesInputItem, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, ResponsesInputItem{Role: "assistant", Content: partsJSON})
|
||||
items = append(items, ResponsesInputItem{Type: "message", Role: "assistant", Content: partsJSON})
|
||||
}
|
||||
|
||||
// Emit one function_call item per tool_call.
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// GPT-5 and every later generation are reasoning-only: forwarding temperature or
|
||||
// top_p makes the Responses API reject the request with "Unsupported parameter".
|
||||
// The check is keyed on the generation number so a new family (gpt-6-astra and
|
||||
// whatever follows) does not silently fall into the sampling branch.
|
||||
func TestIsReasoningModelCoversLaterGenerations(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
model string
|
||||
want bool
|
||||
}{
|
||||
{"gpt-6-astra", true},
|
||||
{"gpt-6", true},
|
||||
{"gpt-7-whatever", true},
|
||||
{"gpt-6-sol", true},
|
||||
{"gpt-5.5", true},
|
||||
{"gpt-5.2", true},
|
||||
{"gpt-5", true},
|
||||
{"GPT-6-Astra", true},
|
||||
{" gpt-6-astra ", true},
|
||||
{"gpt-4o", false},
|
||||
{"gpt-4.1", false},
|
||||
{"gpt-image-1", false},
|
||||
{"gpt-audio", false},
|
||||
{"claude-opus-4-6", false},
|
||||
{"gemini-3.1-pro", false},
|
||||
{"", false},
|
||||
} {
|
||||
if got := isReasoningModel(tc.model); got != tc.want {
|
||||
t.Errorf("isReasoningModel(%q) = %v, want %v", tc.model, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIModelGeneration(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
model string
|
||||
wantMajor int
|
||||
wantOK bool
|
||||
}{
|
||||
{"gpt-6-astra", 6, true},
|
||||
{"gpt-5.5", 5, true},
|
||||
{"gpt-4o", 4, true},
|
||||
{"gpt-10-future", 10, true},
|
||||
{"gpt-image-1", 0, false},
|
||||
{"claude-opus-4-6", 0, false},
|
||||
{"", 0, false},
|
||||
} {
|
||||
major, ok := openAIModelGeneration(tc.model)
|
||||
if major != tc.wantMajor || ok != tc.wantOK {
|
||||
t.Errorf("openAIModelGeneration(%q) = (%d, %v), want (%d, %v)",
|
||||
tc.model, major, ok, tc.wantMajor, tc.wantOK)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_TemperatureStrippedForGPT6Astra(t *testing.T) {
|
||||
temp := 0.7
|
||||
req := &AnthropicRequest{
|
||||
Model: "gpt-6-astra",
|
||||
MaxTokens: 1024,
|
||||
Messages: []AnthropicMessage{{Role: "user", Content: json.RawMessage(`"Hello"`)}},
|
||||
Temperature: &temp,
|
||||
TopP: &temp,
|
||||
}
|
||||
|
||||
resp, err := AnthropicToResponses(req)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, resp.Temperature, "gpt-6-astra is reasoning-only: temperature must be stripped")
|
||||
require.Nil(t, resp.TopP, "gpt-6-astra is reasoning-only: top_p must be stripped")
|
||||
}
|
||||
@@ -109,7 +109,23 @@ func TestResponsesToAnthropic_UnknownItemTypeKeepsRecognizableText(t *testing.T)
|
||||
require.NotContains(t, string(messages[0].Content), "drop me")
|
||||
}
|
||||
|
||||
// user 消息的分片全部不可识别时,以前会退化成 content:"",Anthropic 拒收空内容消息。
|
||||
// data URI 形式的 input_file 要变成 Anthropic document,供后续 Gemini inlineData 使用。
|
||||
func TestResponsesToAnthropic_InputFileDataURIBecomesDocument(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"message","role":"user","content":[
|
||||
{"type":"input_text","text":"read this"},
|
||||
{"type":"input_file","filename":"token.pdf","file_data":"data:application/pdf;base64,JVBERi0="}
|
||||
]}
|
||||
]`)
|
||||
|
||||
requireAnthropicMessagesAreSendable(t, messages)
|
||||
require.Len(t, messages, 1)
|
||||
require.Contains(t, string(messages[0].Content), `"type":"document"`)
|
||||
require.Contains(t, string(messages[0].Content), `"media_type":"application/pdf"`)
|
||||
require.Contains(t, string(messages[0].Content), `"data":"JVBERi0="`)
|
||||
}
|
||||
|
||||
// 只有 file_id、没有 data URI 的 input_file 仍然无法转换,整条消息丢掉。
|
||||
func TestResponsesToAnthropic_UserMessageWithOnlyUnknownPartsIsDropped(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"message","role":"user","content":[{"type":"input_file","file_id":"file_1"}]}
|
||||
|
||||
@@ -14,7 +14,9 @@ import (
|
||||
// enables Anthropic platform groups to accept OpenAI Responses API requests
|
||||
// by converting them to the native /v1/messages format before forwarding upstream.
|
||||
func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, error) {
|
||||
system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input, claude.IsOpus55(req.Model))
|
||||
isOpus55 := claude.IsOpus55(req.Model)
|
||||
isSonnet55 := claude.IsSonnet55(req.Model)
|
||||
system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input, isOpus55 || isSonnet55)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -54,9 +56,11 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
out.ToolChoice = tc
|
||||
}
|
||||
|
||||
// Opus 5.5 always uses adaptive thinking. Resolve the upstream model before
|
||||
// conversion, because client aliases need not identify a Claude model.
|
||||
if claude.IsOpus55(req.Model) {
|
||||
// The 5.5 models reject manual thinking and forced tool use. Sonnet 5.5
|
||||
// additionally supports between_tools to disable up-front thinking.
|
||||
// Resolve the upstream model before conversion: client aliases need not
|
||||
// identify a Claude model.
|
||||
if isOpus55 || isSonnet55 {
|
||||
var choice struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
@@ -66,16 +70,36 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
}
|
||||
}
|
||||
if choice.Type == "any" || choice.Type == "tool" {
|
||||
return nil, fmt.Errorf("claude-opus-5-5 does not support forced tool_choice; use auto or none")
|
||||
return nil, fmt.Errorf("%s does not support forced tool_choice; use auto or none", req.Model)
|
||||
}
|
||||
effort := "medium"
|
||||
if isSonnet55 {
|
||||
effort = "high"
|
||||
if req.Temperature != nil && *req.Temperature != 1 {
|
||||
return nil, fmt.Errorf("claude-sonnet-5-5 does not support non-default temperature")
|
||||
}
|
||||
if req.TopP != nil && (*req.TopP < 0.99 || *req.TopP > 1) {
|
||||
return nil, fmt.Errorf("claude-sonnet-5-5 does not support non-default top_p")
|
||||
}
|
||||
}
|
||||
if req.Reasoning != nil && req.Reasoning.Effort != "" {
|
||||
effort = req.Reasoning.Effort
|
||||
}
|
||||
if isSonnet55 && effort == "none" {
|
||||
// OpenAI's no-reasoning request maps to Sonnet 5.5's lowest
|
||||
// thinking mode. between_tools still preserves signed progress
|
||||
// blocks produced during tool use.
|
||||
out.Thinking = &AnthropicThinking{Type: "between_tools"}
|
||||
if out.OutputConfig == nil {
|
||||
out.OutputConfig = &AnthropicOutputConfig{}
|
||||
}
|
||||
out.OutputConfig.Effort = "low"
|
||||
return out, nil
|
||||
}
|
||||
switch effort {
|
||||
case "low", "medium", "high", "xhigh", "max":
|
||||
default:
|
||||
return nil, fmt.Errorf("claude-opus-5-5 does not support reasoning effort %q; use low, medium, high, xhigh or max", effort)
|
||||
return nil, fmt.Errorf("%s does not support reasoning effort %q; use low, medium, high, xhigh or max", req.Model, effort)
|
||||
}
|
||||
out.Thinking = &AnthropicThinking{Type: "adaptive"}
|
||||
out.OutputConfig = &AnthropicOutputConfig{Effort: effort}
|
||||
@@ -530,6 +554,14 @@ func convertResponsesUserToAnthropicContent(raw json.RawMessage) (json.RawMessag
|
||||
Source: src,
|
||||
})
|
||||
}
|
||||
case "input_file":
|
||||
src := dataURIToAnthropicFileSource(p.FileData)
|
||||
if src != nil {
|
||||
blocks = append(blocks, AnthropicContentBlock{
|
||||
Type: "document",
|
||||
Source: src,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -617,6 +649,12 @@ func dataURIToAnthropicImageSource(dataURI string) *AnthropicImageSource {
|
||||
}
|
||||
}
|
||||
|
||||
// dataURIToAnthropicFileSource parses a data URI into a document source.
|
||||
// file_id-only parts are not convertible here and stay dropped.
|
||||
func dataURIToAnthropicFileSource(fileData string) *AnthropicImageSource {
|
||||
return dataURIToAnthropicImageSource(fileData)
|
||||
}
|
||||
|
||||
// mergeConsecutiveMessages merges consecutive messages with the same role
|
||||
// because Anthropic requires alternating user/assistant turns.
|
||||
func mergeConsecutiveMessages(messages []AnthropicMessage) []AnthropicMessage {
|
||||
|
||||
@@ -609,6 +609,24 @@ func (a *BufferedResponseAccumulator) SupplementResponseOutput(resp *ResponsesRe
|
||||
return
|
||||
}
|
||||
|
||||
// The terminal event can carry a non-empty output array whose message has
|
||||
// no usable text. Trusting it as-is silently drops the text that already
|
||||
// streamed: the client gets an empty reply while usage still bills the
|
||||
// terminal output_tokens. Refill it from the accumulated deltas; non-empty
|
||||
// terminal text stays authoritative.
|
||||
if a.text.Len() > 0 && !responsesOutputHasText(resp.Output) {
|
||||
if !fillResponsesOutputText(resp.Output, a.text.String()) {
|
||||
resp.Output = append(resp.Output, ResponsesOutput{
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []ResponsesContentPart{{
|
||||
Type: "output_text",
|
||||
Text: a.text.String(),
|
||||
}},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
for outputIndex := range resp.Output {
|
||||
item := &resp.Output[outputIndex]
|
||||
if item.Type != "function_call" || item.Arguments != "" {
|
||||
@@ -627,3 +645,46 @@ func (a *BufferedResponseAccumulator) SupplementResponseOutput(resp *ResponsesRe
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// responsesOutputHasText reports whether the terminal output already carries
|
||||
// usable message text. Whitespace-only text does not count.
|
||||
func responsesOutputHasText(output []ResponsesOutput) bool {
|
||||
for i := range output {
|
||||
if output[i].Type != "message" {
|
||||
continue
|
||||
}
|
||||
for _, part := range output[i].Content {
|
||||
if part.Type == "output_text" && strings.TrimSpace(part.Text) != "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// fillResponsesOutputText writes text into the first empty output_text part of
|
||||
// the first message item, adding an output_text part when that message has
|
||||
// none. It returns false when the output holds no message item, leaving the
|
||||
// caller to append one.
|
||||
func fillResponsesOutputText(output []ResponsesOutput, text string) bool {
|
||||
for i := range output {
|
||||
if output[i].Type != "message" {
|
||||
continue
|
||||
}
|
||||
for j := range output[i].Content {
|
||||
if output[i].Content[j].Type != "output_text" {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(output[i].Content[j].Text) == "" {
|
||||
output[i].Content[j].Text = text
|
||||
return true
|
||||
}
|
||||
}
|
||||
output[i].Content = append(output[i].Content, ResponsesContentPart{
|
||||
Type: "output_text",
|
||||
Text: text,
|
||||
})
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// When the terminal event carries a non-empty output array whose message has
|
||||
// no text, the text that already streamed must be restored. Otherwise the
|
||||
// client gets an empty reply while usage still bills the terminal
|
||||
// output_tokens. The accumulator previously rebuilt output only when the
|
||||
// terminal output array was empty, so this upstream shape silently lost text.
|
||||
func TestSupplementResponseOutput_RecoversTextWhenTerminalMessageIsEmpty(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{Type: "response.output_text.delta", Delta: "Hello, "})
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{Type: "response.output_text.delta", Delta: "world"})
|
||||
|
||||
resp := &ResponsesResponse{
|
||||
ID: "resp_empty_msg",
|
||||
Status: "completed",
|
||||
Output: []ResponsesOutput{{
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []ResponsesContentPart{},
|
||||
}},
|
||||
Usage: &ResponsesUsage{InputTokens: 13168, OutputTokens: 3329},
|
||||
}
|
||||
|
||||
acc.SupplementResponseOutput(resp)
|
||||
|
||||
chat := ResponsesToChatCompletions(resp, "gpt-5.5")
|
||||
require.Len(t, chat.Choices, 1)
|
||||
require.NotNil(t, chat.Choices[0].Message.Content, "the client must not receive an empty reply")
|
||||
|
||||
var got string
|
||||
require.NoError(t, json.Unmarshal(chat.Choices[0].Message.Content, &got))
|
||||
assert.Equal(t, "Hello, world", got)
|
||||
}
|
||||
|
||||
// Whitespace-only terminal text counts as missing too.
|
||||
func TestSupplementResponseOutput_RecoversTextWhenTerminalTextIsBlank(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{Type: "response.output_text.delta", Delta: "real text"})
|
||||
|
||||
resp := &ResponsesResponse{
|
||||
Status: "completed",
|
||||
Output: []ResponsesOutput{{
|
||||
Type: "message",
|
||||
Content: []ResponsesContentPart{{Type: "output_text", Text: " "}},
|
||||
}},
|
||||
}
|
||||
|
||||
acc.SupplementResponseOutput(resp)
|
||||
|
||||
chat := ResponsesToChatCompletions(resp, "m")
|
||||
var got string
|
||||
require.NotNil(t, chat.Choices[0].Message.Content)
|
||||
require.NoError(t, json.Unmarshal(chat.Choices[0].Message.Content, &got))
|
||||
assert.Equal(t, "real text", got)
|
||||
}
|
||||
|
||||
// A terminal output array without any message item gets one appended to carry
|
||||
// the accumulated text; existing items are kept.
|
||||
func TestSupplementResponseOutput_AppendsMessageWhenTerminalHasNoMessage(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{Type: "response.output_text.delta", Delta: "only in the stream"})
|
||||
|
||||
resp := &ResponsesResponse{
|
||||
Status: "completed",
|
||||
Output: []ResponsesOutput{{
|
||||
Type: "function_call",
|
||||
CallID: "call_1",
|
||||
Name: "verify",
|
||||
}},
|
||||
}
|
||||
|
||||
acc.SupplementResponseOutput(resp)
|
||||
|
||||
var text string
|
||||
for _, item := range resp.Output {
|
||||
if item.Type != "message" {
|
||||
continue
|
||||
}
|
||||
for _, p := range item.Content {
|
||||
text += p.Text
|
||||
}
|
||||
}
|
||||
assert.Equal(t, "only in the stream", text)
|
||||
|
||||
require.Len(t, resp.Output, 2)
|
||||
assert.Equal(t, "function_call", resp.Output[0].Type)
|
||||
assert.Equal(t, "verify", resp.Output[0].Name)
|
||||
}
|
||||
|
||||
// Non-empty terminal text stays authoritative and is never overwritten by the
|
||||
// accumulated stream.
|
||||
func TestSupplementResponseOutput_KeepsTerminalTextAuthoritative(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{Type: "response.output_text.delta", Delta: "from the stream"})
|
||||
|
||||
resp := &ResponsesResponse{
|
||||
Status: "completed",
|
||||
Output: []ResponsesOutput{{
|
||||
Type: "message",
|
||||
Content: []ResponsesContentPart{{Type: "output_text", Text: "from the terminal event"}},
|
||||
}},
|
||||
}
|
||||
|
||||
acc.SupplementResponseOutput(resp)
|
||||
|
||||
require.Len(t, resp.Output, 1)
|
||||
assert.Equal(t, "from the terminal event", resp.Output[0].Content[0].Text)
|
||||
}
|
||||
|
||||
// Without streamed text no synthetic message is created.
|
||||
func TestSupplementResponseOutput_NoSyntheticMessageWithoutStreamText(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
|
||||
resp := &ResponsesResponse{
|
||||
Status: "completed",
|
||||
Output: []ResponsesOutput{{
|
||||
Type: "message",
|
||||
Content: []ResponsesContentPart{},
|
||||
}},
|
||||
}
|
||||
|
||||
acc.SupplementResponseOutput(resp)
|
||||
|
||||
require.Len(t, resp.Output, 1)
|
||||
assert.Empty(t, resp.Output[0].Content)
|
||||
}
|
||||
@@ -42,7 +42,7 @@ type AnthropicOutputConfig struct {
|
||||
|
||||
// AnthropicThinking configures extended thinking in the Anthropic API.
|
||||
type AnthropicThinking struct {
|
||||
Type string `json:"type"` // "enabled" | "adaptive" | "disabled"
|
||||
Type string `json:"type"` // "enabled" | "adaptive" | "disabled" | "between_tools"
|
||||
BudgetTokens int `json:"budget_tokens,omitempty"` // max thinking tokens
|
||||
}
|
||||
|
||||
@@ -263,7 +263,7 @@ type ResponsesText struct {
|
||||
// The Type field determines which other fields are populated.
|
||||
type ResponsesInputItem struct {
|
||||
// Common
|
||||
Type string `json:"type,omitempty"` // "" for role-based messages
|
||||
Type string `json:"type,omitempty"` // "message" for role-based messages
|
||||
|
||||
// Role-based messages (developer/system/user/assistant)
|
||||
Role string `json:"role,omitempty"`
|
||||
|
||||
@@ -18,6 +18,8 @@ const (
|
||||
BetaTokenCounting = "token-counting-2024-11-01"
|
||||
BetaContext1M = "context-1m-2025-08-07"
|
||||
BetaFastMode = "fast-mode-2026-02-01"
|
||||
// Legacy structured output compatibility; forwarded only when explicitly requested.
|
||||
BetaStructuredOutputs = "structured-outputs-2025-11-13"
|
||||
|
||||
// 新增(对齐官方 CLI 2.1.9x 以来的流量)
|
||||
BetaPromptCachingScope = "prompt-caching-scope-2026-01-05"
|
||||
@@ -188,6 +190,12 @@ var DefaultModels = []Model{
|
||||
DisplayName: "Claude Opus 5",
|
||||
CreatedAt: "2026-07-25T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-sonnet-5-5",
|
||||
Type: "model",
|
||||
DisplayName: "Claude Sonnet 5.5",
|
||||
CreatedAt: "2026-09-28T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-sonnet-5",
|
||||
Type: "model",
|
||||
|
||||
@@ -30,3 +30,16 @@ func TestDefaultModelsContainsOpus55(t *testing.T) {
|
||||
}
|
||||
t.Fatal("claude-opus-5-5 missing")
|
||||
}
|
||||
|
||||
func TestDefaultModelsContainsSonnet55(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
for _, model := range DefaultModels {
|
||||
if model.ID == "claude-sonnet-5-5" {
|
||||
require.Equal(t, "Claude Sonnet 5.5", model.DisplayName)
|
||||
require.Equal(t, "2026-09-28T00:00:00Z", model.CreatedAt)
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatal("claude-sonnet-5-5 missing from DefaultModels")
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ var effortFamilies = []struct {
|
||||
{family: "claude-mythos-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-fable-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-sonnet-4-6", levels: effortLowMediumHighMax},
|
||||
{family: "claude-sonnet-5-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-sonnet-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-opus-4-8", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-opus-4-7", levels: effortLowMediumHighXHighMax},
|
||||
@@ -45,14 +46,30 @@ func IsOpus55(model string) bool {
|
||||
return normalizeEffortModelID(model) == "claude-opus-5-5"
|
||||
}
|
||||
|
||||
// IsSonnet55 identifies the fixed Sonnet 5.5 ID after provider/local suffix normalization.
|
||||
func IsSonnet55(model string) bool {
|
||||
return normalizeEffortModelID(model) == "claude-sonnet-5-5"
|
||||
}
|
||||
|
||||
func normalizeEffortModelID(model string) string {
|
||||
id := strings.ToLower(strings.TrimSpace(model))
|
||||
id = strings.TrimPrefix(id, "models/")
|
||||
if slash := strings.IndexByte(id, '/'); slash >= 0 {
|
||||
id = strings.TrimPrefix(strings.TrimSpace(id[slash+1:]), "models/")
|
||||
}
|
||||
for _, prefix := range []string{"us.", "eu.", "apac.", "jp.", "au.", "us-gov.", "global."} {
|
||||
id = strings.TrimPrefix(id, prefix)
|
||||
}
|
||||
id = strings.TrimPrefix(id, "anthropic.")
|
||||
id = strings.TrimSuffix(id, "-thinking")
|
||||
// OpenRouter uses dotted minor versions for some models. Normalize them
|
||||
// before effort, thinking, and billing family lookups.
|
||||
if id == "claude-opus-5.5" {
|
||||
id = "claude-opus-5-5"
|
||||
}
|
||||
if id == "claude-sonnet-5.5" {
|
||||
id = "claude-sonnet-5-5"
|
||||
}
|
||||
if mapped, ok := ModelIDReverseOverrides[id]; ok {
|
||||
id = mapped
|
||||
}
|
||||
|
||||
@@ -16,6 +16,9 @@ func TestEffortLevelsForModel(t *testing.T) {
|
||||
{model: "claude-opus-4-6", want: []string{"low", "medium", "high", "max"}},
|
||||
{model: "anthropic/claude-sonnet-4-6", want: []string{"low", "medium", "high", "max"}},
|
||||
{model: "claude-opus-5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "anthropic/claude-opus-5.5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "claude-sonnet-5-5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "us.anthropic.claude-sonnet-5-5", want: []string{"low", "medium", "high", "xhigh", "max"}},
|
||||
{model: "claude-opus-4-5-20251101", want: []string{"low", "medium", "high"}},
|
||||
{model: "claude-haiku-4-5-20251001", want: nil},
|
||||
{model: "gpt-5.6", want: nil},
|
||||
@@ -27,3 +30,30 @@ func TestEffortLevelsForModel(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsOpus55OpenRouterExactAlias(t *testing.T) {
|
||||
t.Parallel()
|
||||
for _, model := range []string{"claude-opus-5-5", "anthropic/claude-opus-5.5"} {
|
||||
require.True(t, IsOpus55(model), model)
|
||||
}
|
||||
for _, model := range []string{"claude-opus-5", "anthropic/claude-opus-5.6", "anthropic/claude-opus-5.5-preview"} {
|
||||
require.False(t, IsOpus55(model), model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsSonnet55(t *testing.T) {
|
||||
t.Parallel()
|
||||
for _, model := range []string{
|
||||
"claude-sonnet-5-5",
|
||||
"anthropic/claude-sonnet-5.5",
|
||||
"anthropic.claude-sonnet-5-5",
|
||||
"us.anthropic.claude-sonnet-5-5",
|
||||
"us-gov.anthropic.claude-sonnet-5-5",
|
||||
"global.anthropic.claude-sonnet-5-5-thinking",
|
||||
} {
|
||||
require.True(t, IsSonnet55(model), model)
|
||||
}
|
||||
for _, model := range []string{"claude-sonnet-5", "claude-sonnet-5-5-preview", "claude-opus-5-5"} {
|
||||
require.False(t, IsSonnet55(model), model)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,6 +16,8 @@ func HTTPStatusToGoogleStatus(status int) string {
|
||||
return "NOT_FOUND"
|
||||
case http.StatusTooManyRequests:
|
||||
return "RESOURCE_EXHAUSTED"
|
||||
case 499: // client closed request
|
||||
return "CANCELLED"
|
||||
default:
|
||||
if status >= 500 {
|
||||
return "INTERNAL"
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
package googleapi
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestHTTPStatusToGoogleStatus(t *testing.T) {
|
||||
cases := []struct {
|
||||
status int
|
||||
want string
|
||||
}{
|
||||
{http.StatusBadRequest, "INVALID_ARGUMENT"},
|
||||
{http.StatusUnauthorized, "UNAUTHENTICATED"},
|
||||
{http.StatusForbidden, "PERMISSION_DENIED"},
|
||||
{http.StatusNotFound, "NOT_FOUND"},
|
||||
{http.StatusTooManyRequests, "RESOURCE_EXHAUSTED"},
|
||||
{499, "CANCELLED"},
|
||||
{http.StatusBadGateway, "INTERNAL"},
|
||||
{http.StatusServiceUnavailable, "INTERNAL"},
|
||||
{http.StatusConflict, "UNKNOWN"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
if got := HTTPStatusToGoogleStatus(tc.status); got != tc.want {
|
||||
t.Errorf("HTTPStatusToGoogleStatus(%d) = %q, want %q", tc.status, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
const (
|
||||
apiKeyRateLimitKeyPrefix = "apikey:ratelimit:"
|
||||
apiKeyRateLimitDuration = 24 * time.Hour
|
||||
apiKeyCreateCountKeyPrefix = "apikey:create_count:"
|
||||
apiKeyAuthCachePrefix = "apikey:auth:"
|
||||
authCacheInvalidateChannel = "auth:cache:invalidate"
|
||||
)
|
||||
@@ -24,6 +25,11 @@ func apiKeyRateLimitKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeyRateLimitKeyPrefix, userID)
|
||||
}
|
||||
|
||||
// apiKeyCreateCountKey generates the Redis key for per-user API key creation counting.
|
||||
func apiKeyCreateCountKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeyCreateCountKeyPrefix, userID)
|
||||
}
|
||||
|
||||
func apiKeyAuthCacheKey(key string) string {
|
||||
return fmt.Sprintf("%s%s", apiKeyAuthCachePrefix, key)
|
||||
}
|
||||
@@ -54,9 +60,18 @@ func (c *apiKeyCache) IncrementCreateAttemptCount(ctx context.Context, userID in
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *apiKeyCache) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
key := apiKeyRateLimitKey(userID)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
// IncrementCreateCount 在固定窗口内累加创建次数并返回累加后的值。
|
||||
// ExpireNX 只在首次创建计数键时设置过期,后续创建不会延长窗口;
|
||||
// MULTI 保证 INCR 与 EXPIRE 同时生效,避免计数键丢失 TTL 后永久封禁。
|
||||
func (c *apiKeyCache) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
key := apiKeyCreateCountKey(userID)
|
||||
pipe := c.rdb.TxPipeline()
|
||||
incr := pipe.Incr(ctx, key)
|
||||
pipe.ExpireNX(ctx, key, window)
|
||||
if _, err := pipe.Exec(ctx); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return incr.Val(), nil
|
||||
}
|
||||
|
||||
func (c *apiKeyCache) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
|
||||
@@ -51,19 +51,6 @@ func (s *ApiKeyCacheSuite) TestCreateAttemptCount() {
|
||||
s.AssertTTLWithin(ttl, 1*time.Second, apiKeyRateLimitDuration)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "delete_removes_key",
|
||||
fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
|
||||
userID := int64(1)
|
||||
|
||||
require.NoError(s.T(), cache.IncrementCreateAttemptCount(ctx, userID))
|
||||
require.NoError(s.T(), cache.DeleteCreateAttemptCount(ctx, userID), "DeleteCreateAttemptCount")
|
||||
|
||||
count, err := cache.GetCreateAttemptCount(ctx, userID)
|
||||
require.NoError(s.T(), err, "expected nil error after delete")
|
||||
require.Equal(s.T(), 0, count, "expected zero count after delete")
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyCacheIncrementCreateCountUsesFixedWindow(t *testing.T) {
|
||||
server := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
|
||||
defer func() { _ = client.Close() }()
|
||||
cache := NewAPIKeyCache(client)
|
||||
ctx := context.Background()
|
||||
key := apiKeyCreateCountKey(7)
|
||||
|
||||
count, err := cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
require.Equal(t, time.Hour, server.TTL(key))
|
||||
|
||||
// 后续创建不得延长窗口,否则持续创建会让计数永不过期。
|
||||
server.FastForward(40 * time.Minute)
|
||||
count, err = cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(2), count)
|
||||
require.Equal(t, 20*time.Minute, server.TTL(key))
|
||||
|
||||
// 其他用户独立计数。
|
||||
count, err = cache.IncrementCreateCount(ctx, 8, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
|
||||
// 窗口到期后重新计数。
|
||||
server.FastForward(21 * time.Minute)
|
||||
count, err = cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
}
|
||||
@@ -197,6 +197,8 @@ func (r *contentModerationRepository) CountFlaggedByUserSince(ctx context.Contex
|
||||
return 0, nil
|
||||
}
|
||||
// SQL 中的 'cyber_policy' 字面量须与 service.ContentModerationActionCyberPolicy 保持一致。
|
||||
// 'cyber_log_only' matches service.ContentModerationModeCyberLogOnly; these
|
||||
// events remain evidence but never become penalties after allowlist removal.
|
||||
var count int
|
||||
err := r.db.QueryRowContext(ctx, `
|
||||
WITH last_auto_ban AS (
|
||||
@@ -209,6 +211,8 @@ FROM content_moderation_logs
|
||||
WHERE user_id = $1
|
||||
AND flagged = TRUE
|
||||
AND action <> 'hash_block'
|
||||
AND mode <> 'cyber_log_only'
|
||||
AND mode <> 'risk_control_log_only'
|
||||
AND ($3::bool IS FALSE OR action <> 'cyber_policy')
|
||||
AND created_at >= $2
|
||||
AND created_at > COALESCE((SELECT at FROM last_auto_ban), '-infinity'::timestamptz)
|
||||
|
||||
@@ -102,3 +102,18 @@ func TestContentModerationRepositoryCountFlaggedByUserSince_ExcludesCyberPolicyW
|
||||
require.Equal(t, 3, count)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestContentModerationRepositoryCountFlaggedByUserSince_ExcludesCyberLogOnly(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
repo := NewContentModerationRepository(db)
|
||||
since := time.Now().Add(-time.Hour)
|
||||
mock.ExpectQuery(regexp.QuoteMeta("AND mode <> 'cyber_log_only'\n AND mode <> 'risk_control_log_only'")).
|
||||
WithArgs(int64(12), since, false).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
|
||||
count, err := repo.CountFlaggedByUserSince(context.Background(), 12, since, false)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, count)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
@@ -1034,6 +1034,13 @@ func filterSchedulerExtra(extra map[string]any) map[string]any {
|
||||
"auto_pause_7d_threshold",
|
||||
"auto_pause_5h_disabled",
|
||||
"auto_pause_7d_disabled",
|
||||
// 自动用卡:卡可用的 OpenAI 号在暂停阈值与用卡阈值之间继续调度。
|
||||
// 候选过滤读的是本投影,缺这几个键时放行分支永远不会生效,
|
||||
// 账号会在暂停阈值处被一刀切停调,直到窗口自然重置。
|
||||
service.OpenAIAutoResetCreditEnabledExtraKey,
|
||||
service.OpenAIAutoResetCredit5hThresholdExtraKey,
|
||||
service.OpenAIAutoResetCredit7dThresholdExtraKey,
|
||||
service.OpenAIAutoResetCreditStateExtraKey,
|
||||
"model_rate_limits",
|
||||
service.UpstreamBillingProbeExtraKey,
|
||||
service.GrokMediaEligibleExtraKey,
|
||||
|
||||
@@ -446,6 +446,47 @@ func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) {
|
||||
require.Equal(t, false, got.Extra["auto_pause_7d_disabled"])
|
||||
}
|
||||
|
||||
// 候选过滤读的是 Redis 元数据投影;自动用卡的开关、阈值和卡状态缺失时,
|
||||
// 卡可用的 OpenAI 号会在暂停阈值处被一刀切停调,放行到用卡阈值的逻辑永远不生效。
|
||||
func TestSchedulerMetadataPayload_KeepsOpenAIAutoResetCreditFields(t *testing.T) {
|
||||
checkedAt := time.Now().UTC().Format(time.RFC3339)
|
||||
account := service.Account{
|
||||
ID: 40,
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Extra: map[string]any{
|
||||
"codex_7d_used_percent": 90.0,
|
||||
service.OpenAIAutoResetCreditEnabledExtraKey: true,
|
||||
service.OpenAIAutoResetCredit5hThresholdExtraKey: 0.95,
|
||||
service.OpenAIAutoResetCredit7dThresholdExtraKey: 0.98,
|
||||
service.OpenAIAutoResetCreditStateExtraKey: map[string]any{
|
||||
"status": service.OpenAIAutoResetStatusAvailable,
|
||||
"available_count": 3,
|
||||
"checked_at": checkedAt,
|
||||
"trigger_window": "7d",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
_, metaPayload, err := marshalSchedulerCacheAccount(account)
|
||||
require.NoError(t, err)
|
||||
var cached service.Account
|
||||
require.NoError(t, json.Unmarshal(metaPayload, &cached))
|
||||
|
||||
config := service.ResolveOpenAIAutoResetCreditConfig(&cached)
|
||||
require.True(t, config.Enabled)
|
||||
require.Equal(t, 0.95, config.Threshold5h)
|
||||
require.Equal(t, 0.98, config.Threshold7d)
|
||||
|
||||
state, ok := cached.Extra[service.OpenAIAutoResetCreditStateExtraKey].(map[string]any)
|
||||
require.True(t, ok, "卡状态必须进入调度投影")
|
||||
require.Equal(t, service.OpenAIAutoResetStatusAvailable, state["status"])
|
||||
require.EqualValues(t, 3, state["available_count"])
|
||||
require.Equal(t, checkedAt, state["checked_at"])
|
||||
}
|
||||
|
||||
func TestBuildSchedulerMetadataAccount_KeepsQuotaStateForCachedAccounts(t *testing.T) {
|
||||
now := time.Now().UTC()
|
||||
activeStart := now.Add(-time.Hour).Format(time.RFC3339)
|
||||
|
||||
@@ -1577,11 +1577,32 @@ func (s *UsageLogRepoSuite) TestGetUserUsageTrend() {
|
||||
startTime := base.Add(-1 * time.Hour)
|
||||
endTime := base.Add(48 * time.Hour)
|
||||
|
||||
trend, err := s.repo.GetUserUsageTrend(s.ctx, startTime, endTime, "day", 10)
|
||||
trend, err := s.repo.GetUserUsageTrend(s.ctx, startTime, endTime, "day", 10, "tokens")
|
||||
s.Require().NoError(err, "GetUserUsageTrend")
|
||||
s.Require().GreaterOrEqual(len(trend), 2)
|
||||
}
|
||||
|
||||
func (s *UsageLogRepoSuite) TestGetUserUsageTrend_SelectsTopByMetric() {
|
||||
highTokens := mustCreateUser(s.T(), s.client, &service.User{Email: "tokens@test.com"})
|
||||
highSpend := mustCreateUser(s.T(), s.client, &service.User{Email: "spend@test.com"})
|
||||
tokenKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: highTokens.ID, Key: "sk-trend-tokens", Name: "tokens"})
|
||||
spendKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: highSpend.ID, Key: "sk-trend-spend", Name: "spend"})
|
||||
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-trend-metric"})
|
||||
at := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
|
||||
s.createUsageLog(highTokens, tokenKey, account, 1000, 0, 0.1, at)
|
||||
s.createUsageLog(highSpend, spendKey, account, 10, 0, 5, at)
|
||||
start, end := at.Add(-time.Hour), at.Add(time.Hour)
|
||||
tokens, err := s.repo.GetUserUsageTrend(s.ctx, start, end, "day", 1, "tokens")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(tokens, 1)
|
||||
s.Require().Equal(highTokens.ID, tokens[0].UserID)
|
||||
spend, err := s.repo.GetUserUsageTrend(s.ctx, start, end, "day", 1, "actual_cost")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(spend, 1)
|
||||
s.Require().Equal(highSpend.ID, spend[0].UserID)
|
||||
s.Require().Equal(5.0, spend[0].ActualCost)
|
||||
}
|
||||
|
||||
// --- GetAPIKeyUsageTrend ---
|
||||
|
||||
func (s *UsageLogRepoSuite) TestGetAPIKeyUsageTrend() {
|
||||
|
||||
@@ -82,8 +82,12 @@ func (r *usageLogRepository) GetAPIKeyUsageTrend(ctx context.Context, startTime,
|
||||
}
|
||||
|
||||
// GetUserUsageTrend returns usage trend data grouped by user and date
|
||||
func (r *usageLogRepository) GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) (results []UserUsageTrendPoint, err error) {
|
||||
func (r *usageLogRepository) GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int, metric string) (results []UserUsageTrendPoint, err error) {
|
||||
dateFormat := safeDateFormat(granularity)
|
||||
rankExpr := "SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens)"
|
||||
if metric == "actual_cost" {
|
||||
rankExpr = "SUM(actual_cost)"
|
||||
}
|
||||
|
||||
query := fmt.Sprintf(`
|
||||
WITH top_users AS (
|
||||
@@ -91,7 +95,7 @@ func (r *usageLogRepository) GetUserUsageTrend(ctx context.Context, startTime, e
|
||||
FROM usage_logs
|
||||
WHERE created_at >= $1 AND created_at < $2
|
||||
GROUP BY user_id
|
||||
ORDER BY SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) DESC
|
||||
ORDER BY %s DESC, user_id ASC
|
||||
LIMIT $3
|
||||
)
|
||||
SELECT
|
||||
@@ -109,7 +113,7 @@ func (r *usageLogRepository) GetUserUsageTrend(ctx context.Context, startTime, e
|
||||
AND u.created_at >= $4 AND u.created_at < $5
|
||||
GROUP BY date, u.user_id, us.email, us.username
|
||||
ORDER BY date ASC, tokens DESC
|
||||
`, dateFormat)
|
||||
`, rankExpr, dateFormat)
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, startTime, endTime, limit, startTime, endTime)
|
||||
if err != nil {
|
||||
|
||||
@@ -1004,6 +1004,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"model_plaza_description": "",
|
||||
"plugin_management_enabled": false,
|
||||
"risk_control_enabled": false,
|
||||
"cyber_policy_user_allowlist": "",
|
||||
"cyber_session_block_enabled": false,
|
||||
"cyber_session_block_ttl_seconds": 3600,
|
||||
"affiliate_enabled": false,
|
||||
@@ -1322,6 +1323,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"model_plaza_description": "",
|
||||
"plugin_management_enabled": false,
|
||||
"risk_control_enabled": false,
|
||||
"cyber_policy_user_allowlist": "",
|
||||
"cyber_session_block_enabled": false,
|
||||
"cyber_session_block_ttl_seconds": 3600,
|
||||
"affiliate_enabled": false,
|
||||
@@ -1735,8 +1737,8 @@ func (stubApiKeyCache) IncrementCreateAttemptCount(ctx context.Context, userID i
|
||||
return nil
|
||||
}
|
||||
|
||||
func (stubApiKeyCache) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
return nil
|
||||
func (stubApiKeyCache) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (stubApiKeyCache) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
@@ -2652,7 +2654,7 @@ func (r *stubUsageLogRepo) GetAPIKeyUsageTrend(ctx context.Context, startTime, e
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUsageLogRepo) GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.UserUsageTrendPoint, error) {
|
||||
func (r *stubUsageLogRepo) GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int, metric string) ([]usagestats.UserUsageTrendPoint, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
|
||||
@@ -367,6 +367,7 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAu
|
||||
accounts.GET("/opencode-go-usage/settings", h.Admin.Account.GetOpenCodeGoUsageSettings)
|
||||
accounts.PUT("/opencode-go-usage/settings", h.Admin.Account.UpdateOpenCodeGoUsageSettings)
|
||||
accounts.GET("/:id", h.Admin.Account.GetByID)
|
||||
accounts.GET("/:id/claude/reset-credits", h.Admin.Account.ClaudeResetCredits)
|
||||
accounts.POST("", h.Admin.Account.Create)
|
||||
accounts.POST("/:id/duplicate", h.Admin.Account.Duplicate)
|
||||
accounts.POST("/check-mixed-channel", h.Admin.Account.CheckMixedChannel)
|
||||
|
||||
@@ -271,13 +271,19 @@ func openAIThresholdCandidate(extra map[string]any, window string, now time.Time
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if openAIQuotaWindowReset(extra, window, now) || openAICodexSnapshotStaleForPause(extra, now) {
|
||||
if openAIQuotaWindowReset(extra, window, now) || (openAICodexSnapshotStaleForPause(extra, now) && !openAIQuotaWindowResetPending(extra, window, now)) {
|
||||
return nil
|
||||
}
|
||||
until := parseSchedulingResetAt(extra[resetAtKey])
|
||||
if until == nil {
|
||||
if resetAt, ok := openAICodexWindowResetAt(extra, window); ok {
|
||||
until = &resetAt
|
||||
}
|
||||
}
|
||||
return &accountSchedulingThresholdCandidate{
|
||||
window: window,
|
||||
usedPercent: schedulingPercentValue(usedPercent),
|
||||
until: parseSchedulingResetAt(extra[resetAtKey]),
|
||||
until: until,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -172,7 +172,7 @@ func TestEvaluateAccountSchedulingThreshold_OpenAISkipsStaleSnapshot(t *testing.
|
||||
Extra: map[string]any{
|
||||
"codex_usage_updated_at": now.Add(-2 * time.Hour).Format(time.RFC3339),
|
||||
"codex_5h_used_percent": 100.0,
|
||||
"codex_5h_reset_at": now.Add(3 * time.Hour).Format(time.RFC3339),
|
||||
"codex_5h_reset_at": "invalid",
|
||||
},
|
||||
}
|
||||
|
||||
@@ -181,6 +181,63 @@ func TestEvaluateAccountSchedulingThreshold_OpenAISkipsStaleSnapshot(t *testing.
|
||||
require.False(t, decision.ShouldPause)
|
||||
}
|
||||
|
||||
func TestOpenAIThresholdCandidate_StaleSnapshotFutureReset(t *testing.T) {
|
||||
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
reset map[string]any
|
||||
want bool
|
||||
}{
|
||||
{"absolute future", map[string]any{"codex_5h_reset_at": now.Add(time.Hour).Format(time.RFC3339)}, true},
|
||||
{"relative future", map[string]any{"codex_5h_reset_after_seconds": 4 * 3600}, true},
|
||||
{"missing", nil, false},
|
||||
{"invalid", map[string]any{"codex_5h_reset_at": "invalid"}, false},
|
||||
{"past", map[string]any{"codex_5h_reset_at": now.Add(-time.Minute).Format(time.RFC3339)}, false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
extra := map[string]any{"codex_usage_updated_at": now.Add(-3 * time.Hour).Format(time.RFC3339), "codex_5h_used_percent": 99.0}
|
||||
for key, value := range tc.reset {
|
||||
extra[key] = value
|
||||
}
|
||||
candidate := openAIThresholdCandidate(extra, "5h", now)
|
||||
require.Equal(t, tc.want, candidate != nil)
|
||||
if tc.want {
|
||||
decision := EvaluateAccountSchedulingThreshold(&Account{Platform: PlatformOpenAI, Extra: extra}, map[string]int{PlatformOpenAI: 95}, now)
|
||||
require.True(t, decision.ShouldPause)
|
||||
require.NotNil(t, decision.Until)
|
||||
require.True(t, now.Before(*decision.Until))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOpenAIQuotaUtilization_StaleSnapshotFutureReset(t *testing.T) {
|
||||
now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
reset map[string]any
|
||||
want bool
|
||||
}{
|
||||
{"absolute future", map[string]any{"codex_5h_reset_at": now.Add(time.Hour).Format(time.RFC3339)}, true},
|
||||
{"relative future", map[string]any{"codex_5h_reset_after_seconds": 4 * 3600}, true},
|
||||
{"missing", nil, false},
|
||||
{"invalid", map[string]any{"codex_5h_reset_at": "invalid"}, false},
|
||||
{"past", map[string]any{"codex_5h_reset_at": now.Add(-time.Minute).Format(time.RFC3339)}, false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
extra := map[string]any{"codex_usage_updated_at": now.Add(-3 * time.Hour).Format(time.RFC3339), "codex_5h_used_percent": 99.0}
|
||||
for key, value := range tc.reset {
|
||||
extra[key] = value
|
||||
}
|
||||
utilization, ok := resolveOpenAIQuotaUtilization(extra, "5h", now)
|
||||
require.Equal(t, tc.want, ok)
|
||||
if ok {
|
||||
require.Equal(t, 0.99, utilization)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvaluateAccountSchedulingThreshold_OpenAISkipsResetWindow(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -19,6 +19,8 @@ import (
|
||||
// totalCost 是本次请求的客户计费(倍率前),用于优先级 2。
|
||||
// serviceTier 是最终参与用户计费的 OpenAI 服务层级,用于优先级 3。
|
||||
// pricingAt 与本次客户计费使用同一时刻,避免跨峰谷请求的成本与售价错位。
|
||||
// longContextPricingEnabled 表示上游是否对本次请求收取长上下文费率,用于优先级 3;
|
||||
// 由 accountStatsLongContextPricingEnabled 按账号开关得出,不受分组售价开关影响。
|
||||
// reasoningEffort 是最终转发等级;按账号统计定价中配置的等级倍率计费。
|
||||
func resolveAccountStatsCost(
|
||||
ctx context.Context,
|
||||
@@ -32,6 +34,7 @@ func resolveAccountStatsCost(
|
||||
totalCost float64,
|
||||
serviceTier string,
|
||||
pricingAt time.Time,
|
||||
longContextPricingEnabled bool,
|
||||
reasoningEfforts ...string,
|
||||
) *float64 {
|
||||
reasoningEffort := ""
|
||||
@@ -64,7 +67,7 @@ func resolveAccountStatsCost(
|
||||
|
||||
// 优先级 3:模型定价文件(LiteLLM)默认价格
|
||||
if billingService != nil {
|
||||
return tryModelFilePricing(billingService, upstreamModel, tokens, serviceTier, pricingAt, reasoningEffort)
|
||||
return tryModelFilePricing(billingService, upstreamModel, tokens, serviceTier, pricingAt, longContextPricingEnabled, reasoningEffort)
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -74,11 +77,16 @@ func resolveAccountStatsCost(
|
||||
// 与用户计费共用同一条定价管线,避免这里维护第二份"单价 × token 数"实现后,
|
||||
// 每加一个定价特性都要手工镜像一次。解析器不配置渠道或分组,保持优先级 3 的
|
||||
// 语义:只取模型定价文件,不引入自定义售价。
|
||||
func tryModelFilePricing(billingService *BillingService, model string, tokens UsageTokens, serviceTier string, pricingAt time.Time, reasoningEfforts ...string) *float64 {
|
||||
func tryModelFilePricing(billingService *BillingService, model string, tokens UsageTokens, serviceTier string, pricingAt time.Time, longContextPricingEnabled bool, reasoningEfforts ...string) *float64 {
|
||||
reasoningEffort := ""
|
||||
if len(reasoningEfforts) > 0 {
|
||||
reasoningEffort = reasoningEfforts[0]
|
||||
}
|
||||
resolver := NewModelPricingResolver(nil, billingService)
|
||||
resolved := resolver.Resolve(context.Background(), PricingInput{Model: model})
|
||||
// 无分组的解析结果默认开启长上下文,CostInput.LongContextBillingEnabled=false 无法否决,
|
||||
// 因此直接覆写解析结果。
|
||||
resolved.longContextPricingEnabled = longContextPricingEnabled
|
||||
breakdown, err := billingService.CalculateCostUnified(CostInput{
|
||||
Ctx: context.Background(),
|
||||
Model: model,
|
||||
@@ -87,7 +95,8 @@ func tryModelFilePricing(billingService *BillingService, model string, tokens Us
|
||||
ServiceTier: normalizeBillingServiceTier(serviceTier),
|
||||
ReasoningEffort: reasoningEffort,
|
||||
PricingAt: pricingAt,
|
||||
Resolver: NewModelPricingResolver(nil, billingService),
|
||||
Resolver: resolver,
|
||||
Resolved: resolved,
|
||||
})
|
||||
if err != nil || breakdown == nil || breakdown.TotalCost <= 0 {
|
||||
return nil
|
||||
@@ -263,6 +272,7 @@ func applyAccountStatsCost(
|
||||
tokens UsageTokens,
|
||||
totalCost float64,
|
||||
pricingAt time.Time,
|
||||
longContextPricingEnabled bool,
|
||||
) {
|
||||
model := upstreamModel
|
||||
if model == "" {
|
||||
@@ -281,6 +291,14 @@ func applyAccountStatsCost(
|
||||
reasoningEffort = *usageLog.ReasoningEffort
|
||||
}
|
||||
usageLog.AccountStatsCost = resolveAccountStatsCost(
|
||||
ctx, cs, bs, accountID, groupID, model, tokens, requestCount, totalCost, serviceTier, pricingAt, reasoningEffort,
|
||||
ctx, cs, bs, accountID, groupID, model, tokens, requestCount, totalCost, serviceTier, pricingAt, longContextPricingEnabled, reasoningEffort,
|
||||
)
|
||||
}
|
||||
|
||||
// accountStatsLongContextPricingEnabled 判断账号统计成本是否计入长上下文阶梯。
|
||||
// 账号统计成本反映上游实际成本:OpenAI 账号由开关声明上游是否收取长上下文费率;
|
||||
// 其它平台没有该开关(accountGate 为 nil),按官方阶梯计。分组开关只决定客户售价,
|
||||
// 不参与成本判断。
|
||||
func accountStatsLongContextPricingEnabled(accountGate *bool) bool {
|
||||
return accountGate == nil || *accountGate
|
||||
}
|
||||
|
||||
@@ -472,7 +472,7 @@ func TestTryModelFilePricing_Success(t *testing.T) {
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{}, true)
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
|
||||
require.InDelta(t, 0.2, *result, 1e-12)
|
||||
@@ -483,8 +483,8 @@ func TestTryModelFilePricing_Fable51HasNoImplicitReasoningMultiplier(t *testing.
|
||||
"claude-fable-5-1": {InputPricePerToken: 0.001},
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 100}
|
||||
standard := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", time.Time{}, "xhigh")
|
||||
max := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", time.Time{}, "max")
|
||||
standard := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", time.Time{}, true, "xhigh")
|
||||
max := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", time.Time{}, true, "max")
|
||||
require.NotNil(t, standard)
|
||||
require.NotNil(t, max)
|
||||
require.Equal(t, *standard, *max)
|
||||
@@ -503,7 +503,7 @@ func TestTryModelFilePricing_AppliesLongContextPricing(t *testing.T) {
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 101, OutputTokens: 10, CacheReadTokens: 5}
|
||||
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "", time.Time{})
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "", time.Time{}, true)
|
||||
|
||||
require.NotNil(t, result)
|
||||
// Input and cache-read use the 2x input tier; output uses the 1.5x tier.
|
||||
@@ -542,7 +542,7 @@ func TestTryModelFilePricing_AppliesServiceTierPricing(t *testing.T) {
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, tt.serviceTier, time.Time{})
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, tt.serviceTier, time.Time{}, true)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, tt.want, *result, 1e-12)
|
||||
})
|
||||
@@ -572,7 +572,7 @@ func TestTryModelFilePricing_CombinesPriorityAndLongContextPricing(t *testing.T)
|
||||
CacheReadTokens: 5,
|
||||
}
|
||||
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "priority", time.Time{})
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "priority", time.Time{}, true)
|
||||
|
||||
require.NotNil(t, result)
|
||||
// priority 单价先应用,再叠加长上下文输入 2x、输出 1.5x。
|
||||
@@ -583,7 +583,7 @@ func TestTryModelFilePricing_PricingNotFound(t *testing.T) {
|
||||
// "nonexistent-model" does not match any fallback pattern
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{})
|
||||
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
|
||||
result := tryModelFilePricing(bs, "nonexistent-model", tokens, "", time.Time{})
|
||||
result := tryModelFilePricing(bs, "nonexistent-model", tokens, "", time.Time{}, true)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -593,7 +593,7 @@ func TestTryModelFilePricing_NilFallback(t *testing.T) {
|
||||
"claude-sonnet-4": nil,
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 100}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{}, true)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -605,7 +605,7 @@ func TestTryModelFilePricing_ZeroCost(t *testing.T) {
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{} // all zero tokens → cost = 0 → nil
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{}, true)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -622,7 +622,7 @@ func TestTryModelFilePricing_WithImageOutput(t *testing.T) {
|
||||
OutputTokens: 50,
|
||||
ImageOutputTokens: 10,
|
||||
}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{}, true)
|
||||
require.NotNil(t, result)
|
||||
// ImageOutputTokens 是 OutputTokens 的子集,先扣除再按图片单价计。
|
||||
// 100*0.001 + (50-10)*0.002 + 10*0.01 = 0.1 + 0.08 + 0.1 = 0.28
|
||||
@@ -644,7 +644,7 @@ func TestTryModelFilePricing_WithCacheTokens(t *testing.T) {
|
||||
CacheCreationTokens: 200,
|
||||
CacheReadTokens: 300,
|
||||
}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{}, true)
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 + 200*0.003 + 300*0.0005
|
||||
// = 0.1 + 0.1 + 0.6 + 0.15 = 0.95
|
||||
@@ -692,7 +692,7 @@ func TestTryModelFilePricing_DeepSeekPeakPricing(t *testing.T) {
|
||||
{"sunday", time.Date(2026, time.August, 23, 7, 0, 0, 0, time.UTC), 1},
|
||||
} {
|
||||
t.Run(slot.name, func(t *testing.T) {
|
||||
cost := tryModelFilePricing(bs, model.name, tokens, "", slot.at)
|
||||
cost := tryModelFilePricing(bs, model.name, tokens, "", slot.at, true)
|
||||
require.NotNil(t, cost)
|
||||
require.InDelta(t, baseCost*slot.multiplier, *cost, 1e-12)
|
||||
})
|
||||
@@ -738,7 +738,7 @@ func TestResolveAccountStatsCost_DeepSeekPricingPriority(t *testing.T) {
|
||||
groupID = 99
|
||||
}
|
||||
cost := resolveAccountStatsCost(context.Background(), cs, newTestBillingService(),
|
||||
1, groupID, "deepseek-v4-flash", UsageTokens{InputTokens: 1000}, 1, 0.75, "", peak)
|
||||
1, groupID, "deepseek-v4-flash", UsageTokens{InputTokens: 1000}, 1, 0.75, "", peak, true)
|
||||
if tt.noChannel {
|
||||
require.Nil(t, cost)
|
||||
return
|
||||
@@ -759,7 +759,7 @@ func TestResolveAccountStatsCost_NilChannelService(t *testing.T) {
|
||||
nil, // channelService is nil
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 1, "claude-sonnet-4",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{}, true,
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -775,7 +775,7 @@ func TestResolveAccountStatsCost_EmptyUpstreamModel(t *testing.T) {
|
||||
cs,
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 1, "", // empty upstream model
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{}, true,
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -792,7 +792,7 @@ func TestResolveAccountStatsCost_GetChannelForGroupReturnsNil(t *testing.T) {
|
||||
cs,
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 99, "claude-sonnet-4", // groupID 99 has no channel
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{}, true,
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -823,7 +823,7 @@ func TestResolveAccountStatsCost_HitsCustomRule(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, nil, // billingService not needed when custom rule hits
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 999.0, "priority", time.Time{}, // 自定义账号价格不叠加服务层级倍率
|
||||
tokens, 1, 999.0, "priority", time.Time{}, true, // 自定义账号价格不叠加服务层级倍率
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// 100*0.01 + 50*0.02 = 1.0 + 1.0 = 2.0
|
||||
@@ -845,7 +845,7 @@ func TestResolveAccountStatsCost_ApplyPricingToAccountStats_UsesTotalCost(t *tes
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 0.75, "priority", time.Time{}, // 已完成用户计费,不再重复应用服务层级倍率
|
||||
tokens, 1, 0.75, "priority", time.Time{}, true, // 已完成用户计费,不再重复应用服务层级倍率
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 0.75, *result, 1e-12)
|
||||
@@ -863,7 +863,7 @@ func TestResolveAccountStatsCost_ApplyPricingToAccountStats_ZeroTotalCost_Return
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
UsageTokens{}, 1, 0.0, "", time.Time{}, // totalCost = 0
|
||||
UsageTokens{}, 1, 0.0, "", time.Time{}, true, // totalCost = 0
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -890,7 +890,7 @@ func TestResolveAccountStatsCost_FallsBackToLiteLLM(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 999.0, "", time.Time{}, // totalCost ignored
|
||||
tokens, 1, 999.0, "", time.Time{}, true, // totalCost ignored
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
|
||||
@@ -911,7 +911,7 @@ func TestResolveAccountStatsCost_FallbackHonorsAnthropicFast(t *testing.T) {
|
||||
context.Background(), cs, bs,
|
||||
1, 10, "claude-opus-5",
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000},
|
||||
1, 0, "fast", time.Time{},
|
||||
1, 0, "fast", time.Time{}, true,
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 60, *result, 1e-12)
|
||||
@@ -930,7 +930,7 @@ func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "gemini-3.6-flash-low",
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "", time.Time{},
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "", time.Time{}, true,
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 9.15, *result, 1e-12)
|
||||
@@ -954,7 +954,7 @@ func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "totally-unknown-model",
|
||||
tokens, 1, 0.0, "", time.Time{},
|
||||
tokens, 1, 0.0, "", time.Time{}, true,
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -971,7 +971,7 @@ func TestResolveAccountStatsCost_NilBillingService_SkipsLiteLLM(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, nil, // billingService is nil
|
||||
1, 10, "claude-sonnet-4",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.0, "", time.Time{},
|
||||
UsageTokens{InputTokens: 100}, 1, 0.0, "", time.Time{}, true,
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -1004,7 +1004,7 @@ func TestResolveAccountStatsCost_CustomRulePriorityOverApplyPricing(t *testing.T
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 99.0, "", time.Time{}, // totalCost = 99.0 (would be used if ApplyPricing wins)
|
||||
tokens, 1, 99.0, "", time.Time{}, true, // totalCost = 99.0 (would be used if ApplyPricing wins)
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// Custom rule: 100*0.05 = 5.0 (NOT 99.0 from totalCost)
|
||||
@@ -1032,13 +1032,111 @@ func TestApplyAccountStatsCost_UsesUsageLogServiceTier(t *testing.T) {
|
||||
applyAccountStatsCost(
|
||||
context.Background(), usageLog, cs, bs,
|
||||
1, 10, "gpt-5.6-sol", "gpt-5.6-sol",
|
||||
UsageTokens{InputTokens: 100, OutputTokens: 50}, 999, time.Time{},
|
||||
UsageTokens{InputTokens: 100, OutputTokens: 50}, 999, time.Time{}, true,
|
||||
)
|
||||
|
||||
require.NotNil(t, usageLog.AccountStatsCost)
|
||||
require.InDelta(t, 0.4, *usageLog.AccountStatsCost, 1e-12)
|
||||
}
|
||||
|
||||
func TestApplyAccountStatsCost_LongContextFollowsAccountGate(t *testing.T) {
|
||||
// 渠道售价不参与优先级 3:结果只取模型定价文件。
|
||||
channel := &Channel{
|
||||
ID: 1,
|
||||
Status: StatusActive,
|
||||
ModelPricing: []ChannelModelPricing{{
|
||||
Models: []string{"gpt-5.6-sol"}, InputPrice: testPtrFloat64(0.01),
|
||||
}},
|
||||
}
|
||||
cs := newTestChannelServiceForStats(t, channel, 10, PlatformOpenAI)
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
|
||||
"gpt-5.6-sol": {
|
||||
InputPricePerToken: 0.001,
|
||||
InputPricePerTokenPriority: 0.002,
|
||||
OutputPricePerToken: 0.002,
|
||||
OutputPricePerTokenPriority: 0.004,
|
||||
CacheReadPricePerToken: 0.0001,
|
||||
CacheReadPricePerTokenPriority: 0.0002,
|
||||
LongContextInputThreshold: 100,
|
||||
LongContextInputMultiplier: 2,
|
||||
LongContextOutputMultiplier: 1.5,
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 101, OutputTokens: 10, CacheReadTokens: 5}
|
||||
accountOff, accountOn := false, true
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
gate *bool
|
||||
tier string
|
||||
wantCost float64
|
||||
}{
|
||||
{name: "account_off", gate: &accountOff, wantCost: 0.1215},
|
||||
{name: "account_off_priority", gate: &accountOff, tier: "priority", wantCost: 0.243},
|
||||
{name: "account_on_priority", gate: &accountOn, tier: "priority", wantCost: 0.466},
|
||||
// 非 OpenAI 平台没有账号开关,按官方阶梯计。
|
||||
{name: "no_gate", wantCost: 0.233},
|
||||
{name: "no_gate_priority", tier: "priority", wantCost: 0.466},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
usageLog := &UsageLog{ServiceTier: &tt.tier}
|
||||
applyAccountStatsCost(context.Background(), usageLog, cs, bs,
|
||||
1, 10, "gpt-5.6-sol", "gpt-5.6-sol", tokens, 999, time.Time{},
|
||||
accountStatsLongContextPricingEnabled(tt.gate))
|
||||
require.NotNil(t, usageLog.AccountStatsCost)
|
||||
require.InDelta(t, tt.wantCost, *usageLog.AccountStatsCost, 1e-12)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 分组开关只决定客户售价;账号统计成本按账号开关判断上游是否收取长上下文费率。
|
||||
func TestOpenAIGatewayServiceRecordUsage_AccountStatsLongContextFollowsAccountGate(t *testing.T) {
|
||||
baseCost := 300000*2.5e-6 + 2000*15e-6
|
||||
longContextCost := 300000*2.5e-6*2 + 2000*15e-6*1.5
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
groupLongContext bool
|
||||
accountExtra map[string]any
|
||||
wantTotalCost float64
|
||||
wantAccountCost float64
|
||||
}{
|
||||
{name: "group_on_account_off", groupLongContext: true, wantTotalCost: longContextCost, wantAccountCost: baseCost},
|
||||
{name: "group_off_account_off", wantTotalCost: baseCost, wantAccountCost: baseCost},
|
||||
{
|
||||
name: "group_off_account_on",
|
||||
accountExtra: map[string]any{"openai_long_context_billing_enabled": true},
|
||||
wantTotalCost: longContextCost,
|
||||
wantAccountCost: longContextCost,
|
||||
},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
||||
svc := newOpenAIRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{}, nil)
|
||||
swapInOpenAILadderCatalog(t, svc)
|
||||
svc.channelService = newTestChannelServiceForStats(t, &Channel{ID: 1, Status: StatusActive}, 1, PlatformOpenAI)
|
||||
apiKey := openAIRecordUsageAPIKeyWithGroup(svc, 1015, tt.groupLongContext)
|
||||
apiKey.GroupID = i64p(1)
|
||||
|
||||
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
|
||||
Result: &OpenAIForwardResult{
|
||||
RequestID: "resp_account_stats_long_context_" + tt.name,
|
||||
Usage: OpenAIUsage{InputTokens: 300000, OutputTokens: 2000},
|
||||
Model: "gpt-5.4-2026-03-05",
|
||||
Duration: time.Second,
|
||||
},
|
||||
APIKey: apiKey,
|
||||
User: &User{ID: 2015},
|
||||
Account: &Account{ID: 3015, Platform: PlatformOpenAI, Extra: tt.accountExtra},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usageRepo.lastLog)
|
||||
require.InDelta(t, tt.wantTotalCost, usageRepo.lastLog.TotalCost, 1e-10)
|
||||
require.NotNil(t, usageRepo.lastLog.AccountStatsCost)
|
||||
require.InDelta(t, tt.wantAccountCost, *usageRepo.lastLog.AccountStatsCost, 1e-10)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// helpers for resolveAccountStatsCost tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -53,7 +53,7 @@ type UsageLogRepository interface {
|
||||
GetUserBreakdownStats(ctx context.Context, startTime, endTime time.Time, dim usagestats.UserBreakdownDimension, limit int) ([]usagestats.UserBreakdownItem, error)
|
||||
GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error)
|
||||
GetAPIKeyUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.APIKeyUsageTrendPoint, error)
|
||||
GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.UserUsageTrendPoint, error)
|
||||
GetUserUsageTrend(ctx context.Context, startTime, endTime time.Time, granularity string, limit int, metric string) ([]usagestats.UserUsageTrendPoint, error)
|
||||
GetUserSpendingRanking(ctx context.Context, startTime, endTime time.Time, limit int) (*usagestats.UserSpendingRankingResponse, error)
|
||||
GetBatchUserUsageStats(ctx context.Context, userIDs []int64, startTime, endTime time.Time) (map[int64]*usagestats.BatchUserUsageStats, error)
|
||||
GetBatchAPIKeyUsageStats(ctx context.Context, apiKeyIDs []int64, startTime, endTime time.Time) (map[int64]*usagestats.BatchAPIKeyUsageStats, error)
|
||||
|
||||
@@ -76,7 +76,7 @@ func (r *usageBatchLogRepoStub) GetAllGroupUsageSummary(context.Context, time.Ti
|
||||
func (r *usageBatchLogRepoStub) GetAPIKeyUsageTrend(context.Context, time.Time, time.Time, string, int) ([]usagestats.APIKeyUsageTrendPoint, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *usageBatchLogRepoStub) GetUserUsageTrend(context.Context, time.Time, time.Time, string, int) ([]usagestats.UserUsageTrendPoint, error) {
|
||||
func (r *usageBatchLogRepoStub) GetUserUsageTrend(context.Context, time.Time, time.Time, string, int, string) ([]usagestats.UserUsageTrendPoint, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *usageBatchLogRepoStub) GetUserSpendingRanking(context.Context, time.Time, time.Time, int) (*usagestats.UserSpendingRankingResponse, error) {
|
||||
|
||||
@@ -29,7 +29,7 @@ func TestAdminService_CreateGroup_RejectsEmptyEnabledModelAllowlist(t *testing.T
|
||||
require.Nil(t, repo.created, "拒绝时不得落库")
|
||||
}
|
||||
|
||||
func TestAdminService_CreateGroup_RejectsInvalidAllowlistWildcard(t *testing.T) {
|
||||
func TestAdminService_CreateGroup_AcceptsInteriorAllowlistWildcard(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{createID: 51}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
@@ -40,11 +40,9 @@ func TestAdminService_CreateGroup_RejectsInvalidAllowlistWildcard(t *testing.T)
|
||||
ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-*-5.4"}},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
appErr := infraerrors.FromError(err)
|
||||
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
|
||||
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
|
||||
require.Nil(t, repo.created, "拒绝时不得落库")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, repo.created)
|
||||
require.Equal(t, []string{"gpt-*-5.4"}, repo.created.ModelAllowlist.Models)
|
||||
}
|
||||
|
||||
func TestAdminService_CreateGroup_NormalizesModelAllowlist(t *testing.T) {
|
||||
@@ -83,7 +81,7 @@ func TestAdminService_UpdateGroup_RejectsEmptyEnabledModelAllowlist(t *testing.T
|
||||
require.Nil(t, repo.updated, "拒绝时不得落库")
|
||||
}
|
||||
|
||||
func TestAdminService_UpdateGroup_RejectsInvalidAllowlistWildcard(t *testing.T) {
|
||||
func TestAdminService_UpdateGroup_AcceptsInteriorAllowlistWildcard(t *testing.T) {
|
||||
existing := &Group{ID: 1, Name: "existing", Platform: PlatformOpenAI, Status: StatusActive}
|
||||
repo := &groupRepoStubForAdmin{getByID: existing}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
@@ -92,11 +90,9 @@ func TestAdminService_UpdateGroup_RejectsInvalidAllowlistWildcard(t *testing.T)
|
||||
ModelAllowlist: &GroupModelAllowlist{Enabled: true, Models: []string{"foo-*bar"}},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
appErr := infraerrors.FromError(err)
|
||||
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
|
||||
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
|
||||
require.Nil(t, repo.updated, "拒绝时不得落库")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, repo.updated)
|
||||
require.Equal(t, []string{"foo-*bar"}, repo.updated.ModelAllowlist.Models)
|
||||
}
|
||||
|
||||
func TestAdminService_UpdateGroup_NormalizesAndResetsModelAllowlist(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAnthropicChatStreamAuthoritativeUsage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, adapter := range []string{"anthropic", "native"} {
|
||||
for _, options := range []string{"", `,"stream_options":{"include_usage":true}`, `,"stream_options":{"include_usage":false}`} {
|
||||
for _, tc := range []struct {
|
||||
name, start, delta, repeatDelta string
|
||||
want bool
|
||||
prompt, billableInput, totalInput, output, cached, created int
|
||||
}{
|
||||
{"cache", `,"usage":{"input_tokens":308,"cache_read_input_tokens":241,"cache_creation_input_tokens":17}`, `,"usage":{"output_tokens":49}`, "", true, 566, 308, 566, 49, 241, 17},
|
||||
{"deepseek_partial_cache", "", `,"usage":{"input_tokens":1200,"output_tokens":30,"prompt_cache_hit_tokens":800,"prompt_cache_miss_tokens":400}`, "", true, 1200, 400, 1200, 30, 800, 0},
|
||||
{"openai_partial_cache", "", `,"usage":{"prompt_tokens":1200,"output_tokens":30,"prompt_tokens_details":{"cached_tokens":800}}`, "", true, 1200, 400, 1200, 30, 800, 0},
|
||||
{"repeated_authoritative_cache_bucket", `,"usage":{"input_tokens":1200,"prompt_tokens":1200}`, `,"usage":{"output_tokens":30,"prompt_tokens":1200,"cache_read_input_tokens":800}`, `,"usage":{"output_tokens":30,"prompt_tokens":1200,"cache_read_input_tokens":800}`, true, 1200, 400, 1200, 30, 800, 0},
|
||||
{"independent_input_with_late_cache", `,"usage":{"input_tokens":308,"cache_read_input_tokens":241,"cache_creation_input_tokens":17}`, `,"usage":{"input_tokens":0,"output_tokens":49,"cache_read_input_tokens":300,"cache_creation_input_tokens":17}`, "", true, 625, 308, 625, 49, 300, 17},
|
||||
{"kimi_full_cache", `,"usage":{"input_tokens":173306,"prompt_tokens":173306}`, `,"usage":{"input_tokens":0,"output_tokens":49,"prompt_tokens":173306,"cache_read_input_tokens":173306}`, "", true, 173306, 0, 173306, 49, 173306, 0},
|
||||
{"kimi_full_cache_cached_tokens", `,"usage":{"input_tokens":173306,"prompt_tokens":173306}`, `,"usage":{"input_tokens":0,"output_tokens":49,"prompt_tokens":173306,"cached_tokens":173306}`, "", true, 173306, 0, 173306, 49, 173306, 0},
|
||||
{"kimi_full_cache_prompt_details", `,"usage":{"input_tokens":173306,"prompt_tokens":173306}`, `,"usage":{"input_tokens":0,"output_tokens":49,"prompt_tokens":173306,"prompt_tokens_details":{"cached_tokens":173306}}`, "", true, 173306, 0, 173306, 49, 173306, 0},
|
||||
{"zero_start", `,"usage":{"input_tokens":0,"output_tokens":0}`, "", "", true, 0, 0, 0, 0, 0, 0},
|
||||
{"zero_delta", "", `,"usage":{"input_tokens":0,"output_tokens":0}`, "", true, 0, 0, 0, 0, 0, 0},
|
||||
{"absent", "", "", "", false, 0, 0, 0, 0, 0, 0},
|
||||
{"null", `,"usage":null`, `,"usage":null`, "", false, 0, 0, 0, 0, 0, 0},
|
||||
} {
|
||||
for _, terminal := range []string{"stop", "eof"} {
|
||||
t.Run(adapter+"/"+options+"/"+tc.name+"/"+terminal, func(t *testing.T) {
|
||||
sse := fmt.Sprintf("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"k3\",\"content\":[]%s}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}%s}\n\n", tc.start, tc.delta)
|
||||
if tc.repeatDelta != "" {
|
||||
sse += fmt.Sprintf("event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}%s}\n\n", tc.repeatDelta)
|
||||
}
|
||||
if terminal == "stop" {
|
||||
// Repeated terminal events and EOF must not duplicate usage.
|
||||
sse += "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(sse)),
|
||||
}}
|
||||
body := []byte(`{"model":"k3","messages":[{"role":"user","content":"hi"}],"stream":true` + options + "}")
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(string(body)))
|
||||
if adapter == "native" {
|
||||
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
|
||||
result, err := svc.forwardChatCompletionsViaNativeAnthropic(context.Background(), c, nativeAnthropicTestAccount(), body, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.totalInput, result.Usage.InputTokens)
|
||||
require.Equal(t, tc.output, result.Usage.OutputTokens)
|
||||
require.Equal(t, tc.cached, result.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, tc.created, result.Usage.CacheCreationInputTokens)
|
||||
} else {
|
||||
svc := &GatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
|
||||
account := &Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk-test"}}
|
||||
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, nil)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.billableInput, result.Usage.InputTokens)
|
||||
require.Equal(t, tc.output, result.Usage.OutputTokens)
|
||||
require.Equal(t, tc.cached, result.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, tc.created, result.Usage.CacheCreationInputTokens)
|
||||
}
|
||||
count, done := 0, 0
|
||||
var last *apicompat.ChatCompletionsChunk
|
||||
for _, line := range strings.Split(rec.Body.String(), "\n") {
|
||||
if !strings.HasPrefix(line, "data: ") {
|
||||
continue
|
||||
}
|
||||
payload := strings.TrimPrefix(line, "data: ")
|
||||
if payload == "[DONE]" {
|
||||
done++
|
||||
if tc.want {
|
||||
require.NotNil(t, last)
|
||||
require.NotNil(t, last.Usage)
|
||||
}
|
||||
continue
|
||||
}
|
||||
require.Zero(t, done, "chunk after DONE")
|
||||
var chunk apicompat.ChatCompletionsChunk
|
||||
require.NoError(t, json.Unmarshal([]byte(payload), &chunk))
|
||||
last = &chunk
|
||||
if chunk.Usage == nil {
|
||||
continue
|
||||
}
|
||||
count++
|
||||
require.Empty(t, chunk.Choices)
|
||||
require.Equal(t, tc.prompt, chunk.Usage.PromptTokens)
|
||||
require.Equal(t, tc.output, chunk.Usage.CompletionTokens)
|
||||
require.Equal(t, tc.prompt+tc.output, chunk.Usage.TotalTokens)
|
||||
if tc.cached > 0 {
|
||||
require.NotNil(t, chunk.Usage.PromptTokensDetails)
|
||||
require.Equal(t, tc.cached, chunk.Usage.PromptTokensDetails.CachedTokens)
|
||||
require.Equal(t, tc.created, chunk.Usage.PromptTokensDetails.CacheCreationTokens)
|
||||
}
|
||||
}
|
||||
require.Equal(t, 1, done)
|
||||
if tc.want {
|
||||
require.Equal(t, 1, count)
|
||||
} else {
|
||||
require.Zero(t, count)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -131,7 +131,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
|
||||
}
|
||||
// 区分客户端取消和真正的上游失败,返回更准确的错误消息
|
||||
if c.Request.Context().Err() != nil {
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, "client_disconnected", "Client disconnected before upstream response")
|
||||
return nil, s.writeClaudeError(c, antigravityStatusClientClosed, "client_disconnected", "Client disconnected before upstream response")
|
||||
}
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries")
|
||||
}
|
||||
|
||||
@@ -374,7 +374,7 @@ func (s *AntigravityGatewayService) handleAntigravityCompatTransportError(c *gin
|
||||
}
|
||||
}
|
||||
if c.Request.Context().Err() != nil {
|
||||
return s.writeAntigravityCompatError(c, http.StatusBadGateway, "client_disconnected", "Client disconnected before upstream response")
|
||||
return s.writeAntigravityCompatError(c, antigravityStatusClientClosed, "client_disconnected", "Client disconnected before upstream response")
|
||||
}
|
||||
return s.writeAntigravityCompatError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries")
|
||||
}
|
||||
|
||||
@@ -107,16 +107,24 @@ type antigravityCompatScanEvent struct {
|
||||
}
|
||||
|
||||
type antigravityCompatStreamSession struct {
|
||||
processor *antigravity.StreamingProcessor
|
||||
adapter antigravityCompatStreamAdapter
|
||||
writer *antigravityClientWriter
|
||||
usage *ClaudeUsage
|
||||
pendingEvents []apicompat.AnthropicStreamEvent
|
||||
firstTokenMs *int
|
||||
startTime time.Time
|
||||
meaningfulData bool
|
||||
processor *antigravity.StreamingProcessor
|
||||
adapter antigravityCompatStreamAdapter
|
||||
writer *antigravityClientWriter
|
||||
usage *ClaudeUsage
|
||||
pendingEvents []apicompat.AnthropicStreamEvent
|
||||
firstTokenMs *int
|
||||
startTime time.Time
|
||||
meaningfulData bool
|
||||
preContentKeepaliveSent bool
|
||||
preContentKeepaliveInterval time.Duration
|
||||
}
|
||||
|
||||
const (
|
||||
antigravityCompatPreContentKeepaliveInterval = 15 * time.Second
|
||||
// Allow slow first tokens, but do not let comment-only streams reset the idle timer forever.
|
||||
antigravityCompatPreContentMaxWait = 2 * time.Minute
|
||||
)
|
||||
|
||||
func newAntigravityCompatStreamSession(
|
||||
model string,
|
||||
startTime time.Time,
|
||||
@@ -124,11 +132,12 @@ func newAntigravityCompatStreamSession(
|
||||
writer *antigravityClientWriter,
|
||||
) *antigravityCompatStreamSession {
|
||||
return &antigravityCompatStreamSession{
|
||||
processor: antigravity.NewStreamingProcessor(model),
|
||||
adapter: adapter,
|
||||
writer: writer,
|
||||
usage: &ClaudeUsage{},
|
||||
startTime: startTime,
|
||||
processor: antigravity.NewStreamingProcessor(model),
|
||||
adapter: adapter,
|
||||
writer: writer,
|
||||
usage: &ClaudeUsage{},
|
||||
startTime: startTime,
|
||||
preContentKeepaliveInterval: antigravityCompatPreContentKeepaliveInterval,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -141,15 +150,33 @@ func (s *antigravityCompatStreamSession) consume(line string) {
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) hasMeaningfulData() bool {
|
||||
return s.meaningfulData
|
||||
return s.meaningfulData || s.processor.HasContent()
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) finish() *antigravityStreamResult {
|
||||
func (s *antigravityCompatStreamSession) writePreContentKeepalive(now time.Time) {
|
||||
// Intentional tradeoff: this commits HTTP 200 after 15s, so later upstream
|
||||
// failures cannot fail over; report them as SSE errors to the client instead.
|
||||
if s.hasMeaningfulData() || s.writer.Disconnected() || now.Sub(s.startTime) < s.preContentKeepaliveInterval {
|
||||
return
|
||||
}
|
||||
if s.writer.Write([]byte(": ping\n\n")) {
|
||||
s.preContentKeepaliveSent = true
|
||||
}
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) finish() (*antigravityStreamResult, error) {
|
||||
finalEvents, usage := s.processor.Finish()
|
||||
mergeAntigravityCompatUsage(s.usage, usage)
|
||||
s.consumeClaudeEvents(finalEvents)
|
||||
if !s.hasMeaningfulData() && !s.writer.Disconnected() {
|
||||
if s.preContentKeepaliveSent {
|
||||
s.adapter.WriteError(s.writer, "empty_stream")
|
||||
return s.result(false), errors.New("empty Antigravity compatibility stream after keepalive")
|
||||
}
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
}
|
||||
s.adapter.Finalize(s.writer)
|
||||
return s.result(s.writer.Disconnected())
|
||||
return s.result(s.writer.Disconnected()), nil
|
||||
}
|
||||
|
||||
func (s *antigravityCompatStreamSession) collectResult(clientDisconnect bool) *antigravityStreamResult {
|
||||
@@ -220,24 +247,18 @@ func isMeaningfulAntigravityCompatEvent(event *apicompat.AnthropicStreamEvent) b
|
||||
if event == nil {
|
||||
return false
|
||||
}
|
||||
if event.Type == "message_stop" {
|
||||
return true
|
||||
}
|
||||
if event.ContentBlock != nil {
|
||||
block := event.ContentBlock
|
||||
return block.Type == "tool_use" ||
|
||||
block.Text != "" ||
|
||||
block.Thinking != "" ||
|
||||
block.Signature != "" ||
|
||||
block.Source != nil
|
||||
}
|
||||
if event.Delta != nil {
|
||||
delta := event.Delta
|
||||
return delta.Text != "" ||
|
||||
delta.PartialJSON != "" ||
|
||||
delta.Thinking != "" ||
|
||||
delta.Signature != "" ||
|
||||
delta.StopReason != ""
|
||||
delta.Thinking != ""
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -260,6 +281,13 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
|
||||
originalModel string,
|
||||
adapter antigravityCompatStreamAdapter,
|
||||
prefix string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
return s.handleAntigravityCompatStreamWithKeepaliveInterval(c, resp, startTime, originalModel, adapter, prefix, antigravityCompatPreContentKeepaliveInterval, antigravityCompatPreContentMaxWait)
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleAntigravityCompatStreamWithKeepaliveInterval(
|
||||
c *gin.Context, resp *http.Response, startTime time.Time, originalModel string,
|
||||
adapter antigravityCompatStreamAdapter, prefix string, interval, maxPreContentWait time.Duration,
|
||||
) (*antigravityStreamResult, error) {
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
@@ -275,6 +303,7 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
|
||||
c.Status(http.StatusOK)
|
||||
}
|
||||
session := newAntigravityCompatStreamSession(originalModel, startTime, adapter, writer)
|
||||
session.preContentKeepaliveInterval = interval
|
||||
events, stopScanner, maxLineSize := s.startAntigravityCompatScanner(resp.Body)
|
||||
defer stopScanner()
|
||||
|
||||
@@ -287,15 +316,38 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
|
||||
if keepaliveTicker != nil {
|
||||
defer keepaliveTicker.Stop()
|
||||
}
|
||||
preContentDelay := time.Until(startTime.Add(interval))
|
||||
if preContentDelay < 0 {
|
||||
preContentDelay = 0
|
||||
}
|
||||
preContentTimer := time.NewTimer(preContentDelay)
|
||||
defer preContentTimer.Stop()
|
||||
preContentDeadline := startTime.Add(maxPreContentWait)
|
||||
deadlineDelay := time.Until(preContentDeadline)
|
||||
if deadlineDelay < 0 {
|
||||
deadlineDelay = 0
|
||||
}
|
||||
deadlineTimer := time.NewTimer(deadlineDelay)
|
||||
defer deadlineTimer.Stop()
|
||||
preContentTimeout := func() (*antigravityStreamResult, error) {
|
||||
if session.preContentKeepaliveSent {
|
||||
writeAntigravityCompatStreamError(c, adapter, writer, "stream_timeout")
|
||||
return session.collectResult(false), fmt.Errorf("pre-content stream timeout")
|
||||
}
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case event, open := <-events:
|
||||
if !session.hasMeaningfulData() && !writer.Disconnected() && !time.Now().Before(preContentDeadline) {
|
||||
return preContentTimeout()
|
||||
}
|
||||
if !open {
|
||||
if !session.hasMeaningfulData() && !writer.Disconnected() {
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
return handleAntigravityCompatEmptyStream(c, session)
|
||||
}
|
||||
return session.finish(), nil
|
||||
return session.finish()
|
||||
}
|
||||
if event.err != nil {
|
||||
return s.handleAntigravityCompatReadError(c, session, event.err, maxLineSize, prefix)
|
||||
@@ -309,7 +361,7 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
|
||||
return session.collectResult(true), nil
|
||||
}
|
||||
if !session.hasMeaningfulData() {
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
return preContentTimeout()
|
||||
}
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (%s)", prefix)
|
||||
writeAntigravityCompatStreamError(c, adapter, writer, "stream_timeout")
|
||||
@@ -319,10 +371,30 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
|
||||
if session.hasMeaningfulData() && !writer.Disconnected() {
|
||||
writer.Write([]byte(": ping\n\n"))
|
||||
}
|
||||
case now := <-preContentTimer.C:
|
||||
if !session.hasMeaningfulData() && !writer.Disconnected() && !now.Before(preContentDeadline) {
|
||||
return preContentTimeout()
|
||||
}
|
||||
session.writePreContentKeepalive(now)
|
||||
if !session.hasMeaningfulData() && !writer.Disconnected() {
|
||||
preContentTimer.Reset(interval)
|
||||
}
|
||||
case <-deadlineTimer.C:
|
||||
if !session.hasMeaningfulData() && !writer.Disconnected() {
|
||||
return preContentTimeout()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func handleAntigravityCompatEmptyStream(c *gin.Context, session *antigravityCompatStreamSession) (*antigravityStreamResult, error) {
|
||||
if session.preContentKeepaliveSent {
|
||||
writeAntigravityCompatStreamError(c, session.adapter, session.writer, "empty_stream")
|
||||
return session.collectResult(false), errors.New("empty Antigravity compatibility stream after keepalive")
|
||||
}
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) startAntigravityCompatScanner(
|
||||
body io.Reader,
|
||||
) (<-chan antigravityCompatScanEvent, func(), int) {
|
||||
@@ -407,6 +479,10 @@ func (s *AntigravityGatewayService) handleAntigravityCompatReadError(
|
||||
prefix string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
if !session.hasMeaningfulData() && !session.writer.Disconnected() {
|
||||
if session.preContentKeepaliveSent {
|
||||
writeAntigravityCompatStreamError(c, session.adapter, session.writer, "stream_read_error")
|
||||
return session.collectResult(false), fmt.Errorf("stream read error: %w", err)
|
||||
}
|
||||
return nil, antigravityCompatEmptyStreamError()
|
||||
}
|
||||
if disconnect, handled := handleStreamReadError(err, session.writer.Disconnected(), prefix); handled {
|
||||
|
||||
@@ -28,6 +28,20 @@ type antigravityCompatErrorReader struct {
|
||||
err error
|
||||
}
|
||||
|
||||
type antigravityCompatNotifyingWriter struct {
|
||||
gin.ResponseWriter
|
||||
wrote chan struct{}
|
||||
}
|
||||
|
||||
func (w *antigravityCompatNotifyingWriter) Write(data []byte) (int, error) {
|
||||
n, err := w.ResponseWriter.Write(data)
|
||||
select {
|
||||
case w.wrote <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r *antigravityCompatErrorReader) Read(p []byte) (int, error) {
|
||||
if r.off < len(r.data) {
|
||||
n := copy(p, r.data[r.off:])
|
||||
@@ -786,3 +800,230 @@ func TestAntigravityCompatKeepaliveAfterFirstEvent(t *testing.T) {
|
||||
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
|
||||
require.NoError(t, reader.Close())
|
||||
}
|
||||
|
||||
func TestAntigravityCompatPreContentKeepalive(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
adapter antigravityCompatStreamAdapter
|
||||
want string
|
||||
}{
|
||||
{"chat completions", newAntigravityChatStreamAdapter("gemini-3.1-pro", false), "data: "},
|
||||
{"responses", newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "event: "},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
|
||||
writer := newAntigravityClientWriter(c.Writer, c.Writer, "test")
|
||||
writer.beforeFirstWrite = func() { c.Header("Content-Type", "text/event-stream") }
|
||||
start := time.Now().Add(-20 * time.Second)
|
||||
session := newAntigravityCompatStreamSession("gemini-3.1-pro", start, tt.adapter, writer)
|
||||
session.writePreContentKeepalive(start.Add(14 * time.Second))
|
||||
require.Empty(t, recorder.Body.String())
|
||||
session.consumeClaudeData("message_start", `{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant"}}`)
|
||||
require.False(t, session.hasMeaningfulData())
|
||||
session.writePreContentKeepalive(start.Add(14 * time.Second))
|
||||
require.Empty(t, recorder.Body.String())
|
||||
session.writePreContentKeepalive(start.Add(15 * time.Second))
|
||||
require.Equal(t, ": ping\n\n", recorder.Body.String())
|
||||
require.Nil(t, session.firstTokenMs)
|
||||
session.consumeClaudeData("content_block_delta", `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}`)
|
||||
require.True(t, session.hasMeaningfulData())
|
||||
require.NotNil(t, session.firstTokenMs)
|
||||
require.GreaterOrEqual(t, *session.firstTokenMs, 20000)
|
||||
require.Contains(t, recorder.Body.String(), tt.want)
|
||||
require.Greater(t, strings.Index(recorder.Body.String(), tt.want), strings.Index(recorder.Body.String(), ": ping"))
|
||||
before := recorder.Body.String()
|
||||
session.writePreContentKeepalive(start.Add(30 * time.Second))
|
||||
require.Equal(t, before, recorder.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatHandlerPreContentKeepalive(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
adapter func() antigravityCompatStreamAdapter
|
||||
line string
|
||||
want string
|
||||
}{
|
||||
{"silent chat", func() antigravityCompatStreamAdapter { return newAntigravityChatStreamAdapter("gemini-3.1-pro", false) }, "", `"upstream_error"`},
|
||||
{"signature-only responses", func() antigravityCompatStreamAdapter { return newAntigravityResponsesStreamAdapter("gemini-3.1-pro") }, `data: {"response":{"candidates":[{"content":{"parts":[{"thoughtSignature":"sig","text":""}]}}]}}` + "\n\n", "event: error"},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize, StreamDataIntervalTimeout: 30}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
|
||||
notifier := &antigravityCompatNotifyingWriter{ResponseWriter: c.Writer, wrote: make(chan struct{}, 1)}
|
||||
c.Writer = notifier
|
||||
reader, pipeWriter := io.Pipe()
|
||||
defer func() { _ = reader.Close() }()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := svc.handleAntigravityCompatStream(c, resp, time.Now().Add(-antigravityCompatPreContentKeepaliveInterval+200*time.Millisecond), "gemini-3.1-pro", tt.adapter(), "test")
|
||||
done <- err
|
||||
}()
|
||||
if tt.line != "" {
|
||||
_, err := io.WriteString(pipeWriter, tt.line)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
select {
|
||||
case <-notifier.wrote:
|
||||
case <-time.After(2 * time.Second):
|
||||
_ = pipeWriter.Close()
|
||||
t.Fatal("no pre-content keepalive before read timeout")
|
||||
}
|
||||
require.Equal(t, ": ping\n\n", recorder.Body.String())
|
||||
require.NoError(t, pipeWriter.Close())
|
||||
require.Error(t, <-done)
|
||||
require.Contains(t, recorder.Body.String(), tt.want)
|
||||
require.True(t, IsResponseCommitted(c))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatHandlerRepeatsPreContentKeepalive(t *testing.T) {
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{
|
||||
MaxLineSize: defaultMaxLineSize, StreamDataIntervalTimeout: 30, StreamKeepaliveInterval: 0,
|
||||
}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
|
||||
reader, pipeWriter := io.Pipe()
|
||||
defer func() { _ = reader.Close() }()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := svc.handleAntigravityCompatStreamWithKeepaliveInterval(
|
||||
c, resp, time.Now().Add(-15*time.Millisecond), "gemini-3.1-pro",
|
||||
newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "test", 15*time.Millisecond, time.Second,
|
||||
)
|
||||
done <- err
|
||||
}()
|
||||
time.Sleep(55 * time.Millisecond)
|
||||
require.NoError(t, pipeWriter.Close())
|
||||
require.Error(t, <-done)
|
||||
require.GreaterOrEqual(t, strings.Count(recorder.Body.String(), ": ping\n\n"), 3)
|
||||
require.Contains(t, recorder.Body.String(), "event: error")
|
||||
require.True(t, IsResponseCommitted(c))
|
||||
}
|
||||
|
||||
func TestAntigravityCompatHandlerPreContentDeadlineWithCommentOnlyStream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize, StreamDataIntervalTimeout: 1}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
|
||||
reader, pipeWriter := io.Pipe()
|
||||
defer func() { _ = reader.Close() }()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := svc.handleAntigravityCompatStreamWithKeepaliveInterval(c, resp, time.Now(), "gemini-3.1-pro",
|
||||
newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "test", 10*time.Millisecond, 100*time.Millisecond)
|
||||
done <- err
|
||||
}()
|
||||
stop := make(chan struct{})
|
||||
defer close(stop)
|
||||
go func() {
|
||||
ticker := time.NewTicker(5 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-ticker.C:
|
||||
if _, err := io.WriteString(pipeWriter, ": upstream ping\n\n"); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case err := <-done:
|
||||
require.ErrorContains(t, err, "pre-content stream timeout")
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("comment-only stream exceeded pre-content deadline")
|
||||
}
|
||||
require.Contains(t, recorder.Body.String(), ": ping\n\n")
|
||||
require.Contains(t, recorder.Body.String(), "event: error")
|
||||
require.True(t, IsResponseCommitted(c))
|
||||
_ = pipeWriter.Close()
|
||||
}
|
||||
|
||||
func TestAntigravityCompatExpiredPreContentDeadlineDoesNotCommit(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
|
||||
reader, pipeWriter := io.Pipe()
|
||||
defer func() { _ = reader.Close() }()
|
||||
defer func() { _ = pipeWriter.Close() }()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
|
||||
result, err := svc.handleAntigravityCompatStreamWithKeepaliveInterval(c, resp,
|
||||
time.Now().Add(-time.Minute), "gemini-3.1-pro",
|
||||
newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "test", 10*time.Millisecond, 20*time.Millisecond)
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Empty(t, recorder.Body.String())
|
||||
require.False(t, IsResponseCommitted(c))
|
||||
}
|
||||
|
||||
func TestAntigravityCompatHandlerErrorsAfterPreContentPing(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
wait time.Duration
|
||||
want string
|
||||
}{
|
||||
{"read error", time.Second, "stream_read_error"},
|
||||
{"timeout", 25 * time.Millisecond, "stream_timeout"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
|
||||
reader, pipeWriter := io.Pipe()
|
||||
defer func() { _ = reader.Close() }()
|
||||
defer func() { _ = pipeWriter.Close() }()
|
||||
if tc.name == "read error" {
|
||||
go func() {
|
||||
time.Sleep(15 * time.Millisecond)
|
||||
_ = pipeWriter.CloseWithError(io.ErrUnexpectedEOF)
|
||||
}()
|
||||
}
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: reader}
|
||||
_, err := svc.handleAntigravityCompatStreamWithKeepaliveInterval(c, resp,
|
||||
time.Now().Add(-20*time.Millisecond), "gemini-3.1-pro",
|
||||
newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "test", 10*time.Millisecond, tc.wait)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, recorder.Body.String(), ": ping\n\n")
|
||||
require.Contains(t, recorder.Body.String(), tc.want)
|
||||
require.True(t, IsResponseCommitted(c))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityCompatEmptyAfterKeepaliveReportsStreamError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
adapter antigravityCompatStreamAdapter
|
||||
want string
|
||||
}{
|
||||
{"chat completions", newAntigravityChatStreamAdapter("gemini-3.1-pro", false), `"upstream_error"`},
|
||||
{"responses", newAntigravityResponsesStreamAdapter("gemini-3.1-pro"), "event: error"},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
|
||||
writer := newAntigravityClientWriter(c.Writer, c.Writer, "test")
|
||||
writer.beforeFirstWrite = func() { c.Header("Content-Type", "text/event-stream") }
|
||||
start := time.Now().Add(-20 * time.Second)
|
||||
session := newAntigravityCompatStreamSession("gemini-3.1-pro", start, tt.adapter, writer)
|
||||
session.consumeClaudeData("message_start", `{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant"}}`)
|
||||
session.writePreContentKeepalive(start.Add(15 * time.Second))
|
||||
result, err := handleAntigravityCompatEmptyStream(c, session)
|
||||
require.Error(t, err)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.NotErrorAs(t, err, &failoverErr)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, IsResponseCommitted(c))
|
||||
require.Contains(t, recorder.Body.String(), tt.want)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -180,7 +180,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
|
||||
}
|
||||
// 区分客户端取消和真正的上游失败,返回更准确的错误消息
|
||||
if c.Request.Context().Err() != nil {
|
||||
return nil, s.writeGoogleError(c, http.StatusBadGateway, "Client disconnected before upstream response")
|
||||
return nil, s.writeGoogleError(c, antigravityStatusClientClosed, "Client disconnected before upstream response")
|
||||
}
|
||||
return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries")
|
||||
}
|
||||
|
||||
@@ -686,6 +686,9 @@ func mergeTextPartsToResponse(response map[string]any, textParts []string) map[s
|
||||
return result
|
||||
}
|
||||
|
||||
// antigravityStatusClientClosed 是客户端在上游响应前断开时回写的状态码(499 client closed request)。
|
||||
const antigravityStatusClientClosed = 499
|
||||
|
||||
func (s *AntigravityGatewayService) writeClaudeError(c *gin.Context, status int, errType, message string) error {
|
||||
MarkResponseCommitted(c)
|
||||
c.JSON(status, gin.H{
|
||||
@@ -790,6 +793,8 @@ func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int,
|
||||
statusStr = "NOT_FOUND"
|
||||
case 429:
|
||||
statusStr = "RESOURCE_EXHAUSTED"
|
||||
case antigravityStatusClientClosed:
|
||||
statusStr = "CANCELLED"
|
||||
case 500:
|
||||
statusStr = "INTERNAL"
|
||||
case 502, 503:
|
||||
|
||||
@@ -30,6 +30,8 @@ var (
|
||||
ErrAPIKeyTooShort = infraerrors.BadRequest("API_KEY_TOO_SHORT", "api key must be at least 16 characters")
|
||||
ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens")
|
||||
ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later")
|
||||
ErrAPIKeyCreateLimited = infraerrors.TooManyRequests("API_KEY_CREATE_RATE_LIMITED", "too many api keys created recently, please try again later")
|
||||
ErrAPIKeyCountExceeded = infraerrors.Forbidden("API_KEY_COUNT_EXCEEDED", "api key count limit reached, please delete unused keys first")
|
||||
ErrAPIKeyAuthOverloaded = infraerrors.ServiceUnavailable("API_KEY_AUTH_OVERLOADED", "api key authentication is temporarily overloaded")
|
||||
ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern")
|
||||
// ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key has expired")
|
||||
@@ -48,6 +50,7 @@ const (
|
||||
defaultAuthLookupConcurrency = 64
|
||||
defaultNegativeAuthCacheSize = 16384
|
||||
apiKeyMaxErrorsPerHour = 20
|
||||
apiKeyCreateCountWindow = time.Hour
|
||||
apiKeyLastUsedMinTouch = 30 * time.Second
|
||||
apiKeySortCurrentConcurrency = "current_concurrency"
|
||||
// DB 写失败后的短退避,避免请求路径持续同步重试造成写风暴与高延迟。
|
||||
@@ -172,7 +175,7 @@ type APIKeyQuotaUsageState struct {
|
||||
type APIKeyCache interface {
|
||||
GetCreateAttemptCount(ctx context.Context, userID int64) (int, error)
|
||||
IncrementCreateAttemptCount(ctx context.Context, userID int64) error
|
||||
DeleteCreateAttemptCount(ctx context.Context, userID int64) error
|
||||
IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error)
|
||||
|
||||
IncrementDailyUsage(ctx context.Context, apiKey string) error
|
||||
SetDailyUsageExpiry(ctx context.Context, apiKey string, ttl time.Duration) error
|
||||
@@ -435,6 +438,34 @@ func (s *APIKeyService) checkAPIKeyRateLimit(ctx context.Context, userID int64)
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkAPIKeyCreateLimits 校验创建 API Key 的防滥用限制(对自定义与自动生成的 Key 一视同仁)。
|
||||
// 数量上限按未删除的 Key 计;创建次数按固定窗口累计,删除 Key 不返还次数,
|
||||
// 以阻断"删除后反复新建"的循环。Redis 出错时与自定义 Key 限流一致,不阻止用户操作。
|
||||
func (s *APIKeyService) checkAPIKeyCreateLimits(ctx context.Context, userID int64) error {
|
||||
if s.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
if maxActive := s.cfg.APIKeyCreate.MaxActivePerUser; maxActive > 0 {
|
||||
count, err := s.apiKeyRepo.CountByUserID(ctx, userID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("count api keys: %w", err)
|
||||
}
|
||||
if count >= int64(maxActive) {
|
||||
return ErrAPIKeyCountExceeded
|
||||
}
|
||||
}
|
||||
if maxPerHour := s.cfg.APIKeyCreate.MaxPerUserPerHour; maxPerHour > 0 && s.cache != nil {
|
||||
count, err := s.cache.IncrementCreateCount(ctx, userID, apiKeyCreateCountWindow)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if count > int64(maxPerHour) {
|
||||
return ErrAPIKeyCreateLimited
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// incrementAPIKeyErrorCount 增加用户创建自定义Key的错误计数
|
||||
func (s *APIKeyService) incrementAPIKeyErrorCount(ctx context.Context, userID int64) {
|
||||
if s.cache == nil {
|
||||
@@ -530,6 +561,10 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK
|
||||
}
|
||||
}
|
||||
|
||||
if err := s.checkAPIKeyCreateLimits(ctx, userID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 创建API Key记录
|
||||
apiKey := &APIKey{
|
||||
UserID: userID,
|
||||
@@ -820,10 +855,6 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req
|
||||
if req.Status != nil {
|
||||
apiKey.Status = *req.Status
|
||||
fields.Status = true
|
||||
// 如果状态改变,清除Redis缓存
|
||||
if s.cache != nil {
|
||||
_ = s.cache.DeleteCreateAttemptCount(ctx, apiKey.UserID)
|
||||
}
|
||||
}
|
||||
|
||||
// Update quota fields
|
||||
@@ -931,9 +962,7 @@ func (s *APIKeyService) Delete(ctx context.Context, id int64, userID int64) erro
|
||||
}
|
||||
|
||||
// 删除成功后再清理缓存,避免"缓存已清但删除失败"的竞态。
|
||||
if s.cache != nil {
|
||||
_ = s.cache.DeleteCreateAttemptCount(ctx, userID)
|
||||
}
|
||||
// 注意:不清零创建相关计数,否则"删除后反复新建"即可绕过创建限流。
|
||||
s.InvalidateAuthCacheByKey(ctx, key)
|
||||
s.lastUsedTouchL1.Delete(id)
|
||||
|
||||
|
||||
@@ -138,8 +138,8 @@ func (s *authCacheStub) IncrementCreateAttemptCount(ctx context.Context, userID
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *authCacheStub) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
return nil
|
||||
func (s *authCacheStub) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (s *authCacheStub) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type createLimitAPIKeyRepoStub struct {
|
||||
*apiKeyRepoStub
|
||||
activeCount int64
|
||||
countErr error
|
||||
created []*APIKey
|
||||
}
|
||||
|
||||
func (s *createLimitAPIKeyRepoStub) CountByUserID(ctx context.Context, userID int64) (int64, error) {
|
||||
return s.activeCount, s.countErr
|
||||
}
|
||||
|
||||
func (s *createLimitAPIKeyRepoStub) ExistsByKey(ctx context.Context, key string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (s *createLimitAPIKeyRepoStub) Create(ctx context.Context, key *APIKey) error {
|
||||
s.created = append(s.created, key)
|
||||
return nil
|
||||
}
|
||||
|
||||
type createLimitCacheStub struct {
|
||||
*apiKeyCacheStub
|
||||
createCounts map[int64]int64
|
||||
incrErr error
|
||||
windows []time.Duration
|
||||
}
|
||||
|
||||
func (s *createLimitCacheStub) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
s.windows = append(s.windows, window)
|
||||
if s.incrErr != nil {
|
||||
return 0, s.incrErr
|
||||
}
|
||||
s.createCounts[userID]++
|
||||
return s.createCounts[userID], nil
|
||||
}
|
||||
|
||||
func newCreateLimitService(repo *createLimitAPIKeyRepoStub, cache *createLimitCacheStub, maxActive, maxPerHour int) *APIKeyService {
|
||||
cfg := &config.Config{}
|
||||
cfg.APIKeyCreate.MaxActivePerUser = maxActive
|
||||
cfg.APIKeyCreate.MaxPerUserPerHour = maxPerHour
|
||||
return &APIKeyService{
|
||||
apiKeyRepo: repo,
|
||||
userRepo: &userRepoStub{user: &User{ID: 7}},
|
||||
cache: cache,
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
func newCreateLimitStubs() (*createLimitAPIKeyRepoStub, *createLimitCacheStub) {
|
||||
return &createLimitAPIKeyRepoStub{apiKeyRepoStub: &apiKeyRepoStub{}},
|
||||
&createLimitCacheStub{apiKeyCacheStub: &apiKeyCacheStub{}, createCounts: map[int64]int64{}}
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_RateLimitsGeneratedKeys(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
svc := newCreateLimitService(repo, cache, 0, 3)
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCreateLimited)
|
||||
require.Len(t, repo.created, 3)
|
||||
require.Equal(t, apiKeyCreateCountWindow, cache.windows[0])
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_RateLimitsCustomKeys(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
svc := newCreateLimitService(repo, cache, 0, 1)
|
||||
custom := "sk-custom-key-0123456789"
|
||||
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k", CustomKey: &custom})
|
||||
require.NoError(t, err)
|
||||
_, err = svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k", CustomKey: &custom})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCreateLimited)
|
||||
require.Len(t, repo.created, 1)
|
||||
}
|
||||
|
||||
// 删除 Key 释放数量名额,但不返还创建次数,删建循环仍受每小时次数限制。
|
||||
func TestAPIKeyServiceCreate_DeleteDoesNotRefundCreateCount(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
repo.apiKey = &APIKey{ID: 1, UserID: 7, Key: "k"}
|
||||
svc := newCreateLimitService(repo, cache, 1, 2)
|
||||
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
repo.activeCount = 1
|
||||
_, err = svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCountExceeded)
|
||||
|
||||
require.NoError(t, svc.Delete(context.Background(), 1, 7))
|
||||
repo.activeCount = 0
|
||||
_, err = svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
|
||||
repo.activeCount = 0
|
||||
_, err = svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCreateLimited)
|
||||
require.Len(t, repo.created, 2)
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_ActiveLimitDoesNotConsumeCreateCount(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
repo.activeCount = 5
|
||||
svc := newCreateLimitService(repo, cache, 5, 10)
|
||||
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCountExceeded)
|
||||
require.Zero(t, cache.createCounts[7])
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_ZeroLimitsDisableChecks(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
repo.countErr = errors.New("must not count")
|
||||
svc := newCreateLimitService(repo, cache, 0, 0)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Empty(t, cache.windows)
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_CountErrorFailsClosed(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
repo.countErr = errors.New("db down")
|
||||
svc := newCreateLimitService(repo, cache, 5, 0)
|
||||
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.Error(t, err)
|
||||
require.Empty(t, repo.created)
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_RedisErrorFailsOpen(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
cache.incrErr = errors.New("redis down")
|
||||
svc := newCreateLimitService(repo, cache, 0, 1)
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
@@ -234,11 +234,7 @@ func (s *apiKeyRepoStub) GetRateLimitData(ctx context.Context, id int64) (*APIKe
|
||||
|
||||
// apiKeyCacheStub 是 APIKeyCache 接口的测试桩实现。
|
||||
// 用于验证删除操作时缓存清理逻辑是否被正确调用。
|
||||
//
|
||||
// 设计说明:
|
||||
// - invalidated: 记录被清除缓存的用户 ID 列表
|
||||
type apiKeyCacheStub struct {
|
||||
invalidated []int64 // 记录调用 DeleteCreateAttemptCount 时传入的用户 ID
|
||||
deleteAuthKeys []string // 记录调用 DeleteAuthCache 时传入的缓存 key
|
||||
}
|
||||
|
||||
@@ -252,11 +248,9 @@ func (s *apiKeyCacheStub) IncrementCreateAttemptCount(ctx context.Context, userI
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteCreateAttemptCount 记录被清除缓存的用户 ID。
|
||||
// 删除 API Key 时会调用此方法清除用户的创建尝试计数缓存。
|
||||
func (s *apiKeyCacheStub) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
s.invalidated = append(s.invalidated, userID)
|
||||
return nil
|
||||
// IncrementCreateCount 空实现,本测试不验证此行为
|
||||
func (s *apiKeyCacheStub) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// IncrementDailyUsage 空实现,本测试不验证此行为
|
||||
@@ -306,8 +300,7 @@ func TestApiKeyService_Delete_OwnerMismatch(t *testing.T) {
|
||||
|
||||
err := svc.Delete(context.Background(), 10, 2) // API Key ID=10, 调用者 userID=2
|
||||
require.ErrorIs(t, err, ErrInsufficientPerms)
|
||||
require.Empty(t, repo.deletedIDs) // 验证删除操作未被调用
|
||||
require.Empty(t, cache.invalidated) // 验证缓存未被清除
|
||||
require.Empty(t, repo.deletedIDs) // 验证删除操作未被调用
|
||||
require.Empty(t, cache.deleteAuthKeys)
|
||||
}
|
||||
|
||||
@@ -316,7 +309,7 @@ func TestApiKeyService_Delete_OwnerMismatch(t *testing.T) {
|
||||
// - GetKeyAndOwnerID 返回所有者 ID 为 7
|
||||
// - 调用者 userID 为 7(匹配)
|
||||
// - Delete 成功执行
|
||||
// - 缓存被正确清除(使用 ownerID)
|
||||
// - 认证缓存被正确清除
|
||||
// - 返回 nil 错误
|
||||
func TestApiKeyService_Delete_Success(t *testing.T) {
|
||||
repo := &apiKeyRepoStub{
|
||||
@@ -328,8 +321,7 @@ func TestApiKeyService_Delete_Success(t *testing.T) {
|
||||
|
||||
err := svc.Delete(context.Background(), 42, 7) // API Key ID=42, 调用者 userID=7
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []int64{42}, repo.deletedIDs) // 验证正确的 API Key 被删除
|
||||
require.Equal(t, []int64{7}, cache.invalidated) // 验证所有者的缓存被清除
|
||||
require.Equal(t, []int64{42}, repo.deletedIDs) // 验证正确的 API Key 被删除
|
||||
require.Equal(t, []string{svc.authCacheKey("k")}, cache.deleteAuthKeys)
|
||||
_, exists := svc.lastUsedTouchL1.Load(int64(42))
|
||||
require.False(t, exists, "delete should clear touch debounce cache")
|
||||
@@ -349,7 +341,6 @@ func TestApiKeyService_Delete_NotFound(t *testing.T) {
|
||||
err := svc.Delete(context.Background(), 99, 1)
|
||||
require.ErrorIs(t, err, ErrAPIKeyNotFound)
|
||||
require.Empty(t, repo.deletedIDs)
|
||||
require.Empty(t, cache.invalidated)
|
||||
require.Empty(t, cache.deleteAuthKeys)
|
||||
}
|
||||
|
||||
@@ -496,6 +487,5 @@ func TestApiKeyService_Delete_DeleteFails(t *testing.T) {
|
||||
require.Error(t, err)
|
||||
require.ErrorContains(t, err, "delete api key")
|
||||
require.Equal(t, []int64{3}, repo.deletedIDs) // 验证 DeleteWithAudit 被调用
|
||||
require.Empty(t, cache.invalidated) // 验证删除失败时缓存未被清除(新顺序:先删后清)
|
||||
require.Empty(t, cache.deleteAuthKeys) // 验证删除失败时 auth 缓存未被清除
|
||||
}
|
||||
|
||||
@@ -42,8 +42,8 @@ func (s *quotaStateCacheStub) IncrementCreateAttemptCount(context.Context, int64
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *quotaStateCacheStub) DeleteCreateAttemptCount(context.Context, int64) error {
|
||||
return nil
|
||||
func (s *quotaStateCacheStub) IncrementCreateCount(context.Context, int64, time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (s *quotaStateCacheStub) IncrementDailyUsage(context.Context, string) error {
|
||||
|
||||
@@ -67,8 +67,10 @@ func (*authInvalidationCacheStub) GetCreateAttemptCount(context.Context, int64)
|
||||
func (*authInvalidationCacheStub) IncrementCreateAttemptCount(context.Context, int64) error {
|
||||
return nil
|
||||
}
|
||||
func (*authInvalidationCacheStub) DeleteCreateAttemptCount(context.Context, int64) error { return nil }
|
||||
func (*authInvalidationCacheStub) IncrementDailyUsage(context.Context, string) error { return nil }
|
||||
func (*authInvalidationCacheStub) IncrementCreateCount(context.Context, int64, time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (*authInvalidationCacheStub) IncrementDailyUsage(context.Context, string) error { return nil }
|
||||
func (*authInvalidationCacheStub) SetDailyUsageExpiry(context.Context, string, time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -127,8 +127,17 @@ func normalizeBedrockModelID(modelID string) (normalized string, shouldAdjustReg
|
||||
return "", false, false
|
||||
}
|
||||
if mapped, exists := domain.DefaultBedrockModelMapping[modelID]; exists {
|
||||
// Sonnet 5.5 currently has only a global inference profile on
|
||||
// bedrock-runtime. A caller's AWS region selects the endpoint, but must
|
||||
// not rewrite the profile ID to a regional one that does not exist.
|
||||
if mapped == "global.anthropic.claude-sonnet-5-5" {
|
||||
return mapped, false, true
|
||||
}
|
||||
return mapped, true, true
|
||||
}
|
||||
if modelID == "global.anthropic.claude-sonnet-5-5" {
|
||||
return modelID, false, true
|
||||
}
|
||||
if isRegionalBedrockModelID(modelID) {
|
||||
return modelID, true, true
|
||||
}
|
||||
@@ -181,7 +190,7 @@ func BuildBedrockURL(region, modelID string, stream bool) string {
|
||||
// PrepareBedrockRequestBody 处理请求体以适配 Bedrock API
|
||||
// 1. 注入 anthropic_version
|
||||
// 2. 注入 anthropic_beta(从客户端 anthropic-beta 头解析)
|
||||
// 3. 移除 Bedrock 不支持的字段(model, stream, output_format, output_config)
|
||||
// 3. 移除 Bedrock 不支持的字段(model, stream, output_format);Sonnet 5.5 保留 output_config.effort
|
||||
// 4. 移除工具定义中的 custom 字段(Claude Code 会发送 custom: {defer_loading: true})
|
||||
// 5. 清理 cache_control 中 Bedrock 不支持的字段(scope, ttl)
|
||||
// 6. 修复 thinking 字段兼容性(Opus 4.7 仅支持 adaptive,enabled 需要 budget_tokens)
|
||||
@@ -240,10 +249,20 @@ func PrepareBedrockRequestBodyWithTokens(body []byte, modelID string, betaTokens
|
||||
// 参考 litellm: _convert_output_format_to_inline_schema()
|
||||
body = convertOutputFormatToInlineSchema(body)
|
||||
|
||||
// 移除 output_config 字段(Bedrock Invoke 不支持)
|
||||
body, err = sjson.DeleteBytes(body, "output_config")
|
||||
// InvokeModel accepts output_config.effort for Sonnet 5.5. Keep just that
|
||||
// field; output_config.format has already been inlined above, and older
|
||||
// models retain the existing output_config stripping behavior.
|
||||
if claude.IsSonnet55(modelID) {
|
||||
if effort := gjson.GetBytes(body, "output_config.effort"); effort.Exists() {
|
||||
body, err = sjson.SetRawBytes(body, "output_config", []byte(`{"effort":`+effort.Raw+`}`))
|
||||
} else {
|
||||
body, err = sjson.DeleteBytes(body, "output_config")
|
||||
}
|
||||
} else {
|
||||
body, err = sjson.DeleteBytes(body, "output_config")
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("remove output_config field: %w", err)
|
||||
return nil, fmt.Errorf("normalize output_config field: %w", err)
|
||||
}
|
||||
|
||||
// 移除工具定义中的 custom 字段
|
||||
@@ -270,17 +289,21 @@ func ResolveBedrockBetaTokens(betaHeader string, body []byte, modelID string) []
|
||||
return filterBedrockBetaTokens(betaTokens)
|
||||
}
|
||||
|
||||
// convertOutputFormatToInlineSchema 将 output_format 中的 JSON schema 内联到最后一条 user message
|
||||
// convertOutputFormatToInlineSchema 将结构化输出的 JSON schema 内联到最后一条 user message
|
||||
// Bedrock Invoke 不支持 output_format 参数,litellm 的做法是将 schema 追加到用户消息中
|
||||
// 参考: litellm AmazonAnthropicClaudeMessagesConfig._convert_output_format_to_inline_schema()
|
||||
func convertOutputFormatToInlineSchema(body []byte) []byte {
|
||||
outputFormat := gjson.GetBytes(body, "output_format")
|
||||
outputFormat := gjson.GetBytes(body, "output_config.format")
|
||||
if !outputFormat.Exists() {
|
||||
outputFormat = gjson.GetBytes(body, "output_format")
|
||||
}
|
||||
if !outputFormat.Exists() || !outputFormat.IsObject() {
|
||||
return body
|
||||
}
|
||||
|
||||
// 先从请求体中移除 output_format
|
||||
// 先从请求体中移除两个版本的结构化输出字段。
|
||||
body, _ = sjson.DeleteBytes(body, "output_format")
|
||||
body, _ = sjson.DeleteBytes(body, "output_config.format")
|
||||
|
||||
schema := outputFormat.Get("schema")
|
||||
if !schema.Exists() {
|
||||
@@ -707,6 +730,7 @@ func isBedrockFable5(modelID string) bool {
|
||||
const defaultThinkingBudgetTokens = 10000
|
||||
|
||||
// sanitizeBedrockThinking 修复 thinking 字段的 Bedrock 兼容性问题:
|
||||
// - Sonnet 5.5: enabled 改为 adaptive;disabled 改为 between_tools
|
||||
// - Fable 5: 仅使用 always-on adaptive thinking,不支持手动 budget_tokens
|
||||
// - Opus 4.7+: 仅支持 "adaptive",将 "enabled" 转换为 "adaptive" 并移除 budget_tokens
|
||||
// - 其他模型: "enabled" 必须带 budget_tokens,缺失时补充默认值
|
||||
@@ -731,6 +755,17 @@ func sanitizeBedrockThinking(body []byte, modelID string) []byte {
|
||||
return body
|
||||
}
|
||||
|
||||
if claude.IsSonnet55(modelID) {
|
||||
switch thinkingType {
|
||||
case "enabled":
|
||||
body, _ = sjson.SetBytes(body, "thinking.type", "adaptive")
|
||||
body, _ = sjson.DeleteBytes(body, "thinking.budget_tokens")
|
||||
case "disabled":
|
||||
body, _ = sjson.SetBytes(body, "thinking.type", "between_tools")
|
||||
}
|
||||
return body
|
||||
}
|
||||
|
||||
if isBedrockOpus47OrNewer(modelID) {
|
||||
if thinkingType == "enabled" {
|
||||
body, _ = sjson.SetBytes(body, "thinking.type", "adaptive")
|
||||
|
||||
@@ -482,6 +482,37 @@ func TestBedrockCrossRegionPrefix(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestResolveBedrockModelID(t *testing.T) {
|
||||
t.Run("sonnet 5.5 uses the global inference profile in every region", func(t *testing.T) {
|
||||
for _, region := range []string{"us-east-1", "eu-west-1", "ap-southeast-2"} {
|
||||
account := &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeBedrock,
|
||||
Credentials: map[string]any{
|
||||
"aws_region": region,
|
||||
},
|
||||
}
|
||||
modelID, ok := ResolveBedrockModelID(account, "claude-sonnet-5-5")
|
||||
require.True(t, ok, region)
|
||||
assert.Equal(t, "global.anthropic.claude-sonnet-5-5", modelID, region)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("explicit global sonnet 5.5 profile is not rewritten", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeBedrock,
|
||||
Credentials: map[string]any{
|
||||
"aws_region": "eu-west-1",
|
||||
"model_mapping": map[string]any{
|
||||
"public-sonnet": "global.anthropic.claude-sonnet-5-5",
|
||||
},
|
||||
},
|
||||
}
|
||||
modelID, ok := ResolveBedrockModelID(account, "public-sonnet")
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "global.anthropic.claude-sonnet-5-5", modelID)
|
||||
})
|
||||
|
||||
t.Run("default alias resolves and adjusts region", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
@@ -779,6 +810,19 @@ func TestIsBedrockOpus47OrNewer(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestSanitizeBedrockThinking(t *testing.T) {
|
||||
t.Run("sonnet 5.5 converts legacy enabled to adaptive", func(t *testing.T) {
|
||||
input := `{"thinking":{"type":"enabled","budget_tokens":10000},"messages":[]}`
|
||||
result := sanitizeBedrockThinking([]byte(input), "global.anthropic.claude-sonnet-5-5")
|
||||
assert.Equal(t, "adaptive", gjson.GetBytes(result, "thinking.type").String())
|
||||
assert.False(t, gjson.GetBytes(result, "thinking.budget_tokens").Exists())
|
||||
})
|
||||
|
||||
t.Run("sonnet 5.5 converts legacy disabled to between_tools", func(t *testing.T) {
|
||||
input := `{"thinking":{"type":"disabled"},"messages":[]}`
|
||||
result := sanitizeBedrockThinking([]byte(input), "global.anthropic.claude-sonnet-5-5")
|
||||
assert.Equal(t, "between_tools", gjson.GetBytes(result, "thinking.type").String())
|
||||
})
|
||||
|
||||
t.Run("Fable 5 将 enabled 转换为 adaptive 并移除预算", func(t *testing.T) {
|
||||
input := `{"thinking":{"type":"enabled","budget_tokens":10000},"messages":[]}`
|
||||
result := sanitizeBedrockThinking([]byte(input), "anthropic.claude-fable-5")
|
||||
@@ -1000,6 +1044,20 @@ func TestPrepareBedrockRequestBodyWithTokens_CCCompat(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestPrepareBedrockSonnet55PreservesEffort(t *testing.T) {
|
||||
input := []byte(`{"model":"claude-sonnet-5-5","max_tokens":1024,"output_config":{"effort":"medium","format":{"type":"json_schema","schema":{"type":"object"}}},"messages":[{"role":"user","content":"hello"}]}`)
|
||||
result, err := PrepareBedrockRequestBodyWithTokens(input, "global.anthropic.claude-sonnet-5-5", nil, false)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "medium", gjson.GetBytes(result, "output_config.effort").String())
|
||||
assert.False(t, gjson.GetBytes(result, "output_config.format").Exists())
|
||||
assert.Contains(t, gjson.GetBytes(result, "messages.0.content.1.text").String(), `"type":"object"`, "schema should still be included in the final user message")
|
||||
assert.False(t, gjson.GetBytes(result, "model").Exists())
|
||||
|
||||
oldModel, err := PrepareBedrockRequestBodyWithTokens(input, "us.anthropic.claude-sonnet-4-6", nil, false)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, gjson.GetBytes(oldModel, "output_config").Exists())
|
||||
}
|
||||
|
||||
func TestSanitizeBedrockCCFields(t *testing.T) {
|
||||
t.Run("removes service_tier and interface_geo", func(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-opus-4-6","service_tier":"standard","interface_geo":"us","messages":[]}`)
|
||||
|
||||
@@ -430,7 +430,15 @@ func (s *BillingService) initFallbackPricing() {
|
||||
CacheCreation1hPrice: 8e-6,
|
||||
SupportsCacheBreakdown: true,
|
||||
}
|
||||
|
||||
s.fallbackPrices["claude-sonnet-5-5"] = &ModelPricing{
|
||||
InputPricePerToken: 2e-6,
|
||||
OutputPricePerToken: 10e-6,
|
||||
CacheCreationPricePerToken: 2.5e-6,
|
||||
CacheReadPricePerToken: 0.2e-6,
|
||||
CacheCreation5mPrice: 2.5e-6,
|
||||
CacheCreation1hPrice: 4e-6,
|
||||
SupportsCacheBreakdown: true,
|
||||
}
|
||||
// Claude Fable 5.x uses the same input/output and cache-write prices, while
|
||||
// Fable 5.1 reduces cache reads from $1 to $0.25 per MTok.
|
||||
s.fallbackPrices["claude-fable-5"] = &ModelPricing{
|
||||
@@ -978,6 +986,9 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
|
||||
if claude.IsOpus55(modelLower) {
|
||||
return s.fallbackPrices["claude-opus-5-5"]
|
||||
}
|
||||
if claude.IsSonnet55(modelLower) {
|
||||
return s.fallbackPrices["claude-sonnet-5-5"]
|
||||
}
|
||||
if strings.Contains(modelLower, "opus") {
|
||||
// "opus-5" 必须先判:不能用裸 "5" 匹配,否则 claude-opus-4-5 会被误判。
|
||||
if strings.Contains(modelLower, "opus-5") || strings.Contains(modelLower, "opus5") {
|
||||
@@ -1350,8 +1361,8 @@ func (s *BillingService) getModelPricingAt(model string, pricingAt time.Time) (*
|
||||
return nil, fmt.Errorf("%w for model: %s", ErrModelPricingUnavailable, model)
|
||||
}
|
||||
|
||||
// GetModelPricingWithChannel 获取模型定价,渠道配置的价格覆盖默认值
|
||||
// 渠道存在时,未配置的图片输出价格归零(不回退到 LiteLLM)
|
||||
// GetModelPricingWithChannel 获取模型定价,渠道配置的价格覆盖默认值。
|
||||
// 与其他 token 字段一致,渠道留空的图片输入/输出价沿用目录价,见 applyChannelImagePriceOverrides。
|
||||
func (s *BillingService) GetModelPricingWithChannel(model string, channelPricing *ChannelModelPricing) (*ModelPricing, error) {
|
||||
pricing, err := s.GetModelPricing(model)
|
||||
if err != nil {
|
||||
@@ -1367,13 +1378,7 @@ func (s *BillingService) GetModelPricingWithChannel(model string, channelPricing
|
||||
pricing.FastMultiplier = channelPricing.FastMultiplier
|
||||
pricing.FlexMultiplier = channelPricing.FlexMultiplier
|
||||
pricing.ReasoningEffortMultipliers = maps.Clone(channelPricing.ReasoningEffortMultipliers)
|
||||
if channelPricing.ImageOutputPrice != nil {
|
||||
pricing.ImageOutputPricePerToken = *channelPricing.ImageOutputPrice
|
||||
} else {
|
||||
pricing.ImageOutputPricePerToken = 0
|
||||
}
|
||||
pricing.ImageOutputPriceExplicit = true
|
||||
applyChannelImageInputPrice(channelPricing, pricing)
|
||||
applyChannelImagePriceOverrides(channelPricing, pricing)
|
||||
return pricing, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -1874,19 +1874,49 @@ func TestGetModelPricingWithChannel_UnknownModelReturnsError(t *testing.T) {
|
||||
require.Contains(t, err.Error(), "pricing not found")
|
||||
}
|
||||
|
||||
func TestGetModelPricingWithChannel_NilImageOutputPriceZerosAndMarksExplicit(t *testing.T) {
|
||||
svc := newTestBillingService()
|
||||
func TestGetModelPricingWithChannel_NilImagePricesInheritCatalog(t *testing.T) {
|
||||
svc := NewBillingService(&config.Config{}, newStubPricingServiceFromMap(map[string]*LiteLLMModelPricing{
|
||||
"gpt-image-2": {
|
||||
Mode: "image_generation",
|
||||
InputCostPerToken: 5e-6,
|
||||
OutputCostPerToken: 10e-6,
|
||||
InputCostPerImageToken: 8e-6,
|
||||
OutputCostPerImageToken: 30e-6,
|
||||
},
|
||||
}))
|
||||
|
||||
chPricing := &ChannelModelPricing{
|
||||
InputPrice: testPtrFloat64(10e-6),
|
||||
OutputPrice: testPtrFloat64(20e-6),
|
||||
// ImageOutputPrice intentionally nil
|
||||
InputPrice: testPtrFloat64(6e-6),
|
||||
OutputPrice: testPtrFloat64(12e-6),
|
||||
// ImageInputPrice / ImageOutputPrice intentionally nil
|
||||
}
|
||||
pricing, err := svc.GetModelPricingWithChannel("claude-sonnet-4", chPricing)
|
||||
pricing, err := svc.GetModelPricingWithChannel("gpt-image-2", chPricing)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.InDelta(t, 6e-6, pricing.InputPricePerToken, 1e-12)
|
||||
require.InDelta(t, 30e-6, pricing.ImageOutputPricePerToken, 1e-12)
|
||||
require.False(t, pricing.ImageOutputPriceExplicit)
|
||||
require.InDelta(t, 8e-6, pricing.ImageInputPricePerToken, 1e-12)
|
||||
}
|
||||
|
||||
func TestGetModelPricingWithChannel_ExplicitImagePricesOverrideCatalog(t *testing.T) {
|
||||
svc := NewBillingService(&config.Config{}, newStubPricingServiceFromMap(map[string]*LiteLLMModelPricing{
|
||||
"gpt-image-2": {
|
||||
Mode: "image_generation",
|
||||
InputCostPerToken: 5e-6,
|
||||
InputCostPerImageToken: 8e-6,
|
||||
OutputCostPerImageToken: 30e-6,
|
||||
},
|
||||
}))
|
||||
|
||||
pricing, err := svc.GetModelPricingWithChannel("gpt-image-2", &ChannelModelPricing{
|
||||
ImageInputPrice: testPtrFloat64(9e-6),
|
||||
ImageOutputPrice: testPtrFloat64(0),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0.0, pricing.ImageOutputPricePerToken)
|
||||
require.True(t, pricing.ImageOutputPriceExplicit)
|
||||
require.True(t, pricing.ImageOutputPriceExplicit, "显式 0 仍表示图片输出免费")
|
||||
require.InDelta(t, 9e-6, pricing.ImageInputPricePerToken, 1e-12)
|
||||
}
|
||||
|
||||
func TestComputeTokenBreakdown_ExplicitZeroImagePrice_NoFallback(t *testing.T) {
|
||||
@@ -1976,18 +2006,40 @@ func TestNewModelPricingCatalogFallbackAndContext(t *testing.T) {
|
||||
}
|
||||
})
|
||||
}
|
||||
t.Run(source+"/opus", func(t *testing.T) {
|
||||
tokens := UsageTokens{InputTokens: 300000, OutputTokens: 500, CacheReadTokens: 1000, CacheCreationTokens: 1000, CacheCreation5mTokens: 400, CacheCreation1hTokens: 600}
|
||||
for tier, mult := range map[string]float64{"": 1, "fast": 2} {
|
||||
cost, err := svc.CalculateCostWithServiceTier("claude-opus-5-5", tokens, 1, tier)
|
||||
for _, model := range []string{"claude-opus-5-5", "anthropic/claude-opus-5.5"} {
|
||||
t.Run(source+"/"+model, func(t *testing.T) {
|
||||
tokens := UsageTokens{InputTokens: 300000, OutputTokens: 500, CacheReadTokens: 1000, CacheCreationTokens: 1000, CacheCreation5mTokens: 400, CacheCreation1hTokens: 600}
|
||||
for tier, mult := range map[string]float64{"": 1, "fast": 2} {
|
||||
cost, err := svc.CalculateCostWithServiceTier(model, tokens, 1, tier)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 1.2*mult, cost.InputCost, 1e-10)
|
||||
require.InDelta(t, (400*5e-6+600*8e-6)*mult, cost.CacheCreationCost, 1e-10)
|
||||
require.InDelta(t, 1000*0.2e-6*mult, cost.CacheReadCost, 1e-10)
|
||||
require.InDelta(t, 500*20e-6*mult, cost.OutputCost, 1e-10)
|
||||
require.False(t, cost.LongContextBillingApplied)
|
||||
}
|
||||
})
|
||||
}
|
||||
for _, model := range []string{
|
||||
"claude-sonnet-5-5",
|
||||
"anthropic/claude-sonnet-5.5",
|
||||
"us.anthropic.claude-sonnet-5-5",
|
||||
} {
|
||||
t.Run(source+"/"+model, func(t *testing.T) {
|
||||
tokens := UsageTokens{
|
||||
InputTokens: 100_000, OutputTokens: 500,
|
||||
CacheReadTokens: 1000, CacheCreationTokens: 1000,
|
||||
CacheCreation5mTokens: 400, CacheCreation1hTokens: 600,
|
||||
}
|
||||
cost, err := svc.CalculateCost(model, tokens, 1)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 1.2*mult, cost.InputCost, 1e-10)
|
||||
require.InDelta(t, (400*5e-6+600*8e-6)*mult, cost.CacheCreationCost, 1e-10)
|
||||
require.InDelta(t, 1000*0.2e-6*mult, cost.CacheReadCost, 1e-10)
|
||||
require.InDelta(t, 500*20e-6*mult, cost.OutputCost, 1e-10)
|
||||
require.InDelta(t, 100_000*2e-6, cost.InputCost, 1e-10)
|
||||
require.InDelta(t, 400*2.5e-6+600*4e-6, cost.CacheCreationCost, 1e-10)
|
||||
require.InDelta(t, 1000*0.2e-6, cost.CacheReadCost, 1e-10)
|
||||
require.InDelta(t, 500*10e-6, cost.OutputCost, 1e-10)
|
||||
require.False(t, cost.LongContextBillingApplied)
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
|
||||
)
|
||||
|
||||
var claudeResetGrantIDPattern = regexp.MustCompile(`^[a-z0-9_-]{1,40}$`)
|
||||
|
||||
const claudeResetUsageURL = "https://api.anthropic.com/api/oauth/usage?cedar_ember=1&skip_spend=1"
|
||||
|
||||
// ClaudeResetCredit is deliberately free of upstream grant and organization IDs.
|
||||
type ClaudeResetCredit struct {
|
||||
Label string `json:"label"`
|
||||
ResetsLeft int `json:"resets_left"`
|
||||
StartsAt *time.Time `json:"starts_at,omitempty"`
|
||||
ExpiresAt *time.Time `json:"expires_at,omitempty"`
|
||||
Clears []string `json:"clears"`
|
||||
PercentUsed map[string]float64 `json:"percent_used"`
|
||||
Blocking []string `json:"blocking"`
|
||||
UseRequiresLimit bool `json:"use_requires_limit"`
|
||||
Redeemable bool `json:"redeemable"`
|
||||
}
|
||||
|
||||
type ClaudeResetCredits struct {
|
||||
Eligible bool `json:"eligible"`
|
||||
AvailableCount int `json:"available_count"`
|
||||
Credits []ClaudeResetCredit `json:"credits"`
|
||||
CooldownUntil *time.Time `json:"cooldown_until,omitempty"`
|
||||
WeeklyResetsAt *time.Time `json:"weekly_resets_at,omitempty"`
|
||||
FetchedAt time.Time `json:"fetched_at"`
|
||||
}
|
||||
|
||||
type claudeResetGrant struct {
|
||||
ID string `json:"id"`
|
||||
Label string `json:"label"`
|
||||
ResetsLeft int `json:"resets_left"`
|
||||
StartsAt *time.Time `json:"starts_at"`
|
||||
EndsAt *time.Time `json:"ends_at"`
|
||||
Clears []string `json:"clears"`
|
||||
Paused bool `json:"paused"`
|
||||
UsableNow bool `json:"usable_now"`
|
||||
UseRequiresLimit *bool `json:"use_requires_limit"`
|
||||
PercentUsed map[string]float64 `json:"percent_used"`
|
||||
Blocking []string `json:"blocking"`
|
||||
}
|
||||
type claudeResetBlock struct {
|
||||
Eligible bool `json:"eligible"`
|
||||
AtLimit bool `json:"at_limit"`
|
||||
Grants []claudeResetGrant `json:"grants"`
|
||||
NextGrantID string `json:"next_grant_id"`
|
||||
CooldownUntil *time.Time `json:"cooldown_until"`
|
||||
WeeklyResetsAt *time.Time `json:"weekly_resets_at"`
|
||||
}
|
||||
|
||||
type claudeResetAccounts interface {
|
||||
GetByID(context.Context, int64) (*Account, error)
|
||||
}
|
||||
type claudeResetTokens interface {
|
||||
GetAccessToken(context.Context, *Account) (string, error)
|
||||
}
|
||||
|
||||
type ClaudeResetCreditService struct {
|
||||
accounts claudeResetAccounts
|
||||
tokens claudeResetTokens
|
||||
proxies ProxyRepository
|
||||
settings *SettingService
|
||||
do func(*http.Request, string) (*http.Response, error)
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
func NewClaudeResetCreditService(accounts AccountRepository, tokens *ClaudeTokenProvider, proxies ProxyRepository, settings *SettingService) *ClaudeResetCreditService {
|
||||
s := &ClaudeResetCreditService{accounts: accounts, tokens: tokens, proxies: proxies, settings: settings, now: time.Now}
|
||||
s.do = func(req *http.Request, proxy string) (*http.Response, error) {
|
||||
client, err := httpclient.GetClient(httpclient.Options{ProxyURL: proxy, Timeout: 25 * time.Second, ValidateResolvedIP: true})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Never forward OAuth credentials across redirects, even to another public host.
|
||||
isolated := *client
|
||||
isolated.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
|
||||
return isolated.Do(req)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *ClaudeResetCreditService) account(ctx context.Context, id int64) (*Account, string, string, error) {
|
||||
a, err := s.accounts.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, "", "", err
|
||||
}
|
||||
if a == nil || a.Platform != PlatformAnthropic || a.Type != AccountTypeOAuth {
|
||||
return nil, "", "", infraerrors.BadRequest("CLAUDE_RESET_OAUTH_REQUIRED", "Claude OAuth account required")
|
||||
}
|
||||
profile := false
|
||||
for _, scope := range strings.Fields(a.GetCredential("scope")) {
|
||||
if scope == "user:profile" {
|
||||
profile = true
|
||||
}
|
||||
}
|
||||
if !profile {
|
||||
return nil, "", "", infraerrors.BadRequest("CLAUDE_RESET_PROFILE_SCOPE_REQUIRED", "user:profile scope required")
|
||||
}
|
||||
proxy := ""
|
||||
if a.ProxyID != nil {
|
||||
p, e := s.proxies.GetByID(ctx, *a.ProxyID)
|
||||
if e != nil || p == nil {
|
||||
return nil, "", "", infraerrors.ServiceUnavailable("CLAUDE_RESET_PROXY_UNAVAILABLE", "account proxy unavailable")
|
||||
}
|
||||
proxy = p.URL()
|
||||
}
|
||||
token, err := s.tokens.GetAccessToken(ctx, a)
|
||||
if err != nil || strings.TrimSpace(token) == "" {
|
||||
return nil, "", "", infraerrors.ServiceUnavailable("CLAUDE_RESET_TOKEN_UNAVAILABLE", "OAuth token unavailable")
|
||||
}
|
||||
return a, token, proxy, nil
|
||||
}
|
||||
|
||||
func (s *ClaudeResetCreditService) headers(ctx context.Context, req *http.Request, token string) {
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("anthropic-beta", "oauth-2025-04-20")
|
||||
req.Header.Set("x-app", "cli")
|
||||
req.Header.Set("User-Agent", "claude-cli/"+s.settings.GetClaudeCodeClientVersion(ctx)+" (external, cli)")
|
||||
}
|
||||
|
||||
func (s *ClaudeResetCreditService) query(ctx context.Context, id int64) (*ClaudeResetCredits, error) {
|
||||
_, token, proxy, err := s.account(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, claudeResetUsageURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.headers(ctx, req, token)
|
||||
resp, err := s.do(req, proxy)
|
||||
if err != nil {
|
||||
return nil, infraerrors.ServiceUnavailable("CLAUDE_RESET_QUERY_FAILED", "reset status request failed")
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "CLAUDE_RESET_QUERY_FAILED", fmt.Sprintf("reset status upstream HTTP %d", resp.StatusCode))
|
||||
}
|
||||
var envelope map[string]json.RawMessage
|
||||
if err = json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&envelope); err != nil || envelope == nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "CLAUDE_RESET_STATUS_INVALID", "invalid reset status")
|
||||
}
|
||||
if _, ok := envelope["error"]; ok {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "CLAUDE_RESET_STATUS_INVALID", "invalid reset status")
|
||||
}
|
||||
var block *claudeResetBlock
|
||||
raw, present := envelope["cedar_ember"]
|
||||
if present && string(raw) != "null" {
|
||||
if err = json.Unmarshal(raw, &block); err != nil || block == nil || block.Grants == nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "CLAUDE_RESET_STATUS_INVALID", "invalid reset grants")
|
||||
}
|
||||
}
|
||||
return projectClaudeResetCredits(block, s.now()), nil
|
||||
}
|
||||
|
||||
func (s *ClaudeResetCreditService) Query(ctx context.Context, id int64) (*ClaudeResetCredits, error) {
|
||||
return s.query(ctx, id)
|
||||
}
|
||||
|
||||
func projectClaudeResetCredits(b *claudeResetBlock, now time.Time) *ClaudeResetCredits {
|
||||
r := &ClaudeResetCredits{Credits: []ClaudeResetCredit{}, FetchedAt: now.UTC()}
|
||||
if b == nil {
|
||||
return r
|
||||
}
|
||||
r.Eligible = b.Eligible
|
||||
if b.CooldownUntil != nil && now.Before(*b.CooldownUntil) {
|
||||
r.CooldownUntil = b.CooldownUntil
|
||||
}
|
||||
r.WeeklyResetsAt = b.WeeklyResetsAt
|
||||
for _, g := range b.Grants {
|
||||
if !claudeResetGrantIDPattern.MatchString(g.ID) || len(g.Clears) == 0 || g.ResetsLeft <= 0 || g.Paused || (g.StartsAt != nil && now.Before(*g.StartsAt)) || (g.EndsAt != nil && !now.Before(*g.EndsAt)) {
|
||||
continue
|
||||
}
|
||||
requires := g.UseRequiresLimit == nil || *g.UseRequiresLimit
|
||||
usable := b.Eligible && g.UsableNow && g.ID == b.NextGrantID && (!requires || b.AtLimit) && len(g.Blocking) == 0 && (b.CooldownUntil == nil || !now.Before(*b.CooldownUntil))
|
||||
used := map[string]float64{}
|
||||
for k, v := range g.PercentUsed {
|
||||
if v >= 0 && v <= 100 {
|
||||
used[k] = v
|
||||
}
|
||||
}
|
||||
r.Credits = append(r.Credits, ClaudeResetCredit{Label: g.Label, ResetsLeft: g.ResetsLeft, StartsAt: g.StartsAt, ExpiresAt: g.EndsAt, Clears: g.Clears, PercentUsed: used, Blocking: g.Blocking, UseRequiresLimit: requires, Redeemable: usable})
|
||||
if usable {
|
||||
r.AvailableCount += g.ResetsLeft
|
||||
}
|
||||
}
|
||||
return r
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"github.com/stretchr/testify/require"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type resetAccountStub struct{ account *Account }
|
||||
|
||||
func (s resetAccountStub) GetByID(context.Context, int64) (*Account, error) { return s.account, nil }
|
||||
|
||||
type resetTokenStub struct{}
|
||||
|
||||
func (resetTokenStub) GetAccessToken(context.Context, *Account) (string, error) {
|
||||
return "synthetic-token", nil
|
||||
}
|
||||
func TestClaudeResetStatusNativeContract(t *testing.T) {
|
||||
now := time.Date(2026, 9, 25, 0, 0, 0, 0, time.UTC)
|
||||
s := &ClaudeResetCreditService{accounts: resetAccountStub{&Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Credentials: map[string]any{"scope": "user:profile user:inference"}}}, tokens: resetTokenStub{}, now: func() time.Time { return now }}
|
||||
s.do = func(r *http.Request, p string) (*http.Response, error) {
|
||||
require.Equal(t, claudeResetUsageURL, r.URL.String())
|
||||
require.Equal(t, "GET", r.Method)
|
||||
require.Equal(t, "Bearer synthetic-token", r.Header.Get("Authorization"))
|
||||
require.Contains(t, r.Header.Get("User-Agent"), "claude-cli/")
|
||||
return &http.Response{StatusCode: 200, Body: io.NopCloser(strings.NewReader(`{"cedar_ember":{"eligible":true,"at_limit":false,"next_grant_id":"launch","grants":[{"id":"launch","resets_left":2,"usable_now":true,"use_requires_limit":false,"ends_at":"2026-10-22T00:00:00Z","clears":["five_hour"],"percent_used":{"five_hour":65},"blocking":[]},{"id":"later","clears":["five_hour"],"resets_left":1,"usable_now":true,"use_requires_limit":false}]}}`))}, nil
|
||||
}
|
||||
r, e := s.Query(context.Background(), 1)
|
||||
require.NoError(t, e)
|
||||
require.Equal(t, 2, r.AvailableCount)
|
||||
require.True(t, r.Credits[0].Redeemable)
|
||||
require.False(t, r.Credits[1].Redeemable)
|
||||
b, e := json.Marshal(r)
|
||||
require.NoError(t, e)
|
||||
require.NotContains(t, string(b), `"id"`)
|
||||
require.NotContains(t, string(b), "launch")
|
||||
require.NotContains(t, string(b), "synthetic-token")
|
||||
require.NotContains(t, string(b), "selection_token")
|
||||
}
|
||||
func TestClaudeResetPastCooldownIsCleared(t *testing.T) {
|
||||
now := time.Now()
|
||||
past := now.Add(-time.Minute)
|
||||
no := false
|
||||
g := claudeResetGrant{ID: "grant", ResetsLeft: 1, UsableNow: true, UseRequiresLimit: &no, Clears: []string{"five_hour"}}
|
||||
r := projectClaudeResetCredits(&claudeResetBlock{Eligible: true, NextGrantID: g.ID, CooldownUntil: &past, Grants: []claudeResetGrant{g}}, now)
|
||||
require.Nil(t, r.CooldownUntil)
|
||||
require.Equal(t, 1, r.AvailableCount)
|
||||
future := now.Add(time.Hour)
|
||||
r = projectClaudeResetCredits(&claudeResetBlock{Eligible: true, NextGrantID: g.ID, CooldownUntil: &future, Grants: []claudeResetGrant{g}}, now)
|
||||
require.Equal(t, &future, r.CooldownUntil)
|
||||
}
|
||||
func TestClaudeResetEligibilityFailClosed(t *testing.T) {
|
||||
now := time.Now()
|
||||
past := now.Add(-time.Minute)
|
||||
future := now.Add(time.Hour)
|
||||
no := false
|
||||
base := claudeResetGrant{ID: "grant", ResetsLeft: 1, UsableNow: true, UseRequiresLimit: &no, Clears: []string{"five_hour"}}
|
||||
for _, name := range []string{"paused", "expired", "future", "cooldown", "requires-limit", "blocking", "ineligible", "spent", "not-next"} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
g := base
|
||||
b := &claudeResetBlock{Eligible: true, NextGrantID: g.ID}
|
||||
switch name {
|
||||
case "paused":
|
||||
g.Paused = true
|
||||
case "expired":
|
||||
g.EndsAt = &past
|
||||
case "future":
|
||||
g.StartsAt = &future
|
||||
case "cooldown":
|
||||
b.CooldownUntil = &future
|
||||
case "requires-limit":
|
||||
g.UseRequiresLimit = nil
|
||||
case "blocking":
|
||||
g.Blocking = []string{"seven_day"}
|
||||
case "ineligible":
|
||||
b.Eligible = false
|
||||
case "spent":
|
||||
g.ResetsLeft = 0
|
||||
case "not-next":
|
||||
b.NextGrantID = "other"
|
||||
}
|
||||
b.Grants = []claudeResetGrant{g}
|
||||
r := projectClaudeResetCredits(b, now)
|
||||
require.Zero(t, r.AvailableCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
func TestClaudeResetStatusRejectsMissingScopeBeforeNetwork(t *testing.T) {
|
||||
s := &ClaudeResetCreditService{accounts: resetAccountStub{&Account{Platform: PlatformAnthropic, Type: AccountTypeOAuth}}, do: func(*http.Request, string) (*http.Response, error) { t.Fatal("network called"); return nil, nil }}
|
||||
_, e := s.Query(context.Background(), 1)
|
||||
require.Error(t, e)
|
||||
}
|
||||
func TestClaudeResetMalformedAndAbsent(t *testing.T) {
|
||||
for _, body := range []string{`{}`, `{"five_hour":{},"cedar_ember":null}`, `{"cedar_ember":{"eligible":true}}`, `{"error":{"message":"private upstream data"}}`, `not json`} {
|
||||
t.Run(body, func(t *testing.T) {
|
||||
s := &ClaudeResetCreditService{accounts: resetAccountStub{&Account{Platform: PlatformAnthropic, Type: AccountTypeOAuth, Credentials: map[string]any{"scope": "user:profile"}}}, tokens: resetTokenStub{}, now: time.Now, do: func(*http.Request, string) (*http.Response, error) {
|
||||
return &http.Response{StatusCode: 200, Body: io.NopCloser(strings.NewReader(body))}, nil
|
||||
}}
|
||||
r, e := s.Query(context.Background(), 1)
|
||||
if body == `{}` || strings.Contains(body, `"cedar_ember":null`) {
|
||||
require.NoError(t, e)
|
||||
require.Empty(t, r.Credits)
|
||||
} else {
|
||||
require.Error(t, e)
|
||||
require.NotContains(t, e.Error(), "private upstream data")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user