perf(go): 强化 prompt_cache 同会话粘号，提升多轮缓存命中

Codex 多轮依赖 prompt_cache_key + 同账号；候选池 top-N 可能漏掉粘号账号。
- 按 prompt_cache_key / X-Grok-Conv-Id 注入 sticky 账号到候选窗口
- fingerprint 优先 prompt_cache_key 与会话头
- 配合既有 affinity，减少切号导致的缓存冷启动

Co-Authored-By: Claude <noreply@anthropic.com>

Author: root

internal/proxy/chat.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net/http"
 	"strings"
 	"time"
 
@@ -663,11 +664,23 @@ func upstreamAccounts(chain []pool.Candidate) []grok.Account {
 }
 
 func ChatFingerprint(request ChatRequest) string {
-	for _, key := range []string{"conversation_id", "conversation", "thread_id", "session_id", "prompt_cache_key"} {
+	if request.Raw == nil {
+		return ""
+	}
+	// Explicit sticky keys first (Codex / OpenAI Responses multi-turn).
+	for _, key := range []string{"prompt_cache_key", "conversation_id", "conversation", "thread_id", "session_id"} {
 		if value, _ := request.Raw[key].(string); strings.TrimSpace(value) != "" {
 			return "chat:" + strings.TrimSpace(request.Model) + ":" + key + ":" + strings.TrimSpace(value)
 		}
 	}
+	// Nested metadata (Anthropic / some relays).
+	if meta, _ := request.Raw["metadata"].(map[string]any); meta != nil {
+		for _, key := range []string{"prompt_cache_key", "session_id", "sessionId", "thread_id", "conversation_id", "user_id"} {
+			if value, _ := meta[key].(string); strings.TrimSpace(value) != "" {
+				return "chat:" + strings.TrimSpace(request.Model) + ":meta:" + key + ":" + strings.TrimSpace(value)
+			}
+		}
+	}
 	messages, ok := request.Raw["messages"].([]any)
 	if !ok || len(messages) == 0 {
 		return ""
@@ -680,6 +693,35 @@ func ChatFingerprint(request ChatRequest) string {
 	return "chat:" + strings.TrimSpace(request.Model) + ":messages:" + hex.EncodeToString(sum[:16])
 }
 
+// ChatFingerprintFromHeaders picks sticky keys from common client/proxy headers
+// (Codex X-Grok-Conv-Id, session/thread headers).
+func ChatFingerprintFromHeaders(headers http.Header, model string) string {
+	if headers == nil {
+		return ""
+	}
+	get := func(names ...string) string {
+		for _, name := range names {
+			if v := strings.TrimSpace(headers.Get(name)); v != "" {
+				return v
+			}
+		}
+		return ""
+	}
+	if v := get("X-Grok-Conv-Id", "x-grok-conv-id", "X-Grok2API-Conv-Id"); v != "" {
+		return "chat:" + strings.TrimSpace(model) + ":conv:" + v
+	}
+	if v := get("X-Session-Id", "x-session-id", "Session-Id"); v != "" {
+		return "chat:" + strings.TrimSpace(model) + ":session:" + v
+	}
+	if v := get("X-Thread-Id", "x-thread-id", "Thread-Id"); v != "" {
+		return "chat:" + strings.TrimSpace(model) + ":thread:" + v
+	}
+	if v := get("X-Prompt-Cache-Key", "x-prompt-cache-key"); v != "" {
+		return "chat:" + strings.TrimSpace(model) + ":prompt_cache_key:" + v
+	}
+	return ""
+}
+
 func preferAffinity(ctx context.Context, candidates []pool.Candidate, store AffinityStore, fingerprint string) {
 	accountID, err := store.GetAffinity(ctx, fingerprint)
 	if err != nil || accountID == "" {

internal/server/server.go

@@ -479,7 +479,7 @@ func serveChatCompletions(w http.ResponseWriter, r *http.Request, options Option
 		writeJSON(w, http.StatusBadRequest, map[string]any{"detail": err.Error()})
 		return
 	}
-	candidates, err := listCandidates(r.Context(), options)
+	candidates, err := listCandidatesForRequest(r.Context(), options, chatReq, r.Header)
 	if err != nil {
 		writeJSON(w, http.StatusInternalServerError, map[string]any{"detail": err.Error()})
 		return
@@ -763,7 +763,7 @@ func serveMessages(w http.ResponseWriter, r *http.Request, options Options) {
 	}
 	allowedTools := allowedAnthropicToolNames(body)
 	chatReq := proxy.ChatRequest{Model: model, Stream: false, Raw: body}
-	candidates, err := listCandidates(r.Context(), options)
+	candidates, err := listCandidatesForRequest(r.Context(), options, chatReq, r.Header)
 	if err != nil {
 		writeAnthropicError(w, http.StatusInternalServerError, err.Error(), "api_error")
 		return
@@ -960,15 +960,22 @@ func serveResponses(w http.ResponseWriter, r *http.Request, options Options) {
 		writeOpenAIError(w, http.StatusBadRequest, "input must contain at least one message", "invalid_request_error")
 		return
 	}
-	candidates, err := listCandidates(r.Context(), options)
+	chatReq := proxy.ChatRequest{Model: model, Stream: stream, Raw: body}
+	// Preserve client prompt_cache_key on chatReq for sticky fingerprint (stripped later for upstream).
+	if pck := strings.TrimSpace(stringValue(raw["prompt_cache_key"])); pck != "" {
+		if chatReq.Raw == nil {
+			chatReq.Raw = map[string]any{}
+		}
+		chatReq.Raw["prompt_cache_key"] = pck
+	}
+	candidates, err := listCandidatesForRequest(r.Context(), options, chatReq, r.Header)
 	if err != nil {
 		writeOpenAIError(w, http.StatusInternalServerError, err.Error(), "server_error")
 		return
 	}
 	service := proxy.ChatService{Catalog: modelCatalog(options), Client: upstreamClient(options), PickObserver: options.PickObserver, AffinityStore: options.AffinityStore}
 	started := time.Now()
 	responseID := responses.NewResponseID()
-	chatReq := proxy.ChatRequest{Model: model, Stream: stream, Raw: body}
 	respPolicy := historycompact.ResolveOutboundToolPolicy(
 		"openai_responses",
 		r.UserAgent(),
@@ -2433,6 +2440,25 @@ func upstreamClient(options Options) *grok.Client {
 	return &grok.Client{BaseURL: options.Config.UpstreamBase}
 }
 
+func listCandidatesForRequest(ctx context.Context, options Options, chatReq proxy.ChatRequest, headers http.Header) ([]pool.Candidate, error) {
+	candidates, err := listCandidates(ctx, options)
+	if err != nil {
+		return nil, err
+	}
+	fp := proxy.ChatFingerprint(chatReq)
+	if fp == "" {
+		fp = proxy.ChatFingerprintFromHeaders(headers, chatReq.Model)
+	}
+	// If still empty, try raw prompt_cache_key alone.
+	if fp == "" && chatReq.Raw != nil {
+		if pck, _ := chatReq.Raw["prompt_cache_key"].(string); strings.TrimSpace(pck) != "" {
+			fp = "chat:" + strings.TrimSpace(chatReq.Model) + ":prompt_cache_key:" + strings.TrimSpace(pck)
+		}
+	}
+	sticky := stickyAccountID(ctx, options, fp)
+	return ensureStickyCandidate(ctx, options, candidates, sticky), nil
+}
+
 func listCandidates(ctx context.Context, options Options) ([]pool.Candidate, error) {
 	if len(options.Candidates) > 0 {
 		out := make([]pool.Candidate, len(options.Candidates))
@@ -2445,6 +2471,42 @@ func listCandidates(ctx context.Context, options Options) ([]pool.Candidate, err
 	return options.Store.ListPoolCandidates(ctx)
 }
 
+// ensureStickyCandidate injects an affinity-bound account into the pick window
+// even when it is outside the top-N least_used scan. Critical for prompt-cache
+// multi-turn (Codex prompt_cache_key / X-Grok-Conv-Id) high hit rates.
+func ensureStickyCandidate(ctx context.Context, options Options, candidates []pool.Candidate, stickyID string) []pool.Candidate {
+	stickyID = strings.TrimSpace(stickyID)
+	if stickyID == "" || options.Store == nil {
+		return candidates
+	}
+	for _, c := range candidates {
+		if c.ID == stickyID {
+			return candidates
+		}
+	}
+	extra, err := options.Store.GetPoolCandidate(ctx, stickyID)
+	if err != nil || extra == nil {
+		return candidates
+	}
+	// Put sticky first so prepareChain affinity boost is redundant but safe.
+	out := make([]pool.Candidate, 0, len(candidates)+1)
+	out = append(out, *extra)
+	out = append(out, candidates...)
+	return out
+}
+
+func stickyAccountID(ctx context.Context, options Options, fingerprint string) string {
+	fingerprint = strings.TrimSpace(fingerprint)
+	if fingerprint == "" || options.AffinityStore == nil {
+		return ""
+	}
+	id, err := options.AffinityStore.GetAffinity(ctx, fingerprint)
+	if err != nil {
+		return ""
+	}
+	return strings.TrimSpace(id)
+}
+
 func modelCatalog(options Options) *models.Catalog {
 	if options.Models != nil {
 		return options.Models

internal/store/postgres/pool.go

@@ -21,6 +21,42 @@ var (
 
 const candidateCacheTTL = 400 * time.Millisecond
 
+// GetPoolCandidate loads one account as a pick candidate (sticky TTFT / prompt-cache path).
+func (c *Connector) GetPoolCandidate(ctx context.Context, accountID string) (*pool.Candidate, error) {
+	accountID = strings.TrimSpace(accountID)
+	if accountID == "" {
+		return nil, nil
+	}
+	row := c.Pool.QueryRow(ctx, `
+		SELECT a.id, a.payload, a.email, a.user_id, a.team_id, a.expires_at,
+		       COALESCE(ap.enabled, true), COALESCE(ap.disabled_for_quota, false),
+		       ap.cooldown_until, COALESCE(ap.blocked_models, '{}'::jsonb),
+		       COALESCE(ap.request_count, 0), COALESCE(ap.weight, 1)
+		FROM accounts a
+		LEFT JOIN account_pool ap ON ap.account_id = a.id
+		WHERE a.id = $1
+		LIMIT 1`, accountID)
+	var candidate pool.Candidate
+	var payloadBytes, blockedBytes []byte
+	var email, userID, teamID *string
+	var expiresAt, cooldownUntil *time.Time
+	if err := row.Scan(&candidate.ID, &payloadBytes, &email, &userID, &teamID, &expiresAt, &candidate.Enabled, &candidate.DisabledForQuota, &cooldownUntil, &blockedBytes, &candidate.RequestCount, &candidate.Weight); err != nil {
+		return nil, err
+	}
+	payload := decodeMap(payloadBytes)
+	candidate.Token, _ = firstString(payload, "key", "access_token", "token")
+	candidate.Email = stringValue(email, stringFromMap(payload, "email"))
+	candidate.UserID = stringValue(userID, firstMapString(payload, "user_id", "principal_id"))
+	candidate.TeamID = stringValue(teamID, stringFromMap(payload, "team_id"))
+	candidate.ExpiresAt = expiresAt
+	candidate.CooldownUntil = cooldownUntil
+	candidate.BlockedModels = decodeMap(blockedBytes)
+	if strings.TrimSpace(candidate.Token) == "" {
+		return nil, nil
+	}
+	return &candidate, nil
+}
+
 func (c *Connector) ListPoolCandidates(ctx context.Context) ([]pool.Candidate, error) {
 	candidateCacheMu.Lock()
 	if time.Since(candidateCacheAt) < candidateCacheTTL && len(candidateCacheData) > 0 {