diff --git a/internal/cache/signature_cache.go b/internal/cache/signature_cache.go index 72c3ddeb..75201db2 100644 --- a/internal/cache/signature_cache.go +++ b/internal/cache/signature_cache.go @@ -109,6 +109,7 @@ func purgeExpiredCaches() { return true }) purgeExpiredCodexReasoningReplayCache(now) + purgeExpiredXAIReasoningReplayCache(now) purgeExpiredAntigravityReasoningReplayCache(now) } diff --git a/internal/cache/xai_reasoning_replay_cache.go b/internal/cache/xai_reasoning_replay_cache.go new file mode 100644 index 00000000..49bc0d66 --- /dev/null +++ b/internal/cache/xai_reasoning_replay_cache.go @@ -0,0 +1,337 @@ +package cache + +import ( + "context" + "encoding/json" + "sort" + "strings" + "sync" + "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const ( + // XAIReasoningReplayCacheTTL limits how long encrypted reasoning replay + // items stay in process memory. + XAIReasoningReplayCacheTTL = 1 * time.Hour + + // XAIReasoningReplayCacheMaxEntries bounds process memory for replay + // continuity. Oldest entries are evicted first. + XAIReasoningReplayCacheMaxEntries = 10240 + + // XAIReasoningReplayCacheEvictBatchSize leaves headroom after the cache + // reaches capacity so high write volume does not rescan the map every turn. + XAIReasoningReplayCacheEvictBatchSize = 128 +) + +type xaiReasoningReplayEntry struct { + Items [][]byte + Timestamp time.Time +} + +var ( + xaiReasoningReplayMu sync.Mutex + xaiReasoningReplayEntries = make(map[string]xaiReasoningReplayEntry) +) + +type xaiReasoningReplayKVClient 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) + KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) +} + +var currentXAIReasoningReplayKVClient = func() (xaiReasoningReplayKVClient, bool, error) { + return homekv.CurrentKVClient() +} + +// CacheXAIReasoningReplayItem stores a final Grok reasoning item for stateless +// replay. The stored item is normalized to the minimal shape accepted by +// Responses input replay. +func CacheXAIReasoningReplayItem(modelName, sessionKey string, item []byte) bool { + return CacheXAIReasoningReplayItems(modelName, sessionKey, [][]byte{item}) +} + +// CacheXAIReasoningReplayItems stores the final Grok assistant output items +// needed to replay a stateless next turn. +func CacheXAIReasoningReplayItems(modelName, sessionKey string, items [][]byte) bool { + return CacheXAIReasoningReplayItemsBestEffort(context.Background(), modelName, sessionKey, items) +} + +// CacheXAIReasoningReplayItemsBestEffort stores replay items for completed response paths. +func CacheXAIReasoningReplayItemsBestEffort(ctx context.Context, modelName, sessionKey string, items [][]byte) bool { + key := xaiReasoningReplayCacheKey(modelName, sessionKey) + if key == "" { + return false + } + normalized, ok := normalizeXAIReasoningReplayItems(items) + if !ok { + return false + } + if client, homeMode, errClient := currentXAIReasoningReplayKVClient(); homeMode { + if errClient != nil { + log.Errorf("home kv best-effort xai reasoning replay set failed prefix=cpa:xai:*: %v", errClient) + return false + } + raw, errMarshal := json.Marshal(normalized) + if errMarshal != nil { + log.Errorf("home kv best-effort xai reasoning replay set failed prefix=cpa:xai:*: %v", errMarshal) + return false + } + written, errSet := client.KVSet(ctx, xaiReasoningReplayKVKey(modelName, sessionKey), raw, homekv.KVSetOptions{EX: XAIReasoningReplayCacheTTL}) + if errSet != nil { + log.Errorf("home kv best-effort xai reasoning replay set failed prefix=cpa:xai:*: %v", errSet) + return false + } + return written + } + + cacheCleanupOnce.Do(startCacheCleanup) + now := time.Now() + xaiReasoningReplayMu.Lock() + defer xaiReasoningReplayMu.Unlock() + xaiReasoningReplayEntries[key] = xaiReasoningReplayEntry{ + Items: normalized, + Timestamp: now, + } + if len(xaiReasoningReplayEntries) > XAIReasoningReplayCacheMaxEntries { + evictOldestXAIReasoningReplayEntriesLocked(XAIReasoningReplayCacheEvictBatchSize) + } + return true +} + +// GetXAIReasoningReplayItem retrieves a normalized reasoning replay item. +func GetXAIReasoningReplayItem(modelName, sessionKey string) ([]byte, bool) { + items, ok := GetXAIReasoningReplayItems(modelName, sessionKey) + if !ok || len(items) == 0 { + return nil, false + } + return items[0], true +} + +// GetXAIReasoningReplayItems retrieves normalized assistant output items. +func GetXAIReasoningReplayItems(modelName, sessionKey string) ([][]byte, bool) { + items, ok, err := GetXAIReasoningReplayItemsRequired(context.Background(), modelName, sessionKey) + if err == nil { + return items, ok + } + return nil, false +} + +// GetXAIReasoningReplayItemsRequired retrieves replay items for request-time paths. +func GetXAIReasoningReplayItemsRequired(ctx context.Context, modelName, sessionKey string) ([][]byte, bool, error) { + key := xaiReasoningReplayCacheKey(modelName, sessionKey) + if key == "" { + return nil, false, nil + } + client, homeMode, errClient := currentXAIReasoningReplayKVClient() + if homeMode { + if errClient != nil { + return nil, false, errClient + } + raw, found, errGet := client.KVGet(ctx, xaiReasoningReplayKVKey(modelName, sessionKey)) + if errGet != nil || !found { + return nil, false, errGet + } + var homeItems [][]byte + if errUnmarshal := json.Unmarshal(raw, &homeItems); errUnmarshal != nil { + return nil, false, errUnmarshal + } + if _, errExpire := client.KVExpire(ctx, xaiReasoningReplayKVKey(modelName, sessionKey), XAIReasoningReplayCacheTTL); errExpire != nil { + log.Warnf("home kv xai reasoning replay expire failed prefix=cpa:xai:*: %v", errExpire) + } + return cloneXAIReasoningReplayItems(homeItems), true, nil + } + + cacheCleanupOnce.Do(startCacheCleanup) + now := time.Now() + xaiReasoningReplayMu.Lock() + defer xaiReasoningReplayMu.Unlock() + entry, ok := xaiReasoningReplayEntries[key] + if !ok { + return nil, false, nil + } + if now.Sub(entry.Timestamp) > XAIReasoningReplayCacheTTL { + delete(xaiReasoningReplayEntries, key) + return nil, false, nil + } + entry.Timestamp = now + xaiReasoningReplayEntries[key] = entry + return cloneXAIReasoningReplayItems(entry.Items), true, nil +} + +// DeleteXAIReasoningReplayItem removes one replay item after upstream rejects +// it or the caller otherwise knows it is stale. +func DeleteXAIReasoningReplayItem(modelName, sessionKey string) { + if errDelete := DeleteXAIReasoningReplayItemRequired(context.Background(), modelName, sessionKey); errDelete != nil { + return + } +} + +// DeleteXAIReasoningReplayItemRequired removes one replay item for request-time paths. +func DeleteXAIReasoningReplayItemRequired(ctx context.Context, modelName, sessionKey string) error { + key := xaiReasoningReplayCacheKey(modelName, sessionKey) + if key == "" { + return nil + } + client, homeMode, errClient := currentXAIReasoningReplayKVClient() + if homeMode { + if errClient != nil { + return errClient + } + _, errDel := client.KVDel(ctx, xaiReasoningReplayKVKey(modelName, sessionKey)) + return errDel + } + xaiReasoningReplayMu.Lock() + delete(xaiReasoningReplayEntries, key) + xaiReasoningReplayMu.Unlock() + return nil +} + +// ClearXAIReasoningReplayCache clears all xAI reasoning replay state. +func ClearXAIReasoningReplayCache() { + xaiReasoningReplayMu.Lock() + xaiReasoningReplayEntries = make(map[string]xaiReasoningReplayEntry) + xaiReasoningReplayMu.Unlock() +} + +func xaiReasoningReplayCacheKey(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 xAI credential so auth failover can preserve replay. + return strings.Join([]string{"xai-reasoning-replay", modelName, sessionKey}, "\x00") +} + +func xaiReasoningReplayKVKey(modelName, sessionKey string) string { + return "cpa:xai:reasoning-replay:" + homekv.HashKeyPart(strings.TrimSpace(modelName)) + ":" + homekv.HashKeyPart(strings.TrimSpace(sessionKey)) +} + +func normalizeXAIReasoningReplayItems(items [][]byte) ([][]byte, bool) { + normalized := make([][]byte, 0, len(items)) + for _, item := range items { + normalizedItem, ok := normalizeXAIReasoningReplayItem(item) + if ok { + normalized = append(normalized, normalizedItem) + } + } + return normalized, len(normalized) > 0 +} + +func normalizeXAIReasoningReplayItem(item []byte) ([]byte, bool) { + itemResult := gjson.ParseBytes(item) + switch strings.TrimSpace(itemResult.Get("type").String()) { + case "reasoning": + return normalizeXAIReasoningReplayReasoningItem(itemResult) + case "function_call": + return normalizeXAIReasoningReplayFunctionCallItem(itemResult) + case "custom_tool_call": + return normalizeXAIReasoningReplayCustomToolCallItem(itemResult) + default: + return nil, false + } +} + +func normalizeXAIReasoningReplayReasoningItem(itemResult gjson.Result) ([]byte, bool) { + encryptedContentResult := itemResult.Get("encrypted_content") + if encryptedContentResult.Type != gjson.String { + return nil, false + } + encryptedContent := encryptedContentResult.String() + if encryptedContent != strings.TrimSpace(encryptedContent) { + return nil, false + } + if _, err := signature.InspectGrokEncryptedContent(encryptedContent); err != nil { + return nil, false + } + + normalized := []byte(`{"type":"reasoning","summary":[],"content":null}`) + normalized, _ = sjson.SetBytes(normalized, "encrypted_content", encryptedContent) + return normalized, true +} + +func normalizeXAIReasoningReplayFunctionCallItem(itemResult gjson.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 != gjson.String { + return nil, false + } + + normalized := []byte(`{"type":"function_call"}`) + normalized, _ = sjson.SetBytes(normalized, "call_id", callID) + normalized, _ = sjson.SetBytes(normalized, "name", name) + normalized, _ = sjson.SetBytes(normalized, "arguments", arguments.String()) + return normalized, true +} + +func normalizeXAIReasoningReplayCustomToolCallItem(itemResult gjson.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 + } + + normalized := []byte(`{"type":"custom_tool_call","status":"completed"}`) + if status := strings.TrimSpace(itemResult.Get("status").String()); status != "" { + normalized, _ = sjson.SetBytes(normalized, "status", status) + } + normalized, _ = sjson.SetBytes(normalized, "call_id", callID) + normalized, _ = sjson.SetBytes(normalized, "name", name) + if input.Type == gjson.String { + normalized, _ = sjson.SetBytes(normalized, "input", input.String()) + } else { + normalized, _ = sjson.SetRawBytes(normalized, "input", []byte(input.Raw)) + } + return normalized, true +} + +func cloneXAIReasoningReplayItems(items [][]byte) [][]byte { + cloned := make([][]byte, 0, len(items)) + for _, item := range items { + cloned = append(cloned, append([]byte(nil), item...)) + } + return cloned +} + +func evictOldestXAIReasoningReplayEntriesLocked(count int) { + if count <= 0 || len(xaiReasoningReplayEntries) == 0 { + return + } + type candidate struct { + key string + timestamp time.Time + } + candidates := make([]candidate, 0, len(xaiReasoningReplayEntries)) + for key, entry := range xaiReasoningReplayEntries { + 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(xaiReasoningReplayEntries, candidates[i].key) + } +} + +func purgeExpiredXAIReasoningReplayCache(now time.Time) { + xaiReasoningReplayMu.Lock() + for key, entry := range xaiReasoningReplayEntries { + if now.Sub(entry.Timestamp) > XAIReasoningReplayCacheTTL { + delete(xaiReasoningReplayEntries, key) + } + } + xaiReasoningReplayMu.Unlock() +} diff --git a/internal/cache/xai_reasoning_replay_cache_test.go b/internal/cache/xai_reasoning_replay_cache_test.go new file mode 100644 index 00000000..0de4ac39 --- /dev/null +++ b/internal/cache/xai_reasoning_replay_cache_test.go @@ -0,0 +1,161 @@ +package cache + +import ( + "context" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "testing" + "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/tidwall/gjson" +) + +type fakeXAIReasoningReplayKVClient struct { + values map[string][]byte + getErr error + setErr error + delErr error + expireErr error + getCount int + setCount int + delCount int + expireCount int + lastSetTTL time.Duration + lastExpireTTL time.Duration +} + +func newFakeXAIReasoningReplayKVClient() *fakeXAIReasoningReplayKVClient { + return &fakeXAIReasoningReplayKVClient{values: make(map[string][]byte)} +} + +func (c *fakeXAIReasoningReplayKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) { + c.getCount++ + if c.getErr != nil { + return nil, false, c.getErr + } + value, ok := c.values[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), value...), true, nil +} + +func (c *fakeXAIReasoningReplayKVClient) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) { + c.setCount++ + c.lastSetTTL = opts.EX + if c.setErr != nil { + return false, c.setErr + } + c.values[key] = append([]byte(nil), value...) + return true, nil +} + +func (c *fakeXAIReasoningReplayKVClient) KVDel(_ context.Context, keys ...string) (int64, error) { + c.delCount++ + if c.delErr != nil { + return 0, c.delErr + } + var deleted int64 + for _, key := range keys { + if _, ok := c.values[key]; ok { + delete(c.values, key) + deleted++ + } + } + return deleted, nil +} + +func (c *fakeXAIReasoningReplayKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) { + c.expireCount++ + c.lastExpireTTL = ttl + if c.expireErr != nil { + return false, c.expireErr + } + return true, nil +} + +func useFakeXAIReasoningReplayKVClient(t *testing.T, client *fakeXAIReasoningReplayKVClient, homeMode bool, errClient error) { + t.Helper() + previous := currentXAIReasoningReplayKVClient + currentXAIReasoningReplayKVClient = func() (xaiReasoningReplayKVClient, bool, error) { + return client, homeMode, errClient + } + t.Cleanup(func() { + currentXAIReasoningReplayKVClient = previous + }) +} + +func mustXAIReasoningReplayJSON(t *testing.T, items [][]byte) []byte { + t.Helper() + raw, err := json.Marshal(items) + if err != nil { + t.Fatalf("marshal replay items: %v", err) + } + return raw +} + +func TestXAIReasoningReplayCacheRejectsCodexEncryptedContent(t *testing.T) { + ClearXAIReasoningReplayCache() + t.Cleanup(ClearXAIReasoningReplayCache) + + if CacheXAIReasoningReplayItem("grok-4.3", "claude:xai-cache-test", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"gAAAAABinvalid-gpt-shape"}`)) { + t.Fatal("xAI replay cache should reject GPT/Codex-shaped encrypted_content") + } + if _, ok := GetXAIReasoningReplayItem("grok-4.3", "claude:xai-cache-test"); ok { + t.Fatal("xAI replay cache should not store GPT/Codex-shaped encrypted_content") + } +} + +func TestXAIReasoningReplayCacheStoresGrokEncryptedContent(t *testing.T) { + ClearXAIReasoningReplayCache() + t.Cleanup(ClearXAIReasoningReplayCache) + + encryptedContent := validGrokEncryptedContentForReplayCacheTest() + if !CacheXAIReasoningReplayItem("grok-4.3", "claude:xai-cache-test", []byte(`{"type":"reasoning","summary":[{"type":"summary_text","text":"visible"}],"content":null,"encrypted_content":"`+encryptedContent+`"}`)) { + t.Fatal("xAI replay cache should store valid Grok encrypted_content") + } + item, ok := GetXAIReasoningReplayItem("grok-4.3", "claude:xai-cache-test") + if !ok { + t.Fatal("xAI replay cache item missing after store") + } + if got := gjson.GetBytes(item, "encrypted_content").String(); got != encryptedContent { + t.Fatalf("encrypted_content = %q, want %q; item=%s", got, encryptedContent, string(item)) + } + if got := gjson.GetBytes(item, "summary").Array(); len(got) != 0 { + t.Fatalf("summary length = %d, want normalized empty summary; item=%s", len(got), string(item)) + } +} + +func TestXAIReasoningReplayRequiredHomeExpireFailureReturnsItems(t *testing.T) { + ClearXAIReasoningReplayCache() + t.Cleanup(ClearXAIReasoningReplayCache) + client := newFakeXAIReasoningReplayKVClient() + client.expireErr = errors.New("expire failed") + key := xaiReasoningReplayKVKey("grok-4.3", "session-home") + item := []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"` + validGrokEncryptedContentForReplayCacheTest() + `"}`) + client.values[key] = mustXAIReasoningReplayJSON(t, [][]byte{item}) + useFakeXAIReasoningReplayKVClient(t, client, true, nil) + + items, found, errGet := GetXAIReasoningReplayItemsRequired(context.Background(), "grok-4.3", "session-home") + if errGet != nil { + t.Fatalf("GetXAIReasoningReplayItemsRequired() error = %v", errGet) + } + if !found || len(items) != 1 || string(items[0]) != string(item) { + t.Fatalf("GetXAIReasoningReplayItemsRequired() = %q, %v, want item, true", items, found) + } + if client.expireCount != 1 || client.lastExpireTTL != XAIReasoningReplayCacheTTL { + t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, XAIReasoningReplayCacheTTL) + } +} + +func validGrokEncryptedContentForReplayCacheTest() string { + buf := make([]byte, 0, 256) + for i := 0; len(buf) < 256; i++ { + sum := sha256.Sum256([]byte{byte(i), byte(i >> 8), byte(i >> 16), 99}) + buf = append(buf, sum[:]...) + } + return base64.RawStdEncoding.EncodeToString(buf[:256]) +} diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 8b21f93d..cc730dd3 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -182,6 +182,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req } completedData := xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) completedData = xaiNormalizeReasoningSummaryData(completedData) + cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, completedData) var param any out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, completedData, ¶m) return cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}, nil @@ -665,6 +666,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth } eventData = xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) eventData = xaiNormalizeReasoningSummaryData(eventData) + cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, eventData) normalizedEventName = gjson.GetBytes(eventData, "type").String() } @@ -800,6 +802,7 @@ type xaiPreparedRequest struct { originalPayload []byte body []byte sessionID string + replayScope xaiReasoningReplayScope } func (e *XAIExecutor) prepareResponsesRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) (*xaiPreparedRequest, error) { @@ -835,6 +838,11 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox body, _ = sjson.DeleteBytes(body, "stream_options") body = normalizeXAITools(body) body = normalizeXAIToolChoiceForTools(body) + var replayScope xaiReasoningReplayScope + body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body) + if err != nil { + return nil, err + } body = normalizeXAIInputReasoningItems(body) body = sanitizeXAIInputEncryptedContent(body) body = normalizeCodexInstructions(body) @@ -856,6 +864,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox originalPayload: originalPayload, body: body, sessionID: sessionID, + replayScope: replayScope, }, nil } diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index b360dfbd..ed123f9d 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -5,6 +5,7 @@ import ( "context" "crypto/sha256" "encoding/base64" + "errors" "io" "net/http" "net/http/httptest" @@ -12,6 +13,7 @@ import ( "testing" "github.com/google/uuid" + internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" _ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -1139,6 +1141,205 @@ func TestXAIExecutorDropsInvalidCompactionItem(t *testing.T) { } } +func TestXAIExecutorReasoningReplayCacheStoresFinalDoneAndInjectsNextClaudeRequest(t *testing.T) { + internalcache.ClearXAIReasoningReplayCache() + t.Cleanup(internalcache.ClearXAIReasoningReplayCache) + + addedEncryptedContent := testValidGrokEncryptedContentForSeed(1) + doneEncryptedContent := testValidGrokEncryptedContentForSeed(2) + var bodies [][]byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Fatalf("read body: %v", errRead) + } + bodies = append(bodies, body) + + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(`data: {"type":"response.output_item.added","item":{"id":"rs_added","type":"reasoning","status":"in_progress","summary":[],"encrypted_content":"` + addedEncryptedContent + `"},"output_index":0}` + "\n")) + _, _ = w.Write([]byte(`data: {"type":"response.output_item.done","item":{"id":"rs_done","type":"reasoning","summary":[],"encrypted_content":"` + doneEncryptedContent + `"},"output_index":0}` + "\n")) + _, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","object":"response","created_at":0,"status":"completed","model":"grok-4.3","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` + "\n\n")) + })) + defer server.Close() + + executor := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + ID: "xai-auth-replay-1", + Provider: "xai", + Attributes: map[string]string{ + "base_url": server.URL, + "auth_kind": "oauth", + }, + Metadata: map[string]any{ + "access_token": "xai-token", + }, + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Stream: false, + } + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{"model":"grok-4.3","metadata":{"user_id":"{\"device_id\":\"device-test\",\"account_uuid\":\"\",\"session_id\":\"xai-session-1\"}"},"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`), + }, opts) + if err != nil { + t.Fatalf("first Execute error: %v", err) + } + + _, err = executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{"model":"grok-4.3","metadata":{"user_id":"{\"device_id\":\"device-test\",\"account_uuid\":\"\",\"session_id\":\"xai-session-1\"}"},"messages":[{"role":"user","content":[{"type":"text","text":"next"}]}]}`), + }, opts) + if err != nil { + t.Fatalf("second Execute error: %v", err) + } + + if len(bodies) != 2 { + t.Fatalf("upstream request count = %d, want 2", len(bodies)) + } + secondBody := bodies[1] + if got := gjson.GetBytes(secondBody, "input.0.type").String(); got != "reasoning" { + t.Fatalf("input.0.type = %q, want reasoning; body=%s", got, string(secondBody)) + } + if got := gjson.GetBytes(secondBody, "input.0.encrypted_content").String(); got != doneEncryptedContent { + t.Fatalf("injected encrypted_content = %q, want final done %q; body=%s", got, doneEncryptedContent, string(secondBody)) + } + if got := gjson.GetBytes(secondBody, "input.1.role").String(); got != "user" { + t.Fatalf("input.1.role = %q, want user; body=%s", got, string(secondBody)) + } +} + +func TestApplyXAIReasoningReplayCacheFallsBackWhenReadFails(t *testing.T) { + previous := getXAIReasoningReplayItemsRequired + getXAIReasoningReplayItemsRequired = func(context.Context, string, string) ([][]byte, bool, error) { + return nil, false, errors.New("cache unavailable") + } + t.Cleanup(func() { + getXAIReasoningReplayItemsRequired = previous + }) + + body := []byte(`{"model":"grok-4.3","input":[{"role":"user","content":[{"type":"input_text","text":"hello"}]}]}`) + updated, scope, err := applyXAIReasoningReplayCacheRequired(context.Background(), sdktranslator.FormatClaude, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: body, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "xai-read-error", + }, + }, body) + if err != nil { + t.Fatalf("applyXAIReasoningReplayCacheRequired() error = %v", err) + } + if !scope.valid() { + t.Fatalf("replay scope should remain valid") + } + if string(updated) != string(body) { + t.Fatalf("body changed on cache read error: %s", string(updated)) + } +} + +func TestXAIExecutorReasoningReplayCacheReplaysFunctionCallForClaudeToolResult(t *testing.T) { + internalcache.ClearXAIReasoningReplayCache() + t.Cleanup(internalcache.ClearXAIReasoningReplayCache) + + reasoningEncryptedContent := testValidGrokEncryptedContentForSeed(3) + var bodies [][]byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Fatalf("read body: %v", errRead) + } + bodies = append(bodies, body) + + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(`data: {"type":"response.output_item.done","item":{"id":"rs_1","type":"reasoning","summary":[],"encrypted_content":"` + reasoningEncryptedContent + `"},"output_index":0}` + "\n")) + _, _ = w.Write([]byte(`data: {"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"q\":\"weather\"}","status":"in_progress"},"output_index":1}` + "\n")) + _, _ = w.Write([]byte(`data: {"type":"response.output_item.done","item":{"id":"fc_1","type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"q\":\"weather\"}","status":"completed"},"output_index":1}` + "\n")) + _, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","object":"response","created_at":0,"status":"completed","model":"grok-4.3","output":[]}}` + "\n\n")) + })) + defer server.Close() + + executor := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + ID: "xai-auth-replay-tool", + Provider: "xai", + Attributes: map[string]string{ + "base_url": server.URL, + "auth_kind": "oauth", + }, + Metadata: map[string]any{ + "access_token": "xai-token", + }, + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Stream: false, + } + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{ + "model":"grok-4.3", + "metadata":{"user_id":"{\"device_id\":\"device-test\",\"account_uuid\":\"\",\"session_id\":\"xai-session-tool\"}"}, + "messages":[{"role":"user","content":[{"type":"text","text":"call lookup"}]}], + "tools":[{"name":"lookup","input_schema":{"type":"object","properties":{"q":{"type":"string"}}}}] + }`), + }, opts) + if err != nil { + t.Fatalf("first Execute error: %v", err) + } + + _, err = executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{ + "model":"grok-4.3", + "metadata":{"user_id":"{\"device_id\":\"device-test\",\"account_uuid\":\"\",\"session_id\":\"xai-session-tool\"}"}, + "messages":[ + {"role":"user","content":[{"type":"text","text":"call lookup"}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"call_1","content":"sunny"}]} + ], + "tools":[{"name":"lookup","input_schema":{"type":"object","properties":{"q":{"type":"string"}}}}] + }`), + }, opts) + if err != nil { + t.Fatalf("second Execute error: %v", err) + } + + if len(bodies) != 2 { + t.Fatalf("upstream request count = %d, want 2", len(bodies)) + } + secondBody := bodies[1] + if got := gjson.GetBytes(secondBody, "input.0.type").String(); got != "message" { + t.Fatalf("input.0.type = %q, want initial user message; body=%s", got, string(secondBody)) + } + if got := gjson.GetBytes(secondBody, "input.1.type").String(); got != "reasoning" { + t.Fatalf("input.1.type = %q, want cached reasoning; body=%s", got, string(secondBody)) + } + if got := gjson.GetBytes(secondBody, "input.2.type").String(); got != "function_call" { + t.Fatalf("input.2.type = %q, want cached function_call; body=%s", got, string(secondBody)) + } + if got := gjson.GetBytes(secondBody, "input.2.call_id").String(); got != "call_1" { + t.Fatalf("input.2.call_id = %q, want call_1; body=%s", got, string(secondBody)) + } + if got := gjson.GetBytes(secondBody, "input.3.type").String(); got != "function_call_output" { + t.Fatalf("input.3.type = %q, want function_call_output after cached call; body=%s", got, string(secondBody)) + } + if got := gjson.GetBytes(secondBody, "input.3.call_id").String(); got != "call_1" { + t.Fatalf("input.3.call_id = %q, want call_1; body=%s", got, string(secondBody)) + } +} + +func testValidGrokEncryptedContentForSeed(seed byte) string { + buf := make([]byte, 0, 256) + for i := 0; len(buf) < 256; i++ { + sum := sha256.Sum256([]byte{seed, byte(i), byte(i >> 8), byte(i >> 16)}) + buf = append(buf, sum[:]...) + } + return base64.RawStdEncoding.EncodeToString(buf[:256]) +} + func testValidGrokEncryptedContent() string { buf := make([]byte, 0, 256) for i := 0; len(buf) < 256; i++ { diff --git a/internal/runtime/executor/xai_reasoning_replay.go b/internal/runtime/executor/xai_reasoning_replay.go new file mode 100644 index 00000000..a9f2275f --- /dev/null +++ b/internal/runtime/executor/xai_reasoning_replay.go @@ -0,0 +1,171 @@ +package executor + +import ( + "context" + "strings" + + internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" + "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" + "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" +) + +type xaiReasoningReplayScope struct { + modelName string + sessionKey string +} + +var getXAIReasoningReplayItemsRequired = internalcache.GetXAIReasoningReplayItemsRequired + +func (s xaiReasoningReplayScope) valid() bool { + return strings.TrimSpace(s.modelName) != "" && strings.TrimSpace(s.sessionKey) != "" +} + +func applyXAIReasoningReplayCacheRequired(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) ([]byte, xaiReasoningReplayScope, error) { + scope := xaiReasoningReplayScopeFromRequest(ctx, from, req, opts, body) + if !scope.valid() { + return body, scope, nil + } + items, ok, errReplay := getXAIReasoningReplayItemsRequired(ctx, scope.modelName, scope.sessionKey) + if errReplay != nil { + log.Warnf("xai reasoning replay cache read failed: %v", errReplay) + return body, scope, nil + } + if !ok { + return body, scope, nil + } + items = filterXAIReasoningReplayItemsForInput(body, items) + if len(items) == 0 { + return body, scope, nil + } + updated, ok := insertCodexReasoningReplayItems(body, items) + if !ok { + return body, scope, nil + } + return updated, scope, nil +} + +func xaiReasoningReplayScopeFromRequest(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) xaiReasoningReplayScope { + if !xaiReasoningReplayEnabledForSource(from) { + return xaiReasoningReplayScope{} + } + return xaiReasoningReplayScope{ + modelName: thinking.ParseSuffix(req.Model).ModelName, + sessionKey: codexReasoningReplaySessionKey(ctx, from, req, opts, body), + } +} + +func xaiReasoningReplayEnabledForSource(from sdktranslator.Format) bool { + return sourceFormatEqual(from, sdktranslator.FormatClaude) +} + +func xaiInputHasValidReasoningEncryptedContent(body []byte) bool { + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return false + } + for _, item := range input.Array() { + if strings.TrimSpace(item.Get("type").String()) != "reasoning" { + continue + } + encryptedContent := item.Get("encrypted_content") + if encryptedContent.Type != gjson.String { + continue + } + if _, err := signature.InspectGrokEncryptedContent(encryptedContent.String()); err == nil { + return true + } + } + return false +} + +func filterXAIReasoningReplayItemsForInput(body []byte, items [][]byte) [][]byte { + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return nil + } + + hasInputReasoning := xaiInputHasValidReasoningEncryptedContent(body) + existingCalls := make(map[string]bool) + existingOutputs := make(map[string]bool) + for _, inputItem := range input.Array() { + itemType := strings.TrimSpace(inputItem.Get("type").String()) + if itemType == "function_call_output" || itemType == "custom_tool_call_output" { + callID := strings.TrimSpace(inputItem.Get("call_id").String()) + if callID != "" { + for _, candidate := range codexReplayComparableCallIDs(callID) { + existingOutputs[candidate] = true + } + } + } + for _, key := range codexReplayToolCallKeys(inputItem) { + existingCalls[key] = true + } + } + + filtered := make([][]byte, 0, len(items)) + for _, item := range items { + itemResult := gjson.ParseBytes(item) + switch strings.TrimSpace(itemResult.Get("type").String()) { + case "reasoning": + if hasInputReasoning { + continue + } + case "function_call", "custom_tool_call": + keys := codexReplayToolCallKeys(itemResult) + if len(keys) == 0 || codexReplayAnyToolCallKeyExists(existingCalls, keys) { + continue + } + hasMatchingOutput := false + callID := strings.TrimSpace(itemResult.Get("call_id").String()) + if callID != "" { + for _, candidate := range codexReplayComparableCallIDs(callID) { + if existingOutputs[candidate] { + hasMatchingOutput = true + break + } + } + } + if !hasMatchingOutput { + continue + } + for _, key := range keys { + existingCalls[key] = true + } + default: + continue + } + filtered = append(filtered, item) + } + return filtered +} + +func cacheXAIReasoningReplayFromCompleted(ctx context.Context, scope xaiReasoningReplayScope, completedData []byte) { + if !scope.valid() { + return + } + if ctx == nil { + ctx = context.Background() + } + output := gjson.GetBytes(completedData, "response.output") + if !output.IsArray() { + return + } + items := make([][]byte, 0, len(output.Array())) + for _, item := range output.Array() { + switch strings.TrimSpace(item.Get("type").String()) { + case "reasoning", "function_call", "custom_tool_call": + items = append(items, []byte(item.Raw)) + default: + continue + } + } + if !internalcache.CacheXAIReasoningReplayItemsBestEffort(ctx, scope.modelName, scope.sessionKey, items) { + if errDelete := internalcache.DeleteXAIReasoningReplayItemRequired(ctx, scope.modelName, scope.sessionKey); errDelete != nil { + log.Warnf("xai reasoning replay cache delete failed after completed cache store failed: %v", errDelete) + } + } +} diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index fb8cceb8..c9333a74 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -647,6 +647,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } payload = xaiPatchCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) payload = xaiNormalizeReasoningSummaryData(payload) + cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload) if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { idMapper.state.recordTranscriptTurn(wsReqBody, payload) recordedTranscript = true