fix(grok): route OAuth subscriptions through CLI proxy

This commit is contained in:
Heatherm Huang
2026-07-13 10:11:33 +08:00
parent d9e466ad3a
commit 3375b4ed2b
10 changed files with 115 additions and 9 deletions
@@ -98,7 +98,7 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
require.Contains(t, rec.Body.String(), `"source":"active_probe"`)
require.Contains(t, rec.Body.String(), `"headers_observed":true`)
require.NotContains(t, rec.Body.String(), "access-token")
require.Equal(t, xai.DefaultBaseURL+"/responses", upstream.lastReq.URL.String())
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Contains(t, string(upstream.lastBody), `"store":false`)
require.NotNil(t, repo.updates[42])
+6
View File
@@ -1255,6 +1255,12 @@ func (a *Account) GetGrokBaseURL() string {
return ""
}
baseURL := a.GetCredential("base_url")
if a.IsGrokOAuth() {
normalizedBaseURL := strings.TrimRight(strings.TrimSpace(baseURL), "/")
if normalizedBaseURL == "" || strings.EqualFold(normalizedBaseURL, xai.DefaultBaseURL) {
return xai.DefaultCLIBaseURL
}
}
if baseURL != "" {
return baseURL
}
@@ -4,6 +4,9 @@ package service
import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/stretchr/testify/require"
)
func TestGetBaseURL(t *testing.T) {
@@ -158,3 +161,69 @@ func TestGetGeminiBaseURL(t *testing.T) {
})
}
}
func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
tests := []struct {
name string
account Account
expected string
}{
{
name: "oauth without base_url uses CLI subscription proxy",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{},
},
expected: xai.DefaultCLIBaseURL,
},
{
name: "oauth legacy API default is migrated at runtime to CLI subscription proxy",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": xai.DefaultBaseURL,
},
},
expected: xai.DefaultCLIBaseURL,
},
{
name: "oauth legacy API default with trailing slash is migrated at runtime",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": xai.DefaultBaseURL + "/",
},
},
expected: xai.DefaultCLIBaseURL,
},
{
name: "oauth explicit custom base_url remains supported",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://custom.example.com/v1",
},
},
expected: "https://custom.example.com/v1",
},
{
name: "API key without base_url uses official credit-backed API",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGrok,
Credentials: map[string]any{},
},
expected: xai.DefaultBaseURL,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.expected, tt.account.GetGrokBaseURL())
})
}
}
@@ -71,7 +71,7 @@ func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testin
err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault)
require.NoError(t, err)
require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer grok-access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
@@ -235,7 +235,7 @@ func (s *GrokOAuthService) BuildAccountCredentials(tokenInfo *GrokTokenInfo) map
if tokenInfo.EntitlementStatus != "" {
creds["entitlement_status"] = tokenInfo.EntitlementStatus
}
creds["base_url"] = xai.DefaultBaseURL
creds["base_url"] = xai.DefaultCLIBaseURL
return creds
}
@@ -5,6 +5,7 @@ package service
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/stretchr/testify/require"
@@ -66,3 +67,15 @@ func TestGrokOAuthServiceExchangeCodeRequiresStateForCallbackURLAndConsumesSessi
require.Contains(t, err.Error(), "GROK_OAUTH_SESSION_NOT_FOUND")
require.Zero(t, client.exchangeCalls)
}
func TestGrokOAuthServiceBuildAccountCredentialsDefaultsToSubscriptionProxy(t *testing.T) {
svc := NewGrokOAuthService(nil, &grokOAuthClientStub{})
defer svc.Stop()
credentials := svc.BuildAccountCredentials(&GrokTokenInfo{
AccessToken: "access-token",
ExpiresAt: time.Now().Add(time.Hour).Unix(),
})
require.Equal(t, xai.DefaultCLIBaseURL, credentials["base_url"])
}
@@ -112,7 +112,7 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
require.NotNil(t, result.Snapshot.Requests)
require.EqualValues(t, 10, *result.Snapshot.Requests.Limit)
require.EqualValues(t, 7, *result.Snapshot.Requests.Remaining)
require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())