feat: add Vertex AI storage integration

.gitignore

@@ -34,6 +34,7 @@ electron/dist
 token_estimator_test.go
 skills-lock.json
 .playwright-mcp
+.worktrees/
 
 # Local-only live probes and scratch test workspaces.
 .local-tests/

constant/context_key.go

@@ -27,6 +27,8 @@ const (
 	ContextKeyChannelCreateTime        ContextKey = "channel_create_time"
 	ContextKeyChannelBaseUrl           ContextKey = "base_url"
 	ContextKeyChannelType              ContextKey = "channel_type"
+	ContextKeyRequiredChannelType      ContextKey = "required_channel_type"
+	ContextKeyChannelModels            ContextKey = "channel_models"
 	ContextKeyChannelSetting           ContextKey = "channel_setting"
 	ContextKeyChannelOtherSetting      ContextKey = "channel_other_setting"
 	ContextKeyChannelParamOverride     ContextKey = "param_override"

controller/channel-test.go

@@ -70,6 +70,26 @@ func resolveChannelTestUserID(c *gin.Context) (int, error) {
 }
 
 func testChannel(ctx context.Context, channel *model.Channel, testUserID int, testModel string, endpointType string, isStream bool) testResult {
+	return testChannelWithVertexStorageDependencies(
+		ctx,
+		channel,
+		testUserID,
+		testModel,
+		endpointType,
+		isStream,
+		defaultVertexStorageChannelProbeDependencies(),
+	)
+}
+
+func testChannelWithVertexStorageDependencies(
+	ctx context.Context,
+	channel *model.Channel,
+	testUserID int,
+	testModel string,
+	endpointType string,
+	isStream bool,
+	storageDeps vertexStorageChannelProbeDependencies,
+) testResult {
 	if ctx == nil {
 		ctx = context.Background()
 	}
@@ -106,6 +126,23 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te
 			}
 		}
 	}
+	if isVertexStorageChannelTest(channel, testModel) {
+		c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, relayconstant.VertexStorageRoutePrefix, nil)
+		c.Set("channel", channel.Type)
+		c.Set("base_url", channel.GetBaseURL())
+		newAPIError := middleware.SetupContextForSelectedChannel(c, channel, testModel)
+		if newAPIError != nil {
+			return testResult{
+				context:     c,
+				localErr:    newAPIError,
+				newAPIError: newAPIError,
+			}
+		}
+		return testResult{
+			context:  c,
+			localErr: testVertexStorageChannel(ctx, c, testModel, storageDeps),
+		}
+	}
 
 	endpointType = normalizeChannelTestEndpoint(channel, endpointType)
 
@@ -660,6 +697,12 @@ func shouldUseStreamForAutomaticChannelTest(channel *model.Channel) bool {
 	return channel != nil && channel.Type == constant.ChannelTypeCodex
 }
 
+func isVertexStorageChannelTest(channel *model.Channel, testModel string) bool {
+	return channel != nil &&
+		channel.Type == constant.ChannelTypeVertexAi &&
+		strings.HasPrefix(strings.TrimSpace(testModel), relayconstant.VertexStorageModelPrefix)
+}
+
 func detectErrorMessageFromJSONBytes(jsonBytes []byte) string {
 	if len(jsonBytes) == 0 {
 		return ""

controller/relay.go

@@ -181,13 +181,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		}
 	}()
 
-	retryParam := &service.RetryParam{
-		Ctx:         c,
-		TokenGroup:  relayInfo.TokenGroup,
-		ModelName:   relayInfo.OriginModelName,
-		RequestPath: c.Request.URL.Path,
-		Retry:       common.GetPointer(0),
-	}
+	retryParam := newRelayRetryParam(c, relayInfo)
 	relayInfo.RetryIndex = 0
 	relayInfo.LastError = nil
 
@@ -311,6 +305,9 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
 			AutoBan: &autoBanInt,
 		}, nil
 	}
+	if retryParam.RequiredChannelType == 0 {
+		retryParam.RequiredChannelType = requiredChannelTypeForRelay(c)
+	}
 	channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam)
 	if err != nil {
 		return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
@@ -328,6 +325,27 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
 	return channel, nil
 }
 
+func newRelayRetryParam(c *gin.Context, relayInfo *relaycommon.RelayInfo) *service.RetryParam {
+	return &service.RetryParam{
+		Ctx:                 c,
+		TokenGroup:          relayInfo.TokenGroup,
+		ModelName:           relayInfo.OriginModelName,
+		RequestPath:         c.Request.URL.Path,
+		Retry:               common.GetPointer(0),
+		RequiredChannelType: requiredChannelTypeForRelay(c),
+	}
+}
+
+func requiredChannelTypeForRelay(c *gin.Context) int {
+	if requiredChannelType := common.GetContextKeyInt(c, constant.ContextKeyRequiredChannelType); requiredChannelType != 0 {
+		return requiredChannelType
+	}
+	if relayconstant.IsVertexStoragePath(c.Request.URL.Path) {
+		return constant.ChannelTypeVertexAi
+	}
+	return 0
+}
+
 func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
 	if openaiErr == nil {
 		return false
@@ -513,13 +531,7 @@ func RelayTask(c *gin.Context) {
 		}
 	}()
 
-	retryParam := &service.RetryParam{
-		Ctx:         c,
-		TokenGroup:  relayInfo.TokenGroup,
-		ModelName:   relayInfo.OriginModelName,
-		RequestPath: c.Request.URL.Path,
-		Retry:       common.GetPointer(0),
-	}
+	retryParam := newRelayRetryParam(c, relayInfo)
 
 	for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
 		var channel *model.Channel

controller/relay_channel_type_test.go

@@ -0,0 +1,86 @@
+package controller
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/model"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/gin-gonic/gin"
+	"github.com/glebarez/sqlite"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+func setupVertexStorageRetryTest(t *testing.T) *gorm.DB {
+	t.Helper()
+	originalDB := model.DB
+	originalMemoryCacheEnabled := common.MemoryCacheEnabled
+	dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
+	db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
+	require.NoError(t, err)
+	require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}))
+	model.DB = db
+	common.MemoryCacheEnabled = true
+	t.Cleanup(func() {
+		model.DB = originalDB
+		common.MemoryCacheEnabled = originalMemoryCacheEnabled
+		if originalMemoryCacheEnabled && originalDB != nil {
+			model.InitChannelCache()
+		}
+		sqlDB, sqlErr := db.DB()
+		if sqlErr == nil {
+			require.NoError(t, sqlDB.Close())
+		}
+	})
+	return db
+}
+
+func TestVertexStorageRetryKeepsVertexChannelType(t *testing.T) {
+	db := setupVertexStorageRetryTest(t)
+	modelName := "storage:gs:bucket-a"
+	highPriority := int64(200)
+	middlePriority := int64(100)
+	lowPriority := int64(0)
+	weight := uint(100)
+	channels := []model.Channel{
+		{Id: 6200, Name: "vertex-first", Type: constant.ChannelTypeVertexAi, Key: "vertex-first-key", Status: common.ChannelStatusEnabled, Group: "default", Models: modelName, Priority: &highPriority, Weight: &weight},
+		{Id: 6201, Name: "gemini", Type: constant.ChannelTypeGemini, Key: "gemini-key", Status: common.ChannelStatusEnabled, Group: "default", Models: modelName, Priority: &middlePriority, Weight: &weight},
+		{Id: 6202, Name: "vertex", Type: constant.ChannelTypeVertexAi, Key: "vertex-key", Status: common.ChannelStatusEnabled, Group: "default", Models: modelName, Priority: &lowPriority, Weight: &weight},
+	}
+	require.NoError(t, db.Create(&channels).Error)
+	for _, channel := range channels {
+		require.NoError(t, db.Create(&model.Ability{
+			Group: channel.Group, Model: modelName, ChannelId: channel.Id,
+			Enabled: true, Priority: channel.Priority, Weight: weight,
+		}).Error)
+	}
+	model.InitChannelCache()
+
+	gin.SetMode(gin.TestMode)
+	c, _ := gin.CreateTestContext(httptest.NewRecorder())
+	c.Request = httptest.NewRequest(http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", nil)
+	common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
+
+	relayInfo := &relaycommon.RelayInfo{
+		OriginModelName: modelName,
+		TokenGroup:      "default",
+		ChannelMeta:     &relaycommon.ChannelMeta{},
+	}
+	retryParam := newRelayRetryParam(c, relayInfo)
+	retryParam.SetRetry(1)
+
+	assert.Equal(t, constant.ChannelTypeVertexAi, retryParam.RequiredChannelType)
+	channel, relayErr := getChannel(c, relayInfo, retryParam)
+
+	require.Nil(t, relayErr)
+	require.NotNil(t, channel)
+	assert.Equal(t, 6202, channel.Id)
+	assert.Equal(t, constant.ChannelTypeVertexAi, channel.Type)
+}