internal/cache/xai_reasoning_replay_cache_test.go
Rama main · réplica de lectura del repositorio original.
package cache
import (
"context"
"crypto/sha256"
"errors"
"testing"
"time"
"github.com/MoeForge/moego/base64"
"github.com/MoeForge/moego/ojson"
homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home"
)
type fakeXAIReasoningReplayKVClient struct {
values map[string][]byte
getErr error
setErr error
delErr error
expireErr error
getCount int
setCount int
delCount int
expireCount int
lastSetTTL time.Duration
lastExpireTTL time.Duration
}
func newFakeXAIReasoningReplayKVClient() *fakeXAIReasoningReplayKVClient {
return &fakeXAIReasoningReplayKVClient{values: make(map[string][]byte)}
}
func (c *fakeXAIReasoningReplayKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) {
c.getCount++
if c.getErr != nil {
return nil, false, c.getErr
}
value, ok := c.values[key]
if !ok {
return nil, false, nil
}
return append([]byte(nil), value...), true, nil
}
func (c *fakeXAIReasoningReplayKVClient) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) {
c.setCount++
c.lastSetTTL = opts.EX
if c.setErr != nil {
return false, c.setErr
}
c.values[key] = append([]byte(nil), value...)
return true, nil
}
func (c *fakeXAIReasoningReplayKVClient) KVDel(_ context.Context, keys ...string) (int64, error) {
c.delCount++
if c.delErr != nil {
return 0, c.delErr
}
var deleted int64
for _, key := range keys {
if _, ok := c.values[key]; ok {
delete(c.values, key)
deleted++
}
}
return deleted, nil
}
func (c *fakeXAIReasoningReplayKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) {
c.expireCount++
c.lastExpireTTL = ttl
if c.expireErr != nil {
return false, c.expireErr
}
return true, nil
}
func useFakeXAIReasoningReplayKVClient(t *testing.T, client *fakeXAIReasoningReplayKVClient, homeMode bool, errClient error) {
t.Helper()
previous := currentXAIReasoningReplayKVClient
currentXAIReasoningReplayKVClient = func() (xaiReasoningReplayKVClient, bool, error) {
return client, homeMode, errClient
}
t.Cleanup(func() {
currentXAIReasoningReplayKVClient = previous
})
}
func mustXAIReasoningReplayJSON(t *testing.T, items [][]byte) []byte {
t.Helper()
raw, err := marshalReasoningReplayRawJSONArray(items)
if err != nil {
t.Fatalf("marshal replay items: %v", err)
}
return raw
}
func TestXAIReasoningReplayCacheRejectsCodexEncryptedContent(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
if CacheXAIReasoningReplayItem("grok-4.3", "claude:xai-cache-test", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"gAAAAABinvalid-gpt-shape"}`)) {
t.Fatal("xAI replay cache should reject GPT/Codex-shaped encrypted_content")
}
if _, ok := GetXAIReasoningReplayItem("grok-4.3", "claude:xai-cache-test"); ok {
t.Fatal("xAI replay cache should not store GPT/Codex-shaped encrypted_content")
}
}
func TestXAIReasoningReplayCacheStoresGrokEncryptedContent(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
encryptedContent := validGrokEncryptedContentForReplayCacheTest()
if !CacheXAIReasoningReplayItem("grok-4.3", "claude:xai-cache-test", []byte(`{"type":"reasoning","summary":[{"type":"summary_text","text":"visible"}],"content":null,"encrypted_content":"`+encryptedContent+`"}`)) {
t.Fatal("xAI replay cache should store valid Grok encrypted_content")
}
item, ok := GetXAIReasoningReplayItem("grok-4.3", "claude:xai-cache-test")
if !ok {
t.Fatal("xAI replay cache item missing after store")
}
if got := ojson.GetBytes(item, "encrypted_content").String(); got != encryptedContent {
t.Fatalf("encrypted_content = %q, want %q; item=%s", got, encryptedContent, string(item))
}
if got := ojson.GetBytes(item, "summary").Array(); len(got) != 0 {
t.Fatalf("summary length = %d, want normalized empty summary; item=%s", len(got), string(item))
}
}
func TestXAIReasoningReplayCacheStoresAssistantMessageWithReasoning(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
encryptedContent := validGrokEncryptedContentForReplayCacheTest()
items := [][]byte{
[]byte(`{"id":"rs_1","type":"reasoning","summary":[{"type":"summary_text","text":"visible"}],"encrypted_content":"` + encryptedContent + `"}`),
[]byte(`{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"answer","annotations":[],"logprobs":[]}]}`),
}
if !CacheXAIReasoningReplayItems("grok-4.5", "prompt-cache:session", items) {
t.Fatal("expected reasoning replay items to be cached")
}
got, ok := GetXAIReasoningReplayItems("grok-4.5", "prompt-cache:session")
if !ok || len(got) != 2 {
t.Fatalf("cached items = %q, %v, want two items", got, ok)
}
if ojson.GetBytes(got[0], "encrypted_content").String() != encryptedContent {
t.Fatalf("reasoning encrypted_content not preserved: %s", got[0])
}
if gotText := ojson.GetBytes(got[1], "content.0.text").String(); gotText != "answer" {
t.Fatalf("assistant message text = %q, want answer; item=%s", gotText, got[1])
}
if ojson.GetBytes(got[1], "id").Exists() || ojson.GetBytes(got[1], "status").Exists() {
t.Fatalf("assistant message transport fields were not stripped: %s", got[1])
}
}
func TestXAIReasoningReplayCacheRejectsAssistantMessageWithoutReasoning(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
items := [][]byte{
[]byte(`{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"answer"}]}`),
}
if CacheXAIReasoningReplayItems("grok-4.5", "prompt-cache:message-only", items) {
t.Fatal("message-only replay batch must not be cached")
}
if _, ok := GetXAIReasoningReplayItems("grok-4.5", "prompt-cache:message-only"); ok {
t.Fatal("message-only replay batch unexpectedly exists in cache")
}
}
func TestXAIReasoningReplayCacheStoresToolCallWithoutReasoning(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
tests := []struct {
name string
sessionKey string
item []byte
wantType string
wantPayload string
}{
{
name: "function call",
sessionKey: "prompt-cache:function-call-only",
item: []byte(`{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"q\":\"weather\"}"}`),
wantType: "function_call",
wantPayload: `{"q":"weather"}`,
},
{
name: "custom tool call",
sessionKey: "prompt-cache:custom-tool-call-only",
item: []byte(`{"type":"custom_tool_call","call_id":"call_2","name":"shell","input":"pwd"}`),
wantType: "custom_tool_call",
wantPayload: "pwd",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if !CacheXAIReasoningReplayItems("grok-4.3", tt.sessionKey, [][]byte{tt.item}) {
t.Fatal("tool-call-only replay batch must be cached")
}
items, ok := GetXAIReasoningReplayItems("grok-4.3", tt.sessionKey)
if !ok || len(items) != 1 {
t.Fatalf("cached items = %q, %v, want one item", items, ok)
}
if got := ojson.GetBytes(items[0], "type").String(); got != tt.wantType {
t.Fatalf("cached type = %q, want %q; item=%s", got, tt.wantType, items[0])
}
payloadPath := "arguments"
if tt.wantType == "custom_tool_call" {
payloadPath = "input"
}
if got := ojson.GetBytes(items[0], payloadPath).String(); got != tt.wantPayload {
t.Fatalf("cached %s = %q, want %q; item=%s", payloadPath, got, tt.wantPayload, items[0])
}
})
}
}
func TestXAIReasoningReplayRequiredHomeExpireFailureReturnsItems(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
client := newFakeXAIReasoningReplayKVClient()
client.expireErr = errors.New("expire failed")
key := xaiReasoningReplayKVKey("grok-4.3", "session-home")
item := []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"` + validGrokEncryptedContentForReplayCacheTest() + `"}`)
client.values[key] = mustXAIReasoningReplayJSON(t, [][]byte{item})
useFakeXAIReasoningReplayKVClient(t, client, true, nil)
items, found, errGet := GetXAIReasoningReplayItemsRequired(context.Background(), "grok-4.3", "session-home")
if errGet != nil {
t.Fatalf("GetXAIReasoningReplayItemsRequired() error = %v", errGet)
}
if !found || len(items) != 1 || string(items[0]) != string(item) {
t.Fatalf("GetXAIReasoningReplayItemsRequired() = %q, %v, want item, true", items, found)
}
if client.expireCount != 1 || client.lastExpireTTL != XAIReasoningReplayCacheTTL {
t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, XAIReasoningReplayCacheTTL)
}
}
func validGrokEncryptedContentForReplayCacheTest() string {
buf := make([]byte, 0, 256)
for i := 0; len(buf) < 256; i++ {
sum := sha256.Sum256([]byte{byte(i), byte(i >> 8), byte(i >> 16), 99})
buf = append(buf, sum[:]...)
}
return base64.RawStdEncoding.EncodeToString(buf[:256])
}
func TestXAIReasoningReplayCacheStoresRefusalMessagePart(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
encryptedContent := validGrokEncryptedContentForReplayCacheTest()
items := [][]byte{
[]byte(`{"type":"reasoning","summary":[],"encrypted_content":"` + encryptedContent + `"}`),
[]byte(`{"type":"message","role":"assistant","content":[{"type":"refusal","refusal":"I cannot help with that"}]}`),
}
if !CacheXAIReasoningReplayItems("grok-4.5", "prompt-cache:refusal", items) {
t.Fatal("expected refusal message with reasoning to be cached")
}
got, ok := GetXAIReasoningReplayItems("grok-4.5", "prompt-cache:refusal")
if !ok || len(got) != 2 {
t.Fatalf("cached items = %q, %v, want reasoning + refusal message", got, ok)
}
if ojson.GetBytes(got[1], "content.0.type").String() != "refusal" {
t.Fatalf("message part type = %s, want refusal; item=%s", ojson.GetBytes(got[1], "content.0.type").String(), got[1])
}
if ojson.GetBytes(got[1], "content.0.refusal").String() != "I cannot help with that" {
t.Fatalf("refusal text missing; item=%s", got[1])
}
if ojson.GetBytes(got[1], "content.0.text").Exists() {
t.Fatalf("refusal part should not use text field; item=%s", got[1])
}
}