wehub-resource-sync
282 lines
10 KiB
Go
282 lines
10 KiB
Go
package cache
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
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 := json.Marshal(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 := gjson.GetBytes(item, "encrypted_content").String(); got != encryptedContent {
|
|
t.Fatalf("encrypted_content = %q, want %q; item=%s", got, encryptedContent, string(item))
|
|
}
|
|
if got := gjson.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 gjson.GetBytes(got[0], "encrypted_content").String() != encryptedContent {
|
|
t.Fatalf("reasoning encrypted_content not preserved: %s", got[0])
|
|
}
|
|
if gotText := gjson.GetBytes(got[1], "content.0.text").String(); gotText != "answer" {
|
|
t.Fatalf("assistant message text = %q, want answer; item=%s", gotText, got[1])
|
|
}
|
|
if gjson.GetBytes(got[1], "id").Exists() || gjson.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 := gjson.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 := gjson.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 gjson.GetBytes(got[1], "content.0.type").String() != "refusal" {
|
|
t.Fatalf("message part type = %s, want refusal; item=%s", gjson.GetBytes(got[1], "content.0.type").String(), got[1])
|
|
}
|
|
if gjson.GetBytes(got[1], "content.0.refusal").String() != "I cannot help with that" {
|
|
t.Fatalf("refusal text missing; item=%s", got[1])
|
|
}
|
|
if gjson.GetBytes(got[1], "content.0.text").Exists() {
|
|
t.Fatalf("refusal part should not use text field; item=%s", got[1])
|
|
}
|
|
}
|