internal/cache/codex_reasoning_replay_cache.go
Rama main · réplica de lectura del repositorio original.
package cache
import (
"context"
"sort"
"strings"
"sync"
"time"
"github.com/MoeForge/moego/ojson"
homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home"
"github.com/router-for-me/CLIProxyAPI/v8/internal/signature"
"github.com/router-for-me/CLIProxyAPI/v8/internal/util"
log "github.com/sirupsen/logrus"
)
const (
// CodexReasoningReplayTurnType identifies an internal turn-boundary marker.
CodexReasoningReplayTurnType = "cpa_codex_replay_turn"
// CodexReasoningReplayCacheTTL limits how long encrypted reasoning replay
// items stay in process memory.
CodexReasoningReplayCacheTTL = 1 * time.Hour
// CodexReasoningReplayCacheMaxEntries bounds process memory for replay
// continuity. Oldest entries are evicted first.
CodexReasoningReplayCacheMaxEntries = 10240
// CodexReasoningReplayCacheMaxTurnsPerEntry bounds cumulative state for one agent.
CodexReasoningReplayCacheMaxTurnsPerEntry = 256
// CodexReasoningReplayCacheMaxBytesPerEntry bounds cumulative serialized items for one agent.
CodexReasoningReplayCacheMaxBytesPerEntry = 16 << 20
// CodexReasoningReplayCacheEvictBatchSize leaves headroom after the cache
// reaches capacity so high write volume does not rescan the map every turn.
CodexReasoningReplayCacheEvictBatchSize = 128
)
type codexReasoningReplayEntry struct {
Items [][]byte
Timestamp time.Time
}
var (
codexReasoningReplayMu sync.Mutex
codexReasoningReplayEntries = make(map[string]codexReasoningReplayEntry)
)
type codexReasoningReplayKVClient interface {
KVGet(ctx context.Context, key string) ([]byte, bool, error)
KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error)
KVCompareAndSwap(ctx context.Context, key string, expected []byte, expectedExists bool, value []byte, ttl time.Duration) (bool, error)
KVDel(ctx context.Context, keys ...string) (int64, error)
KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error)
}
var currentCodexReasoningReplayKVClient = func() (codexReasoningReplayKVClient, bool, error) {
return homekv.CurrentKVClient()
}
// CacheCodexReasoningReplayItem stores a final GPT/Codex reasoning item for
// stateless replay. The stored item is normalized to the minimal shape accepted
// by Responses input replay.
func CacheCodexReasoningReplayItem(modelName, sessionKey string, item []byte) bool {
return CacheCodexReasoningReplayItems(modelName, sessionKey, [][]byte{item})
}
// CacheCodexReasoningReplayItems stores the final GPT/Codex assistant output
// items needed to replay a stateless next turn.
func CacheCodexReasoningReplayItems(modelName, sessionKey string, items [][]byte) bool {
return CacheCodexReasoningReplayItemsBestEffort(context.Background(), modelName, sessionKey, items)
}
// CacheCodexReasoningReplayItemsBestEffort stores replay items for completed response paths.
func CacheCodexReasoningReplayItemsBestEffort(ctx context.Context, modelName, sessionKey string, items [][]byte) bool {
key := codexReasoningReplayCacheKey(modelName, sessionKey)
if key == "" {
return false
}
normalized, ok := normalizeCodexReasoningReplayItems(items)
if !ok {
return false
}
if client, homeMode, errClient := currentCodexReasoningReplayKVClient(); homeMode {
if errClient != nil {
log.Errorf("home kv best-effort codex reasoning replay set failed prefix=cpa:codex:*: %v", errClient)
return false
}
raw, errMarshal := marshalReasoningReplayRawJSONArray(normalized)
if errMarshal != nil {
log.Errorf("home kv best-effort codex reasoning replay set failed prefix=cpa:codex:*: %v", errMarshal)
return false
}
written, errSet := client.KVSet(ctx, codexReasoningReplayKVKey(modelName, sessionKey), raw, homekv.KVSetOptions{EX: CodexReasoningReplayCacheTTL})
if errSet != nil {
log.Errorf("home kv best-effort codex reasoning replay set failed prefix=cpa:codex:*: %v", errSet)
return false
}
return written
}
cacheCleanupOnce.Do(startCacheCleanup)
now := time.Now()
codexReasoningReplayMu.Lock()
defer codexReasoningReplayMu.Unlock()
codexReasoningReplayEntries[key] = codexReasoningReplayEntry{
Items: normalized,
Timestamp: now,
}
if len(codexReasoningReplayEntries) > CodexReasoningReplayCacheMaxEntries {
evictOldestCodexReasoningReplayEntries(CodexReasoningReplayCacheEvictBatchSize)
}
return true
}
// AppendCodexReasoningReplayItemsBestEffort appends one completed turn to existing replay state.
func AppendCodexReasoningReplayItemsBestEffort(ctx context.Context, modelName, sessionKey string, items [][]byte) bool {
if ctx == nil {
ctx = context.Background()
}
key := codexReasoningReplayCacheKey(modelName, sessionKey)
if key == "" {
return false
}
normalized, ok := normalizeCodexReasoningReplayItems(items)
if !ok {
return false
}
if client, homeMode, errClient := currentCodexReasoningReplayKVClient(); homeMode {
if errClient != nil {
log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errClient)
return false
}
kvKey := codexReasoningReplayKVKey(modelName, sessionKey)
const maxCASAttempts = 32
for attempt := 0; attempt < maxCASAttempts; attempt++ {
if errContext := ctx.Err(); errContext != nil {
return false
}
existingRaw, found, errGet := client.KVGet(ctx, kvKey)
if errGet != nil {
log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errGet)
return false
}
var existing [][]byte
if found {
var errUnmarshal error
existing, errUnmarshal = unmarshalReasoningReplayRawJSONArray(existingRaw)
if errUnmarshal != nil {
log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:* (please upgrade CLIProxyAPI): %v", errUnmarshal)
return false
}
}
combined := appendCodexReasoningReplayTurn(existing, normalized)
raw, errMarshal := marshalReasoningReplayRawJSONArray(combined)
if errMarshal != nil {
log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errMarshal)
return false
}
written, errCAS := client.KVCompareAndSwap(ctx, kvKey, existingRaw, found, raw, CodexReasoningReplayCacheTTL)
if errCAS != nil {
log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errCAS)
return false
}
if written {
return true
}
}
log.Warn("home kv best-effort codex reasoning replay append exhausted compare-and-swap attempts")
return false
}
cacheCleanupOnce.Do(startCacheCleanup)
now := time.Now()
codexReasoningReplayMu.Lock()
entry := codexReasoningReplayEntries[key]
if now.Sub(entry.Timestamp) > CodexReasoningReplayCacheTTL {
entry.Items = nil
}
entry.Items = appendCodexReasoningReplayTurn(entry.Items, normalized)
entry.Timestamp = now
codexReasoningReplayEntries[key] = entry
if len(codexReasoningReplayEntries) > CodexReasoningReplayCacheMaxEntries {
evictOldestCodexReasoningReplayEntries(CodexReasoningReplayCacheEvictBatchSize)
}
codexReasoningReplayMu.Unlock()
return true
}
func appendCodexReasoningReplayTurn(existing, turn [][]byte) [][]byte {
if len(existing) > 0 && strings.TrimSpace(ojson.GetBytesNoCopy(existing[0], "type").String()) != CodexReasoningReplayTurnType {
existing = nil
}
turnID := ""
if len(turn) > 0 && strings.TrimSpace(ojson.GetBytesNoCopy(turn[0], "type").String()) == CodexReasoningReplayTurnType {
turnID = strings.TrimSpace(ojson.GetBytesNoCopy(turn[0], "id").String())
}
if turnID != "" {
for _, item := range existing {
if strings.TrimSpace(ojson.GetBytesNoCopy(item, "type").String()) == CodexReasoningReplayTurnType &&
strings.TrimSpace(ojson.GetBytesNoCopy(item, "id").String()) == turnID {
return trimCodexReasoningReplayItems(cloneCodexReasoningReplayItems(existing))
}
}
}
combined := make([][]byte, 0, len(existing)+len(turn))
combined = append(combined, cloneCodexReasoningReplayItems(existing)...)
combined = append(combined, cloneCodexReasoningReplayItems(turn)...)
return trimCodexReasoningReplayItems(combined)
}
func trimCodexReasoningReplayItems(items [][]byte) [][]byte {
for {
turnStarts := []int{0}
totalBytes := 0
for index, item := range items {
totalBytes += len(item)
if index > 0 && strings.TrimSpace(ojson.GetBytesNoCopy(item, "type").String()) == CodexReasoningReplayTurnType {
turnStarts = append(turnStarts, index)
}
}
if len(turnStarts) <= CodexReasoningReplayCacheMaxTurnsPerEntry && totalBytes <= CodexReasoningReplayCacheMaxBytesPerEntry {
return items
}
if len(turnStarts) <= 1 {
return nil
}
items = items[turnStarts[1]:]
}
}
// GetCodexReasoningReplayItem retrieves the first normalized upstream replay item.
func GetCodexReasoningReplayItem(modelName, sessionKey string) ([]byte, bool) {
items, ok := GetCodexReasoningReplayItems(modelName, sessionKey)
if !ok {
return nil, false
}
for _, item := range items {
if strings.TrimSpace(ojson.GetBytesNoCopy(item, "type").String()) != CodexReasoningReplayTurnType {
return item, true
}
}
return nil, false
}
// GetCodexReasoningReplayItems retrieves normalized assistant output items.
func GetCodexReasoningReplayItems(modelName, sessionKey string) ([][]byte, bool) {
items, ok, err := GetCodexReasoningReplayItemsRequired(context.Background(), modelName, sessionKey)
if err == nil {
return items, ok
}
return nil, false
}
// GetCodexReasoningReplayItemsRequired retrieves replay items for request-time paths.
func GetCodexReasoningReplayItemsRequired(ctx context.Context, modelName, sessionKey string) ([][]byte, bool, error) {
key := codexReasoningReplayCacheKey(modelName, sessionKey)
if key == "" {
return nil, false, nil
}
client, homeMode, errClient := currentCodexReasoningReplayKVClient()
if homeMode {
if errClient != nil {
return nil, false, errClient
}
raw, found, errGet := client.KVGet(ctx, codexReasoningReplayKVKey(modelName, sessionKey))
if errGet != nil || !found {
return nil, false, errGet
}
homeItems, errUnmarshal := unmarshalReasoningReplayRawJSONArray(raw)
if errUnmarshal != nil {
log.Warnf("codex reasoning replay KV decode failed (please upgrade CLIProxyAPI): %v", errUnmarshal)
return nil, false, errUnmarshal
}
if _, errExpire := client.KVExpire(ctx, codexReasoningReplayKVKey(modelName, sessionKey), CodexReasoningReplayCacheTTL); errExpire != nil {
return nil, false, errExpire
}
return cloneCodexReasoningReplayItems(homeItems), true, nil
}
cacheCleanupOnce.Do(startCacheCleanup)
now := time.Now()
codexReasoningReplayMu.Lock()
defer codexReasoningReplayMu.Unlock()
entry, ok := codexReasoningReplayEntries[key]
if !ok {
return nil, false, nil
}
if now.Sub(entry.Timestamp) > CodexReasoningReplayCacheTTL {
delete(codexReasoningReplayEntries, key)
return nil, false, nil
}
entry.Timestamp = now
codexReasoningReplayEntries[key] = entry
return cloneCodexReasoningReplayItems(entry.Items), true, nil
}
// DeleteCodexReasoningReplayItem removes one replay item after upstream rejects
// it or the caller otherwise knows it is stale.
func DeleteCodexReasoningReplayItem(modelName, sessionKey string) {
if errDelete := DeleteCodexReasoningReplayItemRequired(context.Background(), modelName, sessionKey); errDelete != nil {
return
}
}
// DeleteCodexReasoningReplayItemRequired removes one replay item for request-time paths.
func DeleteCodexReasoningReplayItemRequired(ctx context.Context, modelName, sessionKey string) error {
key := codexReasoningReplayCacheKey(modelName, sessionKey)
if key == "" {
return nil
}
client, homeMode, errClient := currentCodexReasoningReplayKVClient()
if homeMode {
if errClient != nil {
return errClient
}
_, errDel := client.KVDel(ctx, codexReasoningReplayKVKey(modelName, sessionKey))
return errDel
}
codexReasoningReplayMu.Lock()
delete(codexReasoningReplayEntries, key)
codexReasoningReplayMu.Unlock()
return nil
}
// ClearCodexReasoningReplayCache clears all Codex reasoning replay state.
func ClearCodexReasoningReplayCache() {
codexReasoningReplayMu.Lock()
codexReasoningReplayEntries = make(map[string]codexReasoningReplayEntry)
codexReasoningReplayMu.Unlock()
}
func codexReasoningReplayCacheKey(modelName, sessionKey string) string {
modelName = strings.TrimSpace(modelName)
sessionKey = strings.TrimSpace(sessionKey)
if modelName == "" || sessionKey == "" {
return ""
}
// The session key is the continuity boundary. Keep this independent from
// the selected upstream Codex credential so auth failover can preserve replay.
return strings.Join([]string{"codex-reasoning-replay", modelName, sessionKey}, "\x00")
}
func codexReasoningReplayKVKey(modelName, sessionKey string) string {
return "cpa:codex:reasoning-replay:" + homekv.HashKeyPart(strings.TrimSpace(modelName)) + ":" + homekv.HashKeyPart(strings.TrimSpace(sessionKey))
}
func normalizeCodexReasoningReplayItems(items [][]byte) ([][]byte, bool) {
normalized := make([][]byte, 0, len(items))
for _, item := range items {
normalizedItem, ok := normalizeCodexReasoningReplayItem(item)
if ok {
normalized = append(normalized, normalizedItem)
}
}
normalized = trimCodexReasoningReplayItems(normalized)
return normalized, len(normalized) > 0
}
func normalizeCodexReasoningReplayItem(item []byte) ([]byte, bool) {
itemResult := util.ParseGJSONBytesNoCopy(item)
switch strings.TrimSpace(itemResult.Get("type").String()) {
case CodexReasoningReplayTurnType:
return normalizeCodexReasoningReplayTurn(itemResult)
case "reasoning":
return normalizeCodexReasoningReplayReasoningItem(itemResult)
case "function_call":
return normalizeCodexReasoningReplayFunctionCallItem(itemResult)
case "custom_tool_call":
return normalizeCodexReasoningReplayCustomToolCallItem(itemResult)
default:
return nil, false
}
}
func normalizeCodexReasoningReplayTurn(itemResult ojson.Result) ([]byte, bool) {
turnID := strings.TrimSpace(itemResult.Get("id").String())
if turnID == "" {
return nil, false
}
doc := ojson.NewDocString(`{"type":"` + CodexReasoningReplayTurnType + `"}`)
_ = doc.Set("id", turnID)
if fingerprint := strings.TrimSpace(itemResult.Get("assistant_fingerprint").String()); fingerprint != "" {
_ = doc.Set("assistant_fingerprint", fingerprint)
}
if fingerprint := strings.TrimSpace(itemResult.Get("request_fingerprint").String()); fingerprint != "" {
_ = doc.Set("request_fingerprint", fingerprint)
}
callIDs := itemResult.Get("call_ids")
if callIDs.IsArray() {
for _, callIDResult := range callIDs.Array() {
if callID := strings.TrimSpace(callIDResult.String()); callID != "" {
_ = doc.Set("call_ids.-1", callID)
}
}
}
return doc.Bytes(), true
}
func normalizeCodexReasoningReplayReasoningItem(itemResult ojson.Result) ([]byte, bool) {
encryptedContentResult := itemResult.Get("encrypted_content")
if encryptedContentResult.Type != ojson.String {
return nil, false
}
encryptedContent := encryptedContentResult.String()
if encryptedContent != strings.TrimSpace(encryptedContent) {
return nil, false
}
if !signature.IsValidGPTReasoningSignature(encryptedContent) && !signature.IsValidMetaReasoningSignature(encryptedContent) {
return nil, false
}
normalized := []byte(`{"type":"reasoning","summary":[],"content":null}`)
normalized, _ = ojson.SetBytes(normalized, "encrypted_content", encryptedContent)
return normalized, true
}
func normalizeCodexReasoningReplayFunctionCallItem(itemResult ojson.Result) ([]byte, bool) {
callID := strings.TrimSpace(itemResult.Get("call_id").String())
name := strings.TrimSpace(itemResult.Get("name").String())
arguments := itemResult.Get("arguments")
if callID == "" || name == "" || arguments.Type != ojson.String {
return nil, false
}
doc := ojson.NewDocString(`{"type":"function_call"}`)
_ = doc.Set("call_id", callID)
_ = doc.Set("name", name)
_ = doc.Set("arguments", arguments.String())
return doc.Bytes(), true
}
func normalizeCodexReasoningReplayCustomToolCallItem(itemResult ojson.Result) ([]byte, bool) {
callID := strings.TrimSpace(itemResult.Get("call_id").String())
name := strings.TrimSpace(itemResult.Get("name").String())
input := itemResult.Get("input")
if callID == "" || name == "" || !input.Exists() {
return nil, false
}
doc := ojson.NewDocString(`{"type":"custom_tool_call","status":"completed"}`)
if status := strings.TrimSpace(itemResult.Get("status").String()); status != "" {
_ = doc.Set("status", status)
}
_ = doc.Set("call_id", callID)
_ = doc.Set("name", name)
if input.Type == ojson.String {
_ = doc.Set("input", input.String())
} else {
_ = doc.SetRaw("input", input.Raw)
}
return doc.Bytes(), true
}
func cloneCodexReasoningReplayItems(items [][]byte) [][]byte {
cloned := make([][]byte, 0, len(items))
for _, item := range items {
cloned = append(cloned, append([]byte(nil), item...))
}
return cloned
}
func evictOldestCodexReasoningReplayEntries(count int) {
if count <= 0 || len(codexReasoningReplayEntries) == 0 {
return
}
type candidate struct {
key string
timestamp time.Time
}
candidates := make([]candidate, 0, len(codexReasoningReplayEntries))
for key, entry := range codexReasoningReplayEntries {
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++ {
delete(codexReasoningReplayEntries, candidates[i].key)
}
}
func purgeExpiredCodexReasoningReplayCache(now time.Time) {
codexReasoningReplayMu.Lock()
for key, entry := range codexReasoningReplayEntries {
if now.Sub(entry.Timestamp) > CodexReasoningReplayCacheTTL {
delete(codexReasoningReplayEntries, key)
}
}
codexReasoningReplayMu.Unlock()
}