internal/cache/kimi_thinking_replay_cache_test.go

Rama main · réplica de lectura del repositorio original.

package cache

import (
	"bytes"
	"context"
	"sync"
	"testing"
	"time"

	homekv "github.com/router-for-me/CLIProxyAPI/v8/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")
	}
}