diff --git a/backend/internal/service/upstream_models.go b/backend/internal/service/upstream_models.go index 3c5eb2fcd2..c7a392c550 100644 --- a/backend/internal/service/upstream_models.go +++ b/backend/internal/service/upstream_models.go @@ -10,6 +10,7 @@ import ( "net/http" "net/url" "sort" + "strconv" "strings" "time" @@ -25,6 +26,10 @@ const ( UpstreamModelMetadataExtraKey = "upstream_model_metadata" UpstreamModelMetadataIncompleteCode = "upstream_model_metadata_incomplete" UpstreamModelMetadataPartialCode = "upstream_model_metadata_partial" + + vertexPublisherModelsPageSize = 300 + vertexPublisherModelsMaxPages = 20 + vertexGooglePublisherModelPrefix = "publishers/google/models/" ) type UpstreamModelMetadata struct { @@ -741,34 +746,19 @@ func (s *AccountTestService) fetchUpstreamModelList(ctx context.Context, account return nil, nil, newUpstreamModelSyncConfigError("Upstream HTTP client is not configured", nil) } + if account.IsGemini() && account.IsVertexServiceAccount() { + models, err := s.fetchVertexGeminiUpstreamModels(ctx, account) + return models, nil, err + } + req, err := s.buildUpstreamModelsRequest(ctx, account) if err != nil { return nil, nil, err } - proxyURL := upstreamModelsProxyURL(account) - resp, err := s.doUpstreamModelsRequest(req, proxyURL, account) + body, err := s.readUpstreamModelsResponse(req, account) if err != nil { - return nil, nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err) - } - defer func() { _ = resp.Body.Close() }() - - bodyLimit := resolveModelsListReadLimit(s.cfg) - body, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1)) - if err != nil { - return nil, nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err) - } - if int64(len(body)) > bodyLimit { - return nil, nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", bodyLimit)) - } - - if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { - return nil, nil, &UpstreamModelSyncError{ - Kind: UpstreamModelSyncErrorUpstream, - Message: fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode), - StatusCode: resp.StatusCode, - Err: fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode), - } + return nil, nil, err } extractModels := extractUpstreamModelIDs @@ -786,6 +776,117 @@ func (s *AccountTestService) fetchUpstreamModelList(ctx context.Context, account return models, body, nil } +// readUpstreamModelsResponse 发送模型列表请求并读取受限大小的响应体;非 2xx 返回带状态码的同步错误。 +func (s *AccountTestService) readUpstreamModelsResponse(req *http.Request, account *Account) ([]byte, error) { + resp, err := s.doUpstreamModelsRequest(req, upstreamModelsProxyURL(account), account) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err) + } + defer func() { _ = resp.Body.Close() }() + + bodyLimit := resolveModelsListReadLimit(s.cfg) + body, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1)) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err) + } + if int64(len(body)) > bodyLimit { + return nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", bodyLimit)) + } + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return nil, &UpstreamModelSyncError{ + Kind: UpstreamModelSyncErrorUpstream, + Message: fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode), + StatusCode: resp.StatusCode, + Err: fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode), + } + } + return body, nil +} + +// fetchVertexGeminiUpstreamModels 按账号默认区域分页读取 Vertex 的 Google 发布方模型列表,只保留 gemini- 前缀的模型。 +// 区域端点还会列出 Gemma、Veo、AutoML 等非 Gemini 的模型花园条目;global 端点只含 Gemini。 +// 列表接口不接受 API Key。 +func (s *AccountTestService) fetchVertexGeminiUpstreamModels(ctx context.Context, account *Account) ([]string, error) { + if account.IsVertexAPIKey() { + return nil, newUpstreamModelSyncUnsupportedError("Vertex API key accounts do not support upstream model listing", nil) + } + if _, err := parseVertexServiceAccountKey(account); err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex service account credentials", err) + } + location := account.VertexLocation("") + if _, err := buildVertexPublisherModelsURL(location, ""); err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err) + } + if s.geminiTokenProvider == nil { + return nil, newUpstreamModelSyncConfigError("Gemini token provider is not configured", nil) + } + accessToken, err := s.geminiTokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, newUpstreamModelSyncUpstreamError("Failed to get Vertex access token", err) + } + accessToken = strings.TrimSpace(accessToken) + if accessToken == "" { + return nil, newUpstreamModelSyncConfigError("No Vertex access token is available", nil) + } + + models := make([]string, 0) + pageToken := "" + for page := 0; page < vertexPublisherModelsMaxPages; page++ { + listURL, err := buildVertexPublisherModelsURL(location, pageToken) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex location", err) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, listURL, nil) + if err != nil { + return nil, newUpstreamModelSyncConfigError("Invalid Vertex model list URL", err) + } + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+accessToken) + + body, err := s.readUpstreamModelsResponse(req, account) + if err != nil { + // 只有首页的 404/405 表示列表端点不可用;翻页途中的失败不触发配置模型回退。 + var syncErr *UpstreamModelSyncError + if page > 0 && errors.As(err, &syncErr) { + syncErr.StatusCode = 0 + } + return nil, err + } + var response struct { + PublisherModels []struct { + Name string `json:"name"` + } `json:"publisherModels"` + NextPageToken string `json:"nextPageToken"` + } + if err := json.Unmarshal(body, &response); err != nil { + return nil, newUpstreamModelSyncUpstreamError("Upstream model list response was not valid JSON", err) + } + for _, entry := range response.PublisherModels { + name := strings.TrimSpace(entry.Name) + if !strings.HasPrefix(name, vertexGooglePublisherModelPrefix) { + continue + } + modelID := strings.TrimPrefix(name, vertexGooglePublisherModelPrefix) + if strings.HasPrefix(modelID, "gemini-") && len(modelID) > len("gemini-") { + models = append(models, modelID) + } + } + + pageToken = strings.TrimSpace(response.NextPageToken) + if pageToken == "" { + if len(models) == 0 { + return nil, newUpstreamModelSyncUpstreamError("Upstream returned no supported models", nil) + } + return dedupeAndSortModelIDs(models), nil + } + } + return nil, newUpstreamModelSyncUpstreamError( + "Upstream model list pagination did not finish", + fmt.Errorf("more than %d pages", vertexPublisherModelsMaxPages), + ) +} + func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) { switch { case account.Platform == PlatformAntigravity: @@ -1218,6 +1319,26 @@ func buildOpenAIModelsURL(base string) string { return buildOpenAIEndpointURL(base, "/v1/models") } +// buildVertexPublisherModelsURL 构造 Google 发布方模型列表地址。该接口只有 v1beta1 版本,pageSize 上限 300。 +func buildVertexPublisherModelsURL(location, pageToken string) (string, error) { + location = strings.TrimSpace(location) + if location == "" { + location = vertexDefaultLocation + } + if !vertexLocationPattern.MatchString(location) { + return "", fmt.Errorf("invalid vertex location: %s", location) + } + host := location + "-aiplatform.googleapis.com" + if location == "global" { + host = vertexGlobalHost + } + query := url.Values{"pageSize": []string{strconv.Itoa(vertexPublisherModelsPageSize)}} + if pageToken != "" { + query.Set("pageToken", pageToken) + } + return fmt.Sprintf("https://%s/v1beta1/publishers/google/models?%s", host, query.Encode()), nil +} + func buildGeminiModelsURL(base string) string { normalized := strings.TrimRight(strings.TrimSpace(base), "/") if strings.HasSuffix(normalized, "/v1beta/models") { diff --git a/backend/internal/service/upstream_models_vertex_test.go b/backend/internal/service/upstream_models_vertex_test.go new file mode 100644 index 0000000000..bd4678614e --- /dev/null +++ b/backend/internal/service/upstream_models_vertex_test.go @@ -0,0 +1,276 @@ +package service + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type vertexModelsTokenCacheStub struct { + token string + calls int +} + +func (c *vertexModelsTokenCacheStub) GetAccessToken(context.Context, string) (string, error) { + c.calls++ + if c.token == "" { + return "", errors.New("missing token") + } + return c.token, nil +} + +func (c *vertexModelsTokenCacheStub) SetAccessToken(context.Context, string, string, time.Duration) error { + return nil +} + +func (c *vertexModelsTokenCacheStub) DeleteAccessToken(context.Context, string) error { + return nil +} + +func (c *vertexModelsTokenCacheStub) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) { + return false, nil +} + +func (c *vertexModelsTokenCacheStub) ReleaseRefreshLock(context.Context, string) error { + return nil +} + +func vertexModelSyncTestAccount(credentials map[string]any) *Account { + merged := map[string]any{ + "service_account_json": map[string]any{ + "type": "service_account", + "project_id": "proj", + "private_key": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----\n", + "client_email": "svc@proj.iam.gserviceaccount.com", + }, + } + for key, value := range credentials { + merged[key] = value + } + return &Account{ + ID: 173, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: merged, + } +} + +func vertexModelSyncTestService(upstream *httpUpstreamRecorder) *AccountTestService { + return &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, &vertexModelsTokenCacheStub{token: "ya29.vertex-token"}, nil), + } +} + +func vertexModelsJSONResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} + +func TestBuildVertexPublisherModelsURL(t *testing.T) { + t.Parallel() + + u, err := buildVertexPublisherModelsURL("global", "") + require.NoError(t, err) + require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", u) + + u, err = buildVertexPublisherModelsURL("us-central1", "next/page+token") + require.NoError(t, err) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=next%2Fpage%2Btoken", u) + + u, err = buildVertexPublisherModelsURL("", "") + require.NoError(t, err) + require.Equal(t, "https://us-central1-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", u) + + _, err = buildVertexPublisherModelsURL("evil.com/x", "") + require.Error(t, err) +} + +// Scenario: Vertex SA 账号按账号区域分页读取发布方模型列表,只保留 gemini- 前缀模型。 +func TestFetchUpstreamSupportedModelsVertexServiceAccountPaginatesAndKeepsGemini(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[ + {"name":"publishers/google/models/gemini-3.8-flash"}, + {"name":"publishers/google/models/gemma4"}, + {"name":"publishers/google/models/veo-3.1-generate-001"} + ],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[ + {"name":" publishers/google/models/gemini-2.5-flash "}, + {"name":"publishers/google/models/gemini-3.8-flash"}, + {"name":"gemini-without-publisher-prefix"}, + {"name":"publishers/anthropic/models/gemini-other-publisher"}, + {"name":"publishers/google/models/gemini-"} + ]}`), + }} + svc := vertexModelSyncTestService(upstream) + + models, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "location": "europe-west4", + })) + require.NoError(t, err) + require.Equal(t, []string{"gemini-2.5-flash", "gemini-3.8-flash"}, models) + + require.Len(t, upstream.requests, 2) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.requests[0].URL.String()) + require.Equal(t, "https://europe-west4-aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300&pageToken=page-2", upstream.requests[1].URL.String()) + for _, req := range upstream.requests { + require.Equal(t, http.MethodGet, req.Method) + require.Equal(t, "Bearer ya29.vertex-token", req.Header.Get("Authorization")) + } +} + +func TestFetchUpstreamSupportedModelsVertexUpstreamErrorKeepsStatus(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusForbidden, `{"error":{"code":403,"status":"PERMISSION_DENIED"}}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{"location": "global"})) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Equal(t, http.StatusForbidden, syncErr.StatusCode) + require.Equal(t, "https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=300", upstream.lastReq.URL.String()) +} + +func TestFetchUpstreamSupportedModelsVertexWithoutGeminiModelsFails(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemma4"}]}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Equal(t, "Upstream returned no supported models", syncErr.SafeMessage()) +} + +func TestFetchUpstreamSupportedModelsVertexStopsEndlessPagination(t *testing.T) { + t.Parallel() + + responses := make([]*http.Response, 0, vertexPublisherModelsMaxPages) + for i := 0; i < vertexPublisherModelsMaxPages; i++ { + responses = append(responses, vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-3.8-flash"}],"nextPageToken":"again"}`)) + } + upstream := &httpUpstreamRecorder{responses: responses} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, "Upstream model list pagination did not finish", syncErr.SafeMessage()) + require.Len(t, upstream.requests, vertexPublisherModelsMaxPages) +} + +func TestFetchUpstreamSupportedModelsVertexAPIKeyUnsupported(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{} + svc := vertexModelSyncTestService(upstream) + account := &Account{ + ID: 174, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{ + "auth_mode": "apikey", + "api_key": "AQ.test-key", + }, + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), account) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUnsupported, syncErr.Kind) + require.Empty(t, upstream.requests) +} + +func TestFetchUpstreamSupportedModelsVertexInvalidCredentialsIsConfigError(t *testing.T) { + t.Parallel() + + cache := &vertexModelsTokenCacheStub{token: "ya29.vertex-token"} + upstream := &httpUpstreamRecorder{} + svc := &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, cache, nil), + } + account := &Account{ + ID: 175, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{"api_key": "not-a-service-account"}, + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), account) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorConfiguration, syncErr.Kind) + require.Equal(t, "Invalid Vertex service account credentials", syncErr.SafeMessage()) + require.Zero(t, cache.calls) + require.Empty(t, upstream.requests) +} + +func TestFetchUpstreamSupportedModelsVertexInvalidLocationSkipsTokenExchange(t *testing.T) { + t.Parallel() + + cache := &vertexModelsTokenCacheStub{token: "ya29.vertex-token"} + upstream := &httpUpstreamRecorder{} + svc := &AccountTestService{ + httpUpstream: upstream, + cfg: upstreamModelSyncTestConfig(), + geminiTokenProvider: NewGeminiTokenProvider(nil, cache, nil), + } + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "location": "us-central1.evil.com", + })) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorConfiguration, syncErr.Kind) + require.Equal(t, "Invalid Vertex location", syncErr.SafeMessage()) + require.Zero(t, cache.calls) + require.Empty(t, upstream.requests) +} + +// Scenario: 翻页途中的 404 不能被当成列表端点不可用而回退到配置模型。 +func TestSyncUpstreamModelCatalogVertexMidPagination404DoesNotFallback(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + vertexModelsJSONResponse(http.StatusOK, `{"publisherModels":[{"name":"publishers/google/models/gemini-a"}],"nextPageToken":"page-2"}`), + vertexModelsJSONResponse(http.StatusNotFound, `{"error":{"code":404}}`), + }} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.SyncUpstreamModelCatalog(context.Background(), vertexModelSyncTestAccount(map[string]any{ + "model_mapping": map[string]any{"x": "gemini-mapped"}, + })) + var syncErr *UpstreamModelSyncError + require.ErrorAs(t, err, &syncErr) + require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind) + require.Zero(t, syncErr.StatusCode) + require.Len(t, upstream.requests, 2) +} + +func TestFetchUpstreamSupportedModelsVertexFirstPage404KeepsStatus(t *testing.T) { + t.Parallel() + + upstream := &httpUpstreamRecorder{resp: vertexModelsJSONResponse(http.StatusNotFound, `{"error":{"code":404}}`)} + svc := vertexModelSyncTestService(upstream) + + _, err := svc.FetchUpstreamSupportedModels(context.Background(), vertexModelSyncTestAccount(nil)) + require.True(t, upstreamModelListEndpointUnsupported(err)) +}