Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
9.4 kB · 237 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238package cache
import ( "bytes" "context" "sync" "testing" "time"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home")
type fakeKimiThinkingReplayKVClient struct { mu sync.Mutex values map[string][]byte}
func newFakeKimiThinkingReplayKVClient() *fakeKimiThinkingReplayKVClient { return &fakeKimiThinkingReplayKVClient{values: make(map[string][]byte)}}
func (c *fakeKimiThinkingReplayKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) { c.mu.Lock() defer c.mu.Unlock() value, found := c.values[key] return append([]byte(nil), value...), found, nil}
func (c *fakeKimiThinkingReplayKVClient) KVSet(_ context.Context, key string, value []byte, _ homekv.KVSetOptions) (bool, error) { c.mu.Lock() c.values[key] = append([]byte(nil), value...) c.mu.Unlock() return true, nil}
func (c *fakeKimiThinkingReplayKVClient) KVDel(_ context.Context, keys ...string) (int64, error) { c.mu.Lock() defer c.mu.Unlock() var deleted int64 for _, key := range keys { if _, found := c.values[key]; found { delete(c.values, key) deleted++ } } return deleted, nil}
func (c *fakeKimiThinkingReplayKVClient) KVCompareAndSwap(_ context.Context, key string, expected []byte, expectedExists bool, value []byte, _ time.Duration) (bool, error) { c.mu.Lock() defer c.mu.Unlock() current, found := c.values[key] if found != expectedExists || (found && !bytes.Equal(current, expected)) { return false, nil } c.values[key] = append([]byte(nil), value...) return true, nil}
func (c *fakeKimiThinkingReplayKVClient) KVExpire(_ context.Context, key string, _ time.Duration) (bool, error) { c.mu.Lock() defer c.mu.Unlock() _, found := c.values[key] return found, nil}
func useFakeKimiThinkingReplayKVClient(t *testing.T, client *fakeKimiThinkingReplayKVClient) { t.Helper() previous := currentKimiThinkingReplayKVClient currentKimiThinkingReplayKVClient = func() (kimiThinkingReplayKVClient, bool, error) { return client, true, nil } t.Cleanup(func() { currentKimiThinkingReplayKVClient = previous })}
func TestKimiThinkingReplayConditionalDeleteKeepsNewerContent(t *testing.T) { ClearKimiThinkingReplayCache() t.Cleanup(ClearKimiThinkingReplayCache)
const modelFamily = "k3" const sessionKey = "execution:conditional-delete" oldContent := []byte(`[{"type":"thinking","signature":"old"}]`) newContent := []byte(`[{"type":"thinking","signature":"new"}]`) if !CacheKimiThinkingReplayBestEffort(context.Background(), modelFamily, sessionKey, oldContent) { t.Fatal("failed to seed old content") } _, snapshot, found, errGet := GetKimiThinkingReplayWithSnapshotRequired(context.Background(), modelFamily, sessionKey) if errGet != nil || !found { t.Fatalf("GetKimiThinkingReplayWithSnapshotRequired() = found %v, error %v", found, errGet) } if !CacheKimiThinkingReplayBestEffort(context.Background(), modelFamily, sessionKey, newContent) { t.Fatal("failed to write newer content") } if !CacheKimiThinkingReplayBestEffort(context.Background(), modelFamily, sessionKey, oldContent) { t.Fatal("failed to write latest content with repeated bytes") }
deleted, errDelete := DeleteKimiThinkingReplayIfUnchanged(context.Background(), modelFamily, sessionKey, snapshot) if errDelete != nil { t.Fatalf("DeleteKimiThinkingReplayIfUnchanged() error = %v", errDelete) } if deleted { t.Fatal("stale snapshot deleted newer content") } got, found, errGet := GetKimiThinkingReplayRequired(context.Background(), modelFamily, sessionKey) if errGet != nil || !found || !bytes.Equal(got, oldContent) { t.Fatalf("cached content = %s, found %v, error %v; want latest repeated content", got, found, errGet) }}
func TestKimiThinkingReplayConditionalReplaceKeepsConcurrentContent(t *testing.T) { ClearKimiThinkingReplayCache() t.Cleanup(ClearKimiThinkingReplayCache)
const modelFamily = "k3" const sessionKey = "execution:conditional-replace" _, snapshot, found, errGet := GetKimiThinkingReplayWithSnapshotRequired(context.Background(), modelFamily, sessionKey) if errGet != nil || found { t.Fatalf("initial cache read = found %v, error %v; want miss", found, errGet) } newContent := []byte(`[{"type":"thinking","signature":"new"}]`) staleContent := []byte(`[{"type":"thinking","signature":"stale"}]`) if !CacheKimiThinkingReplayBestEffort(context.Background(), modelFamily, sessionKey, newContent) { t.Fatal("failed to write concurrent content") }
replaced, errReplace := ReplaceKimiThinkingReplayIfUnchanged(context.Background(), modelFamily, sessionKey, snapshot, staleContent) if errReplace != nil { t.Fatalf("ReplaceKimiThinkingReplayIfUnchanged() error = %v", errReplace) } if replaced { t.Fatal("stale snapshot replaced concurrent content") } got, found, errGet := GetKimiThinkingReplayRequired(context.Background(), modelFamily, sessionKey) if errGet != nil || !found || !bytes.Equal(got, newContent) { t.Fatalf("cached content = %s, found %v, error %v; want concurrent content", got, found, errGet) }}
func TestKimiThinkingReplayTombstoneFencesConcurrentMiss(t *testing.T) { ClearKimiThinkingReplayCache() t.Cleanup(ClearKimiThinkingReplayCache)
const modelFamily = "k3" const sessionKey = "execution:tombstone-fence" _, firstSnapshot, firstFound, errFirst := GetKimiThinkingReplayWithSnapshotRequired(context.Background(), modelFamily, sessionKey) _, secondSnapshot, secondFound, errSecond := GetKimiThinkingReplayWithSnapshotRequired(context.Background(), modelFamily, sessionKey) if errFirst != nil || errSecond != nil || firstFound || secondFound { t.Fatalf("concurrent misses = %v/%v, errors %v/%v", firstFound, secondFound, errFirst, errSecond) } deleted, errDelete := DeleteKimiThinkingReplayIfUnchanged(context.Background(), modelFamily, sessionKey, firstSnapshot) if errDelete != nil || !deleted { t.Fatalf("first miss delete = %v, error %v", deleted, errDelete) } staleContent := []byte(`[{"type":"thinking","signature":"stale"}]`) replaced, errReplace := ReplaceKimiThinkingReplayIfUnchanged(context.Background(), modelFamily, sessionKey, secondSnapshot, staleContent) if errReplace != nil { t.Fatalf("stale miss replace error = %v", errReplace) } if replaced { t.Fatal("stale miss snapshot crossed a newer tombstone") }}
func TestKimiThinkingReplayHomeGenerationPreventsABADelete(t *testing.T) { client := newFakeKimiThinkingReplayKVClient() useFakeKimiThinkingReplayKVClient(t, client)
const modelFamily = "k3" const sessionKey = "execution:home-aba" contentA := []byte(`[{"type":"thinking","signature":"A"}]`) contentB := []byte(`[{"type":"thinking","signature":"B"}]`) if !CacheKimiThinkingReplayBestEffort(context.Background(), modelFamily, sessionKey, contentA) { t.Fatal("failed to seed Home content A") } _, snapshotA, found, errGet := GetKimiThinkingReplayWithSnapshotRequired(context.Background(), modelFamily, sessionKey) if errGet != nil || !found { t.Fatalf("Home snapshot A = found %v, error %v", found, errGet) } if !CacheKimiThinkingReplayBestEffort(context.Background(), modelFamily, sessionKey, contentB) || !CacheKimiThinkingReplayBestEffort(context.Background(), modelFamily, sessionKey, contentA) { t.Fatal("failed to complete Home A-B-A sequence") } deleted, errDelete := DeleteKimiThinkingReplayIfUnchanged(context.Background(), modelFamily, sessionKey, snapshotA) if errDelete != nil { t.Fatalf("Home stale delete error = %v", errDelete) } if deleted { t.Fatal("Home stale snapshot deleted a newer generation with repeated content") } got, found, errGet := GetKimiThinkingReplayRequired(context.Background(), modelFamily, sessionKey) if errGet != nil || !found || !bytes.Equal(got, contentA) { t.Fatalf("Home cached content = %s, found %v, error %v; want latest A", got, found, errGet) }}
func TestKimiThinkingReplayTracksAggregateLocalBytes(t *testing.T) { ClearKimiThinkingReplayCache() t.Cleanup(ClearKimiThinkingReplayCache)
first := []byte(`[{"type":"thinking","signature":"first"}]`) second := []byte(`[{"type":"thinking","signature":"second"}]`) if !CacheKimiThinkingReplayBestEffort(context.Background(), "k3", "execution:bytes-1", first) || !CacheKimiThinkingReplayBestEffort(context.Background(), "k3", "execution:bytes-2", second) { t.Fatal("failed to seed aggregate byte accounting") } if got, want := kimiThinkingReplayTotalBytes, len(first)+len(second); got != want { t.Fatalf("aggregate bytes = %d, want %d", got, want) } if errDelete := DeleteKimiThinkingReplayRequired(context.Background(), "k3", "execution:bytes-1"); errDelete != nil { t.Fatalf("DeleteKimiThinkingReplayRequired() error = %v", errDelete) } if got, want := kimiThinkingReplayTotalBytes, len(second); got != want { t.Fatalf("aggregate bytes after delete = %d, want %d", got, want) } ClearKimiThinkingReplayCache() if kimiThinkingReplayTotalBytes != 0 { t.Fatalf("aggregate bytes after clear = %d, want 0", kimiThinkingReplayTotalBytes) }}
func TestKimiThinkingReplayRejectsOversizedContent(t *testing.T) { ClearKimiThinkingReplayCache() t.Cleanup(ClearKimiThinkingReplayCache)
content := make([]byte, KimiThinkingReplayCacheMaxBytesPerEntry+1) content[0] = '[' for i := 1; i < len(content)-1; i++ { content[i] = ' ' } content[len(content)-1] = ']' if CacheKimiThinkingReplayBestEffort(context.Background(), "k3", "execution:oversized", content) { t.Fatal("oversized content was cached") }}