feat: add grok subscription support

backend/internal/repository/grok_oauth_client.go

@@ -0,0 +1,113 @@
+package repository
+
+import (
+	"context"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+
+	infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/imroc/req/v3"
+)
+
+type grokOAuthClient struct {
+	tokenURL string
+}
+
+func NewGrokOAuthClient() service.GrokOAuthClient {
+	return &grokOAuthClient{tokenURL: xai.EffectiveTokenURL()}
+}
+
+func (c *grokOAuthClient) ExchangeCode(ctx context.Context, code, codeVerifier, codeChallenge, redirectURI, proxyURL, clientID string) (*xai.TokenResponse, error) {
+	client, err := createGrokReqClient(proxyURL)
+	if err != nil {
+		return nil, infraerrors.Newf(http.StatusBadGateway, "GROK_OAUTH_CLIENT_INIT_FAILED", "create HTTP client: %v", err)
+	}
+
+	clientID = strings.TrimSpace(clientID)
+	if clientID == "" {
+		clientID = xai.EffectiveClientID()
+	}
+
+	formData := url.Values{}
+	formData.Set("grant_type", "authorization_code")
+	formData.Set("client_id", clientID)
+	formData.Set("code", code)
+	formData.Set("redirect_uri", xai.EffectiveRedirectURI(redirectURI))
+	formData.Set("code_verifier", codeVerifier)
+	formData.Set("code_challenge", codeChallenge)
+	formData.Set("code_challenge_method", "S256")
+
+	var tokenResp xai.TokenResponse
+	resp, err := client.R().
+		SetContext(ctx).
+		SetHeader("User-Agent", "sub2api-grok-oauth/1.0").
+		SetFormDataFromValues(formData).
+		SetSuccessResult(&tokenResp).
+		Post(c.tokenURL)
+	if err != nil {
+		return nil, infraerrors.Newf(http.StatusBadGateway, "GROK_OAUTH_REQUEST_FAILED", "request failed: %v", err)
+	}
+	if !resp.IsSuccessState() {
+		return nil, grokOAuthStatusError("GROK_OAUTH_TOKEN_EXCHANGE_FAILED", "token exchange failed", resp)
+	}
+	return &tokenResp, nil
+}
+
+func (c *grokOAuthClient) RefreshToken(ctx context.Context, refreshToken, proxyURL, clientID string) (*xai.TokenResponse, error) {
+	client, err := createGrokReqClient(proxyURL)
+	if err != nil {
+		return nil, infraerrors.Newf(http.StatusBadGateway, "GROK_OAUTH_CLIENT_INIT_FAILED", "create HTTP client: %v", err)
+	}
+
+	clientID = strings.TrimSpace(clientID)
+	if clientID == "" {
+		clientID = xai.EffectiveClientID()
+	}
+
+	formData := url.Values{}
+	formData.Set("grant_type", "refresh_token")
+	formData.Set("client_id", clientID)
+	formData.Set("refresh_token", refreshToken)
+
+	var tokenResp xai.TokenResponse
+	resp, err := client.R().
+		SetContext(ctx).
+		SetHeader("User-Agent", "sub2api-grok-oauth/1.0").
+		SetFormDataFromValues(formData).
+		SetSuccessResult(&tokenResp).
+		Post(c.tokenURL)
+	if err != nil {
+		return nil, infraerrors.Newf(http.StatusBadGateway, "GROK_OAUTH_REQUEST_FAILED", "request failed: %v", err)
+	}
+	if !resp.IsSuccessState() {
+		return nil, grokOAuthStatusError("GROK_OAUTH_TOKEN_REFRESH_FAILED", "token refresh failed", resp)
+	}
+	return &tokenResp, nil
+}
+
+func createGrokReqClient(proxyURL string) (*req.Client, error) {
+	return getSharedReqClient(reqClientOptions{
+		ProxyURL: proxyURL,
+		Timeout:  60 * time.Second,
+	})
+}
+
+func grokOAuthStatusError(code, message string, resp *req.Response) error {
+	statusCode := http.StatusBadGateway
+	errorCode := code
+	upstreamStatus := 0
+	if resp != nil && resp.StatusCode == http.StatusForbidden {
+		statusCode = http.StatusForbidden
+		errorCode = "GROK_OAUTH_ENTITLEMENT_DENIED"
+	}
+	body := ""
+	if resp != nil {
+		upstreamStatus = resp.StatusCode
+		body = resp.String()
+	}
+	return infraerrors.Newf(statusCode, errorCode, "%s: status %d, body: %s", message, upstreamStatus, body)
+}

backend/internal/repository/grok_oauth_client_test.go

@@ -0,0 +1,88 @@
+//go:build unit
+
+package repository
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGrokOAuthClientExchangeAndRefreshUseFormFields(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		require.Equal(t, http.MethodPost, r.Method)
+		require.NoError(t, r.ParseForm())
+		require.Equal(t, "client-id", r.Form.Get("client_id"))
+
+		switch r.Form.Get("grant_type") {
+		case "authorization_code":
+			require.Equal(t, "auth-code", r.Form.Get("code"))
+			require.Equal(t, "http://127.0.0.1:56121/callback", r.Form.Get("redirect_uri"))
+			require.Equal(t, "verifier", r.Form.Get("code_verifier"))
+			require.Equal(t, "challenge", r.Form.Get("code_challenge"))
+			require.Equal(t, "S256", r.Form.Get("code_challenge_method"))
+			_ = json.NewEncoder(w).Encode(map[string]any{
+				"access_token":  "exchange-access",
+				"refresh_token": "exchange-refresh",
+				"token_type":    "Bearer",
+				"expires_in":    3600,
+				"scope":         "openid api:access",
+			})
+		case "refresh_token":
+			require.Equal(t, "refresh-token", r.Form.Get("refresh_token"))
+			_ = json.NewEncoder(w).Encode(map[string]any{
+				"access_token":  "refresh-access",
+				"refresh_token": "refresh-rotated",
+				"token_type":    "Bearer",
+				"expires_in":    7200,
+			})
+		default:
+			http.Error(w, "unexpected grant_type", http.StatusBadRequest)
+		}
+	}))
+	defer server.Close()
+	t.Setenv(xai.EnvTokenURL, server.URL)
+
+	client := NewGrokOAuthClient()
+
+	exchanged, err := client.ExchangeCode(
+		context.Background(),
+		"auth-code",
+		"verifier",
+		"challenge",
+		"http://127.0.0.1:56121/callback",
+		"",
+		"client-id",
+	)
+	require.NoError(t, err)
+	require.Equal(t, "exchange-access", exchanged.AccessToken)
+	require.Equal(t, "exchange-refresh", exchanged.RefreshToken)
+	require.Equal(t, int64(3600), exchanged.ExpiresIn)
+	require.Equal(t, "openid api:access", exchanged.Scope)
+
+	refreshed, err := client.RefreshToken(context.Background(), "refresh-token", "", "client-id")
+	require.NoError(t, err)
+	require.Equal(t, "refresh-access", refreshed.AccessToken)
+	require.Equal(t, "refresh-rotated", refreshed.RefreshToken)
+	require.Equal(t, int64(7200), refreshed.ExpiresIn)
+}
+
+func TestGrokOAuthClientRefreshForbiddenClassifiesEntitlement(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusForbidden)
+		_, _ = w.Write([]byte(`{"error":"subscription required"}`))
+	}))
+	defer server.Close()
+	t.Setenv(xai.EnvTokenURL, server.URL)
+
+	client := NewGrokOAuthClient()
+	_, err := client.RefreshToken(context.Background(), "refresh-token", "", "client-id")
+	require.Error(t, err)
+	require.Contains(t, strings.ToUpper(err.Error()), "GROK_OAUTH_ENTITLEMENT_DENIED")
+}

backend/internal/repository/simple_mode_default_groups.go

@@ -19,6 +19,7 @@ func ensureSimpleModeDefaultGroups(ctx context.Context, client *dbent.Client) er
 		service.PlatformOpenAI:      1,
 		service.PlatformGemini:      1,
 		service.PlatformAntigravity: 2,
+		service.PlatformGrok:        1,
 	}
 
 	for platform, minCount := range requiredByPlatform {

backend/internal/repository/wire.go

@@ -141,6 +141,7 @@ var ProviderSet = wire.NewSet(
 	NewClaudeOAuthClient,
 	NewHTTPUpstream,
 	NewOpenAIOAuthClient,
+	NewGrokOAuthClient,
 	NewGeminiOAuthClient,
 	NewGeminiCliCodeAssistClient,
 	NewGeminiDriveClient,