Files
sub2api/backend/internal/handler/gemini_cli_session_test.go

144 lines
4.3 KiB
Go
Raw Normal View History

feat(gemini): 支持 Gemini CLI 粘性会话与跨账号 thoughtSignature 清理 ## 问题背景 1. Gemini CLI 没有明确的会话标识(如 Claude Code 的 metadata.user_id) 2. thoughtSignature 与具体上游账号强绑定,跨账号使用会导致 400 错误 3. 粘性会话切换账号或 cache 丢失时,旧签名会导致请求失败 ## 解决方案 ### 1. Gemini CLI 会话标识提取 - 从 `x-gemini-api-privileged-user-id` header 和请求体中的 tmp 目录哈希生成会话标识 - 组合策略:SHA256(privileged-user-id + ":" + tmp_dir_hash) - 正则提取:`/\.gemini/tmp/([A-Fa-f0-9]{64})` ### 2. 跨账号 thoughtSignature 清理 实现三种场景的智能清理: 1. **Cache 命中 + 账号切换** - 粘性会话绑定的账号与当前选择的账号不同时清理 2. **同一请求内 failover 切换** - 通过 sessionBoundAccountID 跟踪,检测重试时的账号切换 3. **Gemini CLI + Cache 未命中 + 含签名** - 预防性清理,避免 cache 丢失后首次转发就 400 - 仅对 Gemini CLI 请求且请求体包含 thoughtSignature 时触发 ## 修改内容 ### backend/internal/handler/gemini_v1beta_handler.go - 添加 `extractGeminiCLISessionHash` 函数提取 Gemini CLI 会话标识 - 添加 `isGeminiCLIRequest` 函数识别 Gemini CLI 请求 - 实现账号切换检测与 thoughtSignature 清理逻辑 - 添加 `geminiCLITmpDirRegex` 正则表达式 ### backend/internal/service/gateway_service.go - 添加 `GetCachedSessionAccountID` 方法查询粘性会话绑定的账号 ID ### backend/internal/service/gemini_native_signature_cleaner.go (新增) - 实现 `CleanGeminiNativeThoughtSignatures` 函数 - 递归清理 JSON 中的所有 thoughtSignature 字段 - 支持任意 JSON 顶层类型(object/array) ### backend/internal/handler/gemini_cli_session_test.go (新增) - 测试 Gemini CLI 会话哈希提取逻辑 - 测试 tmp 目录正则匹配 - 覆盖有/无 privileged-user-id 的场景 ## 影响范围 - 修复 Gemini CLI 多轮对话时账号切换导致的 400 错误 - 提高粘性会话的稳定性和容错能力 - 不影响其他客户端(Claude Code 等)的会话标识生成 ## 测试 - 单元测试:go test -tags=unit ./internal/handler -run TestExtractGeminiCLISessionHash - 单元测试:go test -tags=unit ./internal/handler -run TestGeminiCLITmpDirRegex - 编译验证:go build ./cmd/server
2026-01-26 04:40:38 +08:00
//go:build unit
package handler
import (
"crypto/sha256"
"encoding/hex"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestExtractGeminiCLISessionHash(t *testing.T) {
tests := []struct {
name string
body string
privilegedUserID string
wantEmpty bool
wantHash string
}{
{
name: "with privileged-user-id and tmp dir",
body: `{"contents":[{"parts":[{"text":"The project's temporary directory is: /Users/ianshaw/.gemini/tmp/f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740"}]}]}`,
privilegedUserID: "90785f52-8bbe-4b17-b111-a1ddea1636c3",
wantEmpty: false,
wantHash: func() string {
combined := "90785f52-8bbe-4b17-b111-a1ddea1636c3:f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740"
hash := sha256.Sum256([]byte(combined))
return hex.EncodeToString(hash[:])
}(),
},
{
name: "without privileged-user-id but with tmp dir",
body: `{"contents":[{"parts":[{"text":"The project's temporary directory is: /Users/ianshaw/.gemini/tmp/f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740"}]}]}`,
privilegedUserID: "",
wantEmpty: false,
wantHash: "f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740",
},
{
name: "without tmp dir",
body: `{"contents":[{"parts":[{"text":"Hello world"}]}]}`,
privilegedUserID: "90785f52-8bbe-4b17-b111-a1ddea1636c3",
wantEmpty: true,
},
{
name: "empty body",
body: "",
privilegedUserID: "90785f52-8bbe-4b17-b111-a1ddea1636c3",
wantEmpty: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// 创建测试上下文
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/test", nil)
if tt.privilegedUserID != "" {
c.Request.Header.Set("x-gemini-api-privileged-user-id", tt.privilegedUserID)
}
// 调用函数
result := extractGeminiCLISessionHash(c, []byte(tt.body))
// 验证结果
if tt.wantEmpty {
require.Empty(t, result, "expected empty session hash")
} else {
require.NotEmpty(t, result, "expected non-empty session hash")
require.Equal(t, tt.wantHash, result, "session hash mismatch")
}
})
}
}
func TestGeminiCLITmpDirRegex(t *testing.T) {
tests := []struct {
name string
input string
wantMatch bool
wantHash string
}{
{
name: "valid tmp dir path",
input: "/Users/ianshaw/.gemini/tmp/f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740",
wantMatch: true,
wantHash: "f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740",
},
{
name: "valid tmp dir path in text",
input: "The project's temporary directory is: /Users/ianshaw/.gemini/tmp/f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740\nOther text",
wantMatch: true,
wantHash: "f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740",
},
{
name: "invalid hash length",
input: "/Users/ianshaw/.gemini/tmp/abc123",
wantMatch: false,
},
{
name: "no tmp dir",
input: "Hello world",
wantMatch: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
match := geminiCLITmpDirRegex.FindStringSubmatch(tt.input)
if tt.wantMatch {
require.NotNil(t, match, "expected regex to match")
require.Len(t, match, 2, "expected 2 capture groups")
require.Equal(t, tt.wantHash, match[1], "hash mismatch")
} else {
require.Nil(t, match, "expected regex not to match")
}
})
}
}
func TestSafeShortPrefix(t *testing.T) {
tests := []struct {
name string
input string
n int
want string
}{
{name: "空字符串", input: "", n: 8, want: ""},
{name: "长度小于截断值", input: "abc", n: 8, want: "abc"},
{name: "长度等于截断值", input: "12345678", n: 8, want: "12345678"},
{name: "长度大于截断值", input: "1234567890", n: 8, want: "12345678"},
{name: "截断值为0", input: "123456", n: 0, want: "123456"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, safeShortPrefix(tt.input, tt.n))
})
}
}