李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
664 lines
23 KiB
Go
664 lines
23 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
type antigravityCompatTokenCache struct {
|
|
token string
|
|
}
|
|
|
|
type antigravityCompatErrorReader struct {
|
|
data []byte
|
|
off int
|
|
err error
|
|
}
|
|
|
|
func (r *antigravityCompatErrorReader) Read(p []byte) (int, error) {
|
|
if r.off < len(r.data) {
|
|
n := copy(p, r.data[r.off:])
|
|
r.off += n
|
|
return n, nil
|
|
}
|
|
return 0, r.err
|
|
}
|
|
|
|
func (r *antigravityCompatErrorReader) Close() error { return nil }
|
|
|
|
func (c *antigravityCompatTokenCache) GetAccessToken(context.Context, string) (string, error) {
|
|
return c.token, nil
|
|
}
|
|
|
|
func (c *antigravityCompatTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error {
|
|
return nil
|
|
}
|
|
|
|
func (c *antigravityCompatTokenCache) DeleteAccessToken(context.Context, string) error {
|
|
return nil
|
|
}
|
|
|
|
func (c *antigravityCompatTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
|
|
return true, nil
|
|
}
|
|
|
|
func (c *antigravityCompatTokenCache) ReleaseRefreshLock(context.Context, string) error {
|
|
return nil
|
|
}
|
|
|
|
func newAntigravityCompatService(cfg config.GatewayConfig, upstream HTTPUpstream) *AntigravityGatewayService {
|
|
tokenProvider := NewAntigravityTokenProvider(
|
|
nil,
|
|
&antigravityCompatTokenCache{token: "fresh-oauth-token"},
|
|
nil,
|
|
)
|
|
return NewAntigravityGatewayService(
|
|
nil,
|
|
nil,
|
|
nil,
|
|
tokenProvider,
|
|
nil,
|
|
upstream,
|
|
NewSettingService(&antigravitySettingRepoStub{}, &config.Config{Gateway: cfg}),
|
|
nil,
|
|
)
|
|
}
|
|
|
|
func newAntigravityCompatAccount(accountType string) *Account {
|
|
return &Account{
|
|
ID: 3757,
|
|
Name: "antigravity-compat",
|
|
Platform: PlatformAntigravity,
|
|
Type: accountType,
|
|
Status: StatusActive,
|
|
Concurrency: 1,
|
|
Credentials: map[string]any{
|
|
"access_token": "stale-account-token",
|
|
"project_id": "project-3757",
|
|
"model_mapping": map[string]any{
|
|
"gemini-3.1-pro-high": "gemini-3.1-pro-high",
|
|
"claude-sonnet-4-5": "claude-sonnet-4-5",
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func newAntigravityCompatContext(method, path string, body []byte) (*gin.Context, *httptest.ResponseRecorder) {
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(method, path, bytes.NewReader(body))
|
|
return c, recorder
|
|
}
|
|
|
|
func antigravityCompatSuccessResponse() *http.Response {
|
|
body := `data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":3}}}` + "\n\n"
|
|
return &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{
|
|
"Content-Type": []string{"text/event-stream"},
|
|
"X-Request-Id": []string{"request-3757"},
|
|
},
|
|
Body: io.NopCloser(strings.NewReader(body)),
|
|
}
|
|
}
|
|
|
|
func TestAntigravityCompatOAuthUsesNativeTokenAndRoute(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
tests := []struct {
|
|
name string
|
|
path string
|
|
body []byte
|
|
call func(*AntigravityGatewayService, context.Context, *gin.Context, *Account, []byte) (*ForwardResult, error)
|
|
}{
|
|
{
|
|
name: "chat completions",
|
|
path: "/v1/chat/completions",
|
|
body: []byte(`{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"Reply exactly: ok"}]}`),
|
|
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
|
return svc.ForwardAsChatCompletions(ctx, c, account, body, nil)
|
|
},
|
|
},
|
|
{
|
|
name: "responses",
|
|
path: "/v1/responses",
|
|
body: []byte(`{"model":"gemini-3.1-pro-high","input":"Reply exactly: ok"}`),
|
|
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
|
return svc.ForwardAsResponses(ctx, c, account, body, nil)
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var authorization string
|
|
var upstreamPath string
|
|
var upstreamAlt string
|
|
upstream := &queuedHTTPUpstreamStub{
|
|
responses: []*http.Response{antigravityCompatSuccessResponse()},
|
|
onCall: func(req *http.Request, _ *queuedHTTPUpstreamStub) {
|
|
authorization = req.Header.Get("Authorization")
|
|
upstreamPath = req.URL.Path
|
|
upstreamAlt = req.URL.Query().Get("alt")
|
|
},
|
|
}
|
|
svc := newAntigravityCompatService(
|
|
config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
|
upstream,
|
|
)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, tt.path, tt.body)
|
|
|
|
result, err := tt.call(svc, context.Background(), c, newAntigravityCompatAccount(AccountTypeOAuth), tt.body)
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result)
|
|
require.Equal(t, "Bearer fresh-oauth-token", authorization)
|
|
require.Equal(t, "/v1internal:streamGenerateContent", upstreamPath)
|
|
require.Equal(t, "sse", upstreamAlt)
|
|
require.Equal(t, "request-3757", result.RequestID)
|
|
require.Equal(t, 8, result.Usage.InputTokens)
|
|
require.Equal(t, 3, result.Usage.OutputTokens)
|
|
require.Equal(t, http.StatusOK, recorder.Code)
|
|
require.Contains(t, recorder.Body.String(), "ok")
|
|
if tt.name == "chat completions" {
|
|
require.Equal(t, "stop", gjson.Get(recorder.Body.String(), "choices.0.finish_reason").String())
|
|
require.Equal(t, int64(8), gjson.Get(recorder.Body.String(), "usage.prompt_tokens").Int())
|
|
require.Equal(t, int64(3), gjson.Get(recorder.Body.String(), "usage.completion_tokens").Int())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAntigravityCompatRejectsUnsupportedAccountType(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
tests := []struct {
|
|
name string
|
|
path string
|
|
accountType string
|
|
call func(*AntigravityGatewayService, context.Context, *gin.Context, *Account, []byte) (*ForwardResult, error)
|
|
}{
|
|
{
|
|
name: "chat completions upstream",
|
|
path: "/v1/chat/completions",
|
|
accountType: AccountTypeUpstream,
|
|
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
|
return svc.ForwardAsChatCompletions(ctx, c, account, body, nil)
|
|
},
|
|
},
|
|
{
|
|
name: "responses setup token",
|
|
path: "/v1/responses",
|
|
accountType: AccountTypeSetupToken,
|
|
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
|
return svc.ForwardAsResponses(ctx, c, account, body, nil)
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
body := []byte(`{"model":"gemini-3.1-pro-high"}`)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, tt.path, body)
|
|
|
|
result, err := tt.call(&AntigravityGatewayService{}, context.Background(), c, newAntigravityCompatAccount(tt.accountType), body)
|
|
|
|
require.Error(t, err)
|
|
require.Nil(t, result)
|
|
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
|
require.Contains(t, recorder.Body.String(), "native OAuth account required for antigravity compatibility mode")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuildAntigravityCompatGeminiBody_ConfiguresMixedToolInvocations(t *testing.T) {
|
|
svc := &AntigravityGatewayService{}
|
|
tests := []struct {
|
|
name string
|
|
tools string
|
|
wantField bool
|
|
}{
|
|
{
|
|
name: "mixed server and client tools",
|
|
tools: `[{"name":"get_weather","input_schema":{"type":"object"}},{"type":"web_search_20250305","name":"web_search"}]`,
|
|
wantField: true,
|
|
},
|
|
{
|
|
name: "client tools only",
|
|
tools: `[{"name":"get_weather","input_schema":{"type":"object"}}]`,
|
|
},
|
|
{
|
|
name: "server tools only",
|
|
tools: `[{"type":"web_search_20250305","name":"web_search"}]`,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
claudeBody := []byte(`{"messages":[{"role":"user","content":"hello"}],"tools":` + tt.tools + `}`)
|
|
claudeBody = bytes.ReplaceAll(claudeBody, []byte{92}, nil)
|
|
body, err := svc.buildAntigravityCompatGeminiBody(context.Background(), claudeBody, nil, "project-1", "gemini-2.5-flash")
|
|
require.NoError(t, err)
|
|
|
|
var wrapped map[string]any
|
|
require.NoError(t, json.Unmarshal(body, &wrapped))
|
|
request, ok := wrapped["request"].(map[string]any)
|
|
require.True(t, ok)
|
|
toolConfig, exists := request["toolConfig"].(map[string]any)
|
|
if !tt.wantField {
|
|
require.False(t, exists)
|
|
return
|
|
}
|
|
require.True(t, exists)
|
|
require.Equal(t, true, toolConfig["includeServerSideToolInvocations"])
|
|
require.NotContains(t, toolConfig, "include_server_side_tool_invocations")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAntigravityCompatPreservesChatTokenLimit(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
want int64
|
|
}{
|
|
{
|
|
name: "legacy max_tokens below bridge floor",
|
|
body: `{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"ok"}],"max_tokens":8}`,
|
|
want: 8,
|
|
},
|
|
{
|
|
name: "max_completion_tokens takes precedence",
|
|
body: `{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"ok"}],"max_tokens":8,"max_completion_tokens":13}`,
|
|
want: 13,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{antigravityCompatSuccessResponse()}}
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream)
|
|
body := []byte(tt.body)
|
|
c, _ := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", body)
|
|
|
|
result, err := svc.ForwardAsChatCompletions(
|
|
context.Background(),
|
|
c,
|
|
newAntigravityCompatAccount(AccountTypeOAuth),
|
|
body,
|
|
nil,
|
|
)
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result)
|
|
require.Len(t, upstream.requestBodies, 1)
|
|
require.Equal(t, tt.want, gjson.GetBytes(upstream.requestBodies[0], "request.generationConfig.maxOutputTokens").Int())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAntigravityCompatRoutesByMappedModelFamily(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
tests := []struct {
|
|
model string
|
|
wantSessionID bool
|
|
}{
|
|
{model: "gemini-3.1-pro-high", wantSessionID: false},
|
|
{model: "claude-sonnet-4-5", wantSessionID: true},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.model, func(t *testing.T) {
|
|
upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{antigravityCompatSuccessResponse()}}
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream)
|
|
body := []byte(`{"model":"` + tt.model + `","messages":[{"role":"user","content":"ok"}],"max_tokens":8}`)
|
|
c, _ := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", body)
|
|
|
|
result, err := svc.ForwardAsChatCompletions(
|
|
context.Background(),
|
|
c,
|
|
newAntigravityCompatAccount(AccountTypeOAuth),
|
|
body,
|
|
nil,
|
|
)
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result)
|
|
require.Len(t, upstream.requestBodies, 1)
|
|
require.Equal(t, tt.model, gjson.GetBytes(upstream.requestBodies[0], "model").String())
|
|
require.Equal(t, tt.wantSessionID, gjson.GetBytes(upstream.requestBodies[0], "request.sessionId").Exists())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAntigravityCompatUnauthorizedIsCredentialFailure(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{{
|
|
StatusCode: http.StatusUnauthorized,
|
|
Header: http.Header{"X-Request-Id": []string{"auth-3757"}},
|
|
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"Invalid bearer token"}}`)),
|
|
}}}
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream)
|
|
body := []byte(`{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"ok"}]}`)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", body)
|
|
|
|
result, err := svc.ForwardAsChatCompletions(
|
|
context.Background(),
|
|
c,
|
|
newAntigravityCompatAccount(AccountTypeOAuth),
|
|
body,
|
|
nil,
|
|
)
|
|
|
|
require.Nil(t, result)
|
|
var failoverErr *UpstreamFailoverError
|
|
require.ErrorAs(t, err, &failoverErr)
|
|
require.Equal(t, GatewayFailureStageAccountAuth, failoverErr.Stage)
|
|
require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
|
|
require.Equal(t, AntigravityCredentialRejectedReason, failoverErr.Reason)
|
|
require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
|
|
require.Equal(t, http.StatusBadGateway, failoverErr.ClientStatusCode)
|
|
require.Equal(t, AntigravityCredentialRejectedClientMessage, failoverErr.ClientMessage)
|
|
require.Equal(t, "auth-3757", failoverErr.ResponseHeaders.Get("X-Request-Id"))
|
|
require.Empty(t, recorder.Body.String())
|
|
}
|
|
|
|
func TestAntigravityCompatEmptyStreamTriggersFailover(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
tests := []struct {
|
|
name string
|
|
run func(*AntigravityGatewayService, *gin.Context, *http.Response) (*antigravityStreamResult, error)
|
|
}{
|
|
{
|
|
name: "chat completions",
|
|
run: func(svc *AntigravityGatewayService, c *gin.Context, resp *http.Response) (*antigravityStreamResult, error) {
|
|
return svc.handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high", true)
|
|
},
|
|
},
|
|
{
|
|
name: "responses",
|
|
run: func(svc *AntigravityGatewayService, c *gin.Context, resp *http.Response) (*antigravityStreamResult, error) {
|
|
return svc.handleResponsesStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high")
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
|
|
resp := &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader("data: malformed\n\ndata: [DONE]\n\n")),
|
|
}
|
|
|
|
result, err := tt.run(svc, c, resp)
|
|
|
|
require.Nil(t, result)
|
|
var failoverErr *UpstreamFailoverError
|
|
require.ErrorAs(t, err, &failoverErr)
|
|
require.True(t, failoverErr.RetryableOnSameAccount)
|
|
require.Empty(t, recorder.Body.String())
|
|
require.Empty(t, recorder.Header().Get("Content-Type"))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAntigravityCompatUsageOnlyStreamTriggersFailover(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
tests := []struct {
|
|
name string
|
|
run func(*AntigravityGatewayService, *gin.Context, *http.Response) (*antigravityStreamResult, error)
|
|
}{
|
|
{
|
|
name: "chat completions",
|
|
run: func(svc *AntigravityGatewayService, c *gin.Context, resp *http.Response) (*antigravityStreamResult, error) {
|
|
return svc.handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high", true)
|
|
},
|
|
},
|
|
{
|
|
name: "responses",
|
|
run: func(svc *AntigravityGatewayService, c *gin.Context, resp *http.Response) (*antigravityStreamResult, error) {
|
|
return svc.handleResponsesStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high")
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, "/", nil)
|
|
resp := &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader(
|
|
`data: {"response":{"responseId":"resp_3757","usageMetadata":{"promptTokenCount":8}}}` + "\n\n",
|
|
)),
|
|
}
|
|
|
|
result, err := tt.run(svc, c, resp)
|
|
|
|
require.Nil(t, result)
|
|
var failoverErr *UpstreamFailoverError
|
|
require.ErrorAs(t, err, &failoverErr)
|
|
require.True(t, failoverErr.RetryableOnSameAccount)
|
|
require.Empty(t, recorder.Body.String())
|
|
require.Empty(t, recorder.Header().Get("Content-Type"))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAntigravityCompatUsageOnlyNonStreamingTriggersFailover(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
tests := []struct {
|
|
name string
|
|
path string
|
|
body []byte
|
|
call func(*AntigravityGatewayService, context.Context, *gin.Context, *Account, []byte) (*ForwardResult, error)
|
|
}{
|
|
{
|
|
name: "chat completions",
|
|
path: "/v1/chat/completions",
|
|
body: []byte(`{"model":"gemini-3.1-pro-high","messages":[{"role":"user","content":"ok"}]}`),
|
|
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
|
return svc.ForwardAsChatCompletions(ctx, c, account, body, nil)
|
|
},
|
|
},
|
|
{
|
|
name: "responses",
|
|
path: "/v1/responses",
|
|
body: []byte(`{"model":"gemini-3.1-pro-high","input":"ok"}`),
|
|
call: func(svc *AntigravityGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
|
return svc.ForwardAsResponses(ctx, c, account, body, nil)
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader(
|
|
`data: {"response":{"responseId":"resp_3757","usageMetadata":{"promptTokenCount":8}}}` + "\n\n",
|
|
)),
|
|
}}}
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, tt.path, tt.body)
|
|
|
|
result, err := tt.call(
|
|
svc,
|
|
context.Background(),
|
|
c,
|
|
newAntigravityCompatAccount(AccountTypeOAuth),
|
|
tt.body,
|
|
)
|
|
|
|
require.Nil(t, result)
|
|
var failoverErr *UpstreamFailoverError
|
|
require.ErrorAs(t, err, &failoverErr)
|
|
require.True(t, failoverErr.RetryableOnSameAccount)
|
|
require.Empty(t, recorder.Body.String())
|
|
require.Empty(t, recorder.Header().Get("Content-Type"))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAntigravityCompatChatStreamMapsToolCallAndUsage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", nil)
|
|
body := `data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"functionCall":{"id":"call_3757","name":"get_weather","args":{"city":"Tokyo"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":3}}}` + "\n\n"
|
|
resp := &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader(body)),
|
|
}
|
|
|
|
result, err := svc.handleChatCompletionsStreamingFromAntigravity(
|
|
c,
|
|
resp,
|
|
time.Now(),
|
|
"gemini-3.1-pro-high",
|
|
true,
|
|
)
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result)
|
|
require.Equal(t, 8, result.usage.InputTokens)
|
|
require.Equal(t, 3, result.usage.OutputTokens)
|
|
require.Contains(t, recorder.Body.String(), `"tool_calls"`)
|
|
require.Contains(t, recorder.Body.String(), `"get_weather"`)
|
|
require.Contains(t, recorder.Body.String(), `"finish_reason":"tool_calls"`)
|
|
require.Contains(t, recorder.Body.String(), `"prompt_tokens":8`)
|
|
require.Equal(t, 1, strings.Count(recorder.Body.String(), "data: [DONE]"))
|
|
}
|
|
|
|
func TestAntigravityCompatFirstEventTimeoutTriggersFailover(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
svc := newAntigravityCompatService(
|
|
config.GatewayConfig{MaxLineSize: defaultMaxLineSize, StreamDataIntervalTimeout: 1},
|
|
nil,
|
|
)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", nil)
|
|
reader, writer := io.Pipe()
|
|
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: reader}
|
|
type outcome struct {
|
|
result *antigravityStreamResult
|
|
err error
|
|
}
|
|
done := make(chan outcome, 1)
|
|
|
|
go func() {
|
|
result, err := svc.handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high", false)
|
|
done <- outcome{result: result, err: err}
|
|
}()
|
|
|
|
select {
|
|
case got := <-done:
|
|
require.Nil(t, got.result)
|
|
var failoverErr *UpstreamFailoverError
|
|
require.ErrorAs(t, got.err, &failoverErr)
|
|
require.True(t, failoverErr.RetryableOnSameAccount)
|
|
require.Empty(t, recorder.Header().Get("Content-Type"))
|
|
case <-time.After(2 * time.Second):
|
|
_ = writer.Close()
|
|
_ = reader.Close()
|
|
t.Fatal("compat stream ignored StreamDataIntervalTimeout")
|
|
}
|
|
_ = writer.Close()
|
|
_ = reader.Close()
|
|
}
|
|
|
|
func TestAntigravityCompatClientDisconnectDrainsUsage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
|
c, _ := newAntigravityCompatContext(http.MethodPost, "/v1/chat/completions", nil)
|
|
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
|
body := strings.Join([]string{
|
|
`data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"text":"partial"}]}}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":1}}}`,
|
|
"",
|
|
`data: {"response":{"responseId":"resp_3757","candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":15}}}`,
|
|
"",
|
|
}, "\n")
|
|
resp := &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader(body)),
|
|
}
|
|
|
|
result, err := svc.handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high", false)
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result)
|
|
require.True(t, result.clientDisconnect)
|
|
require.Equal(t, 15, result.usage.OutputTokens)
|
|
}
|
|
|
|
func TestAntigravityCompatStreamErrorCommitsSingleTerminalFrame(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, nil)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
|
|
body := []byte(`data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"text":"partial"}]}}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":1}}}` + "\n\n")
|
|
resp := &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: &antigravityCompatErrorReader{
|
|
data: body,
|
|
err: io.ErrUnexpectedEOF,
|
|
},
|
|
}
|
|
|
|
result, err := svc.handleResponsesStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high")
|
|
|
|
require.Error(t, err)
|
|
require.Nil(t, result)
|
|
require.True(t, IsResponseCommitted(c))
|
|
require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: error"))
|
|
}
|
|
|
|
func TestAntigravityCompatKeepaliveAfterFirstEvent(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
svc := newAntigravityCompatService(
|
|
config.GatewayConfig{MaxLineSize: defaultMaxLineSize, StreamKeepaliveInterval: 1},
|
|
nil,
|
|
)
|
|
c, recorder := newAntigravityCompatContext(http.MethodPost, "/v1/responses", nil)
|
|
reader, writer := io.Pipe()
|
|
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: reader}
|
|
done := make(chan error, 1)
|
|
|
|
go func() {
|
|
_, err := svc.handleResponsesStreamingFromAntigravity(c, resp, time.Now(), "gemini-3.1-pro-high")
|
|
done <- err
|
|
}()
|
|
_, err := io.WriteString(
|
|
writer,
|
|
`data: {"response":{"responseId":"resp_3757","candidates":[{"content":{"parts":[{"text":"partial"}]}}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":1}}}`+"\n\n",
|
|
)
|
|
require.NoError(t, err)
|
|
time.Sleep(1200 * time.Millisecond)
|
|
require.NoError(t, writer.Close())
|
|
require.NoError(t, <-done)
|
|
require.Contains(t, recorder.Body.String(), ": ping\n\n")
|
|
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
|
|
require.NoError(t, reader.Close())
|
|
}
|