diff --git a/backend/internal/repository/http_upstream_test.go b/backend/internal/repository/http_upstream_test.go index d34e1df02c..cb8fe70d5b 100644 --- a/backend/internal/repository/http_upstream_test.go +++ b/backend/internal/repository/http_upstream_test.go @@ -15,49 +15,58 @@ import ( "github.com/stretchr/testify/suite" ) -func TestHTTPUpstreamDoAppliesGrokCLIIdentityBeforeRoundTrip(t *testing.T) { +func TestHTTPUpstreamDoAppliesGrokCLIIdentityBeforeOAuthRoundTrip(t *testing.T) { t.Setenv("XAI_GROK_CLI_VERSION", "") - upstream := NewHTTPUpstream(nil) - svc, ok := upstream.(*httpUpstreamService) - require.True(t, ok) + for _, endpoint := range []string{"responses", "chat/completions"} { + t.Run(endpoint, func(t *testing.T) { + upstream := NewHTTPUpstream(nil) + svc, ok := upstream.(*httpUpstreamService) + require.True(t, ok) - const accountID int64 = 4079 - isolation := svc.getIsolationMode() - profile := service.HTTPUpstreamProfileDefault - proxyKey := directProxyKey - protocolMode := svc.resolveProtocolMode(profile, proxyKey, nil) - settings := svc.resolvePoolSettings(isolation, 1) - settings = svc.applyProfilePoolSettings(settings, profile) - cacheKey := buildCacheKey(isolation, proxyKey, accountID, protocolMode) + const accountID int64 = 4084 + isolation := svc.getIsolationMode() + profile := service.HTTPUpstreamProfileDefault + proxyKey := directProxyKey + protocolMode := svc.resolveProtocolMode(profile, proxyKey, nil) + settings := svc.resolvePoolSettings(isolation, 1) + settings = svc.applyProfilePoolSettings(settings, profile) + cacheKey := buildCacheKey(isolation, proxyKey, accountID, protocolMode) - var capturedHeaders http.Header - svc.clients[cacheKey] = &upstreamClientEntry{ - client: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { - capturedHeaders = req.Header.Clone() - return &http.Response{ - StatusCode: http.StatusOK, - Header: make(http.Header), - Body: http.NoBody, - Request: req, - }, nil - })}, - proxyKey: proxyKey, - poolKey: buildPoolKey(settings, protocolMode), - protocolMode: protocolMode, + var capturedHeaders http.Header + svc.clients[cacheKey] = &upstreamClientEntry{ + client: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + capturedHeaders = req.Header.Clone() + statusCode := http.StatusOK + if req.Header.Get("X-XAI-Token-Auth") != "xai-grok-cli" { + statusCode = http.StatusForbidden + } + return &http.Response{ + StatusCode: statusCode, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + }, nil + })}, + proxyKey: proxyKey, + poolKey: buildPoolKey(settings, protocolMode), + protocolMode: protocolMode, + } + + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/"+endpoint, nil) + require.NoError(t, err) + req.Header.Set("User-Agent", "sub2api-grok/1.0") + + resp, err := svc.Do(req, "", accountID, 1) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.NoError(t, resp.Body.Close()) + + require.Equal(t, "0.2.93", capturedHeaders.Get("x-grok-client-version")) + require.Equal(t, "xai-grok-cli", capturedHeaders.Get("X-XAI-Token-Auth")) + require.Equal(t, "xai-grok-workspace/0.2.93", capturedHeaders.Get("User-Agent")) + }) } - - req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil) - require.NoError(t, err) - req.Header.Set("User-Agent", "sub2api-grok/1.0") - - resp, err := svc.Do(req, "", accountID, 1) - require.NoError(t, err) - require.NoError(t, resp.Body.Close()) - - require.Equal(t, "0.2.93", capturedHeaders.Get("x-grok-client-version")) - require.Equal(t, "xai-grok-cli", capturedHeaders.Get("X-XAI-Token-Auth")) - require.Equal(t, "xai-grok-workspace/0.2.93", capturedHeaders.Get("User-Agent")) } func TestApplyGrokCLIProxyHeaders(t *testing.T) {