diff --git a/internal/cache/kimi_thinking_replay_cache.go b/internal/cache/kimi_thinking_replay_cache.go new file mode 100644 index 00000000..c23871bb --- /dev/null +++ b/internal/cache/kimi_thinking_replay_cache.go @@ -0,0 +1,426 @@ +package cache + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + "sync" + "time" + + "github.com/google/uuid" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +const ( + // KimiThinkingReplayCacheTTL limits how long signed assistant content stays replayable. + KimiThinkingReplayCacheTTL = 1 * time.Hour + + // KimiThinkingReplayCacheMaxEntries bounds process memory used for replay continuity. + KimiThinkingReplayCacheMaxEntries = 10240 + + // KimiThinkingReplayCacheEvictBatchSize leaves headroom after reaching capacity. + KimiThinkingReplayCacheEvictBatchSize = 128 + + // KimiThinkingReplayCacheMaxBytesPerEntry bounds one complete assistant content array. + KimiThinkingReplayCacheMaxBytesPerEntry = 8 << 20 + + // KimiThinkingReplayCacheMaxBlocksPerEntry prevents pathological content arrays. + KimiThinkingReplayCacheMaxBlocksPerEntry = 512 + + // KimiThinkingReplayCacheMaxTotalBytes bounds aggregate in-process replay content. + KimiThinkingReplayCacheMaxTotalBytes = 256 << 20 + + kimiThinkingReplayCacheMaxSerializedBytes = KimiThinkingReplayCacheMaxBytesPerEntry + 1024 +) + +type kimiThinkingReplayEntry struct { + Content []byte + Timestamp time.Time + Generation string + Deleted bool +} + +// KimiThinkingReplaySnapshot identifies the exact replay generation read for one request. +type KimiThinkingReplaySnapshot struct { + raw []byte + generation string + loaded bool + found bool +} + +type kimiThinkingReplayHomeValue struct { + Generation string `json:"generation"` + Deleted bool `json:"deleted,omitempty"` + Content json.RawMessage `json:"content,omitempty"` +} + +var ( + kimiThinkingReplayMu sync.Mutex + kimiThinkingReplayEntries = make(map[string]kimiThinkingReplayEntry) + kimiThinkingReplayTotalBytes int +) + +type kimiThinkingReplayKVClient interface { + KVGet(ctx context.Context, key string) ([]byte, bool, error) + KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) + KVDel(ctx context.Context, keys ...string) (int64, error) + KVCompareAndSwap(ctx context.Context, key string, expected []byte, expectedExists bool, value []byte, ttl time.Duration) (bool, error) + KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) +} + +var currentKimiThinkingReplayKVClient = func() (kimiThinkingReplayKVClient, bool, error) { + return homekv.CurrentKVClient() +} + +// CacheKimiThinkingReplayBestEffort stores one complete signed assistant content array. +func CacheKimiThinkingReplayBestEffort(ctx context.Context, modelFamily, sessionKey string, content []byte) bool { + key := kimiThinkingReplayCacheKey(modelFamily, sessionKey) + if key == "" || !validKimiThinkingReplayContent(content) { + return false + } + if ctx == nil { + ctx = context.Background() + } + cloned := append([]byte(nil), content...) + generation := uuid.NewString() + if client, homeMode, errClient := currentKimiThinkingReplayKVClient(); homeMode { + if errClient != nil { + log.Errorf("home kv best-effort kimi thinking replay set failed prefix=cpa:kimi:*: %v", errClient) + return false + } + raw, errMarshal := marshalKimiThinkingReplayHomeValue(generation, false, cloned) + if errMarshal != nil { + log.Errorf("home kv best-effort kimi thinking replay set failed prefix=cpa:kimi:*: %v", errMarshal) + return false + } + written, errSet := client.KVSet(ctx, kimiThinkingReplayKVKey(modelFamily, sessionKey), raw, homekv.KVSetOptions{EX: KimiThinkingReplayCacheTTL}) + if errSet != nil { + log.Errorf("home kv best-effort kimi thinking replay set failed prefix=cpa:kimi:*: %v", errSet) + return false + } + return written + } + + storeKimiThinkingReplayLocal(key, cloned, generation, false, time.Now()) + return true +} + +// GetKimiThinkingReplayRequired retrieves complete assistant content for request-time replay. +func GetKimiThinkingReplayRequired(ctx context.Context, modelFamily, sessionKey string) ([]byte, bool, error) { + content, _, found, errGet := GetKimiThinkingReplayWithSnapshotRequired(ctx, modelFamily, sessionKey) + return content, found, errGet +} + +// GetKimiThinkingReplayWithSnapshotRequired retrieves replay content and the exact cache state read. +func GetKimiThinkingReplayWithSnapshotRequired(ctx context.Context, modelFamily, sessionKey string) ([]byte, KimiThinkingReplaySnapshot, bool, error) { + key := kimiThinkingReplayCacheKey(modelFamily, sessionKey) + if key == "" { + return nil, KimiThinkingReplaySnapshot{}, false, nil + } + if ctx == nil { + ctx = context.Background() + } + client, homeMode, errClient := currentKimiThinkingReplayKVClient() + if homeMode { + if errClient != nil { + return nil, KimiThinkingReplaySnapshot{loaded: true}, false, errClient + } + kvKey := kimiThinkingReplayKVKey(modelFamily, sessionKey) + raw, errRead := readOrReserveKimiThinkingReplayHomeValue(ctx, client, kvKey) + if errRead != nil { + return nil, KimiThinkingReplaySnapshot{loaded: true}, false, errRead + } + snapshot := KimiThinkingReplaySnapshot{raw: append([]byte(nil), raw...), loaded: true, found: true} + content, generation, deleted, okDecode := decodeKimiThinkingReplayHomeValue(raw) + if !okDecode { + return nil, snapshot, false, fmt.Errorf("invalid kimi thinking replay content") + } + snapshot.generation = generation + if _, errExpire := client.KVExpire(ctx, kvKey, KimiThinkingReplayCacheTTL); errExpire != nil { + log.Warnf("home kv kimi thinking replay expire failed prefix=cpa:kimi:*: %v", errExpire) + } + if deleted { + return nil, snapshot, false, nil + } + return content, snapshot, true, nil + } + + cacheCleanupOnce.Do(startCacheCleanup) + now := time.Now() + kimiThinkingReplayMu.Lock() + defer kimiThinkingReplayMu.Unlock() + entry, ok := kimiThinkingReplayEntries[key] + if !ok || now.Sub(entry.Timestamp) > KimiThinkingReplayCacheTTL { + if ok { + kimiThinkingReplayTotalBytes -= len(entry.Content) + delete(kimiThinkingReplayEntries, key) + } + entry = reserveKimiThinkingReplayLocalLocked(key, now) + } + entry.Timestamp = now + kimiThinkingReplayEntries[key] = entry + snapshot := KimiThinkingReplaySnapshot{generation: entry.Generation, loaded: true, found: true} + if entry.Deleted { + return nil, snapshot, false, nil + } + return append([]byte(nil), entry.Content...), snapshot, true, nil +} + +// ReplaceKimiThinkingReplayIfUnchanged stores completed content only if the request snapshot is current. +func ReplaceKimiThinkingReplayIfUnchanged(ctx context.Context, modelFamily, sessionKey string, snapshot KimiThinkingReplaySnapshot, content []byte) (bool, error) { + key := kimiThinkingReplayCacheKey(modelFamily, sessionKey) + if key == "" || !validKimiThinkingReplayContent(content) { + return false, nil + } + if ctx == nil { + ctx = context.Background() + } + if !snapshot.loaded { + return CacheKimiThinkingReplayBestEffort(ctx, modelFamily, sessionKey, content), nil + } + cloned := append([]byte(nil), content...) + generation := uuid.NewString() + client, homeMode, errClient := currentKimiThinkingReplayKVClient() + if homeMode { + if errClient != nil { + return false, errClient + } + raw, errMarshal := marshalKimiThinkingReplayHomeValue(generation, false, cloned) + if errMarshal != nil { + return false, errMarshal + } + return client.KVCompareAndSwap(ctx, kimiThinkingReplayKVKey(modelFamily, sessionKey), snapshot.raw, snapshot.found, raw, KimiThinkingReplayCacheTTL) + } + + cacheCleanupOnce.Do(startCacheCleanup) + kimiThinkingReplayMu.Lock() + defer kimiThinkingReplayMu.Unlock() + entry, found := kimiThinkingReplayEntries[key] + if found != snapshot.found || (found && entry.Generation != snapshot.generation) { + return false, nil + } + kimiThinkingReplayTotalBytes -= len(entry.Content) + kimiThinkingReplayTotalBytes += len(cloned) + kimiThinkingReplayEntries[key] = kimiThinkingReplayEntry{Content: cloned, Timestamp: time.Now(), Generation: generation} + enforceKimiThinkingReplayLimitsLocked() + return true, nil +} + +// DeleteKimiThinkingReplayIfUnchanged clears replay state only if the request snapshot is current. +func DeleteKimiThinkingReplayIfUnchanged(ctx context.Context, modelFamily, sessionKey string, snapshot KimiThinkingReplaySnapshot) (bool, error) { + key := kimiThinkingReplayCacheKey(modelFamily, sessionKey) + if key == "" { + return false, nil + } + if ctx == nil { + ctx = context.Background() + } + if !snapshot.loaded { + return true, DeleteKimiThinkingReplayRequired(ctx, modelFamily, sessionKey) + } + generation := uuid.NewString() + client, homeMode, errClient := currentKimiThinkingReplayKVClient() + if homeMode { + if errClient != nil { + return false, errClient + } + tombstone, errMarshal := marshalKimiThinkingReplayHomeValue(generation, true, nil) + if errMarshal != nil { + return false, errMarshal + } + return client.KVCompareAndSwap(ctx, kimiThinkingReplayKVKey(modelFamily, sessionKey), snapshot.raw, snapshot.found, tombstone, KimiThinkingReplayCacheTTL) + } + + kimiThinkingReplayMu.Lock() + defer kimiThinkingReplayMu.Unlock() + entry, found := kimiThinkingReplayEntries[key] + if found != snapshot.found || (found && entry.Generation != snapshot.generation) { + return false, nil + } + kimiThinkingReplayTotalBytes -= len(entry.Content) + kimiThinkingReplayEntries[key] = kimiThinkingReplayEntry{Timestamp: time.Now(), Generation: generation, Deleted: true} + return true, nil +} + +// DeleteKimiThinkingReplayRequired removes stale replay state unconditionally. +func DeleteKimiThinkingReplayRequired(ctx context.Context, modelFamily, sessionKey string) error { + key := kimiThinkingReplayCacheKey(modelFamily, sessionKey) + if key == "" { + return nil + } + if ctx == nil { + ctx = context.Background() + } + client, homeMode, errClient := currentKimiThinkingReplayKVClient() + if homeMode { + if errClient != nil { + return errClient + } + _, errDelete := client.KVDel(ctx, kimiThinkingReplayKVKey(modelFamily, sessionKey)) + return errDelete + } + kimiThinkingReplayMu.Lock() + if entry, found := kimiThinkingReplayEntries[key]; found { + kimiThinkingReplayTotalBytes -= len(entry.Content) + delete(kimiThinkingReplayEntries, key) + } + kimiThinkingReplayMu.Unlock() + return nil +} + +// ClearKimiThinkingReplayCache clears all in-process Kimi replay state. +func ClearKimiThinkingReplayCache() { + kimiThinkingReplayMu.Lock() + kimiThinkingReplayEntries = make(map[string]kimiThinkingReplayEntry) + kimiThinkingReplayTotalBytes = 0 + kimiThinkingReplayMu.Unlock() +} + +func readOrReserveKimiThinkingReplayHomeValue(ctx context.Context, client kimiThinkingReplayKVClient, key string) ([]byte, error) { + for attempt := 0; attempt < 4; attempt++ { + raw, found, errGet := client.KVGet(ctx, key) + if errGet != nil { + return nil, errGet + } + if found { + if len(raw) > kimiThinkingReplayCacheMaxSerializedBytes { + return nil, fmt.Errorf("kimi thinking replay value exceeds size limit") + } + return raw, nil + } + tombstone, errMarshal := marshalKimiThinkingReplayHomeValue(uuid.NewString(), true, nil) + if errMarshal != nil { + return nil, errMarshal + } + swapped, errReserve := client.KVCompareAndSwap(ctx, key, nil, false, tombstone, KimiThinkingReplayCacheTTL) + if errReserve != nil { + return nil, errReserve + } + if swapped { + return tombstone, nil + } + } + return nil, fmt.Errorf("could not reserve absent kimi thinking replay state") +} + +func marshalKimiThinkingReplayHomeValue(generation string, deleted bool, content []byte) ([]byte, error) { + value := kimiThinkingReplayHomeValue{Generation: generation, Deleted: deleted} + if !deleted { + value.Content = append(json.RawMessage(nil), content...) + } + return json.Marshal(value) +} + +func decodeKimiThinkingReplayHomeValue(raw []byte) ([]byte, string, bool, bool) { + if len(raw) == 0 || len(raw) > kimiThinkingReplayCacheMaxSerializedBytes || !gjson.ValidBytes(raw) { + return nil, "", false, false + } + root := gjson.ParseBytes(raw) + if root.IsArray() { + if !validKimiThinkingReplayContent(raw) { + return nil, "", false, false + } + return append([]byte(nil), raw...), "legacy", false, true + } + var value kimiThinkingReplayHomeValue + if errUnmarshal := json.Unmarshal(raw, &value); errUnmarshal != nil || strings.TrimSpace(value.Generation) == "" { + return nil, "", false, false + } + if value.Deleted { + return nil, value.Generation, true, true + } + if !validKimiThinkingReplayContent(value.Content) { + return nil, "", false, false + } + return append([]byte(nil), value.Content...), value.Generation, false, true +} + +func reserveKimiThinkingReplayLocalLocked(key string, now time.Time) kimiThinkingReplayEntry { + entry := kimiThinkingReplayEntry{Timestamp: now, Generation: uuid.NewString(), Deleted: true} + kimiThinkingReplayEntries[key] = entry + enforceKimiThinkingReplayLimitsLocked() + return entry +} + +func storeKimiThinkingReplayLocal(key string, content []byte, generation string, deleted bool, now time.Time) { + cacheCleanupOnce.Do(startCacheCleanup) + kimiThinkingReplayMu.Lock() + defer kimiThinkingReplayMu.Unlock() + if previous, found := kimiThinkingReplayEntries[key]; found { + kimiThinkingReplayTotalBytes -= len(previous.Content) + } + kimiThinkingReplayTotalBytes += len(content) + kimiThinkingReplayEntries[key] = kimiThinkingReplayEntry{Content: content, Timestamp: now, Generation: generation, Deleted: deleted} + enforceKimiThinkingReplayLimitsLocked() +} + +func kimiThinkingReplayCacheKey(modelFamily, sessionKey string) string { + modelFamily = strings.TrimSpace(modelFamily) + sessionKey = strings.TrimSpace(sessionKey) + if modelFamily == "" || sessionKey == "" { + return "" + } + return strings.Join([]string{"kimi-thinking-replay", modelFamily, sessionKey}, "\x00") +} + +func kimiThinkingReplayKVKey(modelFamily, sessionKey string) string { + return "cpa:kimi:thinking-replay:" + homekv.HashKeyPart(strings.TrimSpace(modelFamily)) + ":" + homekv.HashKeyPart(strings.TrimSpace(sessionKey)) +} + +func validKimiThinkingReplayContent(content []byte) bool { + if len(content) == 0 || len(content) > KimiThinkingReplayCacheMaxBytesPerEntry || !gjson.ValidBytes(content) { + return false + } + root := gjson.ParseBytes(content) + return root.IsArray() && len(root.Array()) > 0 && len(root.Array()) <= KimiThinkingReplayCacheMaxBlocksPerEntry +} + +func enforceKimiThinkingReplayLimitsLocked() { + for len(kimiThinkingReplayEntries) > KimiThinkingReplayCacheMaxEntries || kimiThinkingReplayTotalBytes > KimiThinkingReplayCacheMaxTotalBytes { + if len(kimiThinkingReplayEntries) == 0 { + kimiThinkingReplayTotalBytes = 0 + return + } + evictOldestKimiThinkingReplayEntriesLocked(KimiThinkingReplayCacheEvictBatchSize) + } +} + +func evictOldestKimiThinkingReplayEntriesLocked(count int) { + if count <= 0 || len(kimiThinkingReplayEntries) == 0 { + return + } + type candidate struct { + key string + timestamp time.Time + } + candidates := make([]candidate, 0, len(kimiThinkingReplayEntries)) + for key, entry := range kimiThinkingReplayEntries { + candidates = append(candidates, candidate{key: key, timestamp: entry.Timestamp}) + } + sort.Slice(candidates, func(i, j int) bool { + return candidates[i].timestamp.Before(candidates[j].timestamp) + }) + if count > len(candidates) { + count = len(candidates) + } + for i := 0; i < count; i++ { + entry := kimiThinkingReplayEntries[candidates[i].key] + kimiThinkingReplayTotalBytes -= len(entry.Content) + delete(kimiThinkingReplayEntries, candidates[i].key) + } +} + +func purgeExpiredKimiThinkingReplayCache(now time.Time) { + kimiThinkingReplayMu.Lock() + for key, entry := range kimiThinkingReplayEntries { + if now.Sub(entry.Timestamp) > KimiThinkingReplayCacheTTL { + kimiThinkingReplayTotalBytes -= len(entry.Content) + delete(kimiThinkingReplayEntries, key) + } + } + kimiThinkingReplayMu.Unlock() +} diff --git a/internal/cache/kimi_thinking_replay_cache_test.go b/internal/cache/kimi_thinking_replay_cache_test.go new file mode 100644 index 00000000..4c24f38a --- /dev/null +++ b/internal/cache/kimi_thinking_replay_cache_test.go @@ -0,0 +1,237 @@ +package 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") + } +} diff --git a/internal/cache/signature_cache.go b/internal/cache/signature_cache.go index 75201db2..e9630d5c 100644 --- a/internal/cache/signature_cache.go +++ b/internal/cache/signature_cache.go @@ -111,6 +111,7 @@ func purgeExpiredCaches() { purgeExpiredCodexReasoningReplayCache(now) purgeExpiredXAIReasoningReplayCache(now) purgeExpiredAntigravityReasoningReplayCache(now) + purgeExpiredKimiThinkingReplayCache(now) } // CacheSignature stores a thinking signature for a given model group and text. diff --git a/internal/runtime/executor/kimi_executor.go b/internal/runtime/executor/kimi_executor.go index b9a89425..6593a94b 100644 --- a/internal/runtime/executor/kimi_executor.go +++ b/internal/runtime/executor/kimi_executor.go @@ -86,7 +86,16 @@ func (e *KimiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req from := opts.SourceFormat if from.String() == "claude" { auth.Attributes["base_url"] = kimiauth.KimiAPIBaseURL - return e.ClaudeExecutor.Execute(ctx, auth, req, opts) + preparedReq, replayScope := prepareKimiThinkingReplayRequest(ctx, req, opts) + claudeResp, errExecute := e.ClaudeExecutor.Execute(ctx, auth, preparedReq, opts) + if errExecute != nil { + if replayScope.replayApplied && shouldClearKimiThinkingReplayAfterError(errExecute) { + clearKimiThinkingReplayContent(ctx, replayScope) + } + return claudeResp, errExecute + } + cacheKimiThinkingReplayResponse(ctx, replayScope, claudeResp.Payload) + return claudeResp, nil } responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) @@ -196,7 +205,15 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut from := opts.SourceFormat if from.String() == "claude" { auth.Attributes["base_url"] = kimiauth.KimiAPIBaseURL - return e.ClaudeExecutor.ExecuteStream(ctx, auth, req, opts) + preparedReq, replayScope := prepareKimiThinkingReplayRequest(ctx, req, opts) + claudeResult, errExecute := e.ClaudeExecutor.ExecuteStream(ctx, auth, preparedReq, opts) + if errExecute != nil { + if replayScope.replayApplied && shouldClearKimiThinkingReplayAfterError(errExecute) { + clearKimiThinkingReplayContent(ctx, replayScope) + } + return nil, errExecute + } + return wrapKimiThinkingReplayStream(ctx, claudeResult, replayScope), nil } responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) diff --git a/internal/runtime/executor/kimi_thinking_replay.go b/internal/runtime/executor/kimi_thinking_replay.go new file mode 100644 index 00000000..a670c2d7 --- /dev/null +++ b/internal/runtime/executor/kimi_thinking_replay.go @@ -0,0 +1,477 @@ +package executor + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "sort" + "strings" + + internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +type kimiThinkingReplayScope struct { + modelFamily string + sessionKey string + snapshot internalcache.KimiThinkingReplaySnapshot + cacheReady bool + replayApplied bool +} + +func (s kimiThinkingReplayScope) valid() bool { + return strings.TrimSpace(s.modelFamily) != "" && strings.TrimSpace(s.sessionKey) != "" +} + +func kimiThinkingReplayModelFamily(model string) string { + baseModel := thinking.ParseSuffix(strings.TrimSpace(model)).ModelName + normalized := normalizeKimiUpstreamModel(baseModel) + switch normalized { + case "k3", "k3-256k": + return "k3" + default: + return normalized + } +} + +func kimiThinkingReplayScopeFromRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) kimiThinkingReplayScope { + sessionKey := codexReasoningReplaySessionKey(ctx, sdktranslator.FormatClaude, req, opts, req.Payload) + sessionKey = xaiReasoningReplayIsolateSessionKey(ctx, sessionKey) + return kimiThinkingReplayScope{ + modelFamily: kimiThinkingReplayModelFamily(req.Model), + sessionKey: sessionKey, + } +} + +func prepareKimiThinkingReplayRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Request, kimiThinkingReplayScope) { + scope := kimiThinkingReplayScopeFromRequest(ctx, req, opts) + if !scope.valid() { + return req, scope + } + content, snapshot, found, errGet := internalcache.GetKimiThinkingReplayWithSnapshotRequired(ctx, scope.modelFamily, scope.sessionKey) + scope.snapshot = snapshot + scope.cacheReady = errGet == nil + if errGet != nil { + log.Warnf("kimi thinking replay cache read failed: %v", errGet) + return req, scope + } + if !found { + return req, scope + } + updated, restored := restoreKimiThinkingReplayContent(req.Payload, content) + if restored { + req.Payload = updated + scope.replayApplied = true + } + return req, scope +} + +func cacheKimiThinkingReplayResponse(ctx context.Context, scope kimiThinkingReplayScope, response []byte) { + if !scope.valid() || !scope.cacheReady { + return + } + content := gjson.GetBytes(response, "content") + if !content.IsArray() { + return + } + cacheKimiThinkingReplayContent(ctx, scope, []byte(content.Raw)) +} + +func cacheKimiThinkingReplayContent(ctx context.Context, scope kimiThinkingReplayScope, content []byte) { + if !scope.valid() || !scope.cacheReady { + return + } + if kimiThinkingReplayContentIsReplayable(content) { + if _, errReplace := internalcache.ReplaceKimiThinkingReplayIfUnchanged(ctx, scope.modelFamily, scope.sessionKey, scope.snapshot, content); errReplace != nil { + log.Warnf("kimi thinking replay cache replace failed: %v", errReplace) + } + return + } + clearKimiThinkingReplayContent(ctx, scope) +} + +func shouldClearKimiThinkingReplayAfterError(err error) bool { + if err == nil { + return false + } + var upstreamStatus statusErr + if !errors.As(err, &upstreamStatus) { + return false + } + statusCode := upstreamStatus.StatusCode() + return statusCode == 400 || statusCode == 422 +} + +func clearKimiThinkingReplayContent(ctx context.Context, scope kimiThinkingReplayScope) { + if !scope.valid() || !scope.cacheReady { + return + } + if _, errDelete := internalcache.DeleteKimiThinkingReplayIfUnchanged(ctx, scope.modelFamily, scope.sessionKey, scope.snapshot); errDelete != nil { + log.Warnf("kimi thinking replay cache delete failed: %v", errDelete) + } +} + +func kimiThinkingReplayContentIsReplayable(content []byte) bool { + root := gjson.ParseBytes(content) + if !root.IsArray() { + return false + } + hasSignedThinking := false + hasToolUse := false + for _, part := range root.Array() { + switch strings.TrimSpace(part.Get("type").String()) { + case "thinking": + if strings.TrimSpace(part.Get("signature").String()) != "" { + hasSignedThinking = true + } + case "tool_use": + if strings.TrimSpace(part.Get("id").String()) != "" { + hasToolUse = true + } + } + } + return hasSignedThinking && hasToolUse +} + +func restoreKimiThinkingReplayContent(body, cachedContent []byte) ([]byte, bool) { + cachedParts, cachedOK := kimiNonThinkingContentParts(gjson.ParseBytes(cachedContent)) + if !cachedOK { + return body, false + } + messages := gjson.GetBytes(body, "messages") + if !messages.IsArray() { + return body, false + } + messageItems := messages.Array() + for index := len(messageItems) - 1; index >= 0; index-- { + message := messageItems[index] + if !strings.EqualFold(strings.TrimSpace(message.Get("role").String()), "assistant") { + continue + } + currentContent := message.Get("content") + if kimiJSONEqual([]byte(currentContent.Raw), cachedContent) { + return body, false + } + if kimiContentHasThinking(currentContent) { + continue + } + currentParts, currentOK := kimiNonThinkingContentParts(currentContent) + if !currentOK || !kimiCanonicalPartsEqual(currentParts, cachedParts) { + continue + } + updated, errSet := sjson.SetRawBytes(body, fmt.Sprintf("messages.%d.content", index), cachedContent) + if errSet != nil { + return body, false + } + return updated, true + } + return body, false +} + +func kimiContentHasThinking(content gjson.Result) bool { + if !content.IsArray() { + return false + } + for _, part := range content.Array() { + switch strings.TrimSpace(part.Get("type").String()) { + case "thinking", "redacted_thinking": + return true + } + } + return false +} + +func kimiNonThinkingContentParts(content gjson.Result) ([][]byte, bool) { + if !content.IsArray() { + return nil, false + } + parts := make([][]byte, 0, len(content.Array())) + hasToolUse := false + for _, part := range content.Array() { + switch strings.TrimSpace(part.Get("type").String()) { + case "thinking", "redacted_thinking": + continue + case "tool_use": + if strings.TrimSpace(part.Get("id").String()) == "" { + return nil, false + } + hasToolUse = true + } + canonical, ok := kimiCanonicalJSON([]byte(part.Raw)) + if !ok { + return nil, false + } + parts = append(parts, canonical) + } + return parts, hasToolUse +} + +func kimiCanonicalPartsEqual(left, right [][]byte) bool { + if len(left) != len(right) { + return false + } + for i := range left { + if !bytes.Equal(left[i], right[i]) { + return false + } + } + return true +} + +func kimiJSONEqual(left, right []byte) bool { + canonicalLeft, leftOK := kimiCanonicalJSON(left) + canonicalRight, rightOK := kimiCanonicalJSON(right) + return leftOK && rightOK && bytes.Equal(canonicalLeft, canonicalRight) +} + +func kimiCanonicalJSON(raw []byte) ([]byte, bool) { + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var value any + if errDecode := decoder.Decode(&value); errDecode != nil { + return nil, false + } + canonical, errMarshal := json.Marshal(value) + if errMarshal != nil { + return nil, false + } + return canonical, true +} + +type kimiThinkingReplayStreamBlock struct { + raw []byte + text strings.Builder + thinking strings.Builder + signature strings.Builder + input strings.Builder + textInitialized bool + thinkingInitialized bool + signatureInitialized bool + hasInputDelta bool + finished bool +} + +type kimiThinkingReplayStreamAccumulator struct { + blocks map[int]*kimiThinkingReplayStreamBlock + observed bool + complete bool + upstreamError bool + abandoned bool + bytesUsed int +} + +func newKimiThinkingReplayStreamAccumulator() *kimiThinkingReplayStreamAccumulator { + return &kimiThinkingReplayStreamAccumulator{blocks: make(map[int]*kimiThinkingReplayStreamBlock)} +} + +func (a *kimiThinkingReplayStreamAccumulator) observe(chunk []byte) { + for _, line := range bytes.Split(chunk, []byte("\n")) { + line = bytes.TrimSpace(line) + if !bytes.HasPrefix(line, []byte("data:")) { + continue + } + payload := bytes.TrimSpace(line[len("data:"):]) + if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) { + continue + } + if !gjson.ValidBytes(payload) { + a.abandon() + continue + } + root := gjson.ParseBytes(payload) + switch root.Get("type").String() { + case "message_start": + a.observed = true + case "content_block_start": + if !a.abandoned { + a.observeBlockStart(root) + } + case "content_block_delta": + if !a.abandoned { + a.observeBlockDelta(root) + } + case "content_block_stop": + if !a.abandoned { + a.finishBlock(int(root.Get("index").Int())) + } + case "message_stop": + a.complete = true + case "error": + a.upstreamError = true + a.abandon() + } + } +} + +func (a *kimiThinkingReplayStreamAccumulator) observeBlockStart(root gjson.Result) { + index := int(root.Get("index").Int()) + block := root.Get("content_block") + if !block.IsObject() || len(a.blocks) >= internalcache.KimiThinkingReplayCacheMaxBlocksPerEntry { + a.abandon() + return + } + if _, exists := a.blocks[index]; exists { + a.abandon() + return + } + raw := []byte(block.Raw) + if !a.reserveBytes(len(raw)) { + return + } + a.blocks[index] = &kimiThinkingReplayStreamBlock{raw: append([]byte(nil), raw...)} +} + +func (a *kimiThinkingReplayStreamAccumulator) observeBlockDelta(root gjson.Result) { + index := int(root.Get("index").Int()) + block, ok := a.blocks[index] + if !ok { + a.abandon() + return + } + delta := root.Get("delta") + switch delta.Get("type").String() { + case "text_delta": + a.appendBlockText(block, &block.text, &block.textInitialized, "text", delta.Get("text").String()) + case "thinking_delta": + a.appendBlockText(block, &block.thinking, &block.thinkingInitialized, "thinking", delta.Get("thinking").String()) + case "signature_delta": + a.appendBlockText(block, &block.signature, &block.signatureInitialized, "signature", delta.Get("signature").String()) + case "input_json_delta": + suffix := delta.Get("partial_json").String() + if a.reserveBytes(len(suffix)) { + block.input.WriteString(suffix) + block.hasInputDelta = true + } + default: + a.abandon() + } +} + +func (a *kimiThinkingReplayStreamAccumulator) appendBlockText(block *kimiThinkingReplayStreamBlock, builder *strings.Builder, initialized *bool, path, suffix string) { + if !*initialized { + initial := gjson.GetBytes(block.raw, path).String() + if !a.reserveBytes(len(initial)) { + return + } + builder.WriteString(initial) + *initialized = true + } + if a.reserveBytes(len(suffix)) { + builder.WriteString(suffix) + } +} + +func (a *kimiThinkingReplayStreamAccumulator) finishBlock(index int) { + block, ok := a.blocks[index] + if !ok { + a.abandon() + return + } + if block.hasInputDelta && !gjson.Valid(block.input.String()) { + a.abandon() + return + } + block.finished = true +} + +func (a *kimiThinkingReplayStreamAccumulator) reserveBytes(count int) bool { + if count < 0 || a.bytesUsed > internalcache.KimiThinkingReplayCacheMaxBytesPerEntry-count { + a.abandon() + return false + } + a.bytesUsed += count + return true +} + +func (a *kimiThinkingReplayStreamAccumulator) abandon() { + a.abandoned = true + a.blocks = nil + a.bytesUsed = 0 +} + +func (a *kimiThinkingReplayStreamAccumulator) content() ([]byte, bool) { + if !a.observed || !a.complete || a.upstreamError || a.abandoned { + return nil, false + } + indexes := make([]int, 0, len(a.blocks)) + for index := range a.blocks { + indexes = append(indexes, index) + } + sort.Ints(indexes) + parts := make([][]byte, 0, len(indexes)) + for _, index := range indexes { + block := a.blocks[index] + if !block.finished { + a.abandon() + return nil, false + } + raw := append([]byte(nil), block.raw...) + var errSet error + if block.textInitialized { + raw, errSet = sjson.SetBytes(raw, "text", block.text.String()) + } + if errSet == nil && block.thinkingInitialized { + raw, errSet = sjson.SetBytes(raw, "thinking", block.thinking.String()) + } + if errSet == nil && block.signatureInitialized { + raw, errSet = sjson.SetBytes(raw, "signature", block.signature.String()) + } + if errSet == nil && block.hasInputDelta { + raw, errSet = sjson.SetRawBytes(raw, "input", []byte(block.input.String())) + } + if errSet != nil { + a.abandon() + return nil, false + } + parts = append(parts, raw) + } + content := helps.JoinRawJSONArray(parts) + if len(content) > internalcache.KimiThinkingReplayCacheMaxBytesPerEntry { + a.abandon() + return nil, false + } + return content, true +} + +func wrapKimiThinkingReplayStream(ctx context.Context, result *cliproxyexecutor.StreamResult, scope kimiThinkingReplayScope) *cliproxyexecutor.StreamResult { + if result == nil || !scope.valid() { + return result + } + out := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(out) + accumulator := newKimiThinkingReplayStreamAccumulator() + hasError := false + for chunk := range result.Chunks { + if chunk.Err != nil { + hasError = true + } else { + accumulator.observe(chunk.Payload) + } + select { + case out <- chunk: + case <-ctx.Done(): + return + } + } + if hasError { + return + } + if content, completed := accumulator.content(); completed { + cacheKimiThinkingReplayContent(ctx, scope, content) + return + } + if accumulator.upstreamError && scope.replayApplied { + clearKimiThinkingReplayContent(ctx, scope) + } + }() + return &cliproxyexecutor.StreamResult{Headers: result.Headers.Clone(), Chunks: out} +} diff --git a/internal/runtime/executor/kimi_thinking_replay_test.go b/internal/runtime/executor/kimi_thinking_replay_test.go new file mode 100644 index 00000000..380589c1 --- /dev/null +++ b/internal/runtime/executor/kimi_thinking_replay_test.go @@ -0,0 +1,413 @@ +package executor + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "testing" + + internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "github.com/tidwall/gjson" +) + +type kimiLocalBadRequestError struct{} + +func (kimiLocalBadRequestError) Error() string { return "local validation failed" } +func (kimiLocalBadRequestError) StatusCode() int { return http.StatusBadRequest } + +func TestKimiThinkingReplayModelFamily(t *testing.T) { + cases := []struct { + model string + want string + }{ + {model: "k3", want: "k3"}, + {model: "kimi-k3", want: "k3"}, + {model: "k3-256k", want: "k3"}, + {model: "kimi-k3-256k(high)", want: "k3"}, + {model: "kimi-k2.7-code", want: "k2.7-code"}, + {model: "kimi-k2.7-code-highspeed", want: "k2.7-code-highspeed"}, + } + for _, tc := range cases { + t.Run(tc.model, func(t *testing.T) { + if got := kimiThinkingReplayModelFamily(tc.model); got != tc.want { + t.Fatalf("kimiThinkingReplayModelFamily(%q) = %q, want %q", tc.model, got, tc.want) + } + }) + } +} + +func TestRestoreKimiThinkingReplayContentPreservesCompleteAssistantContent(t *testing.T) { + cached := []byte(`[ + {"type":"thinking","thinking":"full reasoning","signature":"kimi-signature"}, + {"type":"text","text":"I will inspect the file."}, + {"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}} + ]`) + body := []byte(`{"messages":[ + {"role":"user","content":"inspect"}, + {"role":"assistant","content":[ + {"type":"text","text":"I will inspect the file."}, + {"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}} + ]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"ok"}]} + ]}`) + + updated, restored := restoreKimiThinkingReplayContent(body, cached) + if !restored { + t.Fatal("expected cached thinking content to be restored") + } + got := gjson.GetBytes(updated, "messages.1.content") + if !kimiJSONEqual([]byte(got.Raw), cached) { + t.Fatalf("restored content = %s, want complete cached content %s", got.Raw, cached) + } +} + +func TestRestoreKimiThinkingReplayContentDoesNotReplaceExistingThinking(t *testing.T) { + cached := []byte(`[{"type":"thinking","thinking":"cached","signature":"cached-signature"},{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]`) + body := []byte(`{"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"current","signature":"current-signature"},{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]}]}`) + + updated, restored := restoreKimiThinkingReplayContent(body, cached) + if restored { + t.Fatalf("existing thinking must not be replaced: %s", updated) + } + if !kimiJSONEqual(updated, body) { + t.Fatalf("request changed despite existing thinking: got %s want %s", updated, body) + } +} + +func TestPrepareKimiThinkingReplayRequestSharesOnlyK3Variants(t *testing.T) { + internalcache.ClearKimiThinkingReplayCache() + t.Cleanup(internalcache.ClearKimiThinkingReplayCache) + + const sessionID = "family-switch" + const cached = `[{"type":"thinking","thinking":"reasoning","signature":"kimi-signature"},{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]` + if !internalcache.CacheKimiThinkingReplayBestEffort(context.Background(), "k3", "execution:"+sessionID, []byte(cached)) { + t.Fatal("failed to seed K3 thinking replay cache") + } + if !internalcache.CacheKimiThinkingReplayBestEffort(context.Background(), "k2.7-code", "execution:"+sessionID, []byte(cached)) { + t.Fatal("failed to seed K2.7 Code thinking replay cache") + } + + payload := []byte(`{"model":"kimi-k3-256k","messages":[{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]}]}`) + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: sessionID, + }, + } + prepared, scope := prepareKimiThinkingReplayRequest(context.Background(), cliproxyexecutor.Request{Model: "kimi-k3-256k", Payload: payload}, opts) + if scope.modelFamily != "k3" { + t.Fatalf("K3 replay family = %q, want k3", scope.modelFamily) + } + if !gjson.GetBytes(prepared.Payload, "messages.0.content.0.signature").Exists() { + t.Fatalf("K3 variant switch did not restore cached thinking: %s", prepared.Payload) + } + + k27Payload := []byte(`{"model":"kimi-k2.7-code-highspeed","messages":[{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]}]}`) + preparedK27, scopeK27 := prepareKimiThinkingReplayRequest(context.Background(), cliproxyexecutor.Request{Model: "kimi-k2.7-code-highspeed", Payload: k27Payload}, opts) + if scopeK27.modelFamily != "k2.7-code-highspeed" { + t.Fatalf("K2.7 replay family = %q, want k2.7-code-highspeed", scopeK27.modelFamily) + } + if gjson.GetBytes(preparedK27.Payload, "messages.0.content.0.signature").Exists() { + t.Fatalf("K2.7 variants must remain isolated: %s", preparedK27.Payload) + } +} + +func TestKimiThinkingReplayScopeIsolatesClaudeCodeCallers(t *testing.T) { + internalcache.ClearKimiThinkingReplayCache() + t.Cleanup(internalcache.ClearKimiThinkingReplayCache) + + payload := []byte(`{"model":"kimi-k3","metadata":{"user_id":"{\"session_id\":\"claude-session\"}"},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]}]}`) + req := cliproxyexecutor.Request{Model: "kimi-k3", Payload: payload} + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude} + callerAContext := testContextWithAPIKey("caller-a") + callerAScope := kimiThinkingReplayScopeFromRequest(callerAContext, req, opts) + if !callerAScope.valid() || !strings.Contains(callerAScope.sessionKey, ":claude:claude-session:agent:main") { + t.Fatalf("caller A scope = %+v, want isolated Claude Code session", callerAScope) + } + const cached = `[{"type":"thinking","thinking":"reasoning","signature":"kimi-signature"},{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]` + if !internalcache.CacheKimiThinkingReplayBestEffort(callerAContext, callerAScope.modelFamily, callerAScope.sessionKey, []byte(cached)) { + t.Fatal("failed to seed caller A cache") + } + + preparedA, _ := prepareKimiThinkingReplayRequest(callerAContext, req, opts) + if !gjson.GetBytes(preparedA.Payload, "messages.0.content.0.signature").Exists() { + t.Fatalf("caller A did not receive its replay: %s", preparedA.Payload) + } + preparedB, callerBScope := prepareKimiThinkingReplayRequest(testContextWithAPIKey("caller-b"), req, opts) + if callerBScope.sessionKey == callerAScope.sessionKey { + t.Fatal("different downstream API keys shared one replay scope") + } + if gjson.GetBytes(preparedB.Payload, "messages.0.content.0.signature").Exists() { + t.Fatalf("caller B received caller A replay: %s", preparedB.Payload) + } + _, unauthenticatedScope := prepareKimiThinkingReplayRequest(context.Background(), req, opts) + if unauthenticatedScope.valid() { + t.Fatalf("unauthenticated client-controlled session must not enable replay: %+v", unauthenticatedScope) + } +} + +func TestKimiExecutorClaudeNonStreamReplaysThinkingAcrossK3VariantSwitch(t *testing.T) { + internalcache.ClearKimiThinkingReplayCache() + t.Cleanup(internalcache.ClearKimiThinkingReplayCache) + + const cachedContent = `[{"type":"thinking","thinking":"full reasoning","signature":"kimi-signature"},{"type":"text","text":"Inspecting."},{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]` + var upstreamBodies [][]byte + callCount := 0 + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", kimiRoundTripperFunc(func(req *http.Request) (*http.Response, error) { + body, errRead := io.ReadAll(req.Body) + if errRead != nil { + return nil, errRead + } + upstreamBodies = append(upstreamBodies, body) + callCount++ + response := `{"id":"msg_2","type":"message","role":"assistant","model":"k3","content":[{"type":"text","text":"done"}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}` + if callCount == 1 { + response = `{"id":"msg_1","type":"message","role":"assistant","model":"k3-256k","content":` + cachedContent + `,"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}` + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(response)), + }, nil + })) + + executor := NewKimiExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{}, Metadata: map[string]any{"access_token": "test-token"}} + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "nonstream-switch", + }, + } + firstPayload := []byte(`{"model":"kimi-k3-256k","max_tokens":32,"messages":[{"role":"user","content":"inspect"}]}`) + opts.OriginalRequest = firstPayload + if _, errExecute := executor.Execute(ctx, auth, cliproxyexecutor.Request{Model: "kimi-k3-256k", Payload: firstPayload}, opts); errExecute != nil { + t.Fatalf("first Execute() error = %v", errExecute) + } + + secondPayload := []byte(`{"model":"kimi-k3","max_tokens":32,"messages":[{"role":"user","content":"inspect"},{"role":"assistant","content":[{"type":"text","text":"Inspecting."},{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]},{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"ok"}]}]}`) + opts.OriginalRequest = secondPayload + if _, errExecute := executor.Execute(ctx, auth, cliproxyexecutor.Request{Model: "kimi-k3", Payload: secondPayload}, opts); errExecute != nil { + t.Fatalf("second Execute() error = %v", errExecute) + } + if len(upstreamBodies) != 2 { + t.Fatalf("upstream request count = %d, want 2", len(upstreamBodies)) + } + gotContent := gjson.GetBytes(upstreamBodies[1], "messages.1.content") + if !kimiJSONEqual([]byte(gotContent.Raw), []byte(cachedContent)) { + t.Fatalf("second upstream assistant content = %s, want %s", gotContent.Raw, cachedContent) + } + if _, found, errGet := internalcache.GetKimiThinkingReplayRequired(context.Background(), "k3", "execution:nonstream-switch"); errGet != nil || found { + t.Fatalf("unsigned completed turn left stale replay: found %v, error %v", found, errGet) + } +} + +func TestShouldClearKimiThinkingReplayAfterErrorOnlyForUpstreamRequestRejection(t *testing.T) { + if shouldClearKimiThinkingReplayAfterError(errors.New("transport failed")) { + t.Fatal("transport error must not clear valid replay") + } + if shouldClearKimiThinkingReplayAfterError(kimiLocalBadRequestError{}) { + t.Fatal("local bad request must not clear valid replay") + } + if shouldClearKimiThinkingReplayAfterError(statusErr{code: http.StatusInternalServerError}) { + t.Fatal("upstream server error must not clear valid replay") + } + if !shouldClearKimiThinkingReplayAfterError(statusErr{code: http.StatusBadRequest}) { + t.Fatal("upstream bad request should clear applied replay") + } + if !shouldClearKimiThinkingReplayAfterError(statusErr{code: http.StatusUnprocessableEntity}) { + t.Fatal("upstream unprocessable request should clear applied replay") + } +} + +func TestKimiExecutorClaudeErrorClearsAppliedReplay(t *testing.T) { + internalcache.ClearKimiThinkingReplayCache() + t.Cleanup(internalcache.ClearKimiThinkingReplayCache) + + const sessionKey = "execution:error-clears-replay" + const cachedContent = `[{"type":"thinking","thinking":"reasoning","signature":"kimi-signature"},{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]` + if !internalcache.CacheKimiThinkingReplayBestEffort(context.Background(), "k3", sessionKey, []byte(cachedContent)) { + t.Fatal("failed to seed replay cache") + } + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", kimiRoundTripperFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"invalid thinking signature"}}`)), + }, nil + })) + executor := NewKimiExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{}, Metadata: map[string]any{"access_token": "test-token"}} + payload := []byte(`{"model":"kimi-k3-256k","max_tokens":32,"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]},{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"ok"}]}]}`) + _, errExecute := executor.Execute(ctx, auth, cliproxyexecutor.Request{Model: "kimi-k3-256k", Payload: payload}, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "error-clears-replay", + }, + }) + if errExecute == nil { + t.Fatal("Execute() error = nil, want upstream rejection") + } + if _, found, errGet := internalcache.GetKimiThinkingReplayRequired(context.Background(), "k3", sessionKey); errGet != nil || found { + t.Fatalf("rejected replay remained cached: found %v, error %v", found, errGet) + } +} + +func TestKimiExecutorClaudeStreamReplaysThinkingAcrossK3VariantSwitch(t *testing.T) { + internalcache.ClearKimiThinkingReplayCache() + t.Cleanup(internalcache.ClearKimiThinkingReplayCache) + + const firstStream = "event: message_start\n" + + `data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"k3","content":[],"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":0}}}` + "\n\n" + + "event: content_block_start\n" + + `data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"stream reasoning"}}` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"stream-signature"}}` + "\n\n" + + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n" + + "event: content_block_start\n" + + `data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_stream","name":"Read","input":{}}}` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"path\":\"README.md\"}"}}` + "\n\n" + + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":1}` + "\n\n" + + "event: message_delta\n" + + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":1}}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n" + const secondStream = "event: message_start\n" + + `data: {"type":"message_start","message":{"id":"msg_2","type":"message","role":"assistant","model":"k3-256k","content":[],"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":0}}}` + "\n\n" + + "event: content_block_start\n" + + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"done"}}` + "\n\n" + + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n" + + "event: message_delta\n" + + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n" + + var upstreamBodies [][]byte + callCount := 0 + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", kimiRoundTripperFunc(func(req *http.Request) (*http.Response, error) { + body, errRead := io.ReadAll(req.Body) + if errRead != nil { + return nil, errRead + } + upstreamBodies = append(upstreamBodies, body) + callCount++ + stream := firstStream + if callCount == 2 { + stream = secondStream + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(stream)), + }, nil + })) + + executor := NewKimiExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{}, Metadata: map[string]any{"access_token": "test-token"}} + opts := cliproxyexecutor.Options{ + Stream: true, + SourceFormat: sdktranslator.FormatClaude, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "stream-switch", + }, + } + firstPayload := []byte(`{"model":"kimi-k3","max_tokens":32,"stream":true,"messages":[{"role":"user","content":"inspect"}]}`) + opts.OriginalRequest = firstPayload + firstResult, errExecute := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "kimi-k3", Payload: firstPayload}, opts) + if errExecute != nil { + t.Fatalf("first ExecuteStream() error = %v", errExecute) + } + consumeKimiReplayStream(t, firstResult) + + secondPayload := []byte(`{"model":"kimi-k3-256k","max_tokens":32,"stream":true,"messages":[{"role":"user","content":"inspect"},{"role":"assistant","content":[{"type":"tool_use","id":"toolu_stream","name":"Read","input":{"path":"README.md"}}]},{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_stream","content":"ok"}]}]}`) + opts.OriginalRequest = secondPayload + secondResult, errExecute := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "kimi-k3-256k", Payload: secondPayload}, opts) + if errExecute != nil { + t.Fatalf("second ExecuteStream() error = %v", errExecute) + } + consumeKimiReplayStream(t, secondResult) + + if len(upstreamBodies) != 2 { + t.Fatalf("upstream request count = %d, want 2", len(upstreamBodies)) + } + content := gjson.GetBytes(upstreamBodies[1], "messages.1.content") + if got := content.Get("0.thinking").String(); got != "stream reasoning" { + t.Fatalf("replayed stream thinking = %q, want stream reasoning; content=%s", got, content.Raw) + } + if got := content.Get("0.signature").String(); got != "stream-signature" { + t.Fatalf("replayed stream signature = %q, want stream-signature; content=%s", got, content.Raw) + } + if got := content.Get("1.input.path").String(); got != "README.md" { + t.Fatalf("replayed stream tool input path = %q, want README.md; content=%s", got, content.Raw) + } +} + +func TestKimiThinkingReplayUnknownStreamDeltaPreservesPreviousCache(t *testing.T) { + internalcache.ClearKimiThinkingReplayCache() + t.Cleanup(internalcache.ClearKimiThinkingReplayCache) + + const sessionID = "unknown-stream-delta" + const sessionKey = "execution:" + sessionID + cached := []byte(`[{"type":"thinking","thinking":"reasoning","signature":"kimi-signature"},{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]`) + if !internalcache.CacheKimiThinkingReplayBestEffort(context.Background(), "k3", sessionKey, cached) { + t.Fatal("failed to seed replay cache") + } + payload := []byte(`{"model":"kimi-k3-256k","messages":[{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"README.md"}}]}]}`) + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: sessionID, + }, + } + _, scope := prepareKimiThinkingReplayRequest(context.Background(), cliproxyexecutor.Request{Model: "kimi-k3-256k", Payload: payload}, opts) + if !scope.replayApplied { + t.Fatal("expected seeded replay to be applied") + } + + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte( + "event: message_start\n" + + `data: {"type":"message_start","message":{"id":"msg_1","model":"k3"}}` + "\n\n" + + "event: content_block_start\n" + + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"future_delta","value":"new"}}` + "\n\n" + + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n", + )} + close(chunks) + consumeKimiReplayStream(t, wrapKimiThinkingReplayStream(context.Background(), &cliproxyexecutor.StreamResult{Chunks: chunks}, scope)) + + got, found, errGet := internalcache.GetKimiThinkingReplayRequired(context.Background(), "k3", sessionKey) + if errGet != nil || !found || !kimiJSONEqual(got, cached) { + t.Fatalf("unknown successful delta changed previous cache: got %s, found %v, error %v", got, found, errGet) + } +} + +func consumeKimiReplayStream(t *testing.T, result *cliproxyexecutor.StreamResult) { + t.Helper() + if result == nil { + t.Fatal("stream result is nil") + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + } +}