diff --git a/.gitignore b/.gitignore index dc328dd6c80c..7a169e2b517f 100644 --- a/.gitignore +++ b/.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/ diff --git a/constant/context_key.go b/constant/context_key.go index ccb8010f9476..82ab7bf7b45c 100644 --- a/constant/context_key.go +++ b/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" diff --git a/controller/channel-test.go b/controller/channel-test.go index fffc59d24d5a..8b25a76c102c 100644 --- a/controller/channel-test.go +++ b/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 "" diff --git a/controller/relay.go b/controller/relay.go index 8dccfe76dddd..8924dc8cfc6c 100644 --- a/controller/relay.go +++ b/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 diff --git a/controller/relay_channel_type_test.go b/controller/relay_channel_type_test.go new file mode 100644 index 000000000000..ad65c7199daf --- /dev/null +++ b/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) +} diff --git a/controller/vertex_storage_channel_probe.go b/controller/vertex_storage_channel_probe.go new file mode 100644 index 000000000000..aaff9b92d99c --- /dev/null +++ b/controller/vertex_storage_channel_probe.go @@ -0,0 +1,181 @@ +package controller + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/relay/channel/vertex" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" + "github.com/google/uuid" +) + +const vertexStorageChannelTestContent = "new-api Vertex AI Storage channel test\n" + +type vertexStorageChannelProbeDependencies struct { + newObjectName func() string + acquireAccessToken func(vertex.CachedAccessTokenRequest) (string, error) + doProxy func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) +} + +func defaultVertexStorageChannelProbeDependencies() vertexStorageChannelProbeDependencies { + return vertexStorageChannelProbeDependencies{ + newObjectName: func() string { + return ".new-api-channel-test/" + uuid.NewString() + "/test.txt" + }, + acquireAccessToken: vertex.AcquireCachedAccessToken, + doProxy: vertex.DoStorageProxy, + } +} + +func testVertexStorageChannel(ctx context.Context, c *gin.Context, testModel string, deps vertexStorageChannelProbeDependencies) error { + if c == nil { + return errors.New("Vertex storage channel test context is required") + } + if deps.newObjectName == nil || deps.acquireAccessToken == nil || deps.doProxy == nil { + return errors.New("Vertex storage channel test dependencies are incomplete") + } + + modelName := strings.TrimSpace(testModel) + if !strings.HasPrefix(modelName, relayconstant.VertexStorageModelPrefix) { + return errors.New("invalid Vertex storage test model") + } + bucket, err := relayconstant.NormalizeVertexStorageBucket(strings.TrimPrefix(modelName, relayconstant.VertexStorageModelPrefix)) + if err != nil { + return fmt.Errorf("invalid Vertex storage test model: %w", err) + } + if common.GetContextKeyInt(c, constant.ContextKeyChannelType) != constant.ChannelTypeVertexAi { + return errors.New("selected channel is not Vertex AI") + } + if !relayconstant.VertexStorageChannelSupports(common.GetContextKeyStringSlice(c, constant.ContextKeyChannelModels), bucket) { + return fmt.Errorf("selected channel does not allow Cloud Storage bucket %q", bucket) + } + + channelOtherSetting, _ := common.GetContextKeyType[dto.ChannelOtherSettings](c, constant.ContextKeyChannelOtherSetting) + if channelOtherSetting.VertexKeyType == dto.VertexKeyTypeAPIKey { + return errors.New("Vertex storage channel test requires service account JSON") + } + credentials := vertex.Credentials{} + if err := common.Unmarshal([]byte(common.GetContextKeyString(c, constant.ContextKeyChannelKey)), &credentials); err != nil { + return errors.New("selected Vertex AI channel credentials are invalid") + } + + channelSetting, _ := common.GetContextKeyType[dto.ChannelSettings](c, constant.ContextKeyChannelSetting) + accessToken, err := deps.acquireAccessToken(vertex.CachedAccessTokenRequest{ + ChannelID: common.GetContextKeyInt(c, constant.ContextKeyChannelId), + ChannelIsMultiKey: common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey), + ChannelMultiKeyIndex: common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex), + Credentials: credentials, + Proxy: channelSetting.Proxy, + }) + if err != nil { + return errors.New("failed to authorize Vertex storage channel test") + } + + objectName := strings.TrimSpace(deps.newObjectName()) + if objectName == "" { + return errors.New("Vertex storage channel test object name is empty") + } + if ctx == nil { + ctx = context.Background() + } + + var probeErrors []error + uploadQuery := url.Values{} + uploadQuery.Set("uploadType", "media") + uploadQuery.Set("name", objectName) + uploadHeader := make(http.Header) + uploadHeader.Set("Content-Type", "text/plain; charset=utf-8") + _, err = runVertexStorageProbeRequest(ctx, deps.doProxy, vertex.StorageProxyRequest{ + Operation: vertex.StorageOperationUpload, + Method: http.MethodPost, + Bucket: bucket, + RawQuery: uploadQuery.Encode(), + Header: uploadHeader, + Body: strings.NewReader(vertexStorageChannelTestContent), + ContentLength: int64(len(vertexStorageChannelTestContent)), + AccessToken: accessToken, + Proxy: channelSetting.Proxy, + }, 0) + if err != nil { + probeErrors = append(probeErrors, fmt.Errorf("upload temporary object: %w", err)) + } + + downloaded, readErr := runVertexStorageProbeRequest(ctx, deps.doProxy, vertex.StorageProxyRequest{ + Operation: vertex.StorageOperationGet, + Method: http.MethodGet, + Bucket: bucket, + Object: objectName, + RawQuery: "alt=media", + Header: make(http.Header), + ContentLength: 0, + AccessToken: accessToken, + Proxy: channelSetting.Proxy, + }, int64(len(vertexStorageChannelTestContent))) + if readErr != nil { + probeErrors = append(probeErrors, fmt.Errorf("read temporary object: %w", readErr)) + } else if !bytes.Equal(downloaded, []byte(vertexStorageChannelTestContent)) { + probeErrors = append(probeErrors, errors.New("read temporary object: content mismatch")) + } + + cleanupCtx, cancelCleanup := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancelCleanup() + _, deleteErr := runVertexStorageProbeRequest(cleanupCtx, deps.doProxy, vertex.StorageProxyRequest{ + Operation: vertex.StorageOperationDelete, + Method: http.MethodDelete, + Bucket: bucket, + Object: objectName, + Header: make(http.Header), + ContentLength: 0, + AccessToken: accessToken, + Proxy: channelSetting.Proxy, + }, 0) + if deleteErr != nil { + probeErrors = append(probeErrors, fmt.Errorf("delete temporary object %q manually if necessary: %w", objectName, deleteErr)) + } + + return errors.Join(probeErrors...) +} + +func runVertexStorageProbeRequest( + ctx context.Context, + doProxy func(context.Context, vertex.StorageProxyRequest) (*http.Response, error), + input vertex.StorageProxyRequest, + maxResponseBytes int64, +) ([]byte, error) { + response, err := doProxy(ctx, input) + if response != nil && response.Body != nil { + defer service.CloseResponseBodyGracefully(response) + } + if err != nil { + return nil, err + } + if response == nil || response.Body == nil { + return nil, errors.New("Google Cloud Storage returned an empty response") + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return nil, fmt.Errorf("Google Cloud Storage returned status %d", response.StatusCode) + } + if maxResponseBytes <= 0 { + return nil, nil + } + body, err := io.ReadAll(io.LimitReader(response.Body, maxResponseBytes+1)) + if err != nil { + return nil, err + } + if int64(len(body)) > maxResponseBytes { + return nil, errors.New("Google Cloud Storage returned an oversized test object") + } + return body, nil +} diff --git a/controller/vertex_storage_channel_probe_test.go b/controller/vertex_storage_channel_probe_test.go new file mode 100644 index 000000000000..99d00383194b --- /dev/null +++ b/controller/vertex_storage_channel_probe_test.go @@ -0,0 +1,365 @@ +package controller + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relay/channel/vertex" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type vertexStorageProbeBody struct { + io.Reader + closed bool +} + +func (body *vertexStorageProbeBody) Close() error { + body.closed = true + return nil +} + +func TestVertexStorageChannelProbeWritesReadsAndDeletesOnce(t *testing.T) { + c := newVertexStorageChannelProbeContext(t) + objectName := ".new-api-channel-test/probe-id/test.txt" + var operations []vertex.StorageOperation + var responseBodies []*vertexStorageProbeBody + accessTokenCalls := 0 + deps := vertexStorageChannelProbeDependencies{ + newObjectName: func() string { return objectName }, + acquireAccessToken: func(input vertex.CachedAccessTokenRequest) (string, error) { + accessTokenCalls++ + assert.Equal(t, 41, input.ChannelID) + assert.True(t, input.ChannelIsMultiKey) + assert.Equal(t, 2, input.ChannelMultiKeyIndex) + assert.Equal(t, "service@example.com", input.Credentials.ClientEmail) + assert.Equal(t, "http://proxy.example:8080", input.Proxy) + return "access-token", nil + }, + doProxy: func(_ context.Context, input vertex.StorageProxyRequest) (*http.Response, error) { + operations = append(operations, input.Operation) + assert.Equal(t, "bucket-a", input.Bucket) + assert.Equal(t, "access-token", input.AccessToken) + assert.Equal(t, "http://proxy.example:8080", input.Proxy) + + status := http.StatusOK + bodyText := "" + switch input.Operation { + case vertex.StorageOperationUpload: + assert.Equal(t, http.MethodPost, input.Method) + query, err := url.ParseQuery(input.RawQuery) + require.NoError(t, err) + assert.Equal(t, "media", query.Get("uploadType")) + assert.Equal(t, objectName, query.Get("name")) + assert.Equal(t, "text/plain; charset=utf-8", input.Header.Get("Content-Type")) + assert.Equal(t, int64(len(vertexStorageChannelTestContent)), input.ContentLength) + uploaded, err := io.ReadAll(input.Body) + require.NoError(t, err) + assert.Equal(t, vertexStorageChannelTestContent, string(uploaded)) + case vertex.StorageOperationGet: + assert.Equal(t, http.MethodGet, input.Method) + assert.Equal(t, objectName, input.Object) + assert.Equal(t, "alt=media", input.RawQuery) + bodyText = vertexStorageChannelTestContent + case vertex.StorageOperationDelete: + assert.Equal(t, http.MethodDelete, input.Method) + assert.Equal(t, objectName, input.Object) + status = http.StatusNoContent + default: + t.Fatalf("unexpected storage operation %d", input.Operation) + } + + body := &vertexStorageProbeBody{Reader: strings.NewReader(bodyText)} + responseBodies = append(responseBodies, body) + return &http.Response{StatusCode: status, Header: make(http.Header), Body: body}, nil + }, + } + + err := testVertexStorageChannel(context.Background(), c, "storage:gs:bucket-a", deps) + + require.NoError(t, err) + assert.Equal(t, 1, accessTokenCalls) + assert.Equal(t, []vertex.StorageOperation{ + vertex.StorageOperationUpload, + vertex.StorageOperationGet, + vertex.StorageOperationDelete, + }, operations) + require.Len(t, responseBodies, 3) + for _, body := range responseBodies { + assert.True(t, body.closed) + } +} + +func TestVertexStorageChannelProbeContinuesAfterUploadFailure(t *testing.T) { + c := newVertexStorageChannelProbeContext(t) + var operations []vertex.StorageOperation + deps := vertexStorageChannelProbeDependencies{ + newObjectName: func() string { return ".new-api-channel-test/upload-failed/test.txt" }, + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "access-token", nil }, + doProxy: func(_ context.Context, input vertex.StorageProxyRequest) (*http.Response, error) { + operations = append(operations, input.Operation) + status := http.StatusOK + body := "" + if input.Operation == vertex.StorageOperationUpload { + status = http.StatusForbidden + body = `{"error":{"message":"forbidden"}}` + } + if input.Operation == vertex.StorageOperationGet { + body = vertexStorageChannelTestContent + } + if input.Operation == vertex.StorageOperationDelete { + status = http.StatusNoContent + } + return newVertexStorageProbeResponse(status, body), nil + }, + } + + err := testVertexStorageChannel(context.Background(), c, "storage:gs:bucket-a", deps) + + require.Error(t, err) + assert.ErrorContains(t, err, "upload") + assert.Equal(t, []vertex.StorageOperation{ + vertex.StorageOperationUpload, + vertex.StorageOperationGet, + vertex.StorageOperationDelete, + }, operations) +} + +func TestVertexStorageChannelProbeDeletesAfterContentMismatch(t *testing.T) { + c := newVertexStorageChannelProbeContext(t) + deleteCalls := 0 + deps := vertexStorageChannelProbeDependencies{ + newObjectName: func() string { return ".new-api-channel-test/mismatch/test.txt" }, + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "access-token", nil }, + doProxy: func(_ context.Context, input vertex.StorageProxyRequest) (*http.Response, error) { + switch input.Operation { + case vertex.StorageOperationUpload: + return newVertexStorageProbeResponse(http.StatusOK, ""), nil + case vertex.StorageOperationGet: + return newVertexStorageProbeResponse(http.StatusOK, "different content"), nil + case vertex.StorageOperationDelete: + deleteCalls++ + return newVertexStorageProbeResponse(http.StatusNoContent, ""), nil + default: + t.Fatalf("unexpected storage operation %d", input.Operation) + return nil, nil + } + }, + } + + err := testVertexStorageChannel(context.Background(), c, "storage:gs:bucket-a", deps) + + require.Error(t, err) + assert.ErrorContains(t, err, "content mismatch") + assert.Equal(t, 1, deleteCalls) +} + +func TestVertexStorageChannelProbeReportsObjectWhenDeleteFails(t *testing.T) { + c := newVertexStorageChannelProbeContext(t) + objectName := ".new-api-channel-test/delete-failed/test.txt" + deps := vertexStorageChannelProbeDependencies{ + newObjectName: func() string { return objectName }, + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "access-token", nil }, + doProxy: func(_ context.Context, input vertex.StorageProxyRequest) (*http.Response, error) { + switch input.Operation { + case vertex.StorageOperationUpload: + return newVertexStorageProbeResponse(http.StatusOK, ""), nil + case vertex.StorageOperationGet: + return newVertexStorageProbeResponse(http.StatusOK, vertexStorageChannelTestContent), nil + case vertex.StorageOperationDelete: + return newVertexStorageProbeResponse(http.StatusForbidden, `{"error":{"message":"forbidden"}}`), nil + default: + t.Fatalf("unexpected storage operation %d", input.Operation) + return nil, nil + } + }, + } + + err := testVertexStorageChannel(context.Background(), c, "storage:gs:bucket-a", deps) + + require.Error(t, err) + assert.ErrorContains(t, err, "delete") + assert.ErrorContains(t, err, objectName) +} + +func TestVertexStorageChannelProbeUsesIndependentDeleteContext(t *testing.T) { + c := newVertexStorageChannelProbeContext(t) + parent, cancel := context.WithCancel(context.Background()) + cancel() + deleteCalls := 0 + deps := vertexStorageChannelProbeDependencies{ + newObjectName: func() string { return ".new-api-channel-test/canceled/test.txt" }, + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "access-token", nil }, + doProxy: func(ctx context.Context, input vertex.StorageProxyRequest) (*http.Response, error) { + if input.Operation == vertex.StorageOperationDelete { + deleteCalls++ + assert.NoError(t, ctx.Err()) + return newVertexStorageProbeResponse(http.StatusNoContent, ""), nil + } + return nil, ctx.Err() + }, + } + + err := testVertexStorageChannel(parent, c, "storage:gs:bucket-a", deps) + + require.Error(t, err) + assert.Equal(t, 1, deleteCalls) +} + +func TestVertexStorageChannelProbeRejectsInvalidConfigurationBeforeUpstream(t *testing.T) { + tests := []struct { + name string + modelName string + configure func(*gin.Context) + want string + }{ + { + name: "invalid storage model", + modelName: "storage:gs:bucket-a/path", + configure: func(*gin.Context) {}, + want: "invalid", + }, + { + name: "wrong channel type", + modelName: "storage:gs:bucket-a", + configure: func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeGemini) + }, + want: "Vertex AI", + }, + { + name: "bucket not configured", + modelName: "storage:gs:bucket-a", + configure: func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyChannelModels, []string{"storage:gs:bucket-b"}) + }, + want: "does not allow", + }, + { + name: "API key mode", + modelName: "storage:gs:bucket-a", + configure: func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{VertexKeyType: dto.VertexKeyTypeAPIKey}) + }, + want: "service account", + }, + { + name: "invalid credentials", + modelName: "storage:gs:bucket-a", + configure: func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyChannelKey, "not-json") + }, + want: "credentials", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := newVertexStorageChannelProbeContext(t) + tt.configure(c) + deps := vertexStorageChannelProbeDependencies{ + newObjectName: func() string { return ".new-api-channel-test/rejected/test.txt" }, + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { + t.Fatal("OAuth must not run for invalid local configuration") + return "", nil + }, + doProxy: func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) { + t.Fatal("GCS proxy must not run for invalid local configuration") + return nil, nil + }, + } + + err := testVertexStorageChannel(context.Background(), c, tt.modelName, deps) + + require.Error(t, err) + assert.ErrorContains(t, err, tt.want) + }) + } +} + +func TestChannelTestRoutesVertexStorageBeforeBillingAndConsumeLogs(t *testing.T) { + originalDB, originalLogDB := model.DB, model.LOG_DB + model.DB, model.LOG_DB = nil, nil + t.Cleanup(func() { + model.DB, model.LOG_DB = originalDB, originalLogDB + }) + + settingBytes, err := common.Marshal(dto.ChannelSettings{Proxy: "http://proxy.example:8080"}) + require.NoError(t, err) + channel := &model.Channel{ + Id: 41, + Type: constant.ChannelTypeVertexAi, + Key: `{"project_id":"project-a","client_email":"service@example.com","private_key":"private-key"}`, + Models: "gemini-2.5-pro,storage:gs:bucket-a", + Setting: common.GetPointer(string(settingBytes)), + OtherSettings: `{"vertex_key_type":"json"}`, + } + deps := successfulVertexStorageChannelProbeDependencies(t) + + result := testChannelWithVertexStorageDependencies(context.Background(), channel, -1, " storage:gs:bucket-a ", "", false, deps) + + require.NoError(t, result.localErr) + assert.Nil(t, result.newAPIError) + assert.NotNil(t, result.context) +} + +func TestChannelTestOnlyRoutesVertexStorageModelsOnVertexChannels(t *testing.T) { + assert.True(t, isVertexStorageChannelTest(&model.Channel{Type: constant.ChannelTypeVertexAi}, "storage:gs:bucket-a")) + assert.False(t, isVertexStorageChannelTest(&model.Channel{Type: constant.ChannelTypeVertexAi}, "gemini-2.5-pro")) + assert.False(t, isVertexStorageChannelTest(&model.Channel{Type: constant.ChannelTypeGemini}, "storage:gs:bucket-a")) + assert.False(t, isVertexStorageChannelTest(nil, "storage:gs:bucket-a")) +} + +func newVertexStorageChannelProbeContext(t *testing.T) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(c, constant.ContextKeyChannelId, 41) + common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeVertexAi) + common.SetContextKey(c, constant.ContextKeyChannelModels, []string{"gemini-2.5-pro", "storage:gs:bucket-a"}) + common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{VertexKeyType: dto.VertexKeyTypeJSON}) + common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{Proxy: "http://proxy.example:8080"}) + common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true) + common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, 2) + common.SetContextKey(c, constant.ContextKeyChannelKey, `{"project_id":"project-a","client_email":"service@example.com","private_key":"private-key"}`) + return c +} + +func newVertexStorageProbeResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + } +} + +func successfulVertexStorageChannelProbeDependencies(t *testing.T) vertexStorageChannelProbeDependencies { + t.Helper() + return vertexStorageChannelProbeDependencies{ + newObjectName: func() string { return ".new-api-channel-test/routed/test.txt" }, + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "access-token", nil }, + doProxy: func(_ context.Context, input vertex.StorageProxyRequest) (*http.Response, error) { + switch input.Operation { + case vertex.StorageOperationUpload: + return newVertexStorageProbeResponse(http.StatusOK, ""), nil + case vertex.StorageOperationGet: + return newVertexStorageProbeResponse(http.StatusOK, vertexStorageChannelTestContent), nil + case vertex.StorageOperationDelete: + return newVertexStorageProbeResponse(http.StatusNoContent, ""), nil + default: + return nil, errors.New("unexpected storage operation") + } + }, + } +} diff --git a/controller/vertex_storage_proxy.go b/controller/vertex_storage_proxy.go new file mode 100644 index 000000000000..d39262ab69e6 --- /dev/null +++ b/controller/vertex_storage_proxy.go @@ -0,0 +1,151 @@ +package controller + +import ( + "context" + "io" + "net/http" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/relay/channel/vertex" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/gin-gonic/gin" +) + +type vertexStorageProxyDependencies struct { + acquireAccessToken func(vertex.CachedAccessTokenRequest) (string, error) + doProxy func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) +} + +func defaultVertexStorageProxyDependencies() vertexStorageProxyDependencies { + return vertexStorageProxyDependencies{ + acquireAccessToken: vertex.AcquireCachedAccessToken, + doProxy: vertex.DoStorageProxy, + } +} + +func RelayVertexStorageUpload(c *gin.Context) { + relayVertexStorageProxy(c, vertex.StorageOperationUpload, defaultVertexStorageProxyDependencies()) +} + +func RelayVertexStorageList(c *gin.Context) { + relayVertexStorageProxy(c, vertex.StorageOperationList, defaultVertexStorageProxyDependencies()) +} + +func RelayVertexStorageObject(c *gin.Context) { + operation := vertex.StorageOperationGet + if c.Request.Method == http.MethodDelete { + operation = vertex.StorageOperationDelete + } + relayVertexStorageProxy(c, operation, defaultVertexStorageProxyDependencies()) +} + +func relayVertexStorageProxy(c *gin.Context, operation vertex.StorageOperation, deps vertexStorageProxyDependencies) { + bucket, err := relayconstant.NormalizeVertexStorageBucket(c.Param("bucket")) + if err != nil { + respondVertexStorageProxyError(c, http.StatusBadRequest, "invalid_bucket", "invalid Cloud Storage bucket") + return + } + if common.GetContextKeyInt(c, constant.ContextKeyChannelType) != constant.ChannelTypeVertexAi { + respondVertexStorageProxyError(c, http.StatusInternalServerError, "channel_type_mismatch", "selected channel is not Vertex AI") + return + } + if !relayconstant.VertexStorageChannelSupports(common.GetContextKeyStringSlice(c, constant.ContextKeyChannelModels), bucket) { + respondVertexStorageProxyError(c, http.StatusForbidden, "bucket_not_allowed", "selected channel does not allow this Cloud Storage bucket") + return + } + + channelOtherSetting, _ := common.GetContextKeyType[dto.ChannelOtherSettings](c, constant.ContextKeyChannelOtherSetting) + if channelOtherSetting.VertexKeyType == dto.VertexKeyTypeAPIKey { + respondVertexStorageProxyError(c, http.StatusBadRequest, "unsupported_key_type", "Cloud Storage access requires Vertex AI service account JSON") + return + } + + credentials := vertex.Credentials{} + if err = common.Unmarshal([]byte(common.GetContextKeyString(c, constant.ContextKeyChannelKey)), &credentials); err != nil { + respondVertexStorageProxyError(c, http.StatusInternalServerError, "invalid_channel_credentials", "selected Vertex AI channel credentials are invalid") + return + } + + object := strings.TrimPrefix(c.Param("object"), "/") + if (operation == vertex.StorageOperationGet || operation == vertex.StorageOperationDelete) && object == "" { + respondVertexStorageProxyError(c, http.StatusBadRequest, "object_required", "Cloud Storage object is required") + return + } + if (operation == vertex.StorageOperationGet || operation == vertex.StorageOperationDelete) && relayconstant.ValidateVertexStorageObjectName(object) != nil { + respondVertexStorageProxyError(c, http.StatusBadRequest, "invalid_object", "Cloud Storage object contains an invalid path segment") + return + } + + channelSetting, _ := common.GetContextKeyType[dto.ChannelSettings](c, constant.ContextKeyChannelSetting) + accessToken, err := deps.acquireAccessToken(vertex.CachedAccessTokenRequest{ + ChannelID: common.GetContextKeyInt(c, constant.ContextKeyChannelId), + ChannelIsMultiKey: common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey), + ChannelMultiKeyIndex: common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex), + Credentials: credentials, + Proxy: channelSetting.Proxy, + }) + if err != nil { + respondVertexStorageProxyError(c, http.StatusBadGateway, "access_token_failed", "failed to authorize with Vertex AI service account") + return + } + + response, err := deps.doProxy(c.Request.Context(), vertex.StorageProxyRequest{ + Operation: operation, + Method: c.Request.Method, + Bucket: bucket, + Object: object, + RawQuery: c.Request.URL.RawQuery, + Header: c.Request.Header, + Body: c.Request.Body, + ContentLength: c.Request.ContentLength, + AccessToken: accessToken, + Proxy: channelSetting.Proxy, + }) + if response != nil && response.Body != nil { + defer service.CloseResponseBodyGracefully(response) + } + if err != nil || response == nil || response.Body == nil { + respondVertexStorageProxyError(c, http.StatusBadGateway, "upstream_request_failed", "failed to request Google Cloud Storage") + return + } + + responseHeader := vertex.SanitizeStorageResponseHeader(response.Header) + if operation == vertex.StorageOperationUpload && responseHeader.Get("Location") != "" { + rewrittenLocation, rewriteErr := vertex.RewriteStorageResumableLocation(responseHeader.Get("Location"), system_setting.ServerAddress, bucket) + if rewriteErr != nil { + respondVertexStorageProxyError(c, http.StatusBadGateway, "invalid_resumable_location", "Google Cloud Storage returned an invalid resumable upload location") + return + } + responseHeader.Set("Location", rewrittenLocation) + } + if response.StatusCode < http.StatusContinue || response.StatusCode > 599 { + respondVertexStorageProxyError(c, http.StatusBadGateway, "invalid_upstream_status", "Google Cloud Storage returned an invalid response status") + return + } + for name, values := range responseHeader { + if !service.ShouldCopyUpstreamHeader(c, name, values) { + continue + } + for _, value := range values { + c.Writer.Header().Add(name, value) + } + } + c.Status(response.StatusCode) + if _, err = io.Copy(c.Writer, response.Body); err != nil { + logger.LogError(c, "failed to stream Google Cloud Storage response") + } +} + +func respondVertexStorageProxyError(c *gin.Context, status int, code, message string) { + c.JSON(status, gin.H{"error": gin.H{ + "message": message, + "type": "invalid_request_error", + "code": code, + }}) +} diff --git a/controller/vertex_storage_proxy_test.go b/controller/vertex_storage_proxy_test.go new file mode 100644 index 000000000000..0aec3d1243c8 --- /dev/null +++ b/controller/vertex_storage_proxy_test.go @@ -0,0 +1,354 @@ +package controller + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/relay/channel/vertex" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type trackingVertexStorageBody struct { + io.Reader + closed bool +} + +func (body *trackingVertexStorageBody) Close() error { + body.closed = true + return nil +} + +func TestRelayVertexStorageProxyRejectsLocallyInAuthorizationOrder(t *testing.T) { + tests := []struct { + name string + operation vertex.StorageOperation + configure func(*gin.Context) + wantStatus int + wantCode string + }{ + { + name: "bucket before channel type", + operation: vertex.StorageOperationList, + configure: func(c *gin.Context) { + c.Params[0].Value = "bucket-a/path" + common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeOpenAI) + }, + wantStatus: http.StatusBadRequest, + wantCode: "invalid_bucket", + }, + { + name: "channel type before bucket authorization", + operation: vertex.StorageOperationList, + configure: func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeOpenAI) + common.SetContextKey(c, constant.ContextKeyChannelModels, []string{"storage:gs:bucket-b"}) + }, + wantStatus: http.StatusInternalServerError, + wantCode: "channel_type_mismatch", + }, + { + name: "exact bucket authorization before key type", + operation: vertex.StorageOperationList, + configure: func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyChannelModels, []string{"storage:gs:bucket-a-archive"}) + common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{VertexKeyType: dto.VertexKeyTypeAPIKey}) + }, + wantStatus: http.StatusForbidden, + wantCode: "bucket_not_allowed", + }, + { + name: "key type before service account JSON", + operation: vertex.StorageOperationList, + configure: func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{VertexKeyType: dto.VertexKeyTypeAPIKey}) + common.SetContextKey(c, constant.ContextKeyChannelKey, "not-json") + }, + wantStatus: http.StatusBadRequest, + wantCode: "unsupported_key_type", + }, + { + name: "service account JSON before object", + operation: vertex.StorageOperationGet, + configure: func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyChannelKey, "not-json") + }, + wantStatus: http.StatusInternalServerError, + wantCode: "invalid_channel_credentials", + }, + { + name: "object before OAuth", + operation: vertex.StorageOperationDelete, + configure: func(_ *gin.Context) {}, + wantStatus: http.StatusBadRequest, + wantCode: "object_required", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder, c := newVertexStorageProxyTestContext(t, http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", "bucket-a") + tt.configure(c) + deps := rejectingVertexStorageProxyDependencies(t) + + relayVertexStorageProxy(c, tt.operation, deps) + + assert.Equal(t, tt.wantStatus, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"`+tt.wantCode+`"`) + }) + } +} + +func TestRelayVertexStorageProxyStopsAfterOAuthFailure(t *testing.T) { + recorder, c := newVertexStorageProxyTestContext(t, http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", "bucket-a") + deps := vertexStorageProxyDependencies{ + acquireAccessToken: func(input vertex.CachedAccessTokenRequest) (string, error) { + assert.Equal(t, 41, input.ChannelID) + assert.True(t, input.ChannelIsMultiKey) + assert.Equal(t, 3, input.ChannelMultiKeyIndex) + assert.Equal(t, "svc@example.com", input.Credentials.ClientEmail) + assert.Equal(t, "http://proxy.example:8080", input.Proxy) + return "", errors.New("oauth unavailable") + }, + doProxy: func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) { + t.Fatal("GCS proxy must not run when OAuth fails") + return nil, nil + }, + } + + relayVertexStorageProxy(c, vertex.StorageOperationList, deps) + + assert.Equal(t, http.StatusBadGateway, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"access_token_failed"`) + assert.NotContains(t, recorder.Body.String(), "oauth unavailable") +} + +func TestRelayVertexStorageProxyRejectsDotSegmentObjectsBeforeOAuth(t *testing.T) { + tests := []struct { + name string + method string + operation vertex.StorageOperation + object string + }{ + { + name: "get parent segment", + method: http.MethodGet, + operation: vertex.StorageOperationGet, + object: "/../bucket-metadata", + }, + { + name: "delete nested parent segment", + method: http.MethodDelete, + operation: vertex.StorageOperationDelete, + object: "/folder/../object.txt", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder, c := newVertexStorageProxyTestContext(t, tt.method, "/vertexai/storage/v1/b/bucket-a/o"+tt.object, "bucket-a") + c.Params = append(c.Params, gin.Param{Key: "object", Value: tt.object}) + + relayVertexStorageProxy(c, tt.operation, rejectingVertexStorageProxyDependencies(t)) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"invalid_object"`) + }) + } +} + +func TestRelayVertexStorageProxyStreamsRequestAndRangeResponse(t *testing.T) { + incomingBody := &trackingVertexStorageBody{Reader: strings.NewReader("upload-bytes")} + responseBody := &trackingVertexStorageBody{Reader: strings.NewReader("file-bytes")} + deps := vertexStorageProxyDependencies{ + acquireAccessToken: func(input vertex.CachedAccessTokenRequest) (string, error) { + require.Equal(t, vertex.Credentials{ + ProjectID: "project-a", + ClientEmail: "svc@example.com", + PrivateKey: "private-key", + }, input.Credentials) + return "google-token", nil + }, + doProxy: func(_ context.Context, input vertex.StorageProxyRequest) (*http.Response, error) { + assert.Equal(t, vertex.StorageOperationGet, input.Operation) + assert.Equal(t, http.MethodGet, input.Method) + assert.Equal(t, "bucket-a", input.Bucket) + assert.Equal(t, "folder/file.bin", input.Object) + assert.Equal(t, "alt=media&generation=7", input.RawQuery) + assert.Equal(t, "bytes=0-9", input.Header.Get("Range")) + assert.Same(t, incomingBody, input.Body) + assert.Equal(t, int64(12), input.ContentLength) + assert.Equal(t, "google-token", input.AccessToken) + assert.Equal(t, "http://proxy.example:8080", input.Proxy) + return &http.Response{ + StatusCode: http.StatusPartialContent, + Header: http.Header{ + "Content-Type": {"application/octet-stream"}, + "Content-Range": {"bytes 0-9/100"}, + "Etag": {`"etag-1"`}, + "Content-Length": {"10"}, + "Connection": {"keep-alive"}, + }, + Body: responseBody, + }, nil + }, + } + recorder, c := newVertexStorageProxyTestContext(t, http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o/folder%2Ffile.bin?alt=media&generation=7", "bucket-a") + c.Params = append(c.Params, gin.Param{Key: "object", Value: "/folder/file.bin"}) + c.Request.Body = incomingBody + c.Request.ContentLength = 12 + c.Request.Header.Set("Range", "bytes=0-9") + + relayVertexStorageProxy(c, vertex.StorageOperationGet, deps) + + assert.Equal(t, http.StatusPartialContent, recorder.Code) + assert.Equal(t, "file-bytes", recorder.Body.String()) + assert.Equal(t, "bytes 0-9/100", recorder.Header().Get("Content-Range")) + assert.Equal(t, `"etag-1"`, recorder.Header().Get("ETag")) + assert.Empty(t, recorder.Header().Get("Content-Length")) + assert.Empty(t, recorder.Header().Get("Connection")) + assert.True(t, responseBody.closed) +} + +func TestRelayVertexStorageProxyPreservesGCSErrorStatusAndBody(t *testing.T) { + responseBody := &trackingVertexStorageBody{Reader: strings.NewReader(`{"error":{"code":403,"message":"forbidden"}}`)} + deps := vertexStorageProxyDependencies{ + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "google-token", nil }, + doProxy: func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusForbidden, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: responseBody, + }, nil + }, + } + recorder, c := newVertexStorageProxyTestContext(t, http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", "bucket-a") + + relayVertexStorageProxy(c, vertex.StorageOperationList, deps) + + assert.Equal(t, http.StatusForbidden, recorder.Code) + assert.JSONEq(t, `{"error":{"code":403,"message":"forbidden"}}`, recorder.Body.String()) + assert.True(t, responseBody.closed) +} + +func TestRelayVertexStorageProxyClosesResponseBodyWhenProxyReturnsError(t *testing.T) { + responseBody := &trackingVertexStorageBody{Reader: strings.NewReader("must-not-leak")} + deps := vertexStorageProxyDependencies{ + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "google-token", nil }, + doProxy: func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusBadGateway, + Header: http.Header{"X-Upstream-Secret": {"secret"}}, + Body: responseBody, + }, errors.New("redirect rejected") + }, + } + recorder, c := newVertexStorageProxyTestContext(t, http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", "bucket-a") + + relayVertexStorageProxy(c, vertex.StorageOperationList, deps) + + assert.Equal(t, http.StatusBadGateway, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"upstream_request_failed"`) + assert.NotContains(t, recorder.Body.String(), "must-not-leak") + assert.Empty(t, recorder.Header().Get("X-Upstream-Secret")) + assert.True(t, responseBody.closed) +} + +func TestRelayVertexStorageProxyRewritesResumableLocation(t *testing.T) { + previousAddress := system_setting.ServerAddress + system_setting.ServerAddress = "https://api.example.com/base" + t.Cleanup(func() { system_setting.ServerAddress = previousAddress }) + + responseBody := &trackingVertexStorageBody{Reader: strings.NewReader("")} + deps := vertexStorageProxyDependencies{ + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "google-token", nil }, + doProxy: func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Location": { + "https://storage.googleapis.com/upload/storage/v1/b/bucket-a/o?uploadType=resumable&upload_id=session-1", + }}, + Body: responseBody, + }, nil + }, + } + recorder, c := newVertexStorageProxyTestContext(t, http.MethodPost, "/vertexai/upload/storage/v1/b/bucket-a/o?uploadType=resumable&name=file.bin", "bucket-a") + + relayVertexStorageProxy(c, vertex.StorageOperationUpload, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.Equal(t, "https://api.example.com/vertexai/upload/storage/v1/b/bucket-a/o?uploadType=resumable&upload_id=session-1", recorder.Header().Get("Location")) + assert.True(t, responseBody.closed) +} + +func TestRelayVertexStorageProxyDoesNotLeakUnsafeResumableLocation(t *testing.T) { + previousAddress := system_setting.ServerAddress + system_setting.ServerAddress = "" + t.Cleanup(func() { system_setting.ServerAddress = previousAddress }) + + responseBody := &trackingVertexStorageBody{Reader: strings.NewReader("upstream-body")} + deps := vertexStorageProxyDependencies{ + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { return "google-token", nil }, + doProxy: func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Location": { + "https://storage.googleapis.com/upload/storage/v1/b/bucket-a/o?uploadType=resumable&upload_id=session-1", + }}, + Body: responseBody, + }, nil + }, + } + recorder, c := newVertexStorageProxyTestContext(t, http.MethodPost, "/vertexai/upload/storage/v1/b/bucket-a/o?uploadType=resumable&name=file.bin", "bucket-a") + + relayVertexStorageProxy(c, vertex.StorageOperationUpload, deps) + + assert.Equal(t, http.StatusBadGateway, recorder.Code) + assert.Empty(t, recorder.Header().Get("Location")) + assert.NotContains(t, recorder.Body.String(), "storage.googleapis.com") + assert.NotContains(t, recorder.Body.String(), "session-1") + assert.NotContains(t, recorder.Body.String(), "upstream-body") + assert.True(t, responseBody.closed) +} + +func newVertexStorageProxyTestContext(t *testing.T, method, target, bucket string) (*httptest.ResponseRecorder, *gin.Context) { + t.Helper() + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(method, target, nil) + c.Params = gin.Params{{Key: "bucket", Value: bucket}} + common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeVertexAi) + common.SetContextKey(c, constant.ContextKeyChannelModels, []string{"storage:gs:bucket-a"}) + common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{VertexKeyType: dto.VertexKeyTypeJSON}) + common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{Proxy: "http://proxy.example:8080"}) + common.SetContextKey(c, constant.ContextKeyChannelId, 41) + common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true) + common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, 3) + common.SetContextKey(c, constant.ContextKeyChannelKey, `{"project_id":"project-a","client_email":"svc@example.com","private_key":"private-key"}`) + return recorder, c +} + +func rejectingVertexStorageProxyDependencies(t *testing.T) vertexStorageProxyDependencies { + t.Helper() + return vertexStorageProxyDependencies{ + acquireAccessToken: func(vertex.CachedAccessTokenRequest) (string, error) { + t.Fatal("OAuth must not run after local validation rejects the request") + return "", nil + }, + doProxy: func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) { + t.Fatal("GCS proxy must not run after local validation rejects the request") + return nil, nil + }, + } +} diff --git a/docs/openapi/relay.json b/docs/openapi/relay.json index 92a6239074e2..2a0a4cd251c0 100644 --- a/docs/openapi/relay.json +++ b/docs/openapi/relay.json @@ -51,6 +51,9 @@ { "name": "OpenAI音频(Audio)" }, + { + "name": "文件/Vertex AI Cloud Storage" + }, { "name": "重排序(Rerank)" }, @@ -71,6 +74,411 @@ } ], "paths": { + "/vertexai/upload/storage/v1/b/{bucket}/o": { + "post": { + "summary": "上传 Cloud Storage 对象", + "description": "通过 Vertex AI 渠道的服务账号转发 GCS media、multipart 或 resumable 初始化请求。Location 会安全改写为当前网关地址；此接口不计费。", + "operationId": "vertexStorageUploadObject", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "uploadType", + "in": "query", + "required": true, + "schema": { + "type": "string", + "enum": [ + "media", + "multipart", + "resumable" + ] + } + }, + { + "name": "name", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": false, + "content": { + "application/octet-stream": { + "schema": { + "type": "string", + "format": "binary" + } + }, + "application/json": { + "schema": { + "type": "object", + "additionalProperties": true + } + }, + "multipart/related": { + "schema": { + "type": "string", + "format": "binary" + } + } + } + }, + "responses": { + "200": { + "description": "上传成功或 resumable 会话已创建。", + "headers": { + "Location": { + "description": "已改写为当前网关的 resumable 上传地址。", + "schema": { + "type": "string", + "format": "uri" + } + } + } + }, + "201": { + "description": "对象创建成功。" + }, + "400": { + "description": "请求、bucket、上传类型或渠道凭证无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth、GCS 请求或 resumable Location 校验失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + }, + "put": { + "summary": "续传 Cloud Storage 对象分块", + "description": "使用 resumable 初始化返回的网关 Location 上传分块。请求体与响应流式转发；此接口不计费。", + "operationId": "vertexStorageResumeUploadObject", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "uploadType", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "upload_id", + "in": "query", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "name", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "Content-Range", + "in": "header", + "required": false, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "required": false, + "content": { + "application/octet-stream": { + "schema": { + "type": "string", + "format": "binary" + } + } + } + }, + "responses": { + "200": { + "description": "最终分块上传完成。" + }, + "201": { + "description": "对象创建成功。" + }, + "308": { + "description": "分块已接受，resumable 会话尚未完成。" + }, + "400": { + "description": "请求、bucket 或 Content-Range 无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth 或 GCS 请求失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, + "/vertexai/storage/v1/b/{bucket}/o": { + "get": { + "summary": "列举 Cloud Storage 对象", + "description": "列举渠道允许访问的 bucket 中的对象；查询参数按 GCS JSON API 透传。此接口不计费。", + "operationId": "vertexStorageListObjects", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "prefix", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "delimiter", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "pageToken", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Google Cloud Storage 对象列表。" + }, + "400": { + "description": "请求、bucket 或渠道凭证无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth 或 GCS 请求失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, + "/vertexai/storage/v1/b/{bucket}/o/{object}": { + "get": { + "summary": "读取或下载 Cloud Storage 对象", + "description": "默认返回对象元数据；设置 alt=media 时流式下载对象内容。对象名中的 / 应编码为 %2F；此接口不计费。", + "operationId": "vertexStorageGetObject", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "object", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "alt", + "in": "query", + "required": false, + "schema": { + "type": "string", + "enum": [ + "json", + "media" + ] + } + }, + { + "name": "generation", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "Range", + "in": "header", + "required": false, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "对象元数据或完整对象内容。" + }, + "206": { + "description": "Range 下载内容。", + "headers": { + "Content-Range": { + "schema": { + "type": "string" + } + } + } + }, + "400": { + "description": "请求、bucket、object 或渠道凭证无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth 或 GCS 请求失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + }, + "delete": { + "summary": "删除 Cloud Storage 对象", + "description": "删除渠道允许访问的 bucket 中的对象。对象名中的 / 应编码为 %2F；此接口不计费。", + "operationId": "vertexStorageDeleteObject", + "tags": [ + "文件/Vertex AI Cloud Storage" + ], + "parameters": [ + { + "name": "bucket", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "object", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "generation", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "ifGenerationMatch", + "in": "query", + "required": false, + "schema": { + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "对象已删除。" + }, + "400": { + "description": "请求、bucket、object 或渠道凭证无效。" + }, + "403": { + "description": "Token、渠道或 GCS IAM 无权访问该 bucket。" + }, + "502": { + "description": "OAuth 或 GCS 请求失败。" + }, + "default": { + "description": "Google Cloud Storage 原始状态、响应头和响应主体。" + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, "/v1/models": { "get": { "summary": "获取模型列表", diff --git a/docs/superpowers/plans/2026-08-09-vertex-ai-storage-and-channel-test.md b/docs/superpowers/plans/2026-08-09-vertex-ai-storage-and-channel-test.md new file mode 100644 index 000000000000..2927470a945f --- /dev/null +++ b/docs/superpowers/plans/2026-08-09-vertex-ai-storage-and-channel-test.md @@ -0,0 +1,543 @@ +# Vertex AI 文件存储与渠道测试实施计划 + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 为 Vertex AI 渠道增加 GCS bucket 配置、固定 `/vertexai` 文件代理路由和真实写入/读取/删除渠道测试，并提交符合仓库规范的 Pull Request。 + +**Architecture:** 前端继续用 `models` 保存 `storage:gs:`，但将普通模型和 bucket 分开编辑。后端从固定 `/vertexai` 路径提取 bucket，复用模型分发选择类型 41 渠道，再由 Controller 二次授权并通过现有 Vertex 服务账号 OAuth 流式访问固定的 `storage.googleapis.com`。渠道测试直接复用同一 Storage Proxy，不经过公开地址回环，也不进入计费链路。 + +**Tech Stack:** Go 1.22+、Gin、现有 Vertex JWT/OAuth、`net/http`、Testify、React 19、TypeScript、React Hook Form、Base UI、i18next、Vitest/React Testing Library、Bun。 + +## Global Constraints + +- 设计依据：`docs/superpowers/specs/2026-08-09-vertex-ai-storage-and-channel-test-design.md`。 +- 仅 Vertex AI 渠道类型 41 支持 Storage；Gemini 类型 24 和其他渠道不得进入。 +- 对外前缀固定为 `/vertexai`，不得改成 `/v1/rawproxy/vertexai`。 +- 只开放上传、续传、列举、读取、下载和删除对象的五组 method/path，不得增加任意目标通配代理。 +- bucket 授权只接受纯 bucket 名称，并精确匹配渠道 `models` 中的 `storage:gs:`；不支持目录前缀授权。 +- 只支持服务账号 JSON；`dto.VertexKeyTypeAPIKey` 必须在访问 GCS 前拒绝。 +- 上游固定为 `https://storage.googleapis.com`，客户端不得覆盖 Host 或 Authorization。 +- 上传和下载必须流式传输，不得将完整文件读入内存。 +- Resumable `Location` 必须改写到当前服务的 `/vertexai/upload/...`；系统服务地址为空时本地失败，不得泄露 Google Session URL。 +- Storage 文件操作和渠道测试不执行价格查询、quota 预扣、结算、退款或模型消费日志。 +- 不新增数据库字段或迁移，继续兼容 SQLite、MySQL 5.7.8+、PostgreSQL 9.6+。 +- 后端 JSON 编解码使用 `common.*` 包装，不直接调用 `encoding/json` marshal/unmarshal。 +- 新增或大改 Go 测试使用 `require` 做致命断言、`assert` 做非致命断言。 +- 前端新增文案使用 `useTranslation()` 和 `t('English key')`，覆盖 `en`、`zh`、`zh-TW`、`fr`、`ja`、`ru`、`vi`。 +- 前端依赖与命令使用 Bun；修改 TypeScript/TSX 后运行受影响测试、typecheck、涉及文件 lint 和生产构建。 +- 每个任务保留原子提交用于评审；全部测试和最终评审完成后 squash 为一个计划 Commit：`feat: add Vertex AI storage integration`。 +- 创建 PR 前比较当前 Git 用户与历史核心开发者，使用 `.github/PULL_REQUEST_TEMPLATE.md`，必要时在 PR 正文声明 AI 辅助。 + +--- + +## 文件结构 + +- `web/src/features/channels/lib/vertex-storage-models.ts`：bucket 校验、普通模型/Storage 标识拆分与合并。 +- `web/src/features/channels/components/drawers/sections/vertex-storage-buckets-field.tsx`：类型 41 专属 bucket 多值字段。 +- `relay/constant/vertex_storage.go`：`/vertexai` 固定路由、bucket 规范化和 `storage:gs:` 授权 helper。 +- `middleware/distributor.go`：从 Storage 路径构造分发模型并保存渠道模型上下文。 +- `relay/channel/vertex/storage_proxy.go`：固定 GCS URL、请求头过滤、代理请求和 resumable 地址改写。 +- `controller/vertex_storage_proxy.go`：二次授权、本地错误、响应头复制和流式响应。 +- `controller/vertex_storage_channel_probe.go`：唯一临时对象及写入/读取/删除测试。 +- `router/relay-router.go`：注册五组固定 `/vertexai` method/path。 +- `web/src/features/channels/components/data-table-row-actions.tsx`：统一打开渠道测试弹窗。 +- `web/src/features/channels/components/dialogs/channel-test-dialog.tsx`：标记 GCS bucket 测试项。 +- `docs/vertex-ai-storage.md`、`docs/openapi/relay.json`：用户说明和 OpenAPI 契约。 + +### Task 1: 前端 Storage 模型边界与 bucket 字段 + +**Files:** +- Create: `web/src/features/channels/lib/vertex-storage-models.ts` +- Create: `web/src/features/channels/lib/__tests__/vertex-storage-models.test.ts` +- Modify: `web/src/features/channels/lib/index.ts` +- Create: `web/src/features/channels/components/drawers/sections/vertex-storage-buckets-field.tsx` +- Create: `web/src/features/channels/components/drawers/sections/__tests__/vertex-storage-buckets-field.test.tsx` +- Modify: `web/src/features/channels/components/drawers/sections/index.ts` +- Modify: `web/src/features/channels/components/drawers/channel-mutate-drawer.tsx` +- Modify: `web/scripts/add-missing-keys.mjs` +- Modify via i18n flow: `web/src/i18n/locales/{en,zh,zh-TW,fr,ja,ru,vi}.json` + +**Interfaces:** +- Produces: `VERTEX_STORAGE_MODEL_PREFIX`、`normalizeVertexStorageBucket()`、`splitVertexStorageModels()`、`mergeVertexStorageModels()`。 +- Produces: `VertexStorageBucketsField({ channelType, models, onModelsChange })`。 + +- [ ] **Step 1: 写入失败的领域测试** + +```ts +assert.deepEqual( + splitVertexStorageModels([ + 'gemini-2.5-pro', + 'storage:gs:bucket-a', + 'storage:gs:bucket-b', + ]), + { models: ['gemini-2.5-pro'], buckets: ['bucket-a', 'bucket-b'] } +) +assert.equal(normalizeVertexStorageBucket('bucket-a'), 'bucket-a') +assert.equal(normalizeVertexStorageBucket('bucket-a/path'), null) +assert.equal(normalizeVertexStorageBucket('storage:gs:bucket-a'), null) +assert.equal(normalizeVertexStorageBucket('gs://bucket-a'), null) +assert.deepEqual( + mergeVertexStorageModels(['gemini-2.5-pro'], ['bucket-a', 'bucket-a']), + ['gemini-2.5-pro', 'storage:gs:bucket-a'] +) +``` + +- [ ] **Step 2: 运行测试并确认模块缺失** + +```bash +cd web +bun test src/features/channels/lib/__tests__/vertex-storage-models.test.ts +``` + +Expected: FAIL，无法解析 `../vertex-storage-models`。 + +- [ ] **Step 3: 实现纯函数** + +```ts +export const VERTEX_STORAGE_MODEL_PREFIX = 'storage:gs:' + +export function normalizeVertexStorageBucket(value: string): string | null { + const bucket = value.trim() + if (!bucket || bucket === '.' || bucket === '..') return null + if (bucket.startsWith(VERTEX_STORAGE_MODEL_PREFIX)) return null + if (bucket.includes('://') || /[\\/?#]/.test(bucket)) return null + return bucket +} +``` + +`splitVertexStorageModels()` 保序拆分、trim、去重并忽略无效 Storage 项；`mergeVertexStorageModels()` 去掉 models 中旧 Storage 项，仅编码合法且唯一的 bucket。 + +- [ ] **Step 4: 写入组件失败测试** + +```tsx +const nonVertex = await renderField({ channelType: 24, models: [] }) +expect(nonVertex.container).toBeEmptyDOMElement() + +const vertex = await renderControlledField({ + channelType: 41, + models: ['gemini-2.5-pro', 'storage:gs:bucket-a'], +}) +expect(screen.getByText('bucket-a')).toBeInTheDocument() +await userEvent.type(screen.getByLabelText('Storage buckets'), 'bucket-b{enter}') +expect(screen.getByTestId('models-value')).toHaveTextContent( + 'gemini-2.5-pro,storage:gs:bucket-a,storage:gs:bucket-b' +) +``` + +- [ ] **Step 5: 实现字段并接入抽屉** + +字段仅在 `channelType === 41` 渲染，使用现有 `MultiSelect`。普通模型选择器只展示拆分后的 `models`；填充、抓取、移除映射目标和清空普通模型时调用 `mergeVertexStorageModels(newModels, currentBuckets)`，不得静默删除 bucket。 + +- [ ] **Step 6: 按 i18n-translate 流程补齐七语言** + +新增 key：`Storage buckets`、`Configure Google Cloud Storage buckets for this Vertex AI channel.`、`Enter storage bucket names`、`Add storage bucket "{{value}}"`、`Invalid storage bucket name`。 + +- [ ] **Step 7: 验证并提交** + +```bash +cd web +bun test src/features/channels/lib/__tests__/vertex-storage-models.test.ts \ + src/features/channels/components/drawers/sections/__tests__/vertex-storage-buckets-field.test.tsx +bun run typecheck +bunx oxlint -c .oxlintrc.json src/features/channels/lib/vertex-storage-models.ts \ + src/features/channels/components/drawers/sections/vertex-storage-buckets-field.tsx \ + src/features/channels/components/drawers/channel-mutate-drawer.tsx +git add web +git commit -m "feat: add Vertex storage bucket configuration" +``` + +### Task 2: 后端 bucket 契约与类型限定分发 + +**Files:** +- Create: `relay/constant/vertex_storage.go` +- Create: `relay/constant/vertex_storage_test.go` +- Modify: `relay/constant/relay_mode.go` +- Modify: `constant/context_key.go` +- Modify: `middleware/distributor.go` +- Modify: `middleware/distributor_test.go` +- Modify: `service/channel_select.go` +- Modify: `model/ability.go` +- Modify: `model/channel_cache.go` +- Create: `service/channel_select_channel_type_test.go` +- Create: `model/channel_type_selection_test.go` + +**Interfaces:** +- Produces: `VertexStorageRoutePrefix = "/vertexai"`、三条固定 route 常量、bucket/model helper、`ContextKeyChannelModels`、`RelayModeVertexStorage`。 +- Produces: `DistributeByChannelType(requiredChannelType int)`；`service.RetryParam.RequiredChannelType` 向数据库和缓存选择器传递类型限制，值为 0 时保持现有行为。 + +- [ ] **Step 1: 写入失败测试** + +```go +assert.Equal(t, "/vertexai", VertexStorageRoutePrefix) +assert.Equal(t, "/vertexai/upload/storage/v1/b/:bucket/o", VertexStorageUploadRoute) +for _, value := range []string{"", ".", "..", "bucket/path", `bucket\\path`, "storage:gs:bucket", "gs://bucket", "bucket?x=1"} { + _, err := NormalizeVertexStorageBucket(value) + require.Error(t, err, value) +} +assert.True(t, VertexStorageChannelSupports([]string{"storage:gs:bucket-a"}, "bucket-a")) +assert.False(t, VertexStorageChannelSupports([]string{"storage:gs:bucket-ab"}, "bucket-a")) +``` + +- [ ] **Step 2: 运行测试并确认函数缺失** + +```bash +go test ./relay/constant -run 'VertexStorage|NormalizeVertexStorage' -count=1 +``` + +- [ ] **Step 3: 实现常量和 helper** + +```go +const ( + VertexStorageModelPrefix = "storage:gs:" + VertexStorageRoutePrefix = "/vertexai" + VertexStorageUploadRoute = VertexStorageRoutePrefix + "/upload/storage/v1/b/:bucket/o" + VertexStorageListRoute = VertexStorageRoutePrefix + "/storage/v1/b/:bucket/o" + VertexStorageObjectRoute = VertexStorageRoutePrefix + "/storage/v1/b/:bucket/o/*object" +) +``` + +bucket 规范化与前端保持同一拒绝集合；渠道授权逐项 trim 后精确匹配。 + +- [ ] **Step 4: 写入分发失败测试** + +```go +c.Request = httptest.NewRequest(http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", nil) +c.Params = gin.Params{{Key: "bucket", Value: "bucket-a"}} +got, shouldSelect, err := getModelRequest(c) +require.NoError(t, err) +assert.True(t, shouldSelect) +assert.Equal(t, "storage:gs:bucket-a", got.Model) +assert.Equal(t, relayconstant.RelayModeVertexStorage, c.GetInt("relay_mode")) +``` + +增加指定渠道测试：只有类型 41 且精确配置 bucket 才运行下游；类型 24、非法 bucket、未配置 bucket 均提前失败。增加缓存和数据库选择测试：相同 group/model 同时存在类型 24 与 41 时，`RequiredChannelType: 41` 只能返回类型 41；值为 0 时保持原有选择集合。 + +- [ ] **Step 5: 实现分发与上下文** + +`Distribute()` 改为 `return distribute(0)`，新增 `DistributeByChannelType(requiredChannelType)`。类型限制必须覆盖 Token 指定渠道、Affinity 首选渠道、Redis/内存缓存选择和数据库回退选择；`service.RetryParam`、`model.GetChannelByType()` 与 `model.GetRandomSatisfiedChannelByType()` 负责把限制传到底层。`getModelRequest()` 优先识别 `IsVertexStoragePath()`；指定渠道路径增加精确 bucket 检查；`SetupContextForSelectedChannel()` 写入： + +```go +common.SetContextKey(c, constant.ContextKeyChannelModels, channel.GetModels()) +``` + +- [ ] **Step 6: 验证并提交** + +```bash +gofmt -w relay/constant/vertex_storage.go relay/constant/vertex_storage_test.go \ + relay/constant/relay_mode.go constant/context_key.go middleware/distributor.go middleware/distributor_test.go \ + service/channel_select.go model/ability.go model/channel_cache.go +go test ./relay/constant ./middleware ./service ./model -run 'VertexStorage|NormalizeVertexStorage|RequiredChannelType' -count=1 +git add relay/constant constant/context_key.go middleware/distributor.go middleware/distributor_test.go \ + service/channel_select.go service/channel_select_channel_type_test.go model/ability.go model/channel_cache.go model/channel_type_selection_test.go +git commit -m "feat: route Vertex storage buckets" +``` + +### Task 3: 可复用 Vertex OAuth 与固定 GCS Proxy + +**Files:** +- Modify: `relay/channel/vertex/service_account.go` +- Create: `relay/channel/vertex/service_account_test.go` +- Create: `relay/channel/vertex/storage_proxy.go` +- Create: `relay/channel/vertex/storage_proxy_test.go` + +**Interfaces:** +- Produces: `CachedAccessTokenRequest`、`AcquireCachedAccessToken()`、`StorageOperation`、`StorageProxyRequest`、`DoStorageProxy()`、`SanitizeStorageResponseHeader()`、`RewriteStorageResumableLocation()`。 + +- [ ] **Step 1: 写入 OAuth 缓存与固定 URL 失败测试** + +测试同一渠道/多 Key 索引复用缓存、不同索引不共享，以及： + +```go +req, err := buildStorageRequest(context.Background(), StorageProxyRequest{ + Operation: StorageOperationGet, + Method: http.MethodGet, + Bucket: "bucket-a", + Object: "docs/a b.pdf", + RawQuery: "alt=media", + AccessToken: "secret", +}) +require.NoError(t, err) +assert.Equal(t, "storage.googleapis.com", req.URL.Host) +assert.Contains(t, req.URL.EscapedPath(), "docs%2Fa%20b.pdf") +assert.Equal(t, "Bearer secret", req.Header.Get("Authorization")) +``` + +- [ ] **Step 2: 写入 Location 与头部失败测试** + +```go +got, err := RewriteStorageResumableLocation( + "https://storage.googleapis.com/upload/storage/v1/b/bucket-a/o?upload_id=abc", + "https://gateway.example.com", + "bucket-a", +) +require.NoError(t, err) +assert.Equal(t, "https://gateway.example.com/vertexai/upload/storage/v1/b/bucket-a/o?upload_id=abc", got) +``` + +空 ServerAddress、非 Google Location、bucket 不一致必须失败；客户端 Authorization、Host、hop-by-hop headers 必须移除，内容/Range/条件请求/`X-Goog-Meta-*` 必须保留。 + +- [ ] **Step 3: 运行失败测试** + +```bash +go test ./relay/channel/vertex -run 'CachedAccessToken|Storage|RewriteStorage' -count=1 +``` + +- [ ] **Step 4: 实现 OAuth 接口和 Storage Proxy** + +```go +type StorageProxyRequest struct { + Operation StorageOperation + Method string + Bucket string + Object string + RawQuery string + Header http.Header + Body io.Reader + ContentLength int64 + AccessToken string + Proxy string +} +``` + +URL 只由 operation、bucket、object 构造，主机固定；object 使用 `url.PathEscape` 成为单一 segment。`DoStorageProxy` 使用项目现有 Proxy HTTP client、调用方 context 和原始 body，不读取完整 body。 + +- [ ] **Step 5: 验证并提交** + +```bash +gofmt -w relay/channel/vertex/service_account.go relay/channel/vertex/service_account_test.go \ + relay/channel/vertex/storage_proxy.go relay/channel/vertex/storage_proxy_test.go +go test ./relay/channel/vertex -count=1 +git add relay/channel/vertex +git commit -m "feat: add Vertex GCS storage proxy" +``` + +### Task 4: `/vertexai` Controller、路由与公共文档 + +**Files:** +- Create: `controller/vertex_storage_proxy.go` +- Create: `controller/vertex_storage_proxy_test.go` +- Modify: `router/relay-router.go` +- Modify: `router/relay_router_test.go` +- Create: `docs/vertex-ai-storage.md` +- Modify: `docs/openapi/relay.json` + +**Interfaces:** +- Produces: `RelayVertexStorageUpload`、`RelayVertexStorageList`、`RelayVertexStorageObject`。 + +- [ ] **Step 1: 写入 Controller 与精确路由失败测试** + +本地校验失败不得调用上游；合法请求必须透传渠道 ID、多 Key 索引、Proxy、GCS 状态/body/Range/ETag 并关闭 body。路由集合必须精确等于： + +```go +map[string]bool{ + "POST /vertexai/upload/storage/v1/b/:bucket/o": true, + "PUT /vertexai/upload/storage/v1/b/:bucket/o": true, + "GET /vertexai/storage/v1/b/:bucket/o": true, + "GET /vertexai/storage/v1/b/:bucket/o/*object": true, + "DELETE /vertexai/storage/v1/b/:bucket/o/*object": true, +} +``` + +- [ ] **Step 2: 实现 Controller 二次授权** + +依赖边界： + +```go +type vertexStorageProxyDependencies struct { + acquireAccessToken func(vertex.CachedAccessTokenRequest) (string, error) + doProxy func(context.Context, vertex.StorageProxyRequest) (*http.Response, error) +} +``` + +早退顺序：bucket → 渠道类型 → 精确授权 → Key 类型 → 服务账号 JSON → object → Access Token → Storage Proxy。成功响应使用 `io.Copy`；不得建立 RelayInfo 或调用计费服务。 + +- [ ] **Step 3: 注册独立路由组** + +```go +vertexStorageRouter := router.Group("/vertexai") +vertexStorageRouter.Use(middleware.RouteTag("relay")) +vertexStorageRouter.Use(middleware.SystemPerformanceCheck()) +vertexStorageRouter.Use(middleware.TokenAuth()) +vertexStorageRouter.Use(middleware.ModelRequestRateLimit()) +vertexStorageRouter.Use(middleware.DistributeByChannelType(constant.ChannelTypeVertexAi)) +``` + +- [ ] **Step 4: 补充文档与 OpenAPI** + +文档包含 bucket 配置、五组路由、上传/下载/删除 cURL、`fileData.fileUri`、IAM、resumable 和非计费说明。OpenAPI tag 为 `文件/Vertex AI Cloud Storage`，所有路径以 `/vertexai` 开头。 + +- [ ] **Step 5: 验证并提交** + +```bash +gofmt -w controller/vertex_storage_proxy.go controller/vertex_storage_proxy_test.go \ + router/relay-router.go router/relay_router_test.go +go test ./controller ./router -run 'VertexStorage' -count=1 +git add controller/vertex_storage_proxy.go controller/vertex_storage_proxy_test.go \ + router/relay-router.go router/relay_router_test.go docs/vertex-ai-storage.md docs/openapi/relay.json +git commit -m "feat: expose Vertex storage routes" +``` + +### Task 5: Storage 渠道真实读写探测 + +**Files:** +- Create: `controller/vertex_storage_channel_probe.go` +- Create: `controller/vertex_storage_channel_probe_test.go` +- Modify: `controller/channel-test.go` + +**Interfaces:** +- Produces: `vertexStorageChannelProbeDependencies`、`testVertexStorageChannel()`、`testChannelWithVertexStorageDependencies()`。 + +- [ ] **Step 1: 写入三步与失败不中断测试** + +成功用例断言 operation 严格为 Upload、Get、Delete 且各一次。失败用例覆盖写入 403 后仍读取/删除、内容不一致后仍删除、删除失败返回对象路径、非法配置不调用上游、普通模型不进入分支、Storage 分支不产生消费日志。 + +- [ ] **Step 2: 运行失败测试** + +```bash +go test ./controller -run 'VertexStorageChannelProbe|ChannelTestRoutesVertexStorage' -count=1 +``` + +- [ ] **Step 3: 实现唯一对象和三步聚合** + +```go +const vertexStorageChannelTestContent = "new-api Vertex AI Storage channel test\n" + +newObjectName: func() string { + return ".new-api-channel-test/" + uuid.NewString() + "/test.txt" +} +``` + +使用最长 30 秒的独立 cleanup context，各步骤不重试并关闭 response body。写入使用 `uploadType=media&name=...`，读取使用 `alt=media`，删除失败必须附对象路径。 + +- [ ] **Step 4: 在计费前分流** + +```go +if channel.Type == constant.ChannelTypeVertexAi && + strings.HasPrefix(testModel, relayconstant.VertexStorageModelPrefix) { + err := testVertexStorageChannel(ctx, c, testModel, storageDeps) + return storageTestResult(startedAt, err) +} +``` + +分流点必须位于 RelayInfo、价格和日志逻辑之前。 + +- [ ] **Step 5: 验证并提交** + +```bash +gofmt -w controller/channel-test.go controller/vertex_storage_channel_probe.go \ + controller/vertex_storage_channel_probe_test.go +go test ./controller -run 'VertexStorage|ChannelTest' -count=1 +git add controller/channel-test.go controller/vertex_storage_channel_probe.go \ + controller/vertex_storage_channel_probe_test.go +git commit -m "feat: test Vertex storage buckets" +``` + +### Task 6: 统一前端测试入口并标识 Storage 项 + +**Files:** +- Modify: `web/src/features/channels/components/data-table-row-actions.tsx` +- Create: `web/src/features/channels/components/__tests__/channel-test-routing.test.tsx` +- Modify: `web/src/features/channels/components/dialogs/channel-test-dialog.tsx` +- Create: `web/src/features/channels/components/dialogs/__tests__/channel-test-storage.test.tsx` +- Modify: `web/scripts/add-missing-keys.mjs` +- Modify via i18n flow: `web/src/i18n/locales/{en,zh,zh-TW,fr,ja,ru,vi}.json` + +**Interfaces:** +- Consumes: `VERTEX_STORAGE_MODEL_PREFIX`。 +- Produces: 所有单渠道测试按钮统一执行 `setCurrentRow(channel); setOpen('test-channel')`。 + +- [ ] **Step 1: 写入入口一致性失败测试** + +表格 Gauge、卡片按钮、下拉菜单只打开 `test-channel`，不调用 `handleTestChannel`；卡片不得重复显示两个测试按钮。 + +- [ ] **Step 2: 实现统一入口** + +删除直接请求、`isTesting` 和 Loader 状态。Gauge 点击阻止行事件后调用 `handleTest()`；删除卡片额外 PlugZap 按钮。 + +- [ ] **Step 3: 写入 Storage 标识失败测试** + +```tsx +renderDialog({ models: 'gemini-2.5-pro,storage:gs:bucket-a' }) +expect(screen.getByText('storage:gs:bucket-a')).toBeInTheDocument() +expect( + screen.getByText('GCS bucket: writes, reads, and deletes a temporary object') +).toBeInTheDocument() +``` + +- [ ] **Step 4: 实现说明并补七语言** + +Storage 行通过 `model.startsWith(VERTEX_STORAGE_MODEL_PREFIX)` 识别，显示 badge 和 `t('GCS bucket: writes, reads, and deletes a temporary object')`；普通模型保持原样。 + +- [ ] **Step 5: 验证并提交** + +```bash +cd web +bun test src/features/channels/components/__tests__/channel-test-routing.test.tsx \ + src/features/channels/components/dialogs/__tests__/channel-test-storage.test.tsx +bun run typecheck +bunx oxlint -c .oxlintrc.json src/features/channels/components/data-table-row-actions.tsx \ + src/features/channels/components/dialogs/channel-test-dialog.tsx +git add web +git commit -m "feat: expose Vertex bucket tests in channel UI" +``` + +### Task 7: 全量验证、安全评审、单 Commit 与 PR + +**Files:** +- Review: 本计划涉及的全部文件。 +- Read: `.github/PULL_REQUEST_TEMPLATE.md`。 + +**Interfaces:** +- Produces: 一个通过验证的计划 Commit和按模板创建的 PR。 + +- [ ] **Step 1: 运行后端验证** + +```bash +go test ./relay/constant ./middleware ./relay/channel/vertex ./controller ./router -count=1 +go test ./... +go build ./... +``` + +- [ ] **Step 2: 运行前端验证** + +```bash +cd web +bun run i18n:sync +bun run typecheck +bun run lint +bun run build +``` + +- [ ] **Step 3: 安全与规范自检** + +```bash +rg -n 'io\.ReadAll|Authorization|storage\.googleapis\.com|/vertexai|storage:gs:' \ + controller/vertex_storage* relay/channel/vertex/storage_proxy.go \ + relay/constant/vertex_storage.go router/relay-router.go +git diff --check main...HEAD +git status --short +``` + +确认无任意 URL/Host 输入、无 Google Session URL 泄露、无计费调用、无敏感日志、无 `/v1/rawproxy/vertexai` 残留，对象名编码和 bucket 精确授权均有回归测试。 + +- [ ] **Step 4: squash 为一个计划 Commit** + +```bash +base=$(git merge-base HEAD main) +git reset --soft "$base" +git commit -m "feat: add Vertex AI storage integration" +``` + +随后重新运行定向测试和 `git diff --check main...HEAD`。 + +- [ ] **Step 5: 按模板创建 PR** + +```bash +git config user.name +git config user.email +git shortlog -sne --all | head -n 20 +sed -n '1,260p' .github/PULL_REQUEST_TEMPLATE.md +git push -u origin codex/vertexai-storage +``` + +PR 标题使用 `feat: add Vertex AI storage integration`。正文保留模板结构，说明 `/vertexai` 路由、bucket 授权、服务账号限制、流式传输、真实渠道探测、非计费行为和验证结果；若当前 Git 用户不是历史核心开发者，明确说明本变更由 AI 辅助完成。 diff --git a/docs/superpowers/specs/2026-08-09-vertex-ai-storage-and-channel-test-design.md b/docs/superpowers/specs/2026-08-09-vertex-ai-storage-and-channel-test-design.md new file mode 100644 index 000000000000..5c5b03282624 --- /dev/null +++ b/docs/superpowers/specs/2026-08-09-vertex-ai-storage-and-channel-test-design.md @@ -0,0 +1,253 @@ +# Vertex AI 文件存储与渠道测试设计 + +## 背景 + +Vertex AI Gemini 可以通过 `fileData.fileUri` 引用 Google Cloud Storage 对象，例如 `gs://example-bucket/docs/report.pdf`。当前项目已经支持 Vertex AI 推理渠道，但没有提供管理这些对象的受限代理接口，也无法在渠道测试中验证服务账号对存储桶的真实读写权限。 + +本功能参考 `aihub-new` 已有实现，按照当前 new-api 的路由、中间件、渠道表单和测试流程重新适配。功能包括 Vertex AI 存储桶配置、Google Cloud Storage 文件代理，以及存储桶写入/读取/删除测试。 + +本文中的“对象存储”特指 Google Cloud Storage（GCS，URI scheme 为 `gs://`），不引入腾讯云 COS、S3 或通用对象存储抽象。 + +## 目标 + +1. Vertex AI 渠道可以独立配置一个或多个 GCS bucket。 +2. API 客户端可以通过固定 `/vertexai` 前缀路由上传、列举、读取、下载和删除已授权 bucket 中的对象。 +3. 渠道测试可以真实验证服务账号对指定 bucket 的写入、读取和删除权限。 +4. 文件传输采用流式代理，不进入模型计费链路。 +5. 最终变更按项目规范整理为一个计划 Commit，并创建 Pull Request。 + +## 非目标 + +- 不实现 OpenAI `/v1/files` 兼容层。 +- 不新增文件记录、文件 ID、文件归属表或数据库迁移。 +- 不支持 bucket 创建、删除、IAM、ACL、复制、组合或重写接口。 +- 不支持目录前缀级授权；授权粒度固定为整个 bucket。 +- 不支持 Vertex AI API Key 模式访问 GCS。 +- 不新增 Google Cloud Storage SDK。 +- 不建立 S3、腾讯云 COS 或其他对象存储的通用抽象。 +- 不对文件上传、下载或存储桶测试收费。 + +## 渠道配置 + +仅 Vertex AI 渠道类型 41 显示“存储桶”字段。字段支持多值输入，但每一项必须是纯 bucket 名称，例如： + +```text +example-bucket +archive-bucket-01 +``` + +不接受 `gs://example-bucket`、`storage:gs:example-bucket`、`example-bucket/path` 或包含查询字符串、反斜杠、URL scheme 的值。 + +存储桶继续复用渠道现有 `models` 字段持久化： + +```text +storage:gs:example-bucket +``` + +前端加载渠道时，将所有 `storage:gs:` 项从普通模型列表中拆出并去掉前缀回显；保存时重新添加前缀，与普通模型合并、去重后写回 `models`。切换到非 Vertex AI 类型时不主动删除已有存储桶项，避免临时切换造成数据丢失。 + +新增用户界面文案必须使用 `useTranslation()` 和 `t('English key')`，并同步维护 `en`、`zh`、`zh-TW`、`fr`、`ja`、`ru`、`vi` 七种前端语言。 + +## 对外路由 + +采用供应商前缀 `/vertexai`，固定开放以下路径： + +| 方法 | 网关路径 | 上游语义 | +| --- | --- | --- | +| `POST` | `/vertexai/upload/storage/v1/b/:bucket/o` | 简单上传、multipart 上传或初始化 resumable 上传 | +| `PUT` | `/vertexai/upload/storage/v1/b/:bucket/o` | resumable 分片上传或完成上传 | +| `GET` | `/vertexai/storage/v1/b/:bucket/o` | 列举对象 | +| `GET` | `/vertexai/storage/v1/b/:bucket/o/*object` | 获取元数据；`alt=media` 时下载对象内容 | +| `DELETE` | `/vertexai/storage/v1/b/:bucket/o/*object` | 删除对象 | + +路由组依次执行： + +```text +RouteTag("relay") +→ SystemPerformanceCheck +→ TokenAuth +→ ModelRequestRateLimit +→ DistributeByChannelType(Vertex AI) +→ Controller 二次授权 +→ GCS 流式代理 +``` + +不得增加能够选择任意 Google API 路径、任意主机或任意 URL 的通配代理。 + +## 渠道分发与授权 + +中间件从路径参数读取 bucket，并构造分发模型： + +```text +storage:gs: +``` + +该值进入现有模型、用户组、Token、指定渠道和渠道选择流程，候选渠道同时被限制为 Vertex AI 类型 41。Controller 在发送上游请求前再次验证： + +1. bucket 是合法的纯 GCS bucket 名称。 +2. 所选渠道类型是 Vertex AI。 +3. 所选渠道的 `models` 精确包含 `storage:gs:`。 +4. 渠道使用服务账号 JSON 凭证，而不是 Vertex AI API Key。 + +双重校验用于防止未来路由、中间件或上下文变更绕过 bucket 授权。 + +## 上游鉴权与请求代理 + +上游主机固定为： + +```text +https://storage.googleapis.com +``` + +代理复用 Vertex AI 已有服务账号解析、JWT 签名、`cloud-platform` scope、Access Token 缓存和渠道 Proxy 配置。缓存继续按渠道 ID 和多 Key 索引隔离。 + +发送上游请求前必须丢弃客户端 `Authorization`、`Host` 和 hop-by-hop headers，并设置服务端获取的 GCS Bearer Token。允许透传内容、Range、条件请求、`Content-Range`、`X-Goog-Hash` 和 `X-Goog-Meta-*` 等对象操作相关头。 + +对象名来自 `*object`，允许包含目录形式的 `/`，但构造 GCS JSON API URL 时必须将完整对象名编码为单个 path segment，不能将对象名解释成额外的上游路由层级。 + +查询参数在固定主机和固定路径语义下透传，包括 `uploadType`、`name`、`alt`、`prefix`、`delimiter`、`pageToken`、generation 条件和 resumable session 参数。查询参数不得改变 bucket 或上游主机。 + +## 流式传输与 Resumable 上传 + +上传直接使用入站 `Request.Body` 构造上游请求；下载直接从 GCS response body 流式复制到客户端。不得使用 `io.ReadAll` 将完整文件载入内存。请求取消时使用同一 context 取消上游请求，并在所有路径关闭上游 response body。 + +GCS resumable 初始化成功后返回的 `Location` 不能原样暴露。网关必须将其改写为当前服务的： + +```text +/vertexai/upload/storage/v1/b/:bucket/o +``` + +并保留 `upload_id` 等 session 查询参数。绝对地址使用系统配置的服务地址构造，不信任客户端 `Host` 或转发头。后续每个 `PUT` 分片重新经过 Token 鉴权、限流、渠道分发和 bucket 二次授权。 + +## 响应与错误处理 + +GCS 返回的成功状态、4xx/5xx 状态、JSON 错误体、对象元数据、二进制内容和必要响应头原则上原样返回，同时过滤 hop-by-hop headers。 + +以下情况必须在访问 GCS 前返回本地错误： + +- bucket 缺失、非法或包含路径/URL 语义。 +- 没有可用的 Vertex AI 渠道配置目标 bucket。 +- Token、用户组或指定渠道策略无权使用 `storage:gs:`。 +- 所选渠道类型错误或未配置目标 bucket。 +- 渠道使用 Vertex AI API Key。 +- 服务账号 JSON 无法解析或 OAuth Token 获取失败。 +- 对象读取或删除路由缺少对象名。 +- resumable 初始化需要改写绝对地址，但系统服务地址未配置。 + +本地错误沿用项目现有错误响应结构。错误不得包含服务账号 JSON、私钥、Access Token、文件内容或敏感响应头。 + +## 存储桶渠道测试 + +`testChannel` 完成测试项选择和空格清理后，仅在以下条件同时满足时进入 Storage 测试分支： + +1. 渠道类型为 Vertex AI。 +2. 测试项以 `storage:gs:` 开头。 + +普通模型继续走现有推理测试流程。Storage 测试在前置校验通过后生成唯一临时对象： + +```text +.new-api-channel-test/<随机值>/test.txt +``` + +测试使用固定短文本作为内容，并按固定顺序各执行一次： + +1. 使用 GCS media upload 写入临时对象。 +2. 使用 `alt=media` 读取临时对象，并精确比较响应内容。 +3. 删除临时对象。 + +任一步失败仍继续执行剩余步骤，尤其必须尽力执行删除。只有写入成功、读取成功、内容一致且删除成功时测试才通过。删除失败时，响应必须提供临时对象路径供管理员手动清理，但不得暴露凭证。 + +测试直接复用同一套 Access Token 获取和 Storage Proxy，不通过服务公开地址发起 HTTP 回环请求，不自动重试，不计费，也不生成模型消费日志。 + +## 前端测试入口 + +渠道列表仪表按钮、卡片测试按钮和行操作“测试连接”统一打开现有 `ChannelTestDialog`，不再由仪表按钮直接测试默认模型。用户在弹窗中选择普通模型或 `storage:gs:` 测试项。 + +Storage 项需要显示明确的 GCS bucket 类型说明，使用户知道该测试会真实执行写入、读取和删除。单项测试与批量测试复用现有渠道测试 API 和结果区域；普通模型测试行为保持不变。 + +## 计费、日志与数据库 + +Storage Proxy 和 Storage 渠道测试均不得： + +- 查询模型价格； +- 执行 quota 预扣； +- 执行结算或退款； +- 将字节数作为计费乘数； +- 生成模型消费日志。 + +可以记录请求 ID、渠道 ID、bucket、HTTP 方法、上游状态码、耗时和安全处理后的错误类别，用于故障排查。不得记录文件内容、服务账号凭证或 Access Token。 + +本功能不修改数据库结构，继续兼容 SQLite、MySQL 5.7.8+ 和 PostgreSQL 9.6+。 + +## 测试策略 + +### 前端存储桶字段 + +- 测试 `storage:gs:` 拆分、规范化、合并和去重。 +- 测试仅类型 41 显示存储桶字段。 +- 测试多个 bucket 的添加、删除和 `models` 同步。 +- 测试非法 bucket、空值、完整前缀和包含路径的输入不会形成有效配置。 +- 测试切换渠道类型不会静默删除已有 bucket 配置。 + +### 分发与代理 + +- 测试 bucket 映射为 `storage:gs:`。 +- 测试类型限定、精确 bucket 匹配和指定渠道策略。 +- 测试固定生成 `storage.googleapis.com` URL，且对象名中的 `/` 正确编码。 +- 测试客户端不能覆盖主机和 Authorization。 +- 测试请求/响应必要头保留，hop-by-hop headers 移除。 +- 测试上传与下载 body 保持流式语义。 +- 测试 GCS 状态码、错误体、Range 和 `Content-Range` 保持一致。 +- 测试 resumable `Location` 改写，且系统服务地址缺失时不会泄露 Google Session URL。 +- 测试 Vertex AI API Key、非法 bucket 和未授权 bucket 在访问上游前失败。 + +### 渠道测试 + +- 测试 Storage 项进入专用分支，普通模型保持原流程。 +- 测试写入、读取、删除严格按顺序且各执行一次。 +- 测试任一步失败后仍执行剩余步骤。 +- 测试读取内容不一致和删除失败均判定失败。 +- 测试删除失败时返回临时对象路径。 +- 测试成功条件为三步成功且内容一致。 +- 测试不进入计费和模型消费日志路径。 + +### 前端测试弹窗 + +- 测试所有单渠道测试入口统一打开测试弹窗。 +- 测试 Storage 项可以被选择并显示类型说明。 +- 测试单项和批量测试能够展示 Storage 汇总错误。 +- 测试普通模型测试行为不回归。 + +### 验证 + +- 运行受影响 Go 单元测试，并使用 `testify/require` 与 `testify/assert` 编写新增或大幅重写的后端测试。 +- 运行受影响 Vitest/React Testing Library 测试。 +- 在 `web/` 执行 `bun run typecheck`。 +- 对涉及的前端文件执行 lint。 +- 执行 `bun run build` 生产构建 Smoke Test。 +- 执行相关 Go 包测试和根模块构建。 + +## Pull Request 要求 + +实现完成并通过评审后，将本任务产生的提交 squash 为一个计划 Commit，再推送分支并创建 PR。 + +创建 PR 前： + +1. 比较当前 `git config user.name`、`git config user.email` 与仓库历史核心开发者。 +2. 使用 `.github/PULL_REQUEST_TEMPLATE.md` 的结构撰写 PR 内容。 +3. 如果当前 Git 用户不是历史核心开发者，在 PR 正文明确说明代码由 AI 生成或 AI 辅助。 +4. PR 中说明 `/vertexai` 路由、安全边界、非计费行为、验证命令和测试结果。 + +## 验收标准 + +1. Vertex AI 渠道可以配置多个合法纯 bucket 名称，并正确持久化为 `storage:gs:`。 +2. 已授权客户端可以通过固定 `/vertexai` 路由上传、列举、读取、下载和删除对象。 +3. 未配置目标 bucket 的渠道、非 Vertex AI 渠道和 API Key 模式不能访问 GCS。 +4. 对象名可包含目录形式的 `/`，但不会造成上游路径或主机逃逸。 +5. 大文件上传和下载采用流式传输。 +6. Resumable 上传全过程都经过网关鉴权和 bucket 授权，Google Session URL 不会泄露。 +7. Storage 渠道测试真实执行写入、读取校验和删除，失败时仍尽力清理。 +8. Storage 文件操作和渠道测试均不计费、不生成模型消费日志。 +9. 所有新增文案完成七种语言翻译。 +10. Go 测试、前端测试、类型检查、lint、生产构建和相关模块构建通过。 +11. 最终变更整理为一个计划 Commit，并按仓库模板创建 PR。 diff --git a/docs/vertex-ai-storage.md b/docs/vertex-ai-storage.md new file mode 100644 index 000000000000..e027bd4efa78 --- /dev/null +++ b/docs/vertex-ai-storage.md @@ -0,0 +1,181 @@ +# 如何通过 Vertex AI 渠道管理 Cloud Storage 文件 + +## 目标与边界 + +本指南面向需要让 Vertex AI Gemini 通过 `fileData.fileUri` 使用 Google Cloud Storage（GCS）文件的接入工程师。完成配置后，你可以使用 new-api 的固定 `/vertexai` 路由上传、列举、读取、下载和删除指定 bucket 中的对象。 + +这些路由只代理 GCS JSON API 的五组固定操作，不提供任意 Google API、任意主机或任意 URL 的通配代理，也不负责创建 bucket、修改 IAM、管理 ACL、复制或重写对象。授权粒度是整个 bucket，不支持只授权某个对象名前缀。 + +## 前置 IAM 与服务账号 + +1. 在 Google Cloud 项目中创建供 Vertex AI 渠道使用的服务账号。 +2. 为服务账号授予调用 Vertex AI 所需的权限，例如 `roles/aiplatform.user`。 +3. 在目标 bucket 上授予与实际操作匹配的 GCS 权限。需要完整执行本指南的上传、下载、列举和删除操作时，可授予 `roles/storage.objectUser`；生产环境应按最小权限原则拆分只读或只写权限。 +4. 生成服务账号 JSON 密钥，并将完整 JSON 配置到 Vertex AI 渠道的 Key 中。 + +GCS 文件代理不支持 Vertex AI API Key 模式。不要把服务账号 JSON、私钥或网关 Token 写入源码、日志或公开文档。 + +## 配置渠道 bucket + +在管理后台编辑 Vertex AI 渠道，在“存储桶”字段中填写允许访问的纯 bucket 名称，例如： + +```text +example-bucket +archive-bucket-01 +``` + +每个值必须只有 bucket 名称。不要填写 `gs://example-bucket`、`storage:gs:example-bucket`、`example-bucket/path` 或带查询字符串的值。保存后，系统会在渠道 `models` 中以 `storage:gs:` 形式持久化授权；请求只能访问所选渠道精确配置的 bucket。 + +以下示例统一使用脱敏环境变量： + +```bash +export NEW_API_BASE_URL="https://api.example.com" +export NEW_API_TOKEN="" +export GCS_BUCKET="example-bucket" +export GCS_OBJECT="docs%2Freport.pdf" +``` + +## 使用五组固定路由 + +### 1. 使用 POST 上传对象 + +使用 media upload 将本地文件流式上传到指定对象名： + +```bash +curl --fail-with-body \ + -X POST \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + -H "Content-Type: application/pdf" \ + --data-binary @./report.pdf \ + "${NEW_API_BASE_URL}/vertexai/upload/storage/v1/b/${GCS_BUCKET}/o?uploadType=media&name=${GCS_OBJECT}" +``` + +此路由也支持 GCS 的 `multipart` 上传和 resumable 初始化。对象名中的 `/` 应编码为 `%2F`。 + +### 2. 使用 PUT 续传分块 + +先按“Resumable 上传”一节初始化会话，再对返回的网关 `Location` 发送分块： + +```bash +curl --fail-with-body \ + -X PUT \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + -H "Content-Type: application/octet-stream" \ + -H "Content-Range: bytes 0-1048575/2097152" \ + --data-binary @./chunk-000.bin \ + "${RESUMABLE_LOCATION}" +``` + +未完成时 GCS 通常返回 `308 Resume Incomplete`；最终分块完成后返回对象元数据。 + +### 3. 列举对象 + +使用 `prefix` 等 GCS 查询参数筛选列表： + +```bash +curl --fail-with-body \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + "${NEW_API_BASE_URL}/vertexai/storage/v1/b/${GCS_BUCKET}/o?prefix=docs%2F" +``` + +渠道授权仍以整个 bucket 为单位，`prefix` 只筛选响应，不会缩小授权范围。 + +### 4. 读取元数据或下载对象 + +省略 `alt=media` 时读取对象元数据： + +```bash +curl --fail-with-body \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + "${NEW_API_BASE_URL}/vertexai/storage/v1/b/${GCS_BUCKET}/o/${GCS_OBJECT}" +``` + +设置 `alt=media` 时流式下载对象内容： + +```bash +curl --fail-with-body \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + -o ./downloaded-report.pdf \ + "${NEW_API_BASE_URL}/vertexai/storage/v1/b/${GCS_BUCKET}/o/${GCS_OBJECT}?alt=media" +``` + +需要断点下载时可增加 `Range: bytes=0-1048575`。网关会保留 GCS 的状态码、`Content-Range`、`ETag` 和响应体。 + +### 5. 删除对象 + +删除前确认对象名和 bucket，删除操作不可由网关恢复： + +```bash +curl --fail-with-body \ + -X DELETE \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + "${NEW_API_BASE_URL}/vertexai/storage/v1/b/${GCS_BUCKET}/o/${GCS_OBJECT}" +``` + +## 在 `fileData.fileUri` 中引用对象 + +上传完成后，在 Vertex AI Gemini 请求中使用原始 GCS URI，而不是 `/vertexai` 网关下载地址： + +```json +{ + "contents": [ + { + "role": "user", + "parts": [ + { + "fileData": { + "mimeType": "application/pdf", + "fileUri": "gs://example-bucket/docs/report.pdf" + } + }, + { + "text": "请总结这份报告。" + } + ] + } + ] +} +``` + +执行推理的 Vertex AI 服务账号必须能够读取该对象。网关 Token 对 GCS 代理路由的访问权限不会替代 Google Cloud IAM。 + +## 使用 Resumable 上传 + +初始化 resumable 会话，并只从响应头读取 `Location`： + +```bash +RESUMABLE_LOCATION="$({ + curl --silent --show-error --fail-with-body \ + -X POST \ + -D - \ + -o /dev/null \ + -H "Authorization: Bearer ${NEW_API_TOKEN}" \ + -H "Content-Type: application/json; charset=UTF-8" \ + -H "X-Upload-Content-Type: application/octet-stream" \ + --data "{\"name\":\"docs/report.pdf\"}" \ + "${NEW_API_BASE_URL}/vertexai/upload/storage/v1/b/${GCS_BUCKET}/o?uploadType=resumable" +} | awk 'tolower($1) == "location:" { sub(/\r$/, "", $2); print $2 }')" +``` + +返回的 `Location` 必须以 `${NEW_API_BASE_URL}/vertexai/` 开头。网关会校验 Google 返回的 session URL，并使用系统配置的服务地址安全改写；改写失败时返回 `502`，不会把 `storage.googleapis.com` session URL 暴露给客户端。后续每个 `PUT` 都会重新执行 Token 鉴权、限流、渠道分发和 bucket 授权。 + +## 常见错误与非计费说明 + +| 现象 | 常见原因 | 排查方式 | +| --- | --- | --- | +| `400 invalid_bucket` | bucket 包含路径、scheme、查询字符串或完整配置前缀 | 只传纯 bucket 名称 | +| `400 unsupported_key_type` | 渠道使用 Vertex AI API Key | 改用服务账号 JSON | +| `400 object_required` | 读取或删除路由缺少对象名 | 对完整对象名做 URL 编码后放入路径 | +| `403 bucket_not_allowed` | 所选渠道未精确配置目标 bucket | 检查渠道“存储桶”配置及 Token/用户组权限 | +| GCS `403` 响应 | 服务账号缺少 bucket IAM 权限 | 检查 bucket IAM 和服务账号身份 | +| `502 access_token_failed` | 服务账号 JSON、私钥、代理或 Google OAuth 异常 | 检查渠道凭证和 Proxy 配置，不要在日志中输出私钥 | +| `502 invalid_resumable_location` | 系统服务地址缺失，或 Google 返回的 session URL 未通过安全校验 | 配置正确的公开服务地址后重新初始化会话 | + +这些 Storage Proxy 请求采用流式传输，不构造模型 RelayInfo，不执行 quota 预扣或结算，也不写模型消费日志。GCS 本身产生的存储、网络和请求费用仍由你的 Google Cloud 项目承担。 + +## 相关链接 + +- [Google Cloud Storage JSON API](https://cloud.google.com/storage/docs/json_api) +- [Cloud Storage IAM roles](https://cloud.google.com/storage/docs/access-control/iam-roles) +- [Resumable uploads](https://cloud.google.com/storage/docs/performing-resumable-uploads) +- [Vertex AI 文件输入](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/document-understanding) diff --git a/middleware/distributor.go b/middleware/distributor.go index 3f53aa350349..83fcf650ad78 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -31,7 +31,16 @@ type ModelRequest struct { } func Distribute() func(c *gin.Context) { + return distribute(0) +} + +func DistributeByChannelType(requiredChannelType int) func(c *gin.Context) { + return distribute(requiredChannelType) +} + +func distribute(requiredChannelType int) func(c *gin.Context) { return func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyRequiredChannelType, requiredChannelType) var channel *model.Channel channelId, ok := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId) modelRequest, shouldSelectChannel, err := getModelRequest(c) @@ -54,6 +63,14 @@ func Distribute() func(c *gin.Context) { abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled)) return } + if requiredChannelType != 0 && channel.Type != requiredChannelType { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) + return + } + if relayconstant.IsVertexStoragePath(c.Request.URL.Path) && !relayconstant.VertexStorageChannelSupports(channel.GetModels(), c.Param("bucket")) { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) + return + } } else { // Select a channel for the user // check token model mapping @@ -106,6 +123,7 @@ func Distribute() func(c *gin.Context) { affinityUsable := false preferred, err := model.CacheGetChannel(preferredChannelID) if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled && + (requiredChannelType == 0 || preferred.Type == requiredChannelType) && channelSupportsRequestPath(preferred, c.Request.URL.Path, modelRequest.Model) { if usingGroup == "auto" { userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) @@ -134,11 +152,12 @@ func Distribute() func(c *gin.Context) { if channel == nil { channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ - Ctx: c, - ModelName: modelRequest.Model, - TokenGroup: usingGroup, - RequestPath: c.Request.URL.Path, - Retry: common.GetPointer(0), + Ctx: c, + ModelName: modelRequest.Model, + TokenGroup: usingGroup, + RequestPath: c.Request.URL.Path, + Retry: common.GetPointer(0), + RequiredChannelType: requiredChannelType, }) if err != nil { showGroup := usingGroup @@ -254,7 +273,14 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { var modelRequest ModelRequest shouldSelectChannel := true var err error - if strings.Contains(c.Request.URL.Path, "/mj/") { + if relayconstant.IsVertexStoragePath(c.Request.URL.Path) { + modelName, err := relayconstant.VertexStorageModelName(c.Param("bucket")) + if err != nil { + return nil, false, err + } + modelRequest.Model = modelName + c.Set("relay_mode", relayconstant.RelayModeVertexStorage) + } else if strings.Contains(c.Request.URL.Path, "/mj/") { relayMode := relayconstant.Path2RelayModeMidjourney(c.Request.URL.Path) if relayMode == relayconstant.RelayModeMidjourneyTaskFetch || relayMode == relayconstant.RelayModeMidjourneyTaskFetchByCondition || @@ -446,6 +472,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode common.SetContextKey(c, constant.ContextKeyChannelId, channel.Id) common.SetContextKey(c, constant.ContextKeyChannelName, channel.Name) common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type) + common.SetContextKey(c, constant.ContextKeyChannelModels, channel.GetModels()) common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime) common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting()) common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings()) diff --git a/middleware/distributor_test.go b/middleware/distributor_test.go new file mode 100644 index 000000000000..bab8cf20300e --- /dev/null +++ b/middleware/distributor_test.go @@ -0,0 +1,125 @@ +package middleware + +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/i18n" + "github.com/QuantumNous/new-api/model" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupVertexStorageDistributorTest(t *testing.T) *gorm.DB { + t.Helper() + require.NoError(t, i18n.Init()) + 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 TestVertexStoragePathExtractsBucketAsModel(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", nil) + c.Params = gin.Params{{Key: "bucket", Value: "bucket-a"}} + + got, shouldSelect, err := getModelRequest(c) + + require.NoError(t, err) + assert.True(t, shouldSelect) + assert.Equal(t, "storage:gs:bucket-a", got.Model) + assert.Equal(t, relayconstant.RelayModeVertexStorage, c.GetInt("relay_mode")) +} + +func TestVertexStoragePathRejectsInvalidBucket(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", nil) + c.Params = gin.Params{{Key: "bucket", Value: "bucket-a/path"}} + + _, _, err := getModelRequest(c) + + require.Error(t, err) +} + +func TestVertexStoragePinnedChannelRequiresExactVertexBucket(t *testing.T) { + db := setupVertexStorageDistributorTest(t) + channels := []model.Channel{ + {Id: 5101, Name: "vertex-a", Type: constant.ChannelTypeVertexAi, Key: "vertex-a", Status: common.ChannelStatusEnabled, Group: "default", Models: "storage:gs:bucket-a"}, + {Id: 5102, Name: "vertex-b", Type: constant.ChannelTypeVertexAi, Key: "vertex-b", Status: common.ChannelStatusEnabled, Group: "default", Models: "storage:gs:bucket-b"}, + {Id: 5103, Name: "gemini-a", Type: constant.ChannelTypeGemini, Key: "gemini-a", Status: common.ChannelStatusEnabled, Group: "default", Models: "storage:gs:bucket-a"}, + } + require.NoError(t, db.Create(&channels).Error) + priority := int64(0) + for _, channel := range channels { + require.NoError(t, db.Create(&model.Ability{ + Group: channel.Group, Model: channel.Models, ChannelId: channel.Id, + Enabled: true, Priority: &priority, Weight: 100, + }).Error) + } + model.InitChannelCache() + + for _, testCase := range []struct { + name string + channelID string + wantStatus int + wantNext bool + }{ + {name: "matching Vertex bucket", channelID: "5101", wantStatus: http.StatusOK, wantNext: true}, + {name: "different Vertex bucket", channelID: "5102", wantStatus: http.StatusForbidden}, + {name: "Gemini channel", channelID: "5103", wantStatus: http.StatusForbidden}, + } { + t.Run(testCase.name, func(t *testing.T) { + downstreamRan := false + engine := gin.New() + engine.Use(func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyUserGroup, "default") + common.SetContextKey(c, constant.ContextKeyUsingGroup, "default") + common.SetContextKey(c, constant.ContextKeyTokenSpecificChannelId, testCase.channelID) + }) + engine.GET(relayconstant.VertexStorageListRoute, + DistributeByChannelType(constant.ChannelTypeVertexAi), + func(c *gin.Context) { + downstreamRan = true + models := common.GetContextKeyStringSlice(c, constant.ContextKeyChannelModels) + assert.Equal(t, []string{"storage:gs:bucket-a"}, models) + c.Status(http.StatusOK) + }, + ) + + request := httptest.NewRequest(http.MethodGet, "/vertexai/storage/v1/b/bucket-a/o", nil) + response := httptest.NewRecorder() + engine.ServeHTTP(response, request) + + assert.Equal(t, testCase.wantStatus, response.Code) + assert.Equal(t, testCase.wantNext, downstreamRan) + }) + } +} diff --git a/model/ability.go b/model/ability.go index d950a6adbfc4..4822acbd0a72 100644 --- a/model/ability.go +++ b/model/ability.go @@ -60,12 +60,18 @@ func GetAllEnableAbilities() []Ability { return abilities } -func getPriority(group string, model string, retry int) (int, error) { - +func getPriority(group string, model string, retry int, requiredChannelType int) (int, error) { var priorities []int - err := DB.Model(&Ability{}). + priorityQuery := DB.Model(&Ability{}). Select("DISTINCT(priority)"). - Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true). + Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true) + if requiredChannelType != 0 { + enabledChannelIDs := DB.Model(&Channel{}). + Select("id"). + Where("status = ? and type = ?", common.ChannelStatusEnabled, requiredChannelType) + priorityQuery = priorityQuery.Where("channel_id IN (?)", enabledChannelIDs) + } + err := priorityQuery. Order("priority DESC"). // 按优先级降序排序 Pluck("priority", &priorities).Error // Pluck用于将查询的结果直接扫描到一个切片中 @@ -90,15 +96,29 @@ func getPriority(group string, model string, retry int) (int, error) { return priorityToUse, nil } -func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { - maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true) +func getChannelQuery(group string, model string, retry int, requiredChannelType int) (*gorm.DB, error) { + maxPrioritySubQuery := DB.Model(&Ability{}). + Select("MAX(priority)"). + Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true) + var enabledChannelIDs *gorm.DB + if requiredChannelType != 0 { + enabledChannelIDs = DB.Model(&Channel{}). + Select("id"). + Where("status = ? and type = ?", common.ChannelStatusEnabled, requiredChannelType) + maxPrioritySubQuery = maxPrioritySubQuery.Where("channel_id IN (?)", enabledChannelIDs) + } channelQuery := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = (?)", group, model, true, maxPrioritySubQuery) + if enabledChannelIDs != nil { + channelQuery = channelQuery.Where("channel_id IN (?)", enabledChannelIDs) + } if retry != 0 { - priority, err := getPriority(group, model, retry) + priority, err := getPriority(group, model, retry, requiredChannelType) if err != nil { return nil, err - } else { - channelQuery = DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = ?", group, model, true, priority) + } + channelQuery = DB.Where(commonGroupCol+" = ? and model = ? and enabled = ? and priority = ?", group, model, true, priority) + if enabledChannelIDs != nil { + channelQuery = channelQuery.Where("channel_id IN (?)", enabledChannelIDs) } } @@ -106,10 +126,14 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { } func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) { + return GetChannelByType(group, model, retry, requestPath, 0) +} + +func GetChannelByType(group string, model string, retry int, requestPath string, requiredChannelType int) (*Channel, error) { var abilities []Ability var err error = nil - channelQuery, err := getChannelQuery(group, model, retry) + channelQuery, err := getChannelQuery(group, model, retry, requiredChannelType) if err != nil { return nil, err } diff --git a/model/channel_cache.go b/model/channel_cache.go index 86c594384d50..8e8fd68728b4 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -112,21 +112,30 @@ func SyncChannelCache(frequency int) { } func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) { + return GetRandomSatisfiedChannelByType(group, model, retry, requestPath, 0) +} + +func GetRandomSatisfiedChannelByType(group string, model string, retry int, requestPath string, requiredChannelType int) (*Channel, error) { // if memory cache is disabled, get channel directly from database if !common.MemoryCacheEnabled { - return GetChannel(group, model, retry, requestPath) + return GetChannelByType(group, model, retry, requestPath, requiredChannelType) } + return getRandomSatisfiedChannelFromCache(group, model, retry, requestPath, requiredChannelType) +} +func getRandomSatisfiedChannelFromCache(group string, model string, retry int, requestPath string, requiredChannelType int) (*Channel, error) { channelSyncLock.RLock() defer channelSyncLock.RUnlock() // First, try to find channels with the exact model name. channels := filterChannelsByRequestPathAndModel(group2model2channels[group][model], requestPath, model) + channels = filterChannelsByType(channels, requiredChannelType) // If no channels found, try to find channels with the normalized model name. if len(channels) == 0 { normalizedModel := ratio_setting.FormatMatchingModelName(model) channels = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, model) + channels = filterChannelsByType(channels, requiredChannelType) } if len(channels) == 0 { @@ -208,6 +217,20 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat return nil, errors.New("channel not found") } +func filterChannelsByType(channelIDs []int, requiredChannelType int) []int { + if requiredChannelType == 0 { + return channelIDs + } + filtered := make([]int, 0, len(channelIDs)) + for _, channelID := range channelIDs { + channel, ok := channelsIDM[channelID] + if ok && channel.Type == requiredChannelType { + filtered = append(filtered, channelID) + } + } + return filtered +} + // filterChannelsByRequestPathAndModel restricts candidates by request path and // model. Only Advanced Custom (type 58) channels are path-checked: they are kept // only when one of their configured routes matches requestPath and model. All diff --git a/model/channel_type_selection_test.go b/model/channel_type_selection_test.go new file mode 100644 index 000000000000..007d8ad0f448 --- /dev/null +++ b/model/channel_type_selection_test.go @@ -0,0 +1,78 @@ +package model + +import ( + "fmt" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestRequiredChannelTypeFiltersDatabaseCandidates(t *testing.T) { + originalDB := DB + originalGroupCol := commonGroupCol + 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(&Channel{}, &Ability{})) + DB = db + commonGroupCol = "`group`" + common.MemoryCacheEnabled = false + t.Cleanup(func() { + DB = originalDB + commonGroupCol = originalGroupCol + common.MemoryCacheEnabled = originalMemoryCacheEnabled + sqlDB, sqlErr := db.DB() + if sqlErr == nil { + require.NoError(t, sqlDB.Close()) + } + }) + + highPriority := int64(100) + lowPriority := int64(0) + weight := uint(100) + channels := []*Channel{ + // An enabled ability can temporarily outlive a disabled channel. The + // legacy untyped database selector includes that ability, so type 0 must + // preserve its selection semantics. + {Id: 4101, Type: constant.ChannelTypeGemini, Key: "gemini", Status: common.ChannelStatusManuallyDisabled, Name: "gemini", Models: "storage:gs:bucket-a", Group: "default", Priority: &highPriority, Weight: &weight}, + {Id: 4102, Type: constant.ChannelTypeVertexAi, Key: "vertex", Status: common.ChannelStatusEnabled, Name: "vertex", Models: "storage:gs:bucket-a", Group: "default", Priority: &lowPriority, Weight: &weight}, + } + for _, channel := range channels { + require.NoError(t, db.Create(channel).Error) + require.NoError(t, db.Create(&Ability{ + Group: channel.Group, Model: channel.Models, ChannelId: channel.Id, + Enabled: true, Priority: channel.Priority, Weight: weight, + }).Error) + } + + t.Run("limits candidates to the requested channel type", func(t *testing.T) { + selected, err := GetChannelByType( + "default", "storage:gs:bucket-a", 0, + "/vertexai/storage/v1/b/bucket-a/o", constant.ChannelTypeVertexAi, + ) + + require.NoError(t, err) + require.NotNil(t, selected) + assert.Equal(t, 4102, selected.Id) + assert.Equal(t, constant.ChannelTypeVertexAi, selected.Type) + }) + + t.Run("zero preserves the unrestricted candidate set", func(t *testing.T) { + selected, err := GetChannelByType( + "default", "storage:gs:bucket-a", 0, + "/vertexai/storage/v1/b/bucket-a/o", 0, + ) + + require.NoError(t, err) + require.NotNil(t, selected) + assert.Equal(t, 4101, selected.Id) + assert.Equal(t, constant.ChannelTypeGemini, selected.Type) + }) +} diff --git a/relay/channel/vertex/service_account.go b/relay/channel/vertex/service_account.go index 6d5fee0d94fc..cd9f16ccbf2a 100644 --- a/relay/channel/vertex/service_account.go +++ b/relay/channel/vertex/service_account.go @@ -3,21 +3,20 @@ package vertex import ( "crypto/rsa" "crypto/x509" - "encoding/json" "encoding/pem" "errors" + "fmt" "net/http" "net/url" "strings" + "time" + "github.com/QuantumNous/new-api/common" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/service" "github.com/bytedance/gopkg/cache/asynccache" "github.com/golang-jwt/jwt/v5" - - "fmt" - "time" ) type Credentials struct { @@ -37,32 +36,59 @@ var Cache = asynccache.NewAsyncCache(asynccache.Options{ }, }) -func getAccessToken(a *Adaptor, info *relaycommon.RelayInfo) (string, error) { - var cacheKey string - if info.ChannelIsMultiKey { - cacheKey = fmt.Sprintf("access-token-%d-%d", info.ChannelId, info.ChannelMultiKeyIndex) - } else { - cacheKey = fmt.Sprintf("access-token-%d", info.ChannelId) +type CachedAccessTokenRequest struct { + ChannelID int + ChannelIsMultiKey bool + ChannelMultiKeyIndex int + Credentials Credentials + Proxy string +} + +func AcquireCachedAccessToken(input CachedAccessTokenRequest) (string, error) { + return acquireCachedAccessToken(input, func(signedJWT string) (string, error) { + return exchangeJwtForAccessTokenWithProxy(signedJWT, input.Proxy) + }) +} + +func acquireCachedAccessToken(input CachedAccessTokenRequest, exchange func(string) (string, error)) (string, error) { + cacheKey := fmt.Sprintf("access-token-%d", input.ChannelID) + if input.ChannelIsMultiKey { + cacheKey = fmt.Sprintf("access-token-%d-%d", input.ChannelID, input.ChannelMultiKeyIndex) } - val, err := Cache.Get(cacheKey) - if err == nil { - return val.(string), nil + if value, err := Cache.Get(cacheKey); err == nil { + if token, ok := value.(string); ok && token != "" { + return token, nil + } } - signedJWT, err := createSignedJWT(a.AccountCredentials.ClientEmail, a.AccountCredentials.PrivateKey) + signedJWT, err := createSignedJWT(input.Credentials.ClientEmail, input.Credentials.PrivateKey) if err != nil { return "", fmt.Errorf("failed to create signed JWT: %w", err) } - newToken, err := exchangeJwtForAccessToken(signedJWT, info) + newToken, err := exchange(signedJWT) if err != nil { return "", fmt.Errorf("failed to exchange JWT for access token: %w", err) } - if err := Cache.SetDefault(cacheKey, newToken); err { - return newToken, nil + value := Cache.GetOrSet(cacheKey, newToken) + if token, ok := value.(string); ok && token != "" { + return token, nil } return newToken, nil } +func getAccessToken(a *Adaptor, info *relaycommon.RelayInfo) (string, error) { + input := CachedAccessTokenRequest{ + ChannelID: info.ChannelId, + ChannelIsMultiKey: info.ChannelIsMultiKey, + ChannelMultiKeyIndex: info.ChannelMultiKeyIndex, + Credentials: a.AccountCredentials, + Proxy: info.ChannelSetting.Proxy, + } + return acquireCachedAccessToken(input, func(signedJWT string) (string, error) { + return exchangeJwtForAccessToken(signedJWT, info) + }) +} + func createSignedJWT(email, privateKeyPEM string) (string, error) { privateKeyPEM = strings.ReplaceAll(privateKeyPEM, "-----BEGIN PRIVATE KEY-----", "") @@ -123,7 +149,7 @@ func exchangeJwtForAccessToken(signedJWT string, info *relaycommon.RelayInfo) (s defer resp.Body.Close() var result map[string]interface{} - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + if err := common.DecodeJson(resp.Body, &result); err != nil { return "", err } @@ -166,7 +192,7 @@ func exchangeJwtForAccessTokenWithProxy(signedJWT string, proxy string) (string, defer resp.Body.Close() var result map[string]interface{} - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + if err := common.DecodeJson(resp.Body, &result); err != nil { return "", err } diff --git a/relay/channel/vertex/service_account_test.go b/relay/channel/vertex/service_account_test.go new file mode 100644 index 000000000000..4f12b7580dd6 --- /dev/null +++ b/relay/channel/vertex/service_account_test.go @@ -0,0 +1,63 @@ +package vertex + +import ( + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "encoding/pem" + "errors" + "testing" + "time" + + "github.com/bytedance/gopkg/cache/asynccache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAcquireCachedAccessTokenReusesOnlyMatchingMultiKeyIndex(t *testing.T) { + originalCache := Cache + testCache := asynccache.NewAsyncCache(asynccache.Options{ + RefreshDuration: time.Hour, + Fetcher: func(string) (interface{}, error) { + return nil, errors.New("not found") + }, + }) + Cache = testCache + t.Cleanup(func() { + Cache = originalCache + testCache.Close() + }) + + privateKey, err := rsa.GenerateKey(rand.Reader, 1024) + require.NoError(t, err) + privateKeyDER, err := x509.MarshalPKCS8PrivateKey(privateKey) + require.NoError(t, err) + + input := CachedAccessTokenRequest{ + ChannelID: 42, + ChannelIsMultiKey: true, + ChannelMultiKeyIndex: 3, + Credentials: Credentials{ + ClientEmail: "vertex@example.com", + PrivateKey: string(pem.EncodeToMemory(&pem.Block{ + Type: "PRIVATE KEY", + Bytes: privateKeyDER, + })), + }, + Proxy: "://", + } + first, err := acquireCachedAccessToken(input, func(string) (string, error) { + return "cached-token", nil + }) + require.NoError(t, err) + second, err := AcquireCachedAccessToken(input) + require.NoError(t, err) + assert.Equal(t, "cached-token", first) + assert.Equal(t, first, second) + + input.ChannelMultiKeyIndex = 4 + input.Credentials = Credentials{} + _, err = AcquireCachedAccessToken(input) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to create signed JWT") +} diff --git a/relay/channel/vertex/storage_proxy.go b/relay/channel/vertex/storage_proxy.go new file mode 100644 index 000000000000..605cb0240b26 --- /dev/null +++ b/relay/channel/vertex/storage_proxy.go @@ -0,0 +1,207 @@ +package vertex + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/service" +) + +const vertexStorageHost = "storage.googleapis.com" + +type StorageOperation uint8 + +const ( + StorageOperationUpload StorageOperation = iota + 1 + StorageOperationList + StorageOperationGet + StorageOperationDelete +) + +type StorageProxyRequest struct { + Operation StorageOperation + Method string + Bucket string + Object string + RawQuery string + Header http.Header + Body io.Reader + ContentLength int64 + AccessToken string + Proxy string +} + +func buildStorageRequest(ctx context.Context, input StorageProxyRequest) (*http.Request, error) { + if strings.TrimSpace(input.AccessToken) == "" { + return nil, errors.New("Vertex storage access token is required") + } + bucket, err := relayconstant.NormalizeVertexStorageBucket(input.Bucket) + if err != nil { + return nil, err + } + + escapedBucket := url.PathEscape(bucket) + var escapedPath string + switch input.Operation { + case StorageOperationUpload: + if input.Method != http.MethodPost && input.Method != http.MethodPut { + return nil, fmt.Errorf("unsupported Vertex storage upload method %q", input.Method) + } + escapedPath = "/upload/storage/v1/b/" + escapedBucket + "/o" + case StorageOperationList: + if input.Method != http.MethodGet { + return nil, fmt.Errorf("unsupported Vertex storage list method %q", input.Method) + } + escapedPath = "/storage/v1/b/" + escapedBucket + "/o" + case StorageOperationGet, StorageOperationDelete: + expectedMethod := http.MethodGet + if input.Operation == StorageOperationDelete { + expectedMethod = http.MethodDelete + } + if input.Method != expectedMethod { + return nil, fmt.Errorf("unsupported Vertex storage object method %q", input.Method) + } + if err := relayconstant.ValidateVertexStorageObjectName(input.Object); err != nil { + return nil, err + } + escapedPath = "/storage/v1/b/" + escapedBucket + "/o/" + url.PathEscape(input.Object) + default: + return nil, errors.New("unsupported Vertex storage operation") + } + + path, err := url.PathUnescape(escapedPath) + if err != nil { + return nil, fmt.Errorf("build Vertex storage path: %w", err) + } + target := &url.URL{ + Scheme: "https", + Host: vertexStorageHost, + Path: path, + RawPath: escapedPath, + RawQuery: input.RawQuery, + } + request, err := http.NewRequestWithContext(ctx, input.Method, target.String(), input.Body) + if err != nil { + return nil, err + } + request.Header = sanitizeStorageRequestHeader(input.Header) + request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(input.AccessToken)) + if input.ContentLength >= 0 { + request.ContentLength = input.ContentLength + } + return request, nil +} + +func sanitizeStorageRequestHeader(source http.Header) http.Header { + return sanitizeStorageHeader(source, true) +} + +func SanitizeStorageResponseHeader(source http.Header) http.Header { + result := sanitizeStorageHeader(source, false) + result.Set("Cache-Control", "private, no-store") + result.Del("Expires") + result.Del("Age") + return result +} + +func sanitizeStorageHeader(source http.Header, request bool) http.Header { + result := source.Clone() + if result == nil { + result = make(http.Header) + } + remove := map[string]struct{}{ + "connection": {}, + "keep-alive": {}, + "proxy-authenticate": {}, + "proxy-authorization": {}, + "proxy-connection": {}, + "te": {}, + "trailer": {}, + "transfer-encoding": {}, + "upgrade": {}, + } + if request { + remove["authorization"] = struct{}{} + remove["host"] = struct{}{} + remove["content-length"] = struct{}{} + remove["cookie"] = struct{}{} + remove["x-api-key"] = struct{}{} + remove["x-goog-api-key"] = struct{}{} + } + for name, values := range source { + if !strings.EqualFold(name, "Connection") { + continue + } + for _, value := range values { + for _, nominated := range strings.Split(value, ",") { + if nominated = strings.ToLower(strings.TrimSpace(nominated)); nominated != "" { + remove[nominated] = struct{}{} + } + } + } + } + for name := range result { + if _, ok := remove[strings.ToLower(name)]; ok { + delete(result, name) + } + } + return result +} + +func RewriteStorageResumableLocation(location, gatewayBaseURL, bucket string) (string, error) { + normalizedBucket, err := relayconstant.NormalizeVertexStorageBucket(bucket) + if err != nil { + return "", err + } + upstream, err := url.Parse(location) + if err != nil || upstream.Scheme != "https" || upstream.Host != vertexStorageHost || upstream.User != nil || upstream.Fragment != "" { + return "", errors.New("invalid Vertex storage resumable location") + } + expectedPath := "/upload/storage/v1/b/" + url.PathEscape(normalizedBucket) + "/o" + if upstream.EscapedPath() != expectedPath { + return "", errors.New("Vertex storage resumable location bucket does not match") + } + + gateway, err := url.Parse(strings.TrimSpace(gatewayBaseURL)) + if err != nil || (gateway.Scheme != "http" && gateway.Scheme != "https") || gateway.Host == "" || gateway.User != nil { + return "", errors.New("invalid gateway base URL") + } + escapedGatewayPath := relayconstant.VertexStorageRoutePrefix + expectedPath + gatewayPath, err := url.PathUnescape(escapedGatewayPath) + if err != nil { + return "", fmt.Errorf("build Vertex storage resumable gateway path: %w", err) + } + gateway.Path = gatewayPath + gateway.RawPath = escapedGatewayPath + gateway.RawQuery = upstream.RawQuery + gateway.ForceQuery = upstream.ForceQuery + gateway.Fragment = "" + return gateway.String(), nil +} + +func DoStorageProxy(ctx context.Context, input StorageProxyRequest) (*http.Response, error) { + request, err := buildStorageRequest(ctx, input) + if err != nil { + return nil, err + } + client, err := service.GetHttpClientWithProxy(input.Proxy) + if err != nil { + return nil, fmt.Errorf("new proxy http client failed: %w", err) + } + storageClient := *client + storageClient.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { + return http.ErrUseLastResponse + } + response, err := storageClient.Do(request) + if err != nil { + return nil, err + } + response.Header = SanitizeStorageResponseHeader(response.Header) + return response, nil +} diff --git a/relay/channel/vertex/storage_proxy_test.go b/relay/channel/vertex/storage_proxy_test.go new file mode 100644 index 000000000000..f42a62ad9bc9 --- /dev/null +++ b/relay/channel/vertex/storage_proxy_test.go @@ -0,0 +1,315 @@ +package vertex + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "net/url" + "reflect" + "sync/atomic" + "testing" + + "github.com/QuantumNous/new-api/service" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBuildStorageRequestUsesFixedGoogleHostAndEscapesObjectAsOneSegment(t *testing.T) { + req, err := buildStorageRequest(context.Background(), StorageProxyRequest{ + Operation: StorageOperationGet, + Method: http.MethodGet, + Bucket: "bucket-a", + Object: "docs/a b.pdf", + RawQuery: "alt=media", + AccessToken: "secret", + }) + require.NoError(t, err) + assert.Equal(t, "https", req.URL.Scheme) + assert.Equal(t, "storage.googleapis.com", req.URL.Host) + assert.Contains(t, req.URL.EscapedPath(), "docs%2Fa%20b.pdf") + assert.Equal(t, "alt=media", req.URL.RawQuery) + assert.Equal(t, "Bearer secret", req.Header.Get("Authorization")) +} + +func TestBuildStorageRequestRejectsDotSegmentObjects(t *testing.T) { + for _, object := range []string{".", "..", "folder/./file.txt", "folder/../file.txt"} { + t.Run(object, func(t *testing.T) { + _, err := buildStorageRequest(context.Background(), StorageProxyRequest{ + Operation: StorageOperationGet, + Method: http.MethodGet, + Bucket: "bucket-a", + Object: object, + AccessToken: "secret", + }) + require.Error(t, err) + }) + } +} + +func TestBuildStorageRequestFiltersClientCredentialsAndHopByHopHeaders(t *testing.T) { + header := http.Header{ + "Authorization": {"Bearer client-secret"}, + "Host": {"attacker.example.com"}, + "Connection": {"keep-alive, X-Remove-Me"}, + "Keep-Alive": {"timeout=5"}, + "Proxy-Authenticate": {"Basic"}, + "Proxy-Authorization": {"Basic secret"}, + "Proxy-Connection": {"keep-alive"}, + "Te": {"trailers"}, + "Trailer": {"X-Trailer"}, + "Transfer-Encoding": {"chunked"}, + "Upgrade": {"websocket"}, + "X-Remove-Me": {"connection-scoped"}, + "Cookie": {"session=secret"}, + "X-Api-Key": {"api-secret"}, + "X-Goog-Api-Key": {"google-secret"}, + "Content-Type": {"application/pdf"}, + "Content-Range": {"bytes 0-9/10"}, + "Range": {"bytes=0-9"}, + "If-Match": {"etag-a"}, + "If-None-Match": {"etag-b"}, + "X-Goog-Meta-Owner": {"alice"}, + } + + req, err := buildStorageRequest(context.Background(), StorageProxyRequest{ + Operation: StorageOperationUpload, + Method: http.MethodPost, + Bucket: "bucket-a", + RawQuery: "uploadType=resumable&name=docs%2Fa.pdf", + Header: header, + AccessToken: "upstream-token", + }) + require.NoError(t, err) + + assert.Equal(t, "Bearer upstream-token", req.Header.Get("Authorization")) + for _, name := range []string{ + "Host", "Connection", "Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization", + "Proxy-Connection", "Te", "Trailer", "Transfer-Encoding", "Upgrade", "X-Remove-Me", + "Cookie", "X-Api-Key", "X-Goog-Api-Key", + } { + assert.Empty(t, req.Header.Values(name), name) + } + assert.Equal(t, "application/pdf", req.Header.Get("Content-Type")) + assert.Equal(t, "bytes 0-9/10", req.Header.Get("Content-Range")) + assert.Equal(t, "bytes=0-9", req.Header.Get("Range")) + assert.Equal(t, "etag-a", req.Header.Get("If-Match")) + assert.Equal(t, "etag-b", req.Header.Get("If-None-Match")) + assert.Equal(t, "alice", req.Header.Get("X-Goog-Meta-Owner")) + assert.Equal(t, "Bearer client-secret", header.Get("Authorization"), "source header must not be mutated") +} + +func TestSanitizeStorageResponseHeaderRemovesHopByHopHeaders(t *testing.T) { + header := http.Header{ + "Connection": {"X-Internal"}, + "X-Internal": {"secret"}, + "Transfer-Encoding": {"chunked"}, + "Cache-Control": {"public, max-age=3600"}, + "Expires": {"Wed, 12 Aug 2026 00:00:00 GMT"}, + "Age": {"120"}, + "Content-Type": {"application/octet-stream"}, + "Content-Range": {"bytes 0-9/10"}, + "Etag": {"etag-a"}, + "Location": {"https://storage.googleapis.com/upload/storage/v1/b/bucket-a/o?upload_id=abc"}, + } + + got := SanitizeStorageResponseHeader(header) + + assert.Empty(t, got.Values("Connection")) + assert.Empty(t, got.Values("X-Internal")) + assert.Empty(t, got.Values("Transfer-Encoding")) + assert.Equal(t, "private, no-store", got.Get("Cache-Control")) + assert.Empty(t, got.Values("Expires")) + assert.Empty(t, got.Values("Age")) + assert.Equal(t, "application/octet-stream", got.Get("Content-Type")) + assert.Equal(t, "bytes 0-9/10", got.Get("Content-Range")) + assert.Equal(t, "etag-a", got.Get("Etag")) + assert.Equal(t, header.Get("Location"), got.Get("Location")) + assert.Equal(t, "X-Internal", header.Get("Connection"), "source header must not be mutated") + assert.Equal(t, "public, max-age=3600", header.Get("Cache-Control"), "source header must not be mutated") +} + +func TestRewriteStorageResumableLocation(t *testing.T) { + got, err := RewriteStorageResumableLocation( + "https://storage.googleapis.com/upload/storage/v1/b/bucket-a/o?upload_id=abc", + "https://gateway.example.com", + "bucket-a", + ) + require.NoError(t, err) + assert.Equal(t, "https://gateway.example.com/vertexai/upload/storage/v1/b/bucket-a/o?upload_id=abc", got) +} + +func TestRewriteStorageResumableLocationRejectsUnsafeInput(t *testing.T) { + tests := []struct { + name string + location string + gatewayURL string + bucket string + }{ + { + name: "empty server address", + location: "https://storage.googleapis.com/upload/storage/v1/b/bucket-a/o?upload_id=abc", + gatewayURL: "", + bucket: "bucket-a", + }, + { + name: "non google location", + location: "https://attacker.example.com/upload/storage/v1/b/bucket-a/o?upload_id=abc", + gatewayURL: "https://gateway.example.com", + bucket: "bucket-a", + }, + { + name: "google lookalike location", + location: "https://storage.googleapis.com.attacker.example.com/upload/storage/v1/b/bucket-a/o?upload_id=abc", + gatewayURL: "https://gateway.example.com", + bucket: "bucket-a", + }, + { + name: "bucket mismatch", + location: "https://storage.googleapis.com/upload/storage/v1/b/bucket-b/o?upload_id=abc", + gatewayURL: "https://gateway.example.com", + bucket: "bucket-a", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := RewriteStorageResumableLocation(tt.location, tt.gatewayURL, tt.bucket) + require.Error(t, err) + }) + } +} + +type storageCountingReader struct { + reads int +} + +func (r *storageCountingReader) Read([]byte) (int, error) { + r.reads++ + return 0, io.EOF +} + +func TestDoStorageProxyDoesNotBufferRequestBody(t *testing.T) { + body := &storageCountingReader{} + _, err := DoStorageProxy(context.Background(), StorageProxyRequest{ + Operation: StorageOperationUpload, + Method: http.MethodPost, + Bucket: "bucket-a", + RawQuery: "uploadType=media&name=a.txt", + Body: body, + ContentLength: int64(len("contents")), + AccessToken: "secret", + Proxy: "://", + }) + require.Error(t, err) + assert.Zero(t, body.reads) +} + +type storageHostRoutingTransport struct { + upstream *url.URL + crossHostTarget *url.URL + base http.RoundTripper +} + +func (t *storageHostRoutingTransport) RoundTrip(request *http.Request) (*http.Response, error) { + var target *url.URL + switch request.URL.Host { + case vertexStorageHost: + target = t.upstream + case "8.8.8.8": + target = t.crossHostTarget + default: + return t.base.RoundTrip(request) + } + routedRequest := request.Clone(request.Context()) + routedURL := *request.URL + routedURL.Scheme = target.Scheme + routedURL.Host = target.Host + routedRequest.URL = &routedURL + return t.base.RoundTrip(routedRequest) +} + +func TestDoStorageProxyReturnsRedirectWithoutFollowing(t *testing.T) { + service.InitHttpClient() + sharedClient := service.GetHttpClient() + require.NotNil(t, sharedClient) + require.NotNil(t, sharedClient.Transport) + require.NotNil(t, sharedClient.CheckRedirect) + originalRedirectPolicy := reflect.ValueOf(sharedClient.CheckRedirect).Pointer() + + tests := []struct { + name string + redirectLocation string + }{ + { + name: "same host", + redirectLocation: "/redirect-target", + }, + { + name: "cross host", + redirectLocation: "http://8.8.8.8/redirect-target", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var crossHostRequests atomic.Int32 + crossHostTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + crossHostRequests.Add(1) + w.WriteHeader(http.StatusTeapot) + })) + defer crossHostTarget.Close() + + const responseBody = "storage redirect" + var sourceRequests atomic.Int32 + var sourceRedirects atomic.Int32 + source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { + sourceRequests.Add(1) + if request.URL.Path == "/redirect-target" { + sourceRedirects.Add(1) + w.WriteHeader(http.StatusTeapot) + return + } + w.Header().Set("Location", tt.redirectLocation) + w.WriteHeader(http.StatusFound) + _, _ = io.WriteString(w, responseBody) + })) + defer source.Close() + + upstreamURL, err := url.Parse(source.URL) + require.NoError(t, err) + crossHostTargetURL, err := url.Parse(crossHostTarget.URL) + require.NoError(t, err) + originalTransport := sharedClient.Transport + sharedClient.Transport = &storageHostRoutingTransport{ + upstream: upstreamURL, + crossHostTarget: crossHostTargetURL, + base: originalTransport, + } + t.Cleanup(func() { + sharedClient.Transport = originalTransport + }) + + response, err := DoStorageProxy(context.Background(), StorageProxyRequest{ + Operation: StorageOperationGet, + Method: http.MethodGet, + Bucket: "bucket-a", + Object: "docs/a.pdf", + AccessToken: "secret", + }) + require.NoError(t, err) + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + require.NoError(t, err) + + assert.Equal(t, http.StatusFound, response.StatusCode) + assert.Equal(t, tt.redirectLocation, response.Header.Get("Location")) + assert.Equal(t, responseBody, string(body)) + assert.EqualValues(t, 1, sourceRequests.Load()) + assert.Zero(t, sourceRedirects.Load()) + assert.Zero(t, crossHostRequests.Load()) + + assert.Equal(t, originalRedirectPolicy, reflect.ValueOf(sharedClient.CheckRedirect).Pointer(), "the cached client must not be mutated") + }) + } +} diff --git a/relay/constant/relay_mode.go b/relay/constant/relay_mode.go index 5f4b3be5b96f..6c46c8ec0626 100644 --- a/relay/constant/relay_mode.go +++ b/relay/constant/relay_mode.go @@ -54,6 +54,7 @@ const ( RelayModeResponsesCompact RelayModeAlphaSearch + RelayModeVertexStorage ) func Path2RelayMode(path string) int { diff --git a/relay/constant/vertex_storage.go b/relay/constant/vertex_storage.go new file mode 100644 index 000000000000..7193f3222b5d --- /dev/null +++ b/relay/constant/vertex_storage.go @@ -0,0 +1,63 @@ +package constant + +import ( + "errors" + "strings" +) + +const ( + VertexStorageModelPrefix = "storage:gs:" + VertexStorageRoutePrefix = "/vertexai" + VertexStorageUploadRoute = VertexStorageRoutePrefix + "/upload/storage/v1/b/:bucket/o" + VertexStorageListRoute = VertexStorageRoutePrefix + "/storage/v1/b/:bucket/o" + VertexStorageObjectRoute = VertexStorageRoutePrefix + "/storage/v1/b/:bucket/o/*object" +) + +func NormalizeVertexStorageBucket(raw string) (string, error) { + bucket := strings.TrimSpace(raw) + if bucket == "" || bucket == "." || bucket == ".." { + return "", errors.New("invalid Vertex storage bucket") + } + if strings.HasPrefix(bucket, VertexStorageModelPrefix) || + strings.Contains(bucket, "://") || strings.ContainsAny(bucket, `/\?#`) { + return "", errors.New("invalid Vertex storage bucket") + } + return bucket, nil +} + +func VertexStorageModelName(raw string) (string, error) { + bucket, err := NormalizeVertexStorageBucket(raw) + if err != nil { + return "", err + } + return VertexStorageModelPrefix + bucket, nil +} + +func VertexStorageChannelSupports(models []string, rawBucket string) bool { + modelName, err := VertexStorageModelName(rawBucket) + if err != nil { + return false + } + for _, offered := range models { + if strings.TrimSpace(offered) == modelName { + return true + } + } + return false +} + +func ValidateVertexStorageObjectName(object string) error { + if object == "" { + return errors.New("Vertex storage object is required") + } + for _, segment := range strings.Split(object, "/") { + if segment == "." || segment == ".." { + return errors.New("invalid Vertex storage object") + } + } + return nil +} + +func IsVertexStoragePath(path string) bool { + return strings.HasPrefix(path, VertexStorageRoutePrefix+"/") +} diff --git a/relay/constant/vertex_storage_test.go b/relay/constant/vertex_storage_test.go new file mode 100644 index 000000000000..ad24aae668f7 --- /dev/null +++ b/relay/constant/vertex_storage_test.go @@ -0,0 +1,60 @@ +package constant + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestVertexStorageRoutesUseProviderPrefix(t *testing.T) { + assert.Equal(t, "/vertexai", VertexStorageRoutePrefix) + assert.Equal(t, "/vertexai/upload/storage/v1/b/:bucket/o", VertexStorageUploadRoute) + assert.Equal(t, "/vertexai/storage/v1/b/:bucket/o", VertexStorageListRoute) + assert.Equal(t, "/vertexai/storage/v1/b/:bucket/o/*object", VertexStorageObjectRoute) +} + +func TestNormalizeVertexStorageBucketRejectsPathAndURLSyntax(t *testing.T) { + for _, value := range []string{ + "", " ", ".", "..", "bucket/path", `bucket\path`, + "storage:gs:bucket", "gs://bucket", "bucket?x=1", "bucket#fragment", + } { + t.Run(value, func(t *testing.T) { + _, err := NormalizeVertexStorageBucket(value) + require.Error(t, err) + }) + } +} + +func TestVertexStorageModelNameTrimsAndChannelSupportIsExact(t *testing.T) { + modelName, err := VertexStorageModelName(" bucket-a ") + require.NoError(t, err) + assert.Equal(t, "storage:gs:bucket-a", modelName) + + models := []string{"gemini-2.5-pro", " storage:gs:bucket-a ", "storage:gs:bucket-ab"} + assert.True(t, VertexStorageChannelSupports(models, "bucket-a")) + assert.False(t, VertexStorageChannelSupports(models, "bucket")) + assert.False(t, VertexStorageChannelSupports(models, "bucket-a/path")) +} + +func TestValidateVertexStorageObjectNameRejectsDotSegments(t *testing.T) { + for _, value := range []string{ + "", ".", "..", "folder/./file.txt", "folder/../file.txt", "./folder/file.txt", "../folder/file.txt", + } { + t.Run(value, func(t *testing.T) { + require.Error(t, ValidateVertexStorageObjectName(value)) + }) + } + + for _, value := range []string{"file.txt", "folder/file.txt", "folder/.../file.txt", " folder /file.txt "} { + t.Run("valid "+value, func(t *testing.T) { + require.NoError(t, ValidateVertexStorageObjectName(value)) + }) + } +} + +func TestIsVertexStoragePathRequiresSlashAfterPrefix(t *testing.T) { + assert.True(t, IsVertexStoragePath("/vertexai/storage/v1/b/bucket-a/o")) + assert.False(t, IsVertexStoragePath("/vertexai-evil/storage/v1/b/bucket-a/o")) + assert.False(t, IsVertexStoragePath("/v1/rawproxy/vertexai/storage/v1/b/bucket-a/o")) +} diff --git a/router/relay-router.go b/router/relay-router.go index b230a5a8084c..8194a19678f3 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -1,10 +1,13 @@ package router import ( + "strings" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/controller" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/relay" + relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relaykit/types" "github.com/gin-gonic/gin" @@ -79,6 +82,23 @@ func SetRelayRouter(router *gin.Engine) { controller.Relay(c, types.RelayFormatOpenAIRealtime) }) } + + vertexStorageRouter := router.Group(relayconstant.VertexStorageRoutePrefix) + vertexStorageRouter.Use(middleware.RouteTag("relay")) + vertexStorageRouter.Use(middleware.SystemPerformanceCheck()) + vertexStorageRouter.Use(middleware.TokenAuth()) + vertexStorageRouter.Use(middleware.ModelRequestRateLimit()) + vertexStorageRouter.Use(middleware.DistributeByChannelType(constant.ChannelTypeVertexAi)) + { + uploadPath := strings.TrimPrefix(relayconstant.VertexStorageUploadRoute, relayconstant.VertexStorageRoutePrefix) + listPath := strings.TrimPrefix(relayconstant.VertexStorageListRoute, relayconstant.VertexStorageRoutePrefix) + objectPath := strings.TrimPrefix(relayconstant.VertexStorageObjectRoute, relayconstant.VertexStorageRoutePrefix) + vertexStorageRouter.POST(uploadPath, controller.RelayVertexStorageUpload) + vertexStorageRouter.PUT(uploadPath, controller.RelayVertexStorageUpload) + vertexStorageRouter.GET(listPath, controller.RelayVertexStorageList) + vertexStorageRouter.GET(objectPath, controller.RelayVertexStorageObject) + vertexStorageRouter.DELETE(objectPath, controller.RelayVertexStorageObject) + } { //http router httpRouter := relayV1Router.Group("") diff --git a/router/relay_router_test.go b/router/relay_router_test.go index 579bd7bffa25..f81168a4de6c 100644 --- a/router/relay_router_test.go +++ b/router/relay_router_test.go @@ -124,3 +124,126 @@ func setupRelayRouterTestDB(t *testing.T) { } }) } + +func TestVertexStorageRoutesAreExact(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := gin.New() + SetRelayRouter(engine) + + want := map[string]bool{ + "POST /vertexai/upload/storage/v1/b/:bucket/o": true, + "PUT /vertexai/upload/storage/v1/b/:bucket/o": true, + "GET /vertexai/storage/v1/b/:bucket/o": true, + "GET /vertexai/storage/v1/b/:bucket/o/*object": true, + "DELETE /vertexai/storage/v1/b/:bucket/o/*object": true, + } + got := make(map[string]bool, len(want)) + for _, route := range engine.Routes() { + if strings.HasPrefix(route.Path, "/vertexai/") { + got[route.Method+" "+route.Path] = true + } + } + + assert.Equal(t, want, got) +} + +func TestVertexStorageOpenAPIContract(t *testing.T) { + openAPIBytes, err := os.ReadFile("../docs/openapi/relay.json") + require.NoError(t, err) + var document struct { + Tags []map[string]any `json:"tags"` + Paths map[string]map[string]any `json:"paths"` + } + require.NoError(t, common.Unmarshal(openAPIBytes, &document)) + + const tag = "文件/Vertex AI Cloud Storage" + tagFound := false + for _, item := range document.Tags { + if item["name"] == tag { + tagFound = true + break + } + } + assert.True(t, tagFound) + + expected := map[string][]string{ + "/vertexai/upload/storage/v1/b/{bucket}/o": {"post", "put"}, + "/vertexai/storage/v1/b/{bucket}/o": {"get"}, + "/vertexai/storage/v1/b/{bucket}/o/{object}": {"get", "delete"}, + } + vertexPathCount := 0 + for path, pathItem := range document.Paths { + if !strings.HasPrefix(path, "/vertexai/") { + continue + } + vertexPathCount++ + methods, ok := expected[path] + require.True(t, ok, "unexpected Vertex Storage path %s", path) + for _, method := range methods { + operation, ok := pathItem[method].(map[string]any) + require.True(t, ok, "%s %s", method, path) + assert.Equal(t, []any{tag}, operation["tags"]) + security, ok := operation["security"].([]any) + require.True(t, ok) + require.NotEmpty(t, security) + securityItem, ok := security[0].(map[string]any) + require.True(t, ok) + assert.Contains(t, securityItem, "BearerAuth") + + parameters, ok := operation["parameters"].([]any) + require.True(t, ok) + assertOpenAPIRequiredPathParameter(t, parameters, "bucket") + if strings.Contains(path, "{object}") { + assertOpenAPIRequiredPathParameter(t, parameters, "object") + } + if strings.Contains(path, "/upload/") { + assertOpenAPIParameter(t, parameters, "uploadType", "query") + assertOpenAPIParameter(t, parameters, "name", "query") + requestBody, ok := operation["requestBody"].(map[string]any) + require.True(t, ok) + content, ok := requestBody["content"].(map[string]any) + require.True(t, ok) + assert.Contains(t, content, "application/octet-stream") + if method == "put" { + assertOpenAPIParameter(t, parameters, "Content-Range", "header") + } + } + if strings.Contains(path, "{object}") && method == "get" { + assertOpenAPIParameter(t, parameters, "alt", "query") + assertOpenAPIParameter(t, parameters, "generation", "query") + assertOpenAPIParameter(t, parameters, "Range", "header") + } + responses, ok := operation["responses"].(map[string]any) + require.True(t, ok) + for _, status := range []string{"400", "403", "502", "default"} { + assert.Contains(t, responses, status) + } + } + } + assert.Equal(t, len(expected), vertexPathCount) +} + +func assertOpenAPIRequiredPathParameter(t *testing.T, parameters []any, name string) { + t.Helper() + for _, raw := range parameters { + parameter, ok := raw.(map[string]any) + require.True(t, ok) + if parameter["name"] == name && parameter["in"] == "path" { + assert.Equal(t, true, parameter["required"]) + return + } + } + assert.Fail(t, "missing required path parameter", name) +} + +func assertOpenAPIParameter(t *testing.T, parameters []any, name, location string) { + t.Helper() + for _, raw := range parameters { + parameter, ok := raw.(map[string]any) + require.True(t, ok) + if parameter["name"] == name && parameter["in"] == location { + return + } + } + assert.Fail(t, "missing OpenAPI parameter", name+" in "+location) +} diff --git a/service/channel_select.go b/service/channel_select.go index 0ab88dc84ff2..f6ed75c3b95c 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -11,12 +11,13 @@ import ( ) type RetryParam struct { - Ctx *gin.Context - TokenGroup string - ModelName string - RequestPath string - Retry *int - resetNextTry bool + Ctx *gin.Context + TokenGroup string + ModelName string + RequestPath string + Retry *int + RequiredChannelType int + resetNextTry bool } func (p *RetryParam) GetRetry() int { @@ -115,7 +116,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, } logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) - channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath) + channel, _ = model.GetRandomSatisfiedChannelByType(autoGroup, param.ModelName, priorityRetry, param.RequestPath, param.RequiredChannelType) if channel == nil { // Current group has no available channel for this model, try next group // 当前分组没有该模型的可用渠道，尝试下一个分组 @@ -153,7 +154,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, break } } else { - channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath) + channel, err = model.GetRandomSatisfiedChannelByType(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath, param.RequiredChannelType) if err != nil { return nil, param.TokenGroup, err } diff --git a/service/channel_select_channel_type_test.go b/service/channel_select_channel_type_test.go new file mode 100644 index 000000000000..b09ba3cf9407 --- /dev/null +++ b/service/channel_select_channel_type_test.go @@ -0,0 +1,82 @@ +package service + +import ( + "fmt" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func createChannelTypeSelectionFixture( + t *testing.T, + db *gorm.DB, + id int, + channelType int, + priorityValue int64, +) { + t.Helper() + weight := uint(100) + channel := &model.Channel{ + Id: id, + Type: channelType, + Key: fmt.Sprintf("key-%d", id), + Status: common.ChannelStatusEnabled, + Name: fmt.Sprintf("channel-%d", id), + Weight: &weight, + Models: "storage:gs:bucket-a", + Group: "default", + Priority: &priorityValue, + } + require.NoError(t, db.Create(channel).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: "default", + Model: "storage:gs:bucket-a", + ChannelId: id, + Enabled: true, + Priority: &priorityValue, + Weight: weight, + }).Error) +} + +func TestRequiredChannelTypeFiltersCachedCandidates(t *testing.T) { + db := setupChannelSelectAutoGroupsTest(t) + createChannelTypeSelectionFixture(t, db, 3101, constant.ChannelTypeGemini, 100) + createChannelTypeSelectionFixture(t, db, 3102, constant.ChannelTypeVertexAi, 0) + model.InitChannelCache() + + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + for _, testCase := range []struct { + name string + requiredChannelType int + wantID int + wantType int + }{ + {name: "limits candidates to the requested channel type", requiredChannelType: constant.ChannelTypeVertexAi, wantID: 3102, wantType: constant.ChannelTypeVertexAi}, + {name: "zero preserves the unrestricted candidate set", wantID: 3101, wantType: constant.ChannelTypeGemini}, + } { + t.Run(testCase.name, func(t *testing.T) { + selected, group, err := CacheGetRandomSatisfiedChannel(&RetryParam{ + Ctx: ctx, + TokenGroup: "default", + ModelName: "storage:gs:bucket-a", + RequestPath: "/vertexai/storage/v1/b/bucket-a/o", + RequiredChannelType: testCase.requiredChannelType, + }) + + require.NoError(t, err) + require.NotNil(t, selected) + assert.Equal(t, testCase.wantID, selected.Id) + assert.Equal(t, testCase.wantType, selected.Type) + assert.Equal(t, "default", group) + }) + } +} diff --git a/web/src/features/channels/components/__tests__/channel-test-routing.test.tsx b/web/src/features/channels/components/__tests__/channel-test-routing.test.tsx new file mode 100644 index 000000000000..8305d3daa982 --- /dev/null +++ b/web/src/features/channels/components/__tests__/channel-test-routing.test.tsx @@ -0,0 +1,226 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, describe, test } from 'node:test' + +import type { Row } from '@tanstack/react-table' +import { Window } from 'happy-dom' + +import type { Channel } from '../../types' + +const domWindow = new Window({ url: 'https://example.test/channels' }) +for (const key of [ + 'window', + 'document', + 'navigator', + 'localStorage', + 'HTMLElement', + 'HTMLButtonElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'MouseEvent', + 'PointerEvent', + 'KeyboardEvent', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'IntersectionObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +const { act } = await import('react') +const { createRoot } = await import('react-dom/client') +const { QueryClient, QueryClientProvider } = + await import('@tanstack/react-query') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { TooltipProvider } = await import('@/components/ui/tooltip') +const { ChannelRowActionsLayoutContext } = + await import('../channel-row-actions-context') +const { ChannelsProvider, useChannels } = await import('../channels-provider') +const { DataTableRowActions } = await import('../data-table-row-actions') + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { + en: { + translation: { + Edit: 'Edit', + 'Test Connection': 'Test Connection', + 'Test Channel Connection': 'Test Channel Connection', + Enable: 'Enable', + Disable: 'Disable', + 'Open menu': 'Open menu', + }, + }, + }, +}) + +const channel: Channel = { + id: 7, + type: 41, + key: '', + status: 1, + name: 'Vertex storage', + created_time: 0, + test_time: 0, + response_time: 0, + other: '', + balance: 0, + balance_updated_time: 0, + models: 'gemini-2.5-pro,storage:gs:bucket-a', + group: 'default', + used_quota: 0, + other_info: '', + remark: '', + max_input_tokens: 0, + channel_info: { + is_multi_key: false, + multi_key_size: 0, + multi_key_polling_index: 0, + multi_key_mode: 'random', + }, + settings: '{}', +} + +const row = { original: channel } as Row + +function ChannelsStateProbe() { + const channels = useChannels() + return ( + + {channels.open ?? 'closed'}:{channels.currentRow?.id ?? 'none'} + + ) +} + +async function renderActions(layout: 'table' | 'card') { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }) + + await act(async () => { + root.render( + + + + + + + + + + + + + ) + }) + + return { + container, + unmount: async () => { + await act(async () => root.unmount()) + queryClient.clear() + document.body.replaceChildren() + }, + } +} + +async function click(element: Element) { + await act(async () => { + element.dispatchEvent( + new domWindow.MouseEvent('click', { + bubbles: true, + cancelable: true, + }) as unknown as Event + ) + await Promise.resolve() + await domWindow.happyDOM.waitUntilComplete() + }) +} + +function getChannelsState(container: HTMLElement) { + return container.querySelector('[data-testid="channels-state"]')?.textContent +} + +describe('channel test entry routing', () => { + after(() => domWindow.close()) + + test('table gauge opens the channel test dialog for the current channel', async () => { + const rendered = await renderActions('table') + const gauge = rendered.container.querySelector( + 'button[aria-label="Test Connection"]' + ) + assert.ok(gauge) + + await click(gauge) + + assert.equal(getChannelsState(rendered.container), 'test-channel:7') + await rendered.unmount() + }) + + test('card view exposes one test button and opens the same dialog', async () => { + const rendered = await renderActions('card') + const testButtons = rendered.container.querySelectorAll( + 'button[aria-label="Test Connection"], button[aria-label="Test Channel Connection"]' + ) + assert.equal(testButtons.length, 1) + + await click(testButtons[0]) + + assert.equal(getChannelsState(rendered.container), 'test-channel:7') + await rendered.unmount() + }) + + test('dropdown test action opens the same dialog', async () => { + const rendered = await renderActions('table') + const trigger = rendered.container.querySelector( + '[data-slot="dropdown-menu-trigger"]' + ) + assert.ok(trigger) + await click(trigger) + + const menuItem = [ + ...document.querySelectorAll('[data-slot="dropdown-menu-item"]'), + ].find((item) => item.textContent?.includes('Test Connection')) + assert.ok(menuItem) + await click(menuItem) + + assert.equal(getChannelsState(rendered.container), 'test-channel:7') + await rendered.unmount() + }) +}) diff --git a/web/src/features/channels/components/data-table-row-actions.tsx b/web/src/features/channels/components/data-table-row-actions.tsx index 7e15463cc97d..6a744de68e36 100644 --- a/web/src/features/channels/components/data-table-row-actions.tsx +++ b/web/src/features/channels/components/data-table-row-actions.tsx @@ -61,9 +61,7 @@ import { useAuthStore } from '@/stores/auth-store' import { MODEL_FETCHABLE_TYPES } from '../constants' import { - channelsQueryKeys, handleDeleteChannel, - handleTestChannel, handleToggleChannelStatus, isChannelEnabled, isMultiKeyChannel, @@ -85,7 +83,6 @@ export function DataTableRowActions({ row }: DataTableRowActionsProps) { const queryClient = useQueryClient() const currentUser = useAuthStore((s) => s.auth.user) const [deleteConfirmOpen, setDeleteConfirmOpen] = useState(false) - const [isTesting, setIsTesting] = useState(false) const [isTogglingStatus, setIsTogglingStatus] = useState(false) const isEnabled = isChannelEnabled(channel) @@ -106,18 +103,6 @@ export function DataTableRowActions({ row }: DataTableRowActionsProps) { setOpen('test-channel') } - const handleDirectTest = async (e: React.MouseEvent) => { - e.stopPropagation() - setIsTesting(true) - try { - await handleTestChannel(channel.id, { channelName: channel.name }, () => { - queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) - }) - } finally { - setIsTesting(false) - } - } - const handleQueryBalance = () => { setCurrentRow(channel) setOpen('balance-query') @@ -191,42 +176,19 @@ export function DataTableRowActions({ row }: DataTableRowActionsProps) {