diff --git a/backend/internal/repository/http_upstream_test.go b/backend/internal/repository/http_upstream_test.go index 90cbd15ae3..d34e1df02c 100644 --- a/backend/internal/repository/http_upstream_test.go +++ b/backend/internal/repository/http_upstream_test.go @@ -15,6 +15,51 @@ import ( "github.com/stretchr/testify/suite" ) +func TestHTTPUpstreamDoAppliesGrokCLIIdentityBeforeRoundTrip(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", "") + + 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) + + 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, + } + + 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) { t.Run("uses pinned stable version for the CLI proxy", func(t *testing.T) { t.Setenv("XAI_GROK_CLI_VERSION", "")