Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

2026-08-21 18:30:13 +08:00
commit 6d655c9903
3584 changed files with 1270640 additions and 0 deletions
@@ -0,0 +1,360 @@
+//go:build unit
+
+package service
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestOpenAISelectAccountForModelWithExclusions_ChannelMappedRestrictionRejectsEarly(t *testing.T) {
+	t.Parallel()
+
+	channelSvc := newTestChannelService(makeStandardRepo(Channel{
+		ID:                 1,
+		Status:             StatusActive,
+		GroupIDs:           []int64{10},
+		RestrictModels:     true,
+		BillingModelSource: BillingModelSourceChannelMapped,
+		ModelPricing: []ChannelModelPricing{
+			{Platform: PlatformOpenAI, Models: []string{"gpt-4o"}},
+		},
+		ModelMapping: map[string]map[string]string{
+			PlatformOpenAI: {"gpt-4.1": "o3-mini"},
+		},
+	}, map[int64]string{10: PlatformOpenAI}))
+
+	svc := &OpenAIGatewayService{
+		accountRepo: stubOpenAIAccountRepo{accounts: []Account{
+			{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true},
+		}},
+		channelService: channelSvc,
+	}
+
+	groupID := int64(10)
+	_, err := svc.SelectAccountForModelWithExclusions(context.Background(), &groupID, "", "gpt-4.1", nil)
+	require.ErrorIs(t, err, ErrNoAvailableAccounts)
+	require.Contains(t, err.Error(), "channel pricing restriction")
+}
+
+func TestOpenAISelectAccountForModelWithExclusions_UpstreamRestrictionSkipsDisallowedAccount(t *testing.T) {
+	t.Parallel()
+
+	channelSvc := newTestChannelService(makeStandardRepo(Channel{
+		ID:                 1,
+		Status:             StatusActive,
+		GroupIDs:           []int64{10},
+		RestrictModels:     true,
+		BillingModelSource: BillingModelSourceUpstream,
+		ModelPricing: []ChannelModelPricing{
+			{Platform: PlatformOpenAI, Models: []string{"o3-mini"}},
+		},
+	}, map[int64]string{10: PlatformOpenAI}))
+
+	svc := &OpenAIGatewayService{
+		accountRepo: stubOpenAIAccountRepo{accounts: []Account{
+			{
+				ID:          1,
+				Platform:    PlatformOpenAI,
+				Status:      StatusActive,
+				Schedulable: true,
+				Priority:    10,
+				Credentials: map[string]any{
+					"model_mapping": map[string]any{"gpt-4.1": "gpt-4o"},
+				},
+			},
+			{
+				ID:          2,
+				Platform:    PlatformOpenAI,
+				Status:      StatusActive,
+				Schedulable: true,
+				Priority:    20,
+				Credentials: map[string]any{
+					"model_mapping": map[string]any{"gpt-4.1": "o3-mini"},
+				},
+			},
+		}},
+		channelService: channelSvc,
+	}
+
+	groupID := int64(10)
+	account, err := svc.SelectAccountForModelWithExclusions(context.Background(), &groupID, "", "gpt-4.1", nil)
+	require.NoError(t, err)
+	require.NotNil(t, account)
+	require.Equal(t, int64(2), account.ID)
+}
+
+func TestIsUpstreamModelRestrictedByChannel_CompactMappingMatchesForwardPath(t *testing.T) {
+	t.Parallel()
+
+	account := &Account{
+		Platform: PlatformOpenAI,
+		Credentials: map[string]any{
+			"model_mapping":         map[string]any{"gpt-5.4-channel": "gpt-5.4-account"},
+			"compact_model_mapping": map[string]any{"gpt-5.4-account": "gpt-5.4-compact"},
+		},
+	}
+	tests := []struct {
+		name                   string
+		allowedUpstreamModel   string
+		useCompactModelMapping bool
+	}{
+		{
+			name:                   "legacy compact applies compact mapping after channel and account mapping",
+			allowedUpstreamModel:   "gpt-5.4-compact",
+			useCompactModelMapping: true,
+		},
+		{
+			name:                   "native v2 stops after channel and account mapping",
+			allowedUpstreamModel:   "gpt-5.4-account",
+			useCompactModelMapping: false,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			channelSvc := newTestChannelService(makeStandardRepo(Channel{
+				ID:                 1,
+				Status:             StatusActive,
+				GroupIDs:           []int64{10},
+				RestrictModels:     true,
+				BillingModelSource: BillingModelSourceUpstream,
+				ModelPricing: []ChannelModelPricing{
+					{Platform: PlatformOpenAI, Models: []string{tt.allowedUpstreamModel}},
+				},
+				ModelMapping: map[string]map[string]string{
+					PlatformOpenAI: {"gpt-5.4": "gpt-5.4-channel"},
+				},
+			}, map[int64]string{10: PlatformOpenAI}))
+			svc := &OpenAIGatewayService{channelService: channelSvc}
+			mapping := channelSvc.ResolveChannelMapping(context.Background(), 10, "gpt-5.4")
+			require.True(t, mapping.Mapped)
+			require.Equal(t, "gpt-5.4-channel", mapping.MappedModel)
+
+			ctx := WithOpenAIForwardModel(
+				context.Background(),
+				mapping.MappedModel,
+				tt.useCompactModelMapping,
+			)
+			require.False(t, svc.isUpstreamModelRestrictedByChannel(
+				ctx, 10, account, "gpt-5.4", true,
+			))
+			require.True(t, svc.isUpstreamModelRestrictedByChannel(
+				context.Background(), 10, account, "gpt-5.4", true,
+			), "without the forward-model context the restriction check follows a different chain")
+		})
+	}
+}
+
+func TestIsUpstreamModelRestrictedByChannel_PassthroughMatchesForwardPath(t *testing.T) {
+	t.Parallel()
+
+	account := &Account{
+		Platform: PlatformOpenAI,
+		Type:     AccountTypeAPIKey,
+		Credentials: map[string]any{
+			"model_mapping": map[string]any{"gpt-5.4-channel": "gpt-5.4-account"},
+			"compact_model_mapping": map[string]any{
+				"gpt-5.4-channel": "gpt-5.4-compact",
+			},
+		},
+		Extra: map[string]any{"openai_passthrough": true},
+	}
+	tests := []struct {
+		name                   string
+		allowedUpstreamModel   string
+		useCompactModelMapping bool
+	}{
+		{
+			name:                   "native v2 keeps channel-mapped model and ignores normal account mapping",
+			allowedUpstreamModel:   "gpt-5.4-channel",
+			useCompactModelMapping: false,
+		},
+		{
+			name:                   "legacy compact applies compact mapping to channel-mapped model",
+			allowedUpstreamModel:   "gpt-5.4-compact",
+			useCompactModelMapping: true,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			channelSvc := newTestChannelService(makeStandardRepo(Channel{
+				ID:                 1,
+				Status:             StatusActive,
+				GroupIDs:           []int64{10},
+				RestrictModels:     true,
+				BillingModelSource: BillingModelSourceUpstream,
+				ModelPricing: []ChannelModelPricing{
+					{Platform: PlatformOpenAI, Models: []string{tt.allowedUpstreamModel}},
+				},
+				ModelMapping: map[string]map[string]string{
+					PlatformOpenAI: {"gpt-5.4": "gpt-5.4-channel"},
+				},
+			}, map[int64]string{10: PlatformOpenAI}))
+			svc := &OpenAIGatewayService{channelService: channelSvc}
+			mapping := channelSvc.ResolveChannelMapping(context.Background(), 10, "gpt-5.4")
+			require.True(t, mapping.Mapped)
+			require.Equal(t, "gpt-5.4-channel", mapping.MappedModel)
+
+			ctx := WithOpenAIForwardModel(
+				context.Background(),
+				mapping.MappedModel,
+				tt.useCompactModelMapping,
+			)
+			require.False(t, svc.isUpstreamModelRestrictedByChannel(
+				ctx, 10, account, "gpt-5.4", true,
+			))
+		})
+	}
+}
+
+func TestIsUpstreamModelRestrictedByChannel_PassthroughFlagWithRawChatFallbackMatchesForwardPath(t *testing.T) {
+	t.Parallel()
+
+	account := &Account{
+		Platform: PlatformOpenAI,
+		Type:     AccountTypeAPIKey,
+		Credentials: map[string]any{
+			"model_mapping": map[string]any{"gpt-5.4-channel": "gpt-5.4-account"},
+			"compact_model_mapping": map[string]any{
+				"gpt-5.4-account": "gpt-5.4-compact",
+			},
+		},
+		Extra: map[string]any{
+			"openai_passthrough":         true,
+			"openai_responses_supported": false,
+		},
+	}
+
+	for _, useCompactModelMapping := range []bool{false, true} {
+		useCompactModelMapping := useCompactModelMapping
+		name := "native v2"
+		if useCompactModelMapping {
+			name = "legacy compact"
+		}
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			channelSvc := newTestChannelService(makeStandardRepo(Channel{
+				ID:                 1,
+				Status:             StatusActive,
+				GroupIDs:           []int64{10},
+				RestrictModels:     true,
+				BillingModelSource: BillingModelSourceUpstream,
+				ModelPricing: []ChannelModelPricing{
+					{Platform: PlatformOpenAI, Models: []string{"gpt-5.4-account"}},
+				},
+				ModelMapping: map[string]map[string]string{
+					PlatformOpenAI: {"gpt-5.4": "gpt-5.4-channel"},
+				},
+			}, map[int64]string{10: PlatformOpenAI}))
+			svc := &OpenAIGatewayService{channelService: channelSvc}
+			ctx := WithOpenAIForwardModel(
+				context.Background(),
+				"gpt-5.4-channel",
+				useCompactModelMapping,
+			)
+
+			require.False(t, svc.isUpstreamModelRestrictedByChannel(
+				ctx, 10, account, "gpt-5.4", true,
+			))
+		})
+	}
+}
+
+func TestIsUpstreamModelRestrictedByChannel_ForwardModelContextMatchesNormalForwardPath(t *testing.T) {
+	t.Parallel()
+
+	account := &Account{
+		Platform: PlatformOpenAI,
+		Type:     AccountTypeAPIKey,
+		Credentials: map[string]any{
+			"model_mapping": map[string]any{"gpt-5.4-channel": "gpt-5.4-account"},
+		},
+		Extra: map[string]any{
+			"openai_passthrough":         true,
+			"openai_responses_supported": false,
+		},
+	}
+	channelSvc := newTestChannelService(makeStandardRepo(Channel{
+		ID:                 1,
+		Status:             StatusActive,
+		GroupIDs:           []int64{10},
+		RestrictModels:     true,
+		BillingModelSource: BillingModelSourceUpstream,
+		ModelPricing: []ChannelModelPricing{
+			{Platform: PlatformOpenAI, Models: []string{"gpt-5.4-account"}},
+		},
+		ModelMapping: map[string]map[string]string{
+			PlatformOpenAI: {"gpt-5.4": "gpt-5.4-channel"},
+		},
+	}, map[int64]string{10: PlatformOpenAI}))
+	svc := &OpenAIGatewayService{channelService: channelSvc}
+	ctx := WithOpenAIForwardModel(context.Background(), "gpt-5.4-channel", false)
+
+	require.False(t, svc.isUpstreamModelRestrictedByChannel(
+		ctx, 10, account, "gpt-5.4", false,
+	))
+}
+
+func TestOpenAISelectAccountForModelWithExclusions_StickyRestrictedUpstreamFallsBack(t *testing.T) {
+	t.Parallel()
+
+	channelSvc := newTestChannelService(makeStandardRepo(Channel{
+		ID:                 1,
+		Status:             StatusActive,
+		GroupIDs:           []int64{10},
+		RestrictModels:     true,
+		BillingModelSource: BillingModelSourceUpstream,
+		ModelPricing: []ChannelModelPricing{
+			{Platform: PlatformOpenAI, Models: []string{"o3-mini"}},
+		},
+	}, map[int64]string{10: PlatformOpenAI}))
+
+	cache := &stubGatewayCache{
+		sessionBindings: map[string]int64{"openai:sticky-session": 1},
+	}
+	svc := &OpenAIGatewayService{
+		accountRepo: stubOpenAIAccountRepo{accounts: []Account{
+			{
+				ID:          1,
+				Platform:    PlatformOpenAI,
+				Status:      StatusActive,
+				Schedulable: true,
+				Priority:    10,
+				Credentials: map[string]any{
+					"model_mapping": map[string]any{"gpt-4.1": "gpt-4o"},
+				},
+			},
+			{
+				ID:          2,
+				Platform:    PlatformOpenAI,
+				Status:      StatusActive,
+				Schedulable: true,
+				Priority:    20,
+				Credentials: map[string]any{
+					"model_mapping": map[string]any{"gpt-4.1": "o3-mini"},
+				},
+			},
+		}},
+		channelService: channelSvc,
+		cache:          cache,
+	}
+
+	groupID := int64(10)
+	account, err := svc.SelectAccountForModelWithExclusions(context.Background(), &groupID, "sticky-session", "gpt-4.1", nil)
+	require.NoError(t, err)
+	require.NotNil(t, account)
+	require.Equal(t, int64(2), account.ID)
+	require.Equal(t, 1, cache.deletedSessions["openai:sticky-session"])
+	require.Equal(t, int64(2), cache.sessionBindings["openai:sticky-session"])
+}