Merge pull request #543 from slovx2/upstream_main
feat(antigravity): 转发与测试支持 daily/prod 单 URL 切换
This commit is contained in:
@@ -70,6 +70,7 @@ var antigravityPassthroughErrorMessages = []string{
|
|||||||
|
|
||||||
const (
|
const (
|
||||||
antigravityBillingModelEnv = "GATEWAY_ANTIGRAVITY_BILL_WITH_MAPPED_MODEL"
|
antigravityBillingModelEnv = "GATEWAY_ANTIGRAVITY_BILL_WITH_MAPPED_MODEL"
|
||||||
|
antigravityForwardBaseURLEnv = "GATEWAY_ANTIGRAVITY_FORWARD_BASE_URL"
|
||||||
antigravityFallbackSecondsEnv = "GATEWAY_ANTIGRAVITY_FALLBACK_COOLDOWN_SECONDS"
|
antigravityFallbackSecondsEnv = "GATEWAY_ANTIGRAVITY_FALLBACK_COOLDOWN_SECONDS"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -131,6 +132,20 @@ type antigravityRetryLoopResult struct {
|
|||||||
resp *http.Response
|
resp *http.Response
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// resolveAntigravityForwardBaseURL 解析转发用 base URL。
|
||||||
|
// 默认使用 daily(ForwardBaseURLs 的首个地址);当环境变量为 prod 时使用第二个地址。
|
||||||
|
func resolveAntigravityForwardBaseURL() string {
|
||||||
|
baseURLs := antigravity.ForwardBaseURLs()
|
||||||
|
if len(baseURLs) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
mode := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityForwardBaseURLEnv)))
|
||||||
|
if mode == "prod" && len(baseURLs) > 1 {
|
||||||
|
return baseURLs[1]
|
||||||
|
}
|
||||||
|
return baseURLs[0]
|
||||||
|
}
|
||||||
|
|
||||||
// smartRetryAction 智能重试的处理结果
|
// smartRetryAction 智能重试的处理结果
|
||||||
type smartRetryAction int
|
type smartRetryAction int
|
||||||
|
|
||||||
@@ -466,10 +481,11 @@ func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopP
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
availableURLs := antigravity.DefaultURLAvailability.GetAvailableURLs()
|
baseURL := resolveAntigravityForwardBaseURL()
|
||||||
if len(availableURLs) == 0 {
|
if baseURL == "" {
|
||||||
availableURLs = antigravity.BaseURLs
|
return nil, errors.New("no antigravity forward base url configured")
|
||||||
}
|
}
|
||||||
|
availableURLs := []string{baseURL}
|
||||||
|
|
||||||
var resp *http.Response
|
var resp *http.Response
|
||||||
var usedBaseURL string
|
var usedBaseURL string
|
||||||
@@ -907,11 +923,11 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account
|
|||||||
proxyURL = account.Proxy.URL()
|
proxyURL = account.Proxy.URL()
|
||||||
}
|
}
|
||||||
|
|
||||||
// URL fallback 循环
|
baseURL := resolveAntigravityForwardBaseURL()
|
||||||
availableURLs := antigravity.DefaultURLAvailability.GetAvailableURLs()
|
if baseURL == "" {
|
||||||
if len(availableURLs) == 0 {
|
return nil, errors.New("no antigravity forward base url configured")
|
||||||
availableURLs = antigravity.BaseURLs // 所有 URL 都不可用时,重试所有
|
|
||||||
}
|
}
|
||||||
|
availableURLs := []string{baseURL}
|
||||||
|
|
||||||
var lastErr error
|
var lastErr error
|
||||||
for urlIdx, baseURL := range availableURLs {
|
for urlIdx, baseURL := range availableURLs {
|
||||||
|
|||||||
@@ -86,7 +86,9 @@ func (s *stubAntigravityAccountRepo) SetModelRateLimit(ctx context.Context, id i
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestAntigravityRetryLoop_URLFallback_UsesLatestSuccess(t *testing.T) {
|
func TestAntigravityRetryLoop_NoURLFallback_UsesConfiguredBaseURL(t *testing.T) {
|
||||||
|
t.Setenv(antigravityForwardBaseURLEnv, "")
|
||||||
|
|
||||||
oldBaseURLs := append([]string(nil), antigravity.BaseURLs...)
|
oldBaseURLs := append([]string(nil), antigravity.BaseURLs...)
|
||||||
oldAvailability := antigravity.DefaultURLAvailability
|
oldAvailability := antigravity.DefaultURLAvailability
|
||||||
defer func() {
|
defer func() {
|
||||||
@@ -131,15 +133,16 @@ func TestAntigravityRetryLoop_URLFallback_UsesLatestSuccess(t *testing.T) {
|
|||||||
require.NotNil(t, result)
|
require.NotNil(t, result)
|
||||||
require.NotNil(t, result.resp)
|
require.NotNil(t, result.resp)
|
||||||
defer func() { _ = result.resp.Body.Close() }()
|
defer func() { _ = result.resp.Body.Close() }()
|
||||||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
require.Equal(t, http.StatusTooManyRequests, result.resp.StatusCode)
|
||||||
require.False(t, handleErrorCalled)
|
require.True(t, handleErrorCalled)
|
||||||
require.Len(t, upstream.calls, 2)
|
require.Len(t, upstream.calls, antigravityMaxRetries)
|
||||||
require.True(t, strings.HasPrefix(upstream.calls[0], base1))
|
for _, callURL := range upstream.calls {
|
||||||
require.True(t, strings.HasPrefix(upstream.calls[1], base2))
|
require.True(t, strings.HasPrefix(callURL, base1))
|
||||||
|
}
|
||||||
|
|
||||||
available := antigravity.DefaultURLAvailability.GetAvailableURLs()
|
available := antigravity.DefaultURLAvailability.GetAvailableURLs()
|
||||||
require.NotEmpty(t, available)
|
require.NotEmpty(t, available)
|
||||||
require.Equal(t, base2, available[0])
|
require.Equal(t, base1, available[0])
|
||||||
}
|
}
|
||||||
|
|
||||||
// TestHandleUpstreamError_429_ModelRateLimit 测试 429 模型限流场景
|
// TestHandleUpstreamError_429_ModelRateLimit 测试 429 模型限流场景
|
||||||
@@ -915,6 +918,22 @@ func TestIsAntigravityAccountSwitchError(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestResolveAntigravityForwardBaseURL_DefaultDaily(t *testing.T) {
|
||||||
|
t.Setenv(antigravityForwardBaseURLEnv, "")
|
||||||
|
|
||||||
|
oldBaseURLs := append([]string(nil), antigravity.BaseURLs...)
|
||||||
|
defer func() {
|
||||||
|
antigravity.BaseURLs = oldBaseURLs
|
||||||
|
}()
|
||||||
|
|
||||||
|
prodURL := "https://prod.test"
|
||||||
|
dailyURL := "https://daily.test"
|
||||||
|
antigravity.BaseURLs = []string{dailyURL, prodURL}
|
||||||
|
|
||||||
|
resolved := resolveAntigravityForwardBaseURL()
|
||||||
|
require.Equal(t, dailyURL, resolved)
|
||||||
|
}
|
||||||
|
|
||||||
func TestAntigravityAccountSwitchError_Error(t *testing.T) {
|
func TestAntigravityAccountSwitchError_Error(t *testing.T) {
|
||||||
err := &AntigravityAccountSwitchError{
|
err := &AntigravityAccountSwitchError{
|
||||||
OriginalAccountID: 789,
|
OriginalAccountID: 789,
|
||||||
|
|||||||
Reference in New Issue
Block a user