Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/openai_gateway_grok_cache.go

@@ -0,0 +1,549 @@
+package service
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"regexp"
+	"strings"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
+	"github.com/gin-gonic/gin"
+	"github.com/tidwall/gjson"
+	"github.com/tidwall/sjson"
+)
+
+const (
+	grokConversationIDHeader         = "X-Grok-Conv-Id"
+	claudeCodeSessionHeader          = "X-Claude-Code-Session-Id"
+	grokClientToolCacheOptInHeader   = "X-Sub2API-Grok-Client-Tool-Cache"
+	grokFreeCacheNativeToolsJSON     = `[{"type":"web_search"},{"type":"x_search"}]`
+	grokFreeCacheDisabledToolChoice  = "none"
+	grokClientToolCacheOptInExtraKey = "grok_client_tool_cache_enabled"
+)
+
+// Claude Code metadata.user_id often ends with _session_<uuid>.
+var claudeCodeSessionSuffixPattern = regexp.MustCompile(`_session_([a-f0-9-]+)$`)
+
+// extractClaudeCodeSessionID resolves the Claude Code conversation id from
+// headers or Anthropic/OpenAI-compatible payload metadata.
+func extractClaudeCodeSessionID(c *gin.Context, body []byte) string {
+	if c != nil {
+		if seed := strings.TrimSpace(c.GetHeader(claudeCodeSessionHeader)); seed != "" {
+			return seed
+		}
+	}
+	return extractClaudeCodeSessionIDFromPayload(body)
+}
+
+func extractClaudeCodeSessionIDFromPayload(body []byte) string {
+	if len(body) == 0 {
+		return ""
+	}
+	userID := strings.TrimSpace(gjson.GetBytes(body, "metadata.user_id").String())
+	if userID == "" {
+		return ""
+	}
+	if matches := claudeCodeSessionSuffixPattern.FindStringSubmatch(userID); len(matches) >= 2 {
+		return matches[1]
+	}
+	// Claude Code may embed JSON: {"session_id":"..."}
+	if len(userID) > 0 && userID[0] == '{' {
+		if sid := strings.TrimSpace(gjson.Get(userID, "session_id").String()); sid != "" {
+			return sid
+		}
+	}
+	return ""
+}
+
+// resolveGrokCacheIdentity derives one stable, tenant-isolated routing identity
+// for xAI's server-side prompt cache. The returned value is safe to expose to
+// the upstream: it never contains the client's raw session identifier.
+//
+// A valid downstream API key is required. This intentionally fails closed on
+// internal probes and incomplete request contexts instead of creating a cache
+// identity that could be shared by unrelated tenants.
+func resolveGrokCacheIdentity(c *gin.Context, body []byte, explicitKey, upstreamModel string) string {
+	apiKeyID := getAPIKeyIDFromContext(c)
+	if apiKeyID <= 0 {
+		return ""
+	}
+	// /responses/compact rejects tool_choice and does not represent a normal
+	// conversation turn. Keep both cache identity and Free-tier routing
+	// augmentation out of this path.
+	if isOpenAIResponsesCompactPath(c) {
+		return ""
+	}
+
+	model := strings.ToLower(strings.TrimSpace(upstreamModel))
+	if model == "" {
+		return ""
+	}
+
+	seed := explicitGrokCacheSeed(c, body, explicitKey)
+	if seed == "" {
+		seed = deriveOpenAIStablePrefixSessionSeed(body)
+		if seed == "" {
+			// A model alone is too broad for cache routing. Preserve the
+			// existing first-user-derived identity when no reusable prefix is
+			// available so unrelated prompts do not share one tenant-wide key.
+			seed = deriveOpenAIAnchoredContentSessionSeed(body)
+		}
+	}
+	if seed == "" {
+		return ""
+	}
+
+	// generateSessionUUID hashes the whole seed before formatting it as a UUID.
+	// Include a versioned namespace so this identity cannot collide with other
+	// upstream session identifiers derived by sub2api.
+	isolatedSeed := fmt.Sprintf("grok-prompt-cache:v1:%d:%s:%s", apiKeyID, model, seed)
+	return generateSessionUUID(isolatedSeed)
+}
+
+func explicitGrokCacheSeed(c *gin.Context, body []byte, explicitKey string) string {
+	// Claude Code session is the most stable multi-turn identity for
+	// /v1/messages → Grok bridges. Prefer it over generic session headers so
+	// prompt cache routing follows the gateway's existing cache affinity rules.
+	seed := extractClaudeCodeSessionID(c, body)
+	if seed == "" {
+		seed = explicitOpenAIHeaderSessionID(c)
+	}
+	if seed == "" && c != nil {
+		seed = strings.TrimSpace(c.GetHeader(grokConversationIDHeader))
+	}
+	if seed == "" && len(body) > 0 {
+		seed = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String())
+	}
+	if seed == "" {
+		seed = strings.TrimSpace(explicitKey)
+	}
+	// previous_response_id is last-resort: multi-turn Responses without an
+	// explicit session still share one cache identity (model is already in the
+	// isolated seed). Message ids are rejected by the seed helper.
+	if seed == "" && len(body) > 0 {
+		seed = grokPreviousResponseSessionSeed(body)
+	}
+	return seed
+}
+
+func isGrokRequestContext(c *gin.Context) bool {
+	if c == nil {
+		return false
+	}
+	if c.Request != nil {
+		if platform, ok := ResolvedTargetPlatformFromContext(c.Request.Context()); ok {
+			return platform == PlatformGrok
+		}
+	}
+	v, exists := c.Get("api_key")
+	if !exists {
+		return false
+	}
+	apiKey, ok := v.(*APIKey)
+	return ok && apiKey != nil && apiKey.Group != nil && apiKey.Group.Platform == PlatformGrok
+}
+
+// applyGrokResponsesCacheIdentity writes the cache routing identity into an
+// xAI Responses request. Existing client values are deliberately replaced by
+// the tenant-isolated value to prevent collisions on shared OAuth accounts.
+//
+// Free OAuth requests without native search tools are routed by xAI to the
+// non-cacheable build-free model. For otherwise tool-free requests, add the
+// native tools with tool_choice=none: this selects the cache-capable tier
+// without allowing an actual search. Explicit client function tools are handled by
+// applyGrokFreeMessagesFunctionToolCacheRoute (Messages bridge and native Responses).
+func applyGrokResponsesCacheIdentity(body, intentSourceBody []byte, identity string, injectFreeTierTools bool) ([]byte, error) {
+	identity = strings.TrimSpace(identity)
+	if identity == "" {
+		if gjson.GetBytes(body, "prompt_cache_key").Exists() {
+			return sjson.DeleteBytes(body, "prompt_cache_key")
+		}
+		return body, nil
+	}
+	out, err := sjson.SetBytes(body, "prompt_cache_key", identity)
+	if err != nil {
+		return nil, err
+	}
+	if !injectFreeTierTools {
+		return out, nil
+	}
+	// Inspect the pre-sanitization source. patchGrokResponsesBody may remove an
+	// unsupported client tool and its tool_choice; that must not turn an
+	// explicit client tool intent into an eligible native-tool request.
+	if hasGrokResponsesToolIntent(intentSourceBody) {
+		return out, nil
+	}
+	out, err = sjson.SetRawBytes(out, "tools", []byte(grokFreeCacheNativeToolsJSON))
+	if err != nil {
+		return nil, err
+	}
+	return sjson.SetBytes(out, "tool_choice", grokFreeCacheDisabledToolChoice)
+}
+
+func hasGrokResponsesToolIntent(body []byte) bool {
+	if gjson.GetBytes(body, "tools").Exists() || gjson.GetBytes(body, "tool_choice").Exists() {
+		return true
+	}
+	input := gjson.GetBytes(body, "input")
+	if !input.IsArray() {
+		return false
+	}
+	for _, item := range input.Array() {
+		if strings.TrimSpace(item.Get("type").String()) != "additional_tools" {
+			continue
+		}
+		tools := item.Get("tools")
+		if !tools.Exists() || !tools.IsArray() || len(tools.Array()) > 0 {
+			return true
+		}
+	}
+	return false
+}
+
+// applyGrokFreeMessagesFunctionToolCacheRoute enables xAI's cache-capable
+// mixed-tools route only for known Free accounts. Pure client tools default to
+// the cache-capable route so an intermediate sub2api does not need to preserve
+// client-specific opt-in headers. Operators can explicitly disable this per
+// account when native search tools would change the desired behavior (#4486).
+func applyGrokFreeMessagesFunctionToolCacheRoute(body, intentSourceBody []byte, account *Account, cacheIdentity string) ([]byte, error) {
+	allowPureClientTools, _ := grokClientToolCacheAccountPolicy(account)
+	return applyGrokFreeToolCacheRoute(body, intentSourceBody, account, cacheIdentity, allowPureClientTools, true)
+}
+
+// applyGrokFreeRequestToolCacheRoute also accepts a request-scoped opt-in. The
+// sub2api header is consumed locally because buildGrokResponsesRequest only
+// forwards the explicitly supported OpenAI-Beta header from downstream.
+func applyGrokFreeRequestToolCacheRoute(c *gin.Context, body, intentSourceBody []byte, account *Account, cacheIdentity string) ([]byte, error) {
+	allowPureClientTools, accountPolicyExplicit := grokClientToolCacheAccountPolicy(account)
+	requestOptOut := false
+	if c != nil {
+		switch strings.ToLower(strings.TrimSpace(c.GetHeader(grokClientToolCacheOptInHeader))) {
+		case "1", "true", "yes", "on", "prefer-cache":
+			allowPureClientTools = true
+		case "0", "false", "no", "off":
+			allowPureClientTools = false
+			requestOptOut = true
+		}
+	}
+	if !allowPureClientTools && !accountPolicyExplicit && !requestOptOut && isGrokClaudeDesktopResponsesCacheRequest(c) {
+		allowPureClientTools = true
+	}
+	// A function merely named web_search/x_search is still a client function.
+	// Known Free OAuth accounts use the cache route by default; a request-scoped
+	// opt-in may override an account opt-out, while an explicit request opt-out
+	// always wins. The legacy Claude fingerprint remains only as a compatibility
+	// fallback when no account policy has been recorded (#4486).
+	return applyGrokFreeToolCacheRoute(body, intentSourceBody, account, cacheIdentity, allowPureClientTools, allowPureClientTools)
+}
+
+// grokClientToolCacheAccountPolicy is intentionally strict for configured
+// values: only a JSON boolean is accepted. A missing key defaults on solely for
+// accounts positively identified as Grok Free OAuth; paid, API-key, and unknown
+// accounts remain fail-closed.
+func grokClientToolCacheAccountPolicy(account *Account) (enabled, explicit bool) {
+	if !isKnownGrokFreeAccount(account) {
+		return false, false
+	}
+	if account.Extra == nil {
+		return true, false
+	}
+	value, exists := account.Extra[grokClientToolCacheOptInExtraKey]
+	if !exists {
+		return true, false
+	}
+	enabled, valid := value.(bool)
+	if !valid {
+		return false, true
+	}
+	return enabled, true
+}
+
+// isGrokClaudeDesktopResponsesCacheRequest recognizes the strict wire
+// fingerprint emitted when Claude Desktop's local agent is translated by
+// CC Switch into an OpenAI Responses request. Requiring every independent
+// signal prevents a generic Claude-compatible client (or the Chat bridge)
+// from silently opting into the mixed native/client tool route.
+func isGrokClaudeDesktopResponsesCacheRequest(c *gin.Context) bool {
+	if c == nil || c.Request == nil || c.Request.URL == nil || isOpenAIResponsesCompactPath(c) {
+		return false
+	}
+	path := strings.TrimRight(strings.TrimSpace(c.Request.URL.Path), "/")
+	if !strings.HasSuffix(path, "/responses") {
+		return false
+	}
+
+	if !claudeCodeUAPattern.MatchString(strings.TrimSpace(c.GetHeader("User-Agent"))) {
+		return false
+	}
+	switch strings.ToLower(strings.TrimSpace(c.GetHeader("X-App"))) {
+	case "cli", "cli-bg":
+	default:
+		return false
+	}
+	if !strings.EqualFold(strings.TrimSpace(c.GetHeader("anthropic-client-platform")), "desktop_app") {
+		return false
+	}
+	return strings.TrimSpace(c.GetHeader("X-Claude-Code-Session-Id")) != ""
+}
+
+func applyGrokFreeToolCacheRoute(body, intentSourceBody []byte, account *Account, cacheIdentity string, allowPureClientTools, allowFunctionSearch bool) ([]byte, error) {
+	if strings.TrimSpace(cacheIdentity) == "" || !isKnownGrokFreeAccount(account) {
+		return body, nil
+	}
+	intentTools := gjson.GetBytes(intentSourceBody, "tools")
+	intentToolChoice := gjson.GetBytes(intentSourceBody, "tool_choice")
+	if !isGrokFreeCacheFunctionToolIntent(intentTools, intentToolChoice) {
+		return body, nil
+	}
+	if intentToolChoice.Type == gjson.String && strings.TrimSpace(intentToolChoice.String()) == grokFreeCacheDisabledToolChoice {
+		// Adding native cache-routing tools cannot change behavior when the
+		// client has explicitly disabled all tool execution.
+		return appendGrokFreeCacheNativeToolsWithPolicy(body, true, false)
+	}
+	return appendGrokFreeCacheNativeToolsWithPolicy(body, allowPureClientTools, allowFunctionSearch)
+}
+
+// isKnownGrokFreeAccount recognizes free-tier Grok accounts, used for
+// Free cache routing / media free_tier blocks (broader than soft-gate).
+// Soft-gate uses isExplicitGrokFreeOAuthAccount (exact "free" only).
+func isKnownGrokFreeAccount(account *Account) bool {
+	if account == nil || !account.IsGrokOAuth() {
+		return false
+	}
+	// Live access-token JWT wins over stale billing/credential snapshots
+	// so a downgrade to free is visible as soon as the AT is refreshed.
+	if jwtTier := xai.SubscriptionTierFromJWT(account.GetCredential("access_token")); jwtTier != "" {
+		return isGrokFreeSubscriptionTier(jwtTier)
+	}
+	freeSignal := false
+	paidSignal := false
+	inferredFreeSignal := false
+	if billing, err := grokBillingSnapshotFromExtra(account.Extra); err == nil && billing != nil {
+		if tier := strings.TrimSpace(billing.Plan); tier != "" {
+			if isGrokFreeSubscriptionTier(tier) {
+				freeSignal = true
+			} else if !isGrokUnknownSubscriptionTier(tier) {
+				paidSignal = true
+			}
+		}
+		// Usage % or a monthly dollar cap is evidence of a paid plan.
+		if billing.UsagePercent != nil || billing.UsedPercent != nil ||
+			(billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0) {
+			paidSignal = true
+		}
+		// Empty plan + successful monthly observation → inferred free (no paid plan/limit).
+		if strings.TrimSpace(billing.MonthlyUpdatedAt) != "" ||
+			(billing.StatusCode >= http.StatusOK && billing.StatusCode < http.StatusMultipleChoices &&
+				!billing.Partial && len(billing.FailedWindows) == 0) {
+			inferredFreeSignal = true
+		}
+	}
+	if snapshot, err := grokQuotaSnapshotFromExtra(account.Extra); err == nil && snapshot != nil {
+		if tier := strings.TrimSpace(snapshot.SubscriptionTier); tier != "" {
+			if isGrokFreeSubscriptionTier(tier) {
+				freeSignal = true
+			} else if !isGrokUnknownSubscriptionTier(tier) {
+				paidSignal = true
+			}
+		}
+		if snapshot.Tokens != nil && snapshot.Tokens.Limit != nil &&
+			xai.IsGrokFreeRolling24hTokenLimit(*snapshot.Tokens.Limit) {
+			inferredFreeSignal = true
+		}
+	}
+	// Only credentials subscription_tier is authoritative here (not plan_type / extra keys).
+	if tier := strings.TrimSpace(account.GetCredential("subscription_tier")); tier != "" {
+		if isGrokFreeSubscriptionTier(tier) {
+			freeSignal = true
+		} else if !isGrokUnknownSubscriptionTier(tier) {
+			paidSignal = true
+		}
+	}
+	// Explicit paid evidence always wins over an inferred Free signal.
+	return !paidSignal && (freeSignal || inferredFreeSignal)
+}
+
+func isGrokFreeSubscriptionTier(tier string) bool {
+	switch xai.NormalizeSubscriptionTier(tier) {
+	case "free", "x_basic":
+		return true
+	default:
+		return false
+	}
+}
+
+func isGrokUnknownSubscriptionTier(tier string) bool {
+	switch strings.ToLower(strings.TrimSpace(tier)) {
+	case "", "unknown", "n/a", "none":
+		return true
+	default:
+		return false
+	}
+}
+
+func isGrokFreeCacheFunctionToolIntent(tools, toolChoice gjson.Result) bool {
+	if !tools.IsArray() {
+		return false
+	}
+	items := tools.Array()
+	if len(items) == 0 {
+		return false
+	}
+	for _, tool := range items {
+		if !tool.IsObject() {
+			return false
+		}
+		toolType := strings.TrimSpace(tool.Get("type").String())
+		if _, ok := grokResponsesSupportedToolTypes[toolType]; !ok {
+			return false
+		}
+		if toolType == "function" {
+			// Responses function declarations keep name at the top level. Reject
+			// Chat Completions' nested function shape and incomplete declarations.
+			if strings.TrimSpace(tool.Get("name").String()) == "" || tool.Get("function").Exists() {
+				return false
+			}
+		}
+	}
+	if !toolChoice.Exists() {
+		return true
+	}
+	if toolChoice.Type != gjson.String {
+		return false
+	}
+	switch strings.TrimSpace(toolChoice.String()) {
+	case "auto", grokFreeCacheDisabledToolChoice:
+		return true
+	default:
+		return false
+	}
+}
+
+func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) {
+	return appendGrokFreeCacheNativeTools(body, false)
+}
+
+func appendGrokFreeCacheNativeTools(body []byte, allowPureClientTools bool) ([]byte, error) {
+	return appendGrokFreeCacheNativeToolsWithPolicy(body, allowPureClientTools, true)
+}
+
+func appendGrokFreeCacheNativeToolsWithPolicy(body []byte, allowPureClientTools, allowFunctionSearch bool) ([]byte, error) {
+	tools := gjson.GetBytes(body, "tools")
+	if !tools.Exists() || !tools.IsArray() {
+		return body, nil
+	}
+
+	items := tools.Array()
+	if len(items) == 0 {
+		return body, nil
+	}
+	hasNativeSearch := false
+	for _, tool := range items {
+		switch strings.TrimSpace(tool.Get("type").String()) {
+		case "web_search", "x_search":
+			hasNativeSearch = true
+		}
+	}
+	if !allowPureClientTools && !allowFunctionSearch && !hasNativeSearch {
+		return body, nil
+	}
+	merged := make([]json.RawMessage, 0, len(items)+2)
+	present := make(map[string]bool, 2)
+	hasCompanionTool := false
+	for _, tool := range items {
+		toolType := strings.TrimSpace(tool.Get("type").String())
+		switch toolType {
+		case "function":
+			name := strings.TrimSpace(tool.Get("name").String())
+			if !tool.IsObject() || name == "" || tool.Get("function").Exists() {
+				return body, nil
+			}
+			// Grok Build may declare search as function tools. Convert to native
+			// entries so Free OAuth stays cache-capable without duplicate names.
+			if (name == "web_search" || name == "x_search") && allowFunctionSearch {
+				if present[name] {
+					continue
+				}
+				raw, err := json.Marshal(map[string]string{"type": name})
+				if err != nil {
+					return nil, err
+				}
+				merged = append(merged, raw)
+				present[name] = true
+				if allowPureClientTools {
+					hasCompanionTool = true
+				}
+				continue
+			}
+			if name == "web_search" || name == "x_search" {
+				// Keep the client function intact and avoid adding a same-named
+				// native tool unless conversion was explicitly enabled.
+				present[name] = true
+			}
+			hasCompanionTool = true
+			merged = append(merged, json.RawMessage(tool.Raw))
+		case "web_search", "x_search":
+			if present[toolType] {
+				continue
+			}
+			merged = append(merged, json.RawMessage(tool.Raw))
+			present[toolType] = true
+		default:
+			if _, ok := grokResponsesSupportedToolTypes[toolType]; !ok {
+				return body, nil
+			}
+			hasCompanionTool = true
+			merged = append(merged, json.RawMessage(tool.Raw))
+		}
+	}
+	if !hasCompanionTool {
+		return body, nil
+	}
+	// Only complement missing native search tools when the request already contains
+	// at least one search tool (native or function-form). Pure client function tools
+	// (e.g. view_image) must not trigger injection to avoid biasing model tool
+	// selection (#4486).
+	if !allowPureClientTools && !present["web_search"] && !present["x_search"] {
+		return body, nil
+	}
+	for _, toolType := range []string{"web_search", "x_search"} {
+		if present[toolType] {
+			continue
+		}
+		raw, err := json.Marshal(map[string]string{"type": toolType})
+		if err != nil {
+			return nil, err
+		}
+		merged = append(merged, raw)
+	}
+	encoded, err := json.Marshal(merged)
+	if err != nil {
+		return nil, err
+	}
+	return sjson.SetRawBytes(body, "tools", encoded)
+}
+
+// applyGrokCacheHeaders applies the documented Chat Completions conversation
+// routing header. The request is built from a fresh header map, so client
+// supplied x-grok headers cannot override this server-derived value.
+func applyGrokCacheHeaders(headers http.Header, identity string) {
+	if headers == nil {
+		return
+	}
+	identity = strings.TrimSpace(identity)
+	if identity == "" {
+		headers.Del(grokConversationIDHeader)
+		return
+	}
+	headers.Set(grokConversationIDHeader, identity)
+}
+
+// stripGrokChatPromptCacheKey removes the Responses-only body field after it
+// has been used as an identity seed. Chat Completions routes cache by header.
+func stripGrokChatPromptCacheKey(body []byte) ([]byte, error) {
+	if !gjson.GetBytes(body, "prompt_cache_key").Exists() {
+		return body, nil
+	}
+	return sjson.DeleteBytes(body, "prompt_cache_key")
+}