feat: add Anthropic sticky session digest chain matching via Trie
The previous fallback (step 3) in GenerateSessionHash hashed system + all messages together, producing a different hash each round as the conversation grew ([a] -> [a,b] -> [a,b,c]). This made fallback sticky sessions ineffective for multi-turn conversations. Implement per-message Trie digest chain matching (reusing Gemini's Trie infrastructure) so that the previous round's chain is always a prefix of the current round's chain, enabling reliable session affinity.
This commit is contained in:
erio
@@ -22,6 +22,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// GatewayHandler handles API gateway requests
|
||||
@@ -212,6 +213,53 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
if sessionKey != "" {
|
||||
sessionBoundAccountID, _ = h.gatewayService.GetCachedSessionAccountID(c.Request.Context(), apiKey.GroupID, sessionKey)
|
||||
}
|
||||
|
||||
// === Anthropic 内容摘要会话 Fallback 逻辑 ===
|
||||
// 当原有会话标识无效时（sessionBoundAccountID == 0），尝试基于内容摘要链匹配
|
||||
var anthropicDigestChain string
|
||||
var anthropicPrefixHash string
|
||||
var anthropicSessionUUID string
|
||||
useAnthropicDigestFallback := sessionBoundAccountID == 0 && platform != service.PlatformGemini
|
||||
|
||||
if useAnthropicDigestFallback {
|
||||
anthropicDigestChain = service.BuildAnthropicDigestChain(parsedReq)
|
||||
if anthropicDigestChain != "" {
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
anthropicPrefixHash = service.GenerateGeminiPrefixHash(
|
||||
subject.UserID,
|
||||
apiKey.ID,
|
||||
clientIP,
|
||||
userAgent,
|
||||
platform,
|
||||
reqModel,
|
||||
)
|
||||
|
||||
foundUUID, foundAccountID, found := h.gatewayService.FindAnthropicSession(
|
||||
c.Request.Context(),
|
||||
derefGroupID(apiKey.GroupID),
|
||||
anthropicPrefixHash,
|
||||
anthropicDigestChain,
|
||||
)
|
||||
if found {
|
||||
sessionBoundAccountID = foundAccountID
|
||||
anthropicSessionUUID = foundUUID
|
||||
log.Printf("[Anthropic] Digest fallback matched: uuid=%s, accountID=%d, chain=%s",
|
||||
foundUUID[:8], foundAccountID, truncateDigestChain(anthropicDigestChain))
|
||||
|
||||
if sessionKey == "" {
|
||||
sessionKey = service.GenerateAnthropicDigestSessionKey(anthropicPrefixHash, foundUUID)
|
||||
}
|
||||
_ = h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, foundAccountID)
|
||||
} else {
|
||||
anthropicSessionUUID = uuid.New().String()
|
||||
if sessionKey == "" {
|
||||
sessionKey = service.GenerateAnthropicDigestSessionKey(anthropicPrefixHash, anthropicSessionUUID)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 判断是否真的绑定了粘性会话：有 sessionKey 且已经绑定到某个账号
|
||||
hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0
|
||||
|
||||
@@ -540,6 +588,20 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
|
||||
// 保存 Anthropic 内容摘要会话（用于 Fallback 匹配）
|
||||
if useAnthropicDigestFallback && anthropicDigestChain != "" && anthropicPrefixHash != "" {
|
||||
if err := h.gatewayService.SaveAnthropicSession(
|
||||
c.Request.Context(),
|
||||
derefGroupID(apiKey.GroupID),
|
||||
anthropicPrefixHash,
|
||||
anthropicDigestChain,
|
||||
anthropicSessionUUID,
|
||||
account.ID,
|
||||
); err != nil {
|
||||
log.Printf("[Anthropic] Failed to save digest session: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 异步记录使用量（subscription已在函数开头获取）
|
||||
go func(result *service.ForwardResult, usedAccount *service.Account, ua, clientIP string, fcb bool) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
|
||||
@@ -238,3 +238,41 @@ func (c *gatewayCache) SaveGeminiSession(ctx context.Context, groupID int64, pre
|
||||
|
||||
return c.rdb.Eval(ctx, geminiTrieSaveScript, []string{trieKey}, digestChain, value, ttlSeconds).Err()
|
||||
}
|
||||
|
||||
// ============ Anthropic 会话 Fallback 方法 (复用 Trie 实现) ============
|
||||
|
||||
// FindAnthropicSession 查找 Anthropic 会话（复用 Gemini Trie Lua 脚本）
|
||||
func (c *gatewayCache) FindAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
if digestChain == "" {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
trieKey := service.BuildAnthropicTrieKey(groupID, prefixHash)
|
||||
ttlSeconds := int(service.AnthropicSessionTTL().Seconds())
|
||||
|
||||
result, err := c.rdb.Eval(ctx, geminiTrieFindScript, []string{trieKey}, digestChain, ttlSeconds).Result()
|
||||
if err != nil || result == nil {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
value, ok := result.(string)
|
||||
if !ok || value == "" {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
uuid, accountID, ok = service.ParseGeminiSessionValue(value)
|
||||
return uuid, accountID, ok
|
||||
}
|
||||
|
||||
// SaveAnthropicSession 保存 Anthropic 会话（复用 Gemini Trie Lua 脚本）
|
||||
func (c *gatewayCache) SaveAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
if digestChain == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
trieKey := service.BuildAnthropicTrieKey(groupID, prefixHash)
|
||||
value := service.FormatGeminiSessionValue(uuid, accountID)
|
||||
ttlSeconds := int(service.AnthropicSessionTTL().Seconds())
|
||||
|
||||
return c.rdb.Eval(ctx, geminiTrieSaveScript, []string{trieKey}, digestChain, value, ttlSeconds).Err()
|
||||
}
|
||||
|
||||
89
backend/internal/service/anthropic_session.go
Normal file
89
backend/internal/service/anthropic_session.go
Normal file
@@ -0,0 +1,89 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Anthropic 会话 Fallback 相关常量
|
||||
const (
|
||||
// anthropicSessionTTLSeconds Anthropic 会话缓存 TTL（5 分钟）
|
||||
anthropicSessionTTLSeconds = 300
|
||||
|
||||
// anthropicTrieKeyPrefix Anthropic Trie 会话 key 前缀
|
||||
anthropicTrieKeyPrefix = "anthropic:trie:"
|
||||
|
||||
// anthropicDigestSessionKeyPrefix Anthropic 摘要 fallback 会话 key 前缀
|
||||
anthropicDigestSessionKeyPrefix = "anthropic:digest:"
|
||||
)
|
||||
|
||||
// AnthropicSessionTTL 返回 Anthropic 会话缓存 TTL
|
||||
func AnthropicSessionTTL() time.Duration {
|
||||
return anthropicSessionTTLSeconds * time.Second
|
||||
}
|
||||
|
||||
// BuildAnthropicDigestChain 根据 Anthropic 请求生成摘要链
|
||||
// 格式: s:<hash>-u:<hash>-a:<hash>-u:<hash>-...
|
||||
// s = system, u = user, a = assistant
|
||||
func BuildAnthropicDigestChain(parsed *ParsedRequest) string {
|
||||
if parsed == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
var parts []string
|
||||
|
||||
// 1. system prompt
|
||||
if parsed.System != nil {
|
||||
systemData, _ := json.Marshal(parsed.System)
|
||||
if len(systemData) > 0 && string(systemData) != "null" {
|
||||
parts = append(parts, "s:"+shortHash(systemData))
|
||||
}
|
||||
}
|
||||
|
||||
// 2. messages
|
||||
for _, msg := range parsed.Messages {
|
||||
msgMap, ok := msg.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
role, _ := msgMap["role"].(string)
|
||||
prefix := rolePrefix(role)
|
||||
content := msgMap["content"]
|
||||
contentData, _ := json.Marshal(content)
|
||||
parts = append(parts, prefix+":"+shortHash(contentData))
|
||||
}
|
||||
|
||||
return strings.Join(parts, "-")
|
||||
}
|
||||
|
||||
// rolePrefix 将 Anthropic 的 role 映射为单字符前缀
|
||||
func rolePrefix(role string) string {
|
||||
switch role {
|
||||
case "assistant":
|
||||
return "a"
|
||||
default:
|
||||
return "u"
|
||||
}
|
||||
}
|
||||
|
||||
// BuildAnthropicTrieKey 构建 Anthropic Trie Redis key
|
||||
// 格式: anthropic:trie:{groupID}:{prefixHash}
|
||||
func BuildAnthropicTrieKey(groupID int64, prefixHash string) string {
|
||||
return anthropicTrieKeyPrefix + strconv.FormatInt(groupID, 10) + ":" + prefixHash
|
||||
}
|
||||
|
||||
// GenerateAnthropicDigestSessionKey 生成 Anthropic 摘要 fallback 的 sessionKey
|
||||
// 组合 prefixHash 前 8 位 + uuid 前 8 位，确保不同会话产生不同的 sessionKey
|
||||
func GenerateAnthropicDigestSessionKey(prefixHash, uuid string) string {
|
||||
prefix := prefixHash
|
||||
if len(prefixHash) >= 8 {
|
||||
prefix = prefixHash[:8]
|
||||
}
|
||||
uuidPart := uuid
|
||||
if len(uuid) >= 8 {
|
||||
uuidPart = uuid[:8]
|
||||
}
|
||||
return anthropicDigestSessionKeyPrefix + prefix + ":" + uuidPart
|
||||
}
|
||||
357
backend/internal/service/anthropic_session_test.go
Normal file
357
backend/internal/service/anthropic_session_test.go
Normal file
@@ -0,0 +1,357 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) {
|
||||
result := BuildAnthropicDigestChain(nil)
|
||||
if result != "" {
|
||||
t.Errorf("expected empty string for nil request, got: %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
Messages: []any{},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
if result != "" {
|
||||
t.Errorf("expected empty string for empty messages, got: %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 1 {
|
||||
t.Fatalf("expected 1 part, got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "u:") {
|
||||
t.Errorf("expected prefix 'u:', got: %s", parts[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "hi there"},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("expected 2 parts, got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "u:") {
|
||||
t.Errorf("part[0] expected prefix 'u:', got: %s", parts[0])
|
||||
}
|
||||
if !strings.HasPrefix(parts[1], "a:") {
|
||||
t.Errorf("part[1] expected prefix 'a:', got: %s", parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
System: "You are a helpful assistant",
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("expected 2 parts (s + u), got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "s:") {
|
||||
t.Errorf("part[0] expected prefix 's:', got: %s", parts[0])
|
||||
}
|
||||
if !strings.HasPrefix(parts[1], "u:") {
|
||||
t.Errorf("part[1] expected prefix 'u:', got: %s", parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
System: []any{
|
||||
map[string]any{"type": "text", "text": "You are a helpful assistant"},
|
||||
},
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("expected 2 parts (s + u), got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "s:") {
|
||||
t.Errorf("part[0] expected prefix 's:', got: %s", parts[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_ConversationPrefixRelationship(t *testing.T) {
|
||||
// 核心测试：验证对话增长时链的前缀关系
|
||||
// 上一轮的完整链一定是下一轮链的前缀
|
||||
system := "You are a helpful assistant"
|
||||
|
||||
// 第 1 轮: system + user
|
||||
round1 := &ParsedRequest{
|
||||
System: system,
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
chain1 := BuildAnthropicDigestChain(round1)
|
||||
|
||||
// 第 2 轮: system + user + assistant + user
|
||||
round2 := &ParsedRequest{
|
||||
System: system,
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "hi there"},
|
||||
map[string]any{"role": "user", "content": "how are you?"},
|
||||
},
|
||||
}
|
||||
chain2 := BuildAnthropicDigestChain(round2)
|
||||
|
||||
// 第 3 轮: system + user + assistant + user + assistant + user
|
||||
round3 := &ParsedRequest{
|
||||
System: system,
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "hi there"},
|
||||
map[string]any{"role": "user", "content": "how are you?"},
|
||||
map[string]any{"role": "assistant", "content": "I'm doing well"},
|
||||
map[string]any{"role": "user", "content": "great"},
|
||||
},
|
||||
}
|
||||
chain3 := BuildAnthropicDigestChain(round3)
|
||||
|
||||
t.Logf("Chain1: %s", chain1)
|
||||
t.Logf("Chain2: %s", chain2)
|
||||
t.Logf("Chain3: %s", chain3)
|
||||
|
||||
// chain1 是 chain2 的前缀
|
||||
if !strings.HasPrefix(chain2, chain1) {
|
||||
t.Errorf("chain1 should be prefix of chain2:\n chain1: %s\n chain2: %s", chain1, chain2)
|
||||
}
|
||||
|
||||
// chain2 是 chain3 的前缀
|
||||
if !strings.HasPrefix(chain3, chain2) {
|
||||
t.Errorf("chain2 should be prefix of chain3:\n chain2: %s\n chain3: %s", chain2, chain3)
|
||||
}
|
||||
|
||||
// chain1 也是 chain3 的前缀（传递性）
|
||||
if !strings.HasPrefix(chain3, chain1) {
|
||||
t.Errorf("chain1 should be prefix of chain3:\n chain1: %s\n chain3: %s", chain1, chain3)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *testing.T) {
|
||||
parsed1 := &ParsedRequest{
|
||||
System: "System A",
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
parsed2 := &ParsedRequest{
|
||||
System: "System B",
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
|
||||
chain1 := BuildAnthropicDigestChain(parsed1)
|
||||
chain2 := BuildAnthropicDigestChain(parsed2)
|
||||
|
||||
if chain1 == chain2 {
|
||||
t.Error("Different system prompts should produce different chains")
|
||||
}
|
||||
|
||||
// 但 user 部分的 hash 应该相同
|
||||
parts1 := splitChain(chain1)
|
||||
parts2 := splitChain(chain2)
|
||||
if parts1[1] != parts2[1] {
|
||||
t.Error("Same user message should produce same hash regardless of system")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *testing.T) {
|
||||
parsed1 := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "ORIGINAL reply"},
|
||||
map[string]any{"role": "user", "content": "next"},
|
||||
},
|
||||
}
|
||||
parsed2 := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "TAMPERED reply"},
|
||||
map[string]any{"role": "user", "content": "next"},
|
||||
},
|
||||
}
|
||||
|
||||
chain1 := BuildAnthropicDigestChain(parsed1)
|
||||
chain2 := BuildAnthropicDigestChain(parsed2)
|
||||
|
||||
if chain1 == chain2 {
|
||||
t.Error("Different content should produce different chains")
|
||||
}
|
||||
|
||||
parts1 := splitChain(chain1)
|
||||
parts2 := splitChain(chain2)
|
||||
// 第一个 user message hash 应该相同
|
||||
if parts1[0] != parts2[0] {
|
||||
t.Error("First user message hash should be the same")
|
||||
}
|
||||
// assistant reply hash 应该不同
|
||||
if parts1[1] == parts2[1] {
|
||||
t.Error("Assistant reply hash should differ")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
System: "test system",
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "hi"},
|
||||
},
|
||||
}
|
||||
|
||||
chain1 := BuildAnthropicDigestChain(parsed)
|
||||
chain2 := BuildAnthropicDigestChain(parsed)
|
||||
|
||||
if chain1 != chain2 {
|
||||
t.Errorf("BuildAnthropicDigestChain not deterministic: %s vs %s", chain1, chain2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicTrieKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
groupID int64
|
||||
prefixHash string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
groupID: 123,
|
||||
prefixHash: "abcdef12",
|
||||
want: "anthropic:trie:123:abcdef12",
|
||||
},
|
||||
{
|
||||
name: "zero group",
|
||||
groupID: 0,
|
||||
prefixHash: "xyz",
|
||||
want: "anthropic:trie:0:xyz",
|
||||
},
|
||||
{
|
||||
name: "empty prefix",
|
||||
groupID: 1,
|
||||
prefixHash: "",
|
||||
want: "anthropic:trie:1:",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := BuildAnthropicTrieKey(tt.groupID, tt.prefixHash)
|
||||
if got != tt.want {
|
||||
t.Errorf("BuildAnthropicTrieKey(%d, %q) = %q, want %q", tt.groupID, tt.prefixHash, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateAnthropicDigestSessionKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
prefixHash string
|
||||
uuid string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "normal 16 char hash with uuid",
|
||||
prefixHash: "abcdefgh12345678",
|
||||
uuid: "550e8400-e29b-41d4-a716-446655440000",
|
||||
want: "anthropic:digest:abcdefgh:550e8400",
|
||||
},
|
||||
{
|
||||
name: "exactly 8 chars",
|
||||
prefixHash: "12345678",
|
||||
uuid: "abcdefgh",
|
||||
want: "anthropic:digest:12345678:abcdefgh",
|
||||
},
|
||||
{
|
||||
name: "short values",
|
||||
prefixHash: "abc",
|
||||
uuid: "xyz",
|
||||
want: "anthropic:digest:abc:xyz",
|
||||
},
|
||||
{
|
||||
name: "empty values",
|
||||
prefixHash: "",
|
||||
uuid: "",
|
||||
want: "anthropic:digest::",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := GenerateAnthropicDigestSessionKey(tt.prefixHash, tt.uuid)
|
||||
if got != tt.want {
|
||||
t.Errorf("GenerateAnthropicDigestSessionKey(%q, %q) = %q, want %q", tt.prefixHash, tt.uuid, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 验证不同 uuid 产生不同 sessionKey
|
||||
t.Run("different uuid different key", func(t *testing.T) {
|
||||
hash := "sameprefix123456"
|
||||
result1 := GenerateAnthropicDigestSessionKey(hash, "uuid0001-session-a")
|
||||
result2 := GenerateAnthropicDigestSessionKey(hash, "uuid0002-session-b")
|
||||
if result1 == result2 {
|
||||
t.Errorf("Different UUIDs should produce different session keys: %s vs %s", result1, result2)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestAnthropicSessionTTL(t *testing.T) {
|
||||
ttl := AnthropicSessionTTL()
|
||||
if ttl.Seconds() != 300 {
|
||||
t.Errorf("expected 300 seconds, got: %v", ttl.Seconds())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_ContentBlocks(t *testing.T) {
|
||||
// 测试 content 为 content blocks 数组的情况
|
||||
parsed := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": []any{
|
||||
map[string]any{"type": "text", "text": "describe this image"},
|
||||
map[string]any{"type": "image", "source": map[string]any{"type": "base64"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 1 {
|
||||
t.Fatalf("expected 1 part, got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "u:") {
|
||||
t.Errorf("expected prefix 'u:', got: %s", parts[0])
|
||||
}
|
||||
}
|
||||
@@ -232,6 +232,14 @@ func (m *mockGatewayCacheForPlatform) SaveGeminiSession(ctx context.Context, gro
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) FindAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) SaveAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type mockGroupRepoForGateway struct {
|
||||
groups map[int64]*Group
|
||||
getByIDCalls int
|
||||
|
||||
@@ -313,6 +313,14 @@ type GatewayCache interface {
|
||||
// SaveGeminiSession 保存 Gemini 会话
|
||||
// Save Gemini session binding
|
||||
SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error
|
||||
|
||||
// FindAnthropicSession 查找 Anthropic 会话（Trie 匹配）
|
||||
// Find Anthropic session using Trie matching
|
||||
FindAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool)
|
||||
|
||||
// SaveAnthropicSession 保存 Anthropic 会话
|
||||
// Save Anthropic session binding
|
||||
SaveAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error
|
||||
}
|
||||
|
||||
// derefGroupID safely dereferences *int64 to int64, returning 0 if nil
|
||||
@@ -482,23 +490,25 @@ func (s *GatewayService) GenerateSessionHash(parsed *ParsedRequest) string {
|
||||
return s.hashContent(cacheableContent)
|
||||
}
|
||||
|
||||
// 3. Fallback: 使用 system 内容
|
||||
// 3. 最后 fallback: 使用 system + 所有消息的完整摘要串
|
||||
var combined strings.Builder
|
||||
if parsed.System != nil {
|
||||
systemText := s.extractTextFromSystem(parsed.System)
|
||||
if systemText != "" {
|
||||
return s.hashContent(systemText)
|
||||
combined.WriteString(systemText)
|
||||
}
|
||||
}
|
||||
|
||||
// 4. 最后 fallback: 使用第一条消息
|
||||
if len(parsed.Messages) > 0 {
|
||||
if firstMsg, ok := parsed.Messages[0].(map[string]any); ok {
|
||||
msgText := s.extractTextFromContent(firstMsg["content"])
|
||||
for _, msg := range parsed.Messages {
|
||||
if m, ok := msg.(map[string]any); ok {
|
||||
msgText := s.extractTextFromContent(m["content"])
|
||||
if msgText != "" {
|
||||
return s.hashContent(msgText)
|
||||
combined.WriteString(msgText)
|
||||
}
|
||||
}
|
||||
}
|
||||
if combined.Len() > 0 {
|
||||
return s.hashContent(combined.String())
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
@@ -541,6 +551,22 @@ func (s *GatewayService) SaveGeminiSession(ctx context.Context, groupID int64, p
|
||||
return s.cache.SaveGeminiSession(ctx, groupID, prefixHash, digestChain, uuid, accountID)
|
||||
}
|
||||
|
||||
// FindAnthropicSession 查找 Anthropic 会话（基于内容摘要链的 Fallback 匹配）
|
||||
func (s *GatewayService) FindAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
if digestChain == "" || s.cache == nil {
|
||||
return "", 0, false
|
||||
}
|
||||
return s.cache.FindAnthropicSession(ctx, groupID, prefixHash, digestChain)
|
||||
}
|
||||
|
||||
// SaveAnthropicSession 保存 Anthropic 会话
|
||||
func (s *GatewayService) SaveAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
if digestChain == "" || s.cache == nil {
|
||||
return nil
|
||||
}
|
||||
return s.cache.SaveAnthropicSession(ctx, groupID, prefixHash, digestChain, uuid, accountID)
|
||||
}
|
||||
|
||||
func (s *GatewayService) extractCacheableContent(parsed *ParsedRequest) string {
|
||||
if parsed == nil {
|
||||
return ""
|
||||
@@ -1104,7 +1130,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
result.ReleaseFunc() // 释放槽位
|
||||
// 继续到负载感知选择
|
||||
} else {
|
||||
_ = s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL)
|
||||
if s.debugModelRoutingEnabled() {
|
||||
log.Printf("[ModelRoutingDebug] routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), stickyAccountID)
|
||||
}
|
||||
@@ -1258,7 +1283,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
if !s.checkAndRegisterSession(ctx, account, sessionHash) {
|
||||
result.ReleaseFunc() // 释放槽位，继续到 Layer 2
|
||||
} else {
|
||||
_ = s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL)
|
||||
return &AccountSelectionResult{
|
||||
Account: account,
|
||||
Acquired: true,
|
||||
@@ -2163,9 +2187,6 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
if s.debugModelRoutingEnabled() {
|
||||
log.Printf("[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
|
||||
}
|
||||
@@ -2266,9 +2287,6 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
@@ -2377,9 +2395,6 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
||||
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
if s.debugModelRoutingEnabled() {
|
||||
log.Printf("[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
|
||||
}
|
||||
@@ -2482,9 +2497,6 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
||||
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
|
||||
@@ -281,6 +281,14 @@ func (m *mockGatewayCacheForGemini) SaveGeminiSession(ctx context.Context, group
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) FindAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) SaveAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择
|
||||
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
@@ -220,6 +220,14 @@ func (c *stubGatewayCache) SaveGeminiSession(ctx context.Context, groupID int64,
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *stubGatewayCache) FindAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
func (c *stubGatewayCache) SaveAnthropicSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) {
|
||||
now := time.Now()
|
||||
resetAt := now.Add(10 * time.Minute)
|
||||
|
||||
Reference in New Issue
Block a user