feat: add upstream model sync service
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
benjamin
parent
f5bd25bea0
commit
ba676e43fd
474
backend/internal/service/upstream_models.go
Normal file
474
backend/internal/service/upstream_models.go
Normal file
@ -0,0 +1,474 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
|
||||
)
|
||||
|
||||
const upstreamModelsBodyLimit int64 = 8 << 20
|
||||
|
||||
// UpstreamModelSyncErrorKind classifies model sync failures for safe HTTP mapping.
|
||||
type UpstreamModelSyncErrorKind string
|
||||
|
||||
const (
|
||||
// UpstreamModelSyncErrorConfiguration means the account or server configuration cannot perform the sync.
|
||||
UpstreamModelSyncErrorConfiguration UpstreamModelSyncErrorKind = "configuration"
|
||||
// UpstreamModelSyncErrorUnsupported means the account format is intentionally unsupported for live model sync.
|
||||
UpstreamModelSyncErrorUnsupported UpstreamModelSyncErrorKind = "unsupported"
|
||||
// UpstreamModelSyncErrorUpstream means the configured upstream failed or returned an unusable response.
|
||||
UpstreamModelSyncErrorUpstream UpstreamModelSyncErrorKind = "upstream"
|
||||
)
|
||||
|
||||
// UpstreamModelSyncError keeps internal failure details wrapped while exposing a safe client message.
|
||||
type UpstreamModelSyncError struct {
|
||||
Kind UpstreamModelSyncErrorKind
|
||||
Message string
|
||||
Err error
|
||||
}
|
||||
|
||||
func (e *UpstreamModelSyncError) Error() string {
|
||||
if e == nil {
|
||||
return ""
|
||||
}
|
||||
if e.Err == nil {
|
||||
return e.Message
|
||||
}
|
||||
return e.Message + ": " + e.Err.Error()
|
||||
}
|
||||
|
||||
func (e *UpstreamModelSyncError) Unwrap() error {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
return e.Err
|
||||
}
|
||||
|
||||
// SafeMessage returns the sanitized message that can be sent to API clients.
|
||||
func (e *UpstreamModelSyncError) SafeMessage() string {
|
||||
if e == nil || strings.TrimSpace(e.Message) == "" {
|
||||
return "Failed to sync upstream models"
|
||||
}
|
||||
return e.Message
|
||||
}
|
||||
|
||||
func newUpstreamModelSyncConfigError(message string, err error) error {
|
||||
return &UpstreamModelSyncError{Kind: UpstreamModelSyncErrorConfiguration, Message: message, Err: err}
|
||||
}
|
||||
|
||||
func newUpstreamModelSyncUnsupportedError(message string, err error) error {
|
||||
return &UpstreamModelSyncError{Kind: UpstreamModelSyncErrorUnsupported, Message: message, Err: err}
|
||||
}
|
||||
|
||||
func newUpstreamModelSyncUpstreamError(message string, err error) error {
|
||||
return &UpstreamModelSyncError{Kind: UpstreamModelSyncErrorUpstream, Message: message, Err: err}
|
||||
}
|
||||
|
||||
// FetchUpstreamSupportedModels fetches the live model list from the account's upstream API format.
|
||||
func (s *AccountTestService) FetchUpstreamSupportedModels(ctx context.Context, account *Account) ([]string, error) {
|
||||
if s == nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Account test service is not configured", nil)
|
||||
}
|
||||
if account == nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Account is required", nil)
|
||||
}
|
||||
|
||||
if account.Platform == PlatformAntigravity && account.Type != AccountTypeAPIKey {
|
||||
return s.fetchAntigravityOAuthUpstreamModels(ctx, account)
|
||||
}
|
||||
|
||||
if s.httpUpstream == nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Upstream HTTP client is not configured", nil)
|
||||
}
|
||||
|
||||
req, err := s.buildUpstreamModelsRequest(ctx, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
proxyURL := upstreamModelsProxyURL(account)
|
||||
resp, err := s.doUpstreamModelsRequest(req, proxyURL, account)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to request upstream model list", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, upstreamModelsBodyLimit+1))
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to read upstream model list", err)
|
||||
}
|
||||
if int64(len(body)) > upstreamModelsBodyLimit {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Upstream model list response is too large", fmt.Errorf("response exceeds %d bytes", upstreamModelsBodyLimit))
|
||||
}
|
||||
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
return nil, newUpstreamModelSyncUpstreamError(
|
||||
fmt.Sprintf("Upstream model list request failed with HTTP %d", resp.StatusCode),
|
||||
fmt.Errorf("upstream model list returned HTTP %d", resp.StatusCode),
|
||||
)
|
||||
}
|
||||
|
||||
models, err := extractUpstreamModelIDs(body)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Upstream model list response was not valid JSON", err)
|
||||
}
|
||||
if len(models) == 0 {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Upstream returned no supported models", nil)
|
||||
}
|
||||
|
||||
return models, nil
|
||||
}
|
||||
|
||||
func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) {
|
||||
switch {
|
||||
case account.Platform == PlatformAntigravity:
|
||||
return s.buildAntigravityAPIKeyModelsRequest(ctx, account)
|
||||
case account.IsOpenAI():
|
||||
return s.buildOpenAIUpstreamModelsRequest(ctx, account)
|
||||
case account.IsGemini():
|
||||
return s.buildGeminiUpstreamModelsRequest(ctx, account)
|
||||
case account.IsAnthropic():
|
||||
return s.buildAnthropicUpstreamModelsRequest(ctx, account)
|
||||
default:
|
||||
return nil, newUpstreamModelSyncUnsupportedError(
|
||||
fmt.Sprintf("Unsupported platform for upstream model sync: %s", account.Platform), nil,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AccountTestService) buildAnthropicUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) {
|
||||
if account.IsBedrock() || account.Type == AccountTypeServiceAccount {
|
||||
return nil, newUpstreamModelSyncUnsupportedError(
|
||||
fmt.Sprintf("Unsupported Anthropic account type for upstream model sync: %s", account.Type), nil,
|
||||
)
|
||||
}
|
||||
|
||||
baseURL := "https://api.anthropic.com"
|
||||
authHeaderName := ""
|
||||
authHeaderValue := ""
|
||||
betaHeader := ""
|
||||
|
||||
if account.IsOAuth() {
|
||||
accessToken := strings.TrimSpace(account.GetCredential("access_token"))
|
||||
if accessToken == "" && s.claudeTokenProvider != nil {
|
||||
token, tokenErr := s.claudeTokenProvider.GetAccessToken(ctx, account)
|
||||
if tokenErr != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to get Anthropic access token", tokenErr)
|
||||
}
|
||||
accessToken = strings.TrimSpace(token)
|
||||
}
|
||||
if accessToken == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("No Anthropic access token is available", nil)
|
||||
}
|
||||
authHeaderName = "Authorization"
|
||||
authHeaderValue = "Bearer " + accessToken
|
||||
betaHeader = claude.DefaultBetaHeader
|
||||
} else if account.Type == AccountTypeAPIKey {
|
||||
apiKey := strings.TrimSpace(account.GetCredential("api_key"))
|
||||
if apiKey == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("No Anthropic API key is available", nil)
|
||||
}
|
||||
baseURL = account.GetBaseURL()
|
||||
if strings.TrimSpace(baseURL) == "" {
|
||||
baseURL = "https://api.anthropic.com"
|
||||
}
|
||||
authHeaderName = "x-api-key"
|
||||
authHeaderValue = apiKey
|
||||
betaHeader = claude.APIKeyBetaHeader
|
||||
} else {
|
||||
return nil, newUpstreamModelSyncUnsupportedError(
|
||||
fmt.Sprintf("Unsupported Anthropic account type for upstream model sync: %s", account.Type), nil,
|
||||
)
|
||||
}
|
||||
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Anthropic base URL", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, buildV1ModelsURL(normalizedBaseURL), nil)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Anthropic model list URL", err)
|
||||
}
|
||||
for key, value := range claude.DefaultHeaders {
|
||||
req.Header.Set(key, value)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
req.Header.Set("anthropic-beta", betaHeader)
|
||||
req.Header.Set(authHeaderName, authHeaderValue)
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func (s *AccountTestService) buildAntigravityAPIKeyModelsRequest(ctx context.Context, account *Account) (*http.Request, error) {
|
||||
if account.Type != AccountTypeAPIKey {
|
||||
return nil, newUpstreamModelSyncUnsupportedError(
|
||||
fmt.Sprintf("Unsupported Antigravity account type for upstream model sync: %s", account.Type), nil,
|
||||
)
|
||||
}
|
||||
apiKey := strings.TrimSpace(account.GetCredential("api_key"))
|
||||
if apiKey == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("No Antigravity API key is available", nil)
|
||||
}
|
||||
|
||||
baseURL := strings.TrimRight(strings.TrimSpace(account.GetCredential("base_url")), "/")
|
||||
if baseURL == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("Antigravity API-key base URL is required for upstream model sync", nil)
|
||||
}
|
||||
if !strings.HasSuffix(strings.ToLower(baseURL), "/antigravity") {
|
||||
return nil, newUpstreamModelSyncUnsupportedError(
|
||||
"Antigravity API-key upstream model sync requires a compatible gateway base URL ending in /antigravity; use Antigravity OAuth for official Cloud Code upstreams",
|
||||
nil,
|
||||
)
|
||||
}
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Antigravity base URL", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, buildV1ModelsURL(normalizedBaseURL), nil)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Antigravity model list URL", err)
|
||||
}
|
||||
for key, value := range claude.DefaultHeaders {
|
||||
req.Header.Set(key, value)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
req.Header.Set("anthropic-beta", claude.APIKeyBetaHeader)
|
||||
req.Header.Set("x-api-key", apiKey)
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func (s *AccountTestService) buildOpenAIUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) {
|
||||
if account.Type != AccountTypeAPIKey {
|
||||
return nil, newUpstreamModelSyncUnsupportedError(
|
||||
fmt.Sprintf("Unsupported OpenAI account type for upstream model sync: %s", account.Type), nil,
|
||||
)
|
||||
}
|
||||
apiKey := strings.TrimSpace(account.GetOpenAIApiKey())
|
||||
if apiKey == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("No OpenAI API key is available", nil)
|
||||
}
|
||||
|
||||
baseURL := account.GetOpenAIBaseURL()
|
||||
if strings.TrimSpace(baseURL) == "" {
|
||||
baseURL = "https://api.openai.com"
|
||||
}
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid OpenAI base URL", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, buildOpenAIModelsURL(normalizedBaseURL), nil)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid OpenAI model list URL", err)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func (s *AccountTestService) buildGeminiUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) {
|
||||
baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
|
||||
if strings.TrimSpace(baseURL) == "" {
|
||||
baseURL = geminicli.AIStudioBaseURL
|
||||
}
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Gemini base URL", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, buildGeminiModelsURL(normalizedBaseURL), nil)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Invalid Gemini model list URL", err)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
switch account.Type {
|
||||
case AccountTypeAPIKey:
|
||||
apiKey := strings.TrimSpace(account.GetCredential("api_key"))
|
||||
if apiKey == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("No Gemini API key is available", nil)
|
||||
}
|
||||
req.Header.Set("x-goog-api-key", apiKey)
|
||||
case AccountTypeOAuth:
|
||||
if strings.TrimSpace(account.GetCredential("project_id")) != "" {
|
||||
return nil, newUpstreamModelSyncUnsupportedError("Gemini Code Assist model listing is not supported by this sync button", nil)
|
||||
}
|
||||
if s.geminiTokenProvider == nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Gemini token provider is not configured", nil)
|
||||
}
|
||||
accessToken, tokenErr := s.geminiTokenProvider.GetAccessToken(ctx, account)
|
||||
if tokenErr != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to get Gemini access token", tokenErr)
|
||||
}
|
||||
accessToken = strings.TrimSpace(accessToken)
|
||||
if accessToken == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("No Gemini access token is available", nil)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
default:
|
||||
return nil, newUpstreamModelSyncUnsupportedError(
|
||||
fmt.Sprintf("Unsupported Gemini account type for upstream model sync: %s", account.Type), nil,
|
||||
)
|
||||
}
|
||||
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func (s *AccountTestService) fetchAntigravityOAuthUpstreamModels(ctx context.Context, account *Account) ([]string, error) {
|
||||
if s.antigravityGatewayService == nil || s.antigravityGatewayService.GetTokenProvider() == nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Antigravity token provider is not configured", nil)
|
||||
}
|
||||
|
||||
accessToken, err := s.antigravityGatewayService.GetTokenProvider().GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to get Antigravity access token", err)
|
||||
}
|
||||
accessToken = strings.TrimSpace(accessToken)
|
||||
if accessToken == "" {
|
||||
return nil, newUpstreamModelSyncConfigError("No Antigravity access token is available", nil)
|
||||
}
|
||||
|
||||
client, err := antigravity.NewClient(upstreamModelsProxyURL(account))
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncConfigError("Failed to configure Antigravity client", err)
|
||||
}
|
||||
modelsResp, _, err := client.FetchAvailableModels(ctx, accessToken, strings.TrimSpace(account.GetCredential("project_id")))
|
||||
if err != nil {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Failed to fetch Antigravity available models", err)
|
||||
}
|
||||
if modelsResp == nil || len(modelsResp.Models) == 0 {
|
||||
return nil, newUpstreamModelSyncUpstreamError("Upstream returned no supported models", nil)
|
||||
}
|
||||
|
||||
models := make([]string, 0, len(modelsResp.Models))
|
||||
for modelID := range modelsResp.Models {
|
||||
models = append(models, strings.TrimSpace(modelID))
|
||||
}
|
||||
return dedupeAndSortModelIDs(models), nil
|
||||
}
|
||||
|
||||
func (s *AccountTestService) doUpstreamModelsRequest(req *http.Request, proxyURL string, account *Account) (*http.Response, error) {
|
||||
if s.tlsFPProfileService == nil {
|
||||
return s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, nil)
|
||||
}
|
||||
return s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
|
||||
}
|
||||
|
||||
func upstreamModelsProxyURL(account *Account) string {
|
||||
if account != nil && account.ProxyID != nil && account.Proxy != nil {
|
||||
return account.Proxy.URL()
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func buildV1ModelsURL(base string) string {
|
||||
normalized := strings.TrimRight(strings.TrimSpace(base), "/")
|
||||
if strings.HasSuffix(normalized, "/v1/models") {
|
||||
return normalized
|
||||
}
|
||||
if strings.HasSuffix(normalized, "/v1") {
|
||||
return normalized + "/models"
|
||||
}
|
||||
return normalized + "/v1/models"
|
||||
}
|
||||
|
||||
func buildOpenAIModelsURL(base string) string {
|
||||
normalized := strings.TrimRight(strings.TrimSpace(base), "/")
|
||||
if strings.HasSuffix(normalized, "/v1/models") {
|
||||
return normalized
|
||||
}
|
||||
if strings.HasSuffix(normalized, "/v1") {
|
||||
return normalized + "/models"
|
||||
}
|
||||
return normalized + "/v1/models"
|
||||
}
|
||||
|
||||
func buildGeminiModelsURL(base string) string {
|
||||
normalized := strings.TrimRight(strings.TrimSpace(base), "/")
|
||||
if strings.HasSuffix(normalized, "/v1beta/models") {
|
||||
return normalized
|
||||
}
|
||||
if strings.HasSuffix(normalized, "/v1beta") {
|
||||
return normalized + "/models"
|
||||
}
|
||||
return normalized + "/v1beta/models"
|
||||
}
|
||||
|
||||
type upstreamModelEntry struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
func extractUpstreamModelIDs(body []byte) ([]string, error) {
|
||||
var response struct {
|
||||
Data []upstreamModelEntry `json:"data"`
|
||||
Models []upstreamModelEntry `json:"models"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &response); err != nil {
|
||||
var arrayResponse []upstreamModelEntry
|
||||
if arrayErr := json.Unmarshal(body, &arrayResponse); arrayErr != nil {
|
||||
return nil, fmt.Errorf("parse upstream model list: %w", err)
|
||||
}
|
||||
|
||||
models := make([]string, 0, len(arrayResponse))
|
||||
for _, entry := range arrayResponse {
|
||||
models = append(models, upstreamModelEntryID(entry))
|
||||
}
|
||||
return dedupeAndSortModelIDs(models), nil
|
||||
}
|
||||
|
||||
models := make([]string, 0, len(response.Data)+len(response.Models))
|
||||
for _, entry := range response.Data {
|
||||
models = append(models, upstreamModelEntryID(entry))
|
||||
}
|
||||
for _, entry := range response.Models {
|
||||
models = append(models, upstreamModelEntryID(entry))
|
||||
}
|
||||
|
||||
if len(models) == 0 {
|
||||
var arrayResponse []upstreamModelEntry
|
||||
if err := json.Unmarshal(body, &arrayResponse); err == nil {
|
||||
for _, entry := range arrayResponse {
|
||||
models = append(models, upstreamModelEntryID(entry))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return dedupeAndSortModelIDs(models), nil
|
||||
}
|
||||
|
||||
func upstreamModelEntryID(entry upstreamModelEntry) string {
|
||||
modelID := strings.TrimSpace(entry.ID)
|
||||
if modelID == "" {
|
||||
modelID = strings.TrimSpace(entry.Name)
|
||||
}
|
||||
return strings.TrimPrefix(modelID, "models/")
|
||||
}
|
||||
|
||||
func dedupeAndSortModelIDs(models []string) []string {
|
||||
seen := make(map[string]struct{}, len(models))
|
||||
result := make([]string, 0, len(models))
|
||||
for _, model := range models {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := seen[model]; exists {
|
||||
continue
|
||||
}
|
||||
seen[model] = struct{}{}
|
||||
result = append(result, model)
|
||||
}
|
||||
sort.Strings(result)
|
||||
return result
|
||||
}
|
||||
226
backend/internal/service/upstream_models_test.go
Normal file
226
backend/internal/service/upstream_models_test.go
Normal file
@ -0,0 +1,226 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func upstreamModelSyncTestConfig() *config.Config {
|
||||
return &config.Config{
|
||||
Security: config.SecurityConfig{
|
||||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildV1ModelsURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.Equal(t, "https://api.anthropic.com/v1/models", buildV1ModelsURL("https://api.anthropic.com"))
|
||||
require.Equal(t, "https://api.anthropic.com/v1/models", buildV1ModelsURL("https://api.anthropic.com/v1"))
|
||||
require.Equal(t, "https://api.anthropic.com/v1/models", buildV1ModelsURL("https://api.anthropic.com/v1/models"))
|
||||
require.Equal(t, "https://gateway.example.com/antigravity/v1/models", buildV1ModelsURL("https://gateway.example.com/antigravity/"))
|
||||
}
|
||||
|
||||
func TestBuildGeminiModelsURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.Equal(t, "https://generativelanguage.googleapis.com/v1beta/models", buildGeminiModelsURL("https://generativelanguage.googleapis.com"))
|
||||
require.Equal(t, "https://generativelanguage.googleapis.com/v1beta/models", buildGeminiModelsURL("https://generativelanguage.googleapis.com/v1beta"))
|
||||
require.Equal(t, "https://generativelanguage.googleapis.com/v1beta/models", buildGeminiModelsURL("https://generativelanguage.googleapis.com/v1beta/models"))
|
||||
}
|
||||
|
||||
func TestExtractUpstreamModelIDs(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
want []string
|
||||
}{
|
||||
{
|
||||
name: "openai and anthropic data array",
|
||||
body: `{"data":[{"id":"claude-sonnet-4-5"},{"id":"gpt-5"},{"id":"gpt-5"},{"id":""}]}`,
|
||||
want: []string{"claude-sonnet-4-5", "gpt-5"},
|
||||
},
|
||||
{
|
||||
name: "gemini models array strips prefix",
|
||||
body: `{"models":[{"name":"models/gemini-2.5-pro"},{"name":"gemini-2.5-flash"}]}`,
|
||||
want: []string{"gemini-2.5-flash", "gemini-2.5-pro"},
|
||||
},
|
||||
{
|
||||
name: "top level array",
|
||||
body: `[{"id":"z-model"},{"name":"models/a-model"}]`,
|
||||
want: []string{"a-model", "z-model"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got, err := extractUpstreamModelIDs([]byte(tt.body))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildUpstreamModelsRequestsForAPIKeyAccounts(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
svc := &AccountTestService{cfg: upstreamModelSyncTestConfig()}
|
||||
ctx := context.Background()
|
||||
|
||||
anthropicReq, err := svc.buildAnthropicUpstreamModelsRequest(ctx, &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "anthropic-key",
|
||||
"base_url": "https://anthropic.example.com/v1",
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://anthropic.example.com/v1/models", anthropicReq.URL.String())
|
||||
require.Equal(t, "anthropic-key", anthropicReq.Header.Get("x-api-key"))
|
||||
require.Equal(t, "2023-06-01", anthropicReq.Header.Get("anthropic-version"))
|
||||
|
||||
openAIReq, err := svc.buildOpenAIUpstreamModelsRequest(ctx, &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "openai-key",
|
||||
"base_url": "https://openai.example.com",
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://openai.example.com/v1/models", openAIReq.URL.String())
|
||||
require.Equal(t, "Bearer openai-key", openAIReq.Header.Get("Authorization"))
|
||||
|
||||
geminiReq, err := svc.buildGeminiUpstreamModelsRequest(ctx, &Account{
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "gemini-key",
|
||||
"base_url": "https://generativelanguage.googleapis.com/v1beta",
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://generativelanguage.googleapis.com/v1beta/models", geminiReq.URL.String())
|
||||
require.Equal(t, "gemini-key", geminiReq.Header.Get("x-goog-api-key"))
|
||||
|
||||
antigravityReq, err := svc.buildAntigravityAPIKeyModelsRequest(ctx, &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "antigravity-key",
|
||||
"base_url": "https://gateway.example.com/antigravity",
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://gateway.example.com/antigravity/v1/models", antigravityReq.URL.String())
|
||||
require.Equal(t, "antigravity-key", antigravityReq.Header.Get("x-api-key"))
|
||||
}
|
||||
|
||||
func TestBuildAntigravityAPIKeyModelsRequestRejectsOfficialCloudCodeBase(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
svc := &AccountTestService{cfg: upstreamModelSyncTestConfig()}
|
||||
_, err := svc.buildAntigravityAPIKeyModelsRequest(context.Background(), &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "antigravity-key",
|
||||
"base_url": "https://cloudcode-pa.googleapis.com",
|
||||
},
|
||||
})
|
||||
require.Error(t, err)
|
||||
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.True(t, errors.As(err, &syncErr))
|
||||
require.Equal(t, UpstreamModelSyncErrorUnsupported, syncErr.Kind)
|
||||
require.Contains(t, syncErr.SafeMessage(), "compatible gateway")
|
||||
}
|
||||
|
||||
func TestBuildAnthropicUpstreamModelsRequestRejectsBedrock(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
svc := &AccountTestService{cfg: upstreamModelSyncTestConfig()}
|
||||
_, err := svc.buildAnthropicUpstreamModelsRequest(context.Background(), &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeBedrock,
|
||||
})
|
||||
require.Error(t, err)
|
||||
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.True(t, errors.As(err, &syncErr))
|
||||
require.Equal(t, UpstreamModelSyncErrorUnsupported, syncErr.Kind)
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsParsesOpenAIResponse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"data":[{"id":"gpt-5"},{"id":"gpt-5"},{"name":"o3"}]}`)),
|
||||
}}
|
||||
svc := &AccountTestService{
|
||||
httpUpstream: upstream,
|
||||
cfg: upstreamModelSyncTestConfig(),
|
||||
}
|
||||
|
||||
models, err := svc.FetchUpstreamSupportedModels(context.Background(), &Account{
|
||||
ID: 7,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "openai-key",
|
||||
"base_url": "https://openai.example.com/v1",
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"gpt-5", "o3"}, models)
|
||||
require.Equal(t, "https://openai.example.com/v1/models", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "Bearer openai-key", upstream.lastReq.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
func TestFetchUpstreamSupportedModelsDoesNotExposeUpstreamBody(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":"SECRET_TOKEN should not be exposed"}`)),
|
||||
}}
|
||||
svc := &AccountTestService{
|
||||
httpUpstream: upstream,
|
||||
cfg: upstreamModelSyncTestConfig(),
|
||||
}
|
||||
|
||||
_, err := svc.FetchUpstreamSupportedModels(context.Background(), &Account{
|
||||
ID: 8,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "openai-key",
|
||||
"base_url": "https://openai.example.com/v1",
|
||||
},
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.NotContains(t, err.Error(), "SECRET_TOKEN")
|
||||
|
||||
var syncErr *UpstreamModelSyncError
|
||||
require.True(t, errors.As(err, &syncErr))
|
||||
require.Equal(t, UpstreamModelSyncErrorUpstream, syncErr.Kind)
|
||||
require.NotContains(t, syncErr.SafeMessage(), "SECRET_TOKEN")
|
||||
require.Contains(t, syncErr.SafeMessage(), "HTTP 502")
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user