feat(xai): encrypted reasoning replay for Responses/Claude with hardened cache

Preserve include=reasoning.encrypted_content, cache reasoning+assistant message
for store:false multi-turn, and accept Grok thinking signatures on Claude→Codex
for grok targets.

Harden replay cache: blob-level reasoning filter, never double-inject assistant
when history has a last message, isolate session keys by downstream API key,
clear cache on non-replayable completed turns, and normalize refusal parts.
This commit is contained in:
sususu
2026-07-12 22:31:19 +08:00
parent 3586d3e750
commit 041816c2fd
7 changed files with 693 additions and 45 deletions
+84 -8
View File
@@ -63,32 +63,58 @@ func CacheXAIReasoningReplayItems(modelName, sessionKey string, items [][]byte)
return CacheXAIReasoningReplayItemsBestEffort(context.Background(), modelName, sessionKey, items)
}
// XAIReasoningReplayStoreStatus reports why a completed-turn cache write
// succeeded or failed so callers can decide whether to keep prior entries.
type XAIReasoningReplayStoreStatus int
const (
// XAIReasoningReplayStoreInvalidArgs means model/session were empty.
XAIReasoningReplayStoreInvalidArgs XAIReasoningReplayStoreStatus = iota
// XAIReasoningReplayStored means a valid reasoning batch was written.
XAIReasoningReplayStored
// XAIReasoningReplayNoReplayableState means the completed output had no
// cacheable reasoning batch (for example reasoning disabled).
XAIReasoningReplayNoReplayableState
// XAIReasoningReplayStoreBackendError means normalize succeeded but the
// storage backend failed; previous entries should be retained.
XAIReasoningReplayStoreBackendError
)
// CacheXAIReasoningReplayItemsBestEffort stores replay items for completed response paths.
func CacheXAIReasoningReplayItemsBestEffort(ctx context.Context, modelName, sessionKey string, items [][]byte) bool {
return StoreXAIReasoningReplayItems(ctx, modelName, sessionKey, items) == XAIReasoningReplayStored
}
// StoreXAIReasoningReplayItems stores replay items and distinguishes empty
// completed state from backend failures.
func StoreXAIReasoningReplayItems(ctx context.Context, modelName, sessionKey string, items [][]byte) XAIReasoningReplayStoreStatus {
key := xaiReasoningReplayCacheKey(modelName, sessionKey)
if key == "" {
return false
return XAIReasoningReplayStoreInvalidArgs
}
normalized, ok := normalizeXAIReasoningReplayItems(items)
if !ok {
return false
return XAIReasoningReplayNoReplayableState
}
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
return XAIReasoningReplayStoreBackendError
}
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
return XAIReasoningReplayStoreBackendError
}
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 XAIReasoningReplayStoreBackendError
}
return written
if !written {
return XAIReasoningReplayStoreBackendError
}
return XAIReasoningReplayStored
}
cacheCleanupOnce.Do(startCacheCleanup)
@@ -102,7 +128,7 @@ func CacheXAIReasoningReplayItemsBestEffort(ctx context.Context, modelName, sess
if len(xaiReasoningReplayEntries) > XAIReasoningReplayCacheMaxEntries {
evictOldestXAIReasoningReplayEntriesLocked(XAIReasoningReplayCacheEvictBatchSize)
}
return true
return XAIReasoningReplayStored
}
// GetXAIReasoningReplayItem retrieves a normalized reasoning replay item.
@@ -217,13 +243,17 @@ func xaiReasoningReplayKVKey(modelName, sessionKey string) string {
func normalizeXAIReasoningReplayItems(items [][]byte) ([][]byte, bool) {
normalized := make([][]byte, 0, len(items))
hasReasoning := false
for _, item := range items {
normalizedItem, ok := normalizeXAIReasoningReplayItem(item)
if ok {
normalized = append(normalized, normalizedItem)
if strings.TrimSpace(gjson.GetBytes(normalizedItem, "type").String()) == "reasoning" {
hasReasoning = true
}
}
}
return normalized, len(normalized) > 0
return normalized, hasReasoning
}
func normalizeXAIReasoningReplayItem(item []byte) ([]byte, bool) {
@@ -231,6 +261,8 @@ func normalizeXAIReasoningReplayItem(item []byte) ([]byte, bool) {
switch strings.TrimSpace(itemResult.Get("type").String()) {
case "reasoning":
return normalizeXAIReasoningReplayReasoningItem(itemResult)
case "message":
return normalizeXAIReasoningReplayMessageItem(itemResult)
case "function_call":
return normalizeXAIReasoningReplayFunctionCallItem(itemResult)
case "custom_tool_call":
@@ -258,6 +290,50 @@ func normalizeXAIReasoningReplayReasoningItem(itemResult gjson.Result) ([]byte,
return normalized, true
}
func normalizeXAIReasoningReplayMessageItem(itemResult gjson.Result) ([]byte, bool) {
if !strings.EqualFold(strings.TrimSpace(itemResult.Get("role").String()), "assistant") {
return nil, false
}
content := itemResult.Get("content")
if !content.IsArray() || len(content.Array()) == 0 {
return nil, false
}
normalized := []byte(`{"type":"message","role":"assistant","content":[]}`)
for _, part := range content.Array() {
partType := strings.TrimSpace(part.Get("type").String())
var nextPart []byte
switch partType {
case "output_text":
textValue := part.Get("text")
if textValue.Type != gjson.String {
continue
}
nextPart = []byte(`{"type":"output_text","text":""}`)
nextPart, _ = sjson.SetBytes(nextPart, "text", textValue.String())
case "refusal":
// Responses API refusal parts use the "refusal" field, not "text".
refusalValue := part.Get("refusal")
if refusalValue.Type != gjson.String {
continue
}
nextPart = []byte(`{"type":"refusal","refusal":""}`)
nextPart, _ = sjson.SetBytes(nextPart, "refusal", refusalValue.String())
default:
continue
}
updated, errSet := sjson.SetRawBytes(normalized, "content.-1", nextPart)
if errSet != nil {
return nil, false
}
normalized = updated
}
if len(gjson.GetBytes(normalized, "content").Array()) == 0 {
return nil, false
}
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())
+70
View File
@@ -129,6 +129,49 @@ func TestXAIReasoningReplayCacheStoresGrokEncryptedContent(t *testing.T) {
}
}
func TestXAIReasoningReplayCacheStoresAssistantMessageWithReasoning(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
encryptedContent := validGrokEncryptedContentForReplayCacheTest()
items := [][]byte{
[]byte(`{"id":"rs_1","type":"reasoning","summary":[{"type":"summary_text","text":"visible"}],"encrypted_content":"` + encryptedContent + `"}`),
[]byte(`{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"answer","annotations":[],"logprobs":[]}]}`),
}
if !CacheXAIReasoningReplayItems("grok-4.5", "prompt-cache:session", items) {
t.Fatal("expected reasoning replay items to be cached")
}
got, ok := GetXAIReasoningReplayItems("grok-4.5", "prompt-cache:session")
if !ok || len(got) != 2 {
t.Fatalf("cached items = %q, %v, want two items", got, ok)
}
if gjson.GetBytes(got[0], "encrypted_content").String() != encryptedContent {
t.Fatalf("reasoning encrypted_content not preserved: %s", got[0])
}
if gotText := gjson.GetBytes(got[1], "content.0.text").String(); gotText != "answer" {
t.Fatalf("assistant message text = %q, want answer; item=%s", gotText, got[1])
}
if gjson.GetBytes(got[1], "id").Exists() || gjson.GetBytes(got[1], "status").Exists() {
t.Fatalf("assistant message transport fields were not stripped: %s", got[1])
}
}
func TestXAIReasoningReplayCacheRejectsAssistantMessageWithoutReasoning(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
items := [][]byte{
[]byte(`{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"answer"}]}`),
}
if CacheXAIReasoningReplayItems("grok-4.5", "prompt-cache:message-only", items) {
t.Fatal("message-only replay batch must not be cached")
}
if _, ok := GetXAIReasoningReplayItems("grok-4.5", "prompt-cache:message-only"); ok {
t.Fatal("message-only replay batch unexpectedly exists in cache")
}
}
func TestXAIReasoningReplayRequiredHomeExpireFailureReturnsItems(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
@@ -159,3 +202,30 @@ func validGrokEncryptedContentForReplayCacheTest() string {
}
return base64.RawStdEncoding.EncodeToString(buf[:256])
}
func TestXAIReasoningReplayCacheStoresRefusalMessagePart(t *testing.T) {
ClearXAIReasoningReplayCache()
t.Cleanup(ClearXAIReasoningReplayCache)
encryptedContent := validGrokEncryptedContentForReplayCacheTest()
items := [][]byte{
[]byte(`{"type":"reasoning","summary":[],"encrypted_content":"` + encryptedContent + `"}`),
[]byte(`{"type":"message","role":"assistant","content":[{"type":"refusal","refusal":"I cannot help with that"}]}`),
}
if !CacheXAIReasoningReplayItems("grok-4.5", "prompt-cache:refusal", items) {
t.Fatal("expected refusal message with reasoning to be cached")
}
got, ok := GetXAIReasoningReplayItems("grok-4.5", "prompt-cache:refusal")
if !ok || len(got) != 2 {
t.Fatalf("cached items = %q, %v, want reasoning + refusal message", got, ok)
}
if gjson.GetBytes(got[1], "content.0.type").String() != "refusal" {
t.Fatalf("message part type = %s, want refusal; item=%s", gjson.GetBytes(got[1], "content.0.type").String(), got[1])
}
if gjson.GetBytes(got[1], "content.0.refusal").String() != "I cannot help with that" {
t.Fatalf("refusal text missing; item=%s", got[1])
}
if gjson.GetBytes(got[1], "content.0.text").Exists() {
t.Fatalf("refusal part should not use text field; item=%s", got[1])
}
}
-18
View File
@@ -1144,7 +1144,6 @@ func xaiMetadataString(meta map[string]any, key string) string {
}
func sanitizeXAIResponsesBody(body []byte, model string) []byte {
body = removeXAIEncryptedReasoningInclude(body)
if !xaiSupportsReasoningEffort(model) {
if gjson.GetBytes(body, "reasoning.effort").Exists() {
log.Debugf("xai: stripping reasoning.effort for model %s (no thinking levels in model registry)", model)
@@ -1483,23 +1482,6 @@ func appendXAIReasoningSummary(previous json.RawMessage, currentSummary []gjson.
return updated, true
}
func removeXAIEncryptedReasoningInclude(body []byte) []byte {
include := gjson.GetBytes(body, "include")
if !include.Exists() || !include.IsArray() {
return body
}
kept := make([]string, 0, len(include.Array()))
for _, item := range include.Array() {
value := strings.TrimSpace(item.String())
if value == "" || value == "reasoning.encrypted_content" {
continue
}
kept = append(kept, value)
}
body, _ = sjson.SetBytes(body, "include", kept)
return body
}
// xaiSupportsReasoningEffort reports whether the model accepts Responses API
// reasoning.effort. Capability comes from model registry thinking metadata
// (static models.json and dynamic registrations), not a hard-coded name allowlist.
+385 -1
View File
@@ -12,6 +12,7 @@ import (
"strings"
"testing"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai"
internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache"
@@ -24,6 +25,14 @@ import (
"github.com/tidwall/sjson"
)
func testContextWithAPIKey(apiKey string) context.Context {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(rec)
ginCtx.Set("userApiKey", apiKey)
return context.WithValue(context.Background(), "gin", ginCtx)
}
func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) {
var gotPath string
var gotAuth string
@@ -159,11 +168,16 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) {
if !foundNamespaceCustom {
t.Fatalf("namespace custom tool was not moved to top-level tools; body=%s", string(gotBody))
}
foundEncryptedReasoningInclude := false
for _, include := range gjson.GetBytes(gotBody, "include").Array() {
if include.String() == "reasoning.encrypted_content" {
t.Fatalf("xai request must not ask for encrypted reasoning content: %s", string(gotBody))
foundEncryptedReasoningInclude = true
break
}
}
if !foundEncryptedReasoningInclude {
t.Fatalf("xai request must preserve reasoning.encrypted_content include: %s", string(gotBody))
}
}
func TestXAIExecutorComposerSessionIsolation(t *testing.T) {
@@ -1452,6 +1466,376 @@ func TestXAIExecutorReasoningReplayCacheStoresFinalDoneAndInjectsNextClaudeReque
}
}
func TestXAIExecutorResponsesSSEReplaysEncryptedReasoningAndAssistantMessage(t *testing.T) {
internalcache.ClearXAIReasoningReplayCache()
t.Cleanup(internalcache.ClearXAIReasoningReplayCache)
encryptedContent := testValidGrokEncryptedContentForSeed(9)
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")
if len(bodies) == 1 {
_, _ = w.Write([]byte(`data: {"type":"response.output_item.done","item":{"id":"rs_1","type":"reasoning","summary":[],"encrypted_content":"` + encryptedContent + `"},"output_index":0}` + "\n"))
_, _ = w.Write([]byte(`data: {"type":"response.output_item.done","item":{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"first answer"}]},"output_index":1}` + "\n"))
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed","model":"grok-4.5","output":[]}}` + "\n\n"))
return
}
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_2","status":"completed","model":"grok-4.5","output":[]}}` + "\n\n"))
}))
defer server.Close()
executor := NewXAIExecutor(&config.Config{})
auth := &cliproxyauth.Auth{
ID: "xai-auth-responses-sse-replay",
Provider: "xai",
Attributes: map[string]string{
"base_url": server.URL,
},
Metadata: map[string]any{"access_token": "xai-token"},
}
opts := cliproxyexecutor.Options{
SourceFormat: sdktranslator.FormatOpenAIResponse,
ResponseFormat: sdktranslator.FormatOpenAIResponse,
Stream: true,
}
firstPayload := []byte(`{"model":"grok-4.5","stream":true,"store":false,"prompt_cache_key":"codex-sse-session","include":["reasoning.encrypted_content"],"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"first"}]}]}`)
secondPayload := []byte(`{"model":"grok-4.5","stream":true,"store":false,"prompt_cache_key":"codex-sse-session","include":["reasoning.encrypted_content"],"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"second"}]}]}`)
streamedResponses := make([][]byte, 0, 2)
ctx := testContextWithAPIKey("codex-sse-api-key")
for _, payload := range [][]byte{firstPayload, secondPayload} {
result, err := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "grok-4.5", Payload: payload}, opts)
if err != nil {
t.Fatalf("ExecuteStream error: %v", err)
}
var streamed bytes.Buffer
for chunk := range result.Chunks {
if chunk.Err != nil {
t.Fatalf("stream chunk error: %v", chunk.Err)
}
streamed.Write(chunk.Payload)
}
streamedResponses = append(streamedResponses, bytes.Clone(streamed.Bytes()))
}
if len(bodies) != 2 {
t.Fatalf("upstream request count = %d, want 2", len(bodies))
}
if includes := gjson.GetBytes(bodies[0], "include").Array(); len(includes) != 1 || includes[0].String() != "reasoning.encrypted_content" {
t.Fatalf("first request include was not preserved: %s", bodies[0])
}
var downstreamEncryptedContent string
for _, line := range bytes.Split(streamedResponses[0], []byte("\n")) {
if !bytes.HasPrefix(line, xaiDataTag) {
continue
}
eventData := bytes.TrimSpace(line[len(xaiDataTag):])
if gjson.GetBytes(eventData, "type").String() != "response.output_item.done" ||
gjson.GetBytes(eventData, "item.type").String() != "reasoning" {
continue
}
downstreamEncryptedContent = gjson.GetBytes(eventData, "item.encrypted_content").String()
break
}
if downstreamEncryptedContent != encryptedContent {
t.Fatalf("downstream encrypted_content = %q, want upstream Grok blob; stream=%s", downstreamEncryptedContent, streamedResponses[0])
}
if got := gjson.GetBytes(bodies[1], "input.0.type").String(); got != "reasoning" {
t.Fatalf("second input.0.type = %q, want reasoning; body=%s", got, bodies[1])
}
if got := gjson.GetBytes(bodies[1], "input.0.encrypted_content").String(); got != encryptedContent {
t.Fatalf("replayed encrypted_content = %q, want cached Grok blob; body=%s", got, bodies[1])
}
if got := gjson.GetBytes(bodies[1], "input.1.type").String(); got != "message" {
t.Fatalf("second input.1.type = %q, want assistant message; body=%s", got, bodies[1])
}
if got := gjson.GetBytes(bodies[1], "input.1.content.0.text").String(); got != "first answer" {
t.Fatalf("replayed assistant text = %q, want first answer; body=%s", got, bodies[1])
}
if got := gjson.GetBytes(bodies[1], "input.2.content.0.text").String(); got != "second" {
t.Fatalf("new user text = %q, want second; body=%s", got, bodies[1])
}
}
func TestFilterXAIReasoningReplayItemsSkipsMatchingCachedTurn(t *testing.T) {
encryptedContent := testValidGrokEncryptedContentForSeed(10)
body := []byte(`{"input":[{"type":"reasoning","summary":[],"encrypted_content":""},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"first answer"}]},{"type":"message","role":"user","content":[{"type":"input_text","text":"second"}]}]}`)
body, _ = sjson.SetBytes(body, "input.0.encrypted_content", encryptedContent)
items := [][]byte{
[]byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":""}`),
[]byte(`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"first answer"}]}`),
}
items[0], _ = sjson.SetBytes(items[0], "encrypted_content", encryptedContent)
filtered := filterXAIReasoningReplayItemsForInput(body, items)
if len(filtered) != 0 {
t.Fatalf("filtered replay items = %q, want none for client-provided history", filtered)
}
}
func TestFilterXAIReasoningReplayItemsKeepsNewCachedTurnWhenInputHasOlderReasoning(t *testing.T) {
oldEncryptedContent := testValidGrokEncryptedContentForSeed(10)
newEncryptedContent := testValidGrokEncryptedContentForSeed(12)
body := []byte(`{"input":[{"type":"reasoning","summary":[],"encrypted_content":""},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"older answer"}]},{"type":"message","role":"user","content":[{"type":"input_text","text":"next"}]}]}`)
body, _ = sjson.SetBytes(body, "input.0.encrypted_content", oldEncryptedContent)
items := [][]byte{
[]byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":""}`),
[]byte(`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"new answer"}]}`),
}
items[0], _ = sjson.SetBytes(items[0], "encrypted_content", newEncryptedContent)
filtered := filterXAIReasoningReplayItemsForInput(body, items)
if len(filtered) != 1 || gjson.GetBytes(filtered[0], "type").String() != "reasoning" {
t.Fatalf("filtered replay items = %q, want new reasoning only (input already has last assistant)", filtered)
}
if got := gjson.GetBytes(filtered[0], "encrypted_content").String(); got != newEncryptedContent {
t.Fatalf("filtered reasoning encrypted_content = %q, want new cached blob", got)
}
}
func TestFilterXAIReasoningReplayItemsSkipsDuplicateAssistantMessage(t *testing.T) {
encryptedContent := testValidGrokEncryptedContentForSeed(11)
body := []byte(`{"input":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"first answer"}]},{"type":"message","role":"user","content":[{"type":"input_text","text":"second"}]}]}`)
items := [][]byte{
[]byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":""}`),
[]byte(`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"first answer"}]}`),
}
items[0], _ = sjson.SetBytes(items[0], "encrypted_content", encryptedContent)
filtered := filterXAIReasoningReplayItemsForInput(body, items)
if len(filtered) != 1 || gjson.GetBytes(filtered[0], "type").String() != "reasoning" {
t.Fatalf("filtered replay items = %q, want reasoning only", filtered)
}
}
func TestFilterXAIReasoningReplayItemsDoesNotMatchOlderAssistantMessage(t *testing.T) {
encryptedContent := testValidGrokEncryptedContentForSeed(13)
body := []byte(`{"input":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]},{"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"different answer"}]},{"type":"message","role":"user","content":[{"type":"input_text","text":"next"}]}]}`)
items := [][]byte{
[]byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":""}`),
[]byte(`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]}`),
}
items[0], _ = sjson.SetBytes(items[0], "encrypted_content", encryptedContent)
filtered := filterXAIReasoningReplayItemsForInput(body, items)
if len(filtered) != 1 || gjson.GetBytes(filtered[0], "type").String() != "reasoning" {
t.Fatalf("filtered replay items = %q, want reasoning only when a later assistant already exists", filtered)
}
}
// Scenario #3 (partial inject): client already has a last assistant whose text
// drifts from the cached message. Injecting the cached message would produce
// two assistants around the same turn. Expect reasoning-only inject.
func TestFilterXAIReasoningReplayItemsSkipsMessageWhenLastAssistantTextDrifts(t *testing.T) {
encryptedContent := testValidGrokEncryptedContentForSeed(20)
body := []byte(`{"input":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"first answer."}]},{"type":"message","role":"user","content":[{"type":"input_text","text":"second"}]}]}`)
items := [][]byte{
[]byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":""}`),
[]byte(`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"first answer"}]}`),
}
items[0], _ = sjson.SetBytes(items[0], "encrypted_content", encryptedContent)
filtered := filterXAIReasoningReplayItemsForInput(body, items)
if len(filtered) != 1 || gjson.GetBytes(filtered[0], "type").String() != "reasoning" {
t.Fatalf("filtered = %q, want reasoning only for drifted last assistant", filtered)
}
// Full insert path: must not create two assistant messages.
updated, ok := insertCodexReasoningReplayItems(body, filtered)
if !ok {
t.Fatal("insertCodexReasoningReplayItems failed")
}
assistantCount := 0
for _, item := range gjson.GetBytes(updated, "input").Array() {
if item.Get("type").String() == "message" && item.Get("role").String() == "assistant" {
assistantCount++
}
}
if assistantCount != 1 {
t.Fatalf("assistant messages after inject = %d, want 1; body=%s", assistantCount, updated)
}
if gjson.GetBytes(updated, "input.0.type").String() != "reasoning" {
t.Fatalf("expected reasoning inserted before last assistant; body=%s", updated)
}
}
// Scenario #2: Claude multi-turn where the client resends older thinking signature
// but drops the latest turn's signature. Cache holds the latest R(+M); upstream
// must receive the latest encrypted blob, not only the older client-provided one.
func TestXAIExecutorClaudeInjectsLatestCachedReasoningWhenHistoryHasOnlyOlderSignature(t *testing.T) {
internalcache.ClearXAIReasoningReplayCache()
t.Cleanup(internalcache.ClearXAIReasoningReplayCache)
oldEncrypted := testValidGrokEncryptedContentForSeed(21)
latestEncrypted := testValidGrokEncryptedContentForSeed(22)
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")
if len(bodies) == 1 {
_, _ = w.Write([]byte(`data: {"type":"response.output_item.done","item":{"id":"rs_latest","type":"reasoning","summary":[],"encrypted_content":"` + latestEncrypted + `"},"output_index":0}` + "\n"))
_, _ = w.Write([]byte(`data: {"type":"response.output_item.done","item":{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"latest answer"}]},"output_index":1}` + "\n"))
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed","model":"grok-4.5","output":[]}}` + "\n\n"))
return
}
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_2","status":"completed","model":"grok-4.5","output":[]}}` + "\n\n"))
}))
defer server.Close()
executor := NewXAIExecutor(&config.Config{})
auth := &cliproxyauth.Auth{
ID: "xai-auth-claude-missing-latest-sig",
Provider: "xai",
Attributes: map[string]string{"base_url": server.URL},
Metadata: map[string]any{"access_token": "xai-token"},
}
opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude, Stream: false}
ctx := testContextWithAPIKey("claude-missing-sig-key")
// Turn 1: user only -> cache latest R+M
_, err := executor.Execute(ctx, auth, cliproxyexecutor.Request{
Model: "grok-4.5",
Payload: []byte(`{"model":"grok-4.5","metadata":{"user_id":"{\"session_id\":\"claude-missing-latest\"}"},"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`),
}, opts)
if err != nil {
t.Fatalf("first Execute: %v", err)
}
// Turn 2 (actual failure shape): client keeps an OLDER thinking signature and the
// assistant text, but does not resend the latest encrypted/signature blob.
secondPayload := []byte(`{
"model":"grok-4.5",
"metadata":{"user_id":"{\"session_id\":\"claude-missing-latest\"}"},
"messages":[
{"role":"user","content":[{"type":"text","text":"hello"}]},
{"role":"assistant","content":[
{"type":"thinking","thinking":"older summary","signature":""},
{"type":"text","text":"latest answer"}
]},
{"role":"user","content":[{"type":"text","text":"next"}]}
]
}`)
secondPayload, _ = sjson.SetBytes(secondPayload, "messages.1.content.0.signature", oldEncrypted)
_, err = executor.Execute(ctx, auth, cliproxyexecutor.Request{
Model: "grok-4.5",
Payload: secondPayload,
}, opts)
if err != nil {
t.Fatalf("second Execute: %v", err)
}
if len(bodies) != 2 {
t.Fatalf("upstream requests = %d, want 2", len(bodies))
}
// Upstream must include BOTH older client signature (as reasoning) and latest cached blob.
// At minimum the latest cached encrypted_content must be present for continuity.
second := bodies[1]
foundLatest := false
foundOld := false
assistantCount := 0
for _, item := range gjson.GetBytes(second, "input").Array() {
switch item.Get("type").String() {
case "reasoning":
enc := item.Get("encrypted_content").String()
if enc == latestEncrypted {
foundLatest = true
}
if enc == oldEncrypted {
foundOld = true
}
case "message":
if item.Get("role").String() == "assistant" {
assistantCount++
}
}
}
if !foundLatest {
t.Fatalf("latest cached encrypted_content missing from upstream body (broken Claude missing-signature scenario): %s", second)
}
if !foundOld {
t.Fatalf("older client signature/reasoning missing after translate: %s", second)
}
if assistantCount != 1 {
t.Fatalf("assistant messages = %d, want 1 (no partial double-message inject); body=%s", assistantCount, second)
}
}
func TestCacheXAIReasoningReplayFromCompletedClearsPreviousEntryWhenNoReplayableState(t *testing.T) {
internalcache.ClearXAIReasoningReplayCache()
t.Cleanup(internalcache.ClearXAIReasoningReplayCache)
modelName := "grok-4.5"
sessionKey := "prompt-cache:clear-previous"
encryptedContent := testValidGrokEncryptedContentForSeed(14)
previousItems := [][]byte{
[]byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":""}`),
[]byte(`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"previous answer"}]}`),
}
previousItems[0], _ = sjson.SetBytes(previousItems[0], "encrypted_content", encryptedContent)
if !internalcache.CacheXAIReasoningReplayItems(modelName, sessionKey, previousItems) {
t.Fatal("failed to seed xAI reasoning replay cache")
}
cacheXAIReasoningReplayFromCompleted(context.Background(), xaiReasoningReplayScope{
modelName: modelName,
sessionKey: sessionKey,
}, []byte(`{"response":{"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"message without reasoning"}]}]}}`))
if _, ok := internalcache.GetXAIReasoningReplayItems(modelName, sessionKey); ok {
t.Fatal("expected previous replay entry to be cleared after non-replayable completed output")
}
}
func TestXAIReasoningReplayScopeIsolatesOpenAIResponsePromptCacheKeyByAPIKey(t *testing.T) {
payload := []byte(`{"model":"grok-4.5","prompt_cache_key":"shared-session","input":[]}`)
opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}
req := cliproxyexecutor.Request{Model: "grok-4.5", Payload: payload}
scopeA := xaiReasoningReplayScopeFromRequest(testContextWithAPIKey("api-key-a"), sdktranslator.FormatOpenAIResponse, req, opts, payload)
scopeB := xaiReasoningReplayScopeFromRequest(testContextWithAPIKey("api-key-b"), sdktranslator.FormatOpenAIResponse, req, opts, payload)
if !scopeA.valid() || !scopeB.valid() {
t.Fatalf("scopes must be valid with caller api keys: A=%+v B=%+v", scopeA, scopeB)
}
if scopeA.sessionKey == scopeB.sessionKey {
t.Fatalf("session keys must differ across callers, both %q", scopeA.sessionKey)
}
if !strings.HasPrefix(scopeA.sessionKey, "caller:") || !strings.Contains(scopeA.sessionKey, "prompt-cache:shared-session") {
t.Fatalf("session key A = %q, want caller-isolated prompt-cache key", scopeA.sessionKey)
}
scopeNoKey := xaiReasoningReplayScopeFromRequest(context.Background(), sdktranslator.FormatOpenAIResponse, req, opts, payload)
if scopeNoKey.valid() {
t.Fatalf("OpenAI Responses without caller API key must disable replay: %+v", scopeNoKey)
}
}
func TestXAIReasoningReplayScopeSkipsIncrementalWebsocketPreviousResponse(t *testing.T) {
scope := xaiReasoningReplayScopeFromRequest(
cliproxyexecutor.WithDownstreamWebsocket(context.Background()),
sdktranslator.FormatOpenAIResponse,
cliproxyexecutor.Request{
Model: "grok-4.5",
Payload: []byte(`{"model":"grok-4.5","previous_response_id":"resp_1","prompt_cache_key":"codex-ws-session","input":[]}`),
},
cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse},
[]byte(`{"model":"grok-4.5","prompt_cache_key":"codex-ws-session","input":[]}`),
)
if scope.valid() {
t.Fatalf("incremental websocket request must not enable cache replay: %+v", scope)
}
}
func TestApplyXAIReasoningReplayCacheFallsBackWhenReadFails(t *testing.T) {
previous := getXAIReasoningReplayItemsRequired
getXAIReasoningReplayItemsRequired = func(context.Context, string, string) ([][]byte, bool, error) {
+109 -16
View File
@@ -2,10 +2,12 @@ package executor
import (
"context"
"crypto/sha256"
"encoding/hex"
"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/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"
@@ -52,30 +54,61 @@ func xaiReasoningReplayScopeFromRequest(ctx context.Context, from sdktranslator.
if !xaiReasoningReplayEnabledForSource(from) {
return xaiReasoningReplayScope{}
}
// End-to-end WebSocket requests use upstream previous_response_id state.
// Replaying encrypted reasoning as input as well would duplicate the turn.
if cliproxyexecutor.DownstreamWebsocket(ctx) && strings.TrimSpace(gjson.GetBytes(req.Payload, "previous_response_id").String()) != "" {
return xaiReasoningReplayScope{}
}
sessionKey := codexReasoningReplaySessionKey(ctx, from, req, opts, body)
sessionKey = xaiReasoningReplayIsolateSessionKey(ctx, from, sessionKey)
return xaiReasoningReplayScope{
modelName: thinking.ParseSuffix(req.Model).ModelName,
sessionKey: codexReasoningReplaySessionKey(ctx, from, req, opts, body),
sessionKey: sessionKey,
}
}
// xaiReasoningReplayIsolateSessionKey namespaces client-controlled session keys
// by the downstream CPA API key so two callers cannot share encrypted reasoning
// or assistant text by reusing prompt_cache_key / window / session headers.
// Trusted execution session keys keep their existing form. OpenAI Responses
// without a caller API key is disabled rather than storing under a global key.
func xaiReasoningReplayIsolateSessionKey(ctx context.Context, from sdktranslator.Format, sessionKey string) string {
sessionKey = strings.TrimSpace(sessionKey)
if sessionKey == "" {
return ""
}
if strings.HasPrefix(sessionKey, "execution:") {
return sessionKey
}
apiKey := strings.TrimSpace(helps.APIKeyFromContext(ctx))
if apiKey == "" {
if sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse) {
return ""
}
return sessionKey
}
sum := sha256.Sum256([]byte(apiKey))
return "caller:" + hex.EncodeToString(sum[:8]) + ":" + sessionKey
}
func xaiReasoningReplayEnabledForSource(from sdktranslator.Format) bool {
return sourceFormatEqual(from, sdktranslator.FormatClaude)
return sourceFormatEqual(from, sdktranslator.FormatClaude) ||
sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse)
}
func xaiInputHasValidReasoningEncryptedContent(body []byte) bool {
input := gjson.GetBytes(body, "input")
if !input.IsArray() {
func xaiInputHasReasoningEncryptedContent(inputItems []gjson.Result, encryptedContent string) bool {
if encryptedContent == "" {
return false
}
for _, item := range input.Array() {
for _, item := range inputItems {
if strings.TrimSpace(item.Get("type").String()) != "reasoning" {
continue
}
encryptedContent := item.Get("encrypted_content")
if encryptedContent.Type != gjson.String {
inputEncryptedContent := item.Get("encrypted_content")
if inputEncryptedContent.Type != gjson.String {
continue
}
if _, err := signature.InspectGrokEncryptedContent(encryptedContent.String()); err == nil {
if inputEncryptedContent.String() == encryptedContent {
return true
}
}
@@ -88,10 +121,11 @@ func filterXAIReasoningReplayItemsForInput(body []byte, items [][]byte) [][]byte
return nil
}
hasInputReasoning := xaiInputHasValidReasoningEncryptedContent(body)
inputItems := input.Array()
_, hasLastAssistantMessage := xaiInputLastAssistantMessage(inputItems)
existingCalls := make(map[string]bool)
existingOutputs := make(map[string]bool)
for _, inputItem := range input.Array() {
for _, inputItem := range inputItems {
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())
@@ -111,7 +145,15 @@ func filterXAIReasoningReplayItemsForInput(body []byte, items [][]byte) [][]byte
itemResult := gjson.ParseBytes(item)
switch strings.TrimSpace(itemResult.Get("type").String()) {
case "reasoning":
if hasInputReasoning {
if xaiInputHasReasoningEncryptedContent(inputItems, itemResult.Get("encrypted_content").String()) {
continue
}
case "message":
// If the client already has any trailing assistant message, never inject
// another one. Equal content is a true duplicate; unequal content would
// create two adjacent assistants (partial-history / drift), which breaks
// turn pairing. Inject reasoning alone in that case.
if hasLastAssistantMessage {
continue
}
case "function_call", "custom_tool_call":
@@ -143,6 +185,48 @@ func filterXAIReasoningReplayItemsForInput(body []byte, items [][]byte) [][]byte
return filtered
}
func xaiInputLastAssistantMessage(inputItems []gjson.Result) (gjson.Result, bool) {
for i := len(inputItems) - 1; i >= 0; i-- {
inputItem := inputItems[i]
if strings.TrimSpace(inputItem.Get("type").String()) != "message" ||
!strings.EqualFold(strings.TrimSpace(inputItem.Get("role").String()), "assistant") {
continue
}
return inputItem, true
}
return gjson.Result{}, false
}
func xaiAssistantMessageContentEqual(left, right gjson.Result) bool {
if !left.IsArray() || !right.IsArray() {
return false
}
leftParts := left.Array()
rightParts := right.Array()
if len(leftParts) != len(rightParts) {
return false
}
for i := range leftParts {
if strings.TrimSpace(leftParts[i].Get("type").String()) != strings.TrimSpace(rightParts[i].Get("type").String()) ||
xaiAssistantMessagePartValue(leftParts[i]) != xaiAssistantMessagePartValue(rightParts[i]) {
return false
}
}
return true
}
func xaiAssistantMessagePartValue(part gjson.Result) string {
switch strings.TrimSpace(part.Get("type").String()) {
case "refusal":
if refusal := part.Get("refusal"); refusal.Type == gjson.String {
return refusal.String()
}
return part.Get("text").String()
default:
return part.Get("text").String()
}
}
func cacheXAIReasoningReplayFromCompleted(ctx context.Context, scope xaiReasoningReplayScope, completedData []byte) {
if !scope.valid() {
return
@@ -157,15 +241,24 @@ func cacheXAIReasoningReplayFromCompleted(ctx context.Context, scope xaiReasonin
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":
case "reasoning", "message", "function_call", "custom_tool_call":
items = append(items, []byte(item.Raw))
default:
continue
}
}
if !internalcache.CacheXAIReasoningReplayItemsBestEffort(ctx, scope.modelName, scope.sessionKey, items) {
switch internalcache.StoreXAIReasoningReplayItems(ctx, scope.modelName, scope.sessionKey, items) {
case internalcache.XAIReasoningReplayStored:
return
case internalcache.XAIReasoningReplayNoReplayableState:
// Successful completed turn without cacheable reasoning must not leave
// a previous turn's encrypted state to be injected later.
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)
log.Warnf("xai reasoning replay cache delete failed after non-replayable completed output: %v", errDelete)
}
case internalcache.XAIReasoningReplayStoreBackendError:
log.Debug("xai reasoning replay cache store backend error; retaining previous entry")
default:
// Invalid args: nothing to store or clear.
}
}
@@ -139,9 +139,16 @@ func ConvertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool)
return
}
signature, ok := sigcompat.CompatibleSignatureForProvider(sigcompat.SignatureProviderGPT, part.Get("signature").String())
rawSignature := part.Get("signature").String()
signature, ok := sigcompat.CompatibleSignatureForProvider(sigcompat.SignatureProviderGPT, rawSignature)
if !ok {
return
if !codexClaudeTargetAcceptsGrokSignature(modelName) {
return
}
if _, err := sigcompat.InspectGrokEncryptedContent(rawSignature); err != nil {
return
}
signature = rawSignature
}
flushMessage()
@@ -343,6 +350,11 @@ func ConvertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool)
return template
}
func codexClaudeTargetAcceptsGrokSignature(modelName string) bool {
baseModel := strings.ToLower(strings.TrimSpace(thinking.ParseSuffix(modelName).ModelName))
return strings.Contains(baseModel, "grok")
}
func normalizeCodexServiceTier(result gjson.Result) string {
if !result.Exists() || result.Type != gjson.String {
return ""
@@ -6,6 +6,7 @@ import (
"testing"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
func TestConvertClaudeRequestToCodex_SystemMessageScenarios(t *testing.T) {
@@ -480,6 +481,36 @@ func TestConvertClaudeRequestToCodex_AssistantThinkingSignatureToReasoningItem(t
}
}
func TestConvertClaudeRequestToCodex_AssistantGrokSignatureToReasoningItem(t *testing.T) {
signature := "HmlYdr2aCAqCYP/m9mr8PS6KOsdMs72FGDigmydR+Jsmuv8KX97yWPlbOwmXJgWn0CbHaCacdQD3+n5EvpgLfPNmafS3kdICBjRuDf4bzHy7uBiUhNVhqPtp/ee1y9q4imPE4LYgD1VZ4J+bp9mTeqA1+nC9Oue58CiNEMV9SVaGenCD+aBnVuSTzQhD32Y+68i6HLJW0Dx6ifaRfb8hxYtA/sPM+/FTvAMW11nRho5a2BBSkpnzfqqAz/e/vGJ77/bygpXM823QA9wL9i0X"
payload := []byte(`{"model":"grok-4.5","messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"summary","signature":""},{"type":"text","text":"answer"}]},{"role":"user","content":"next"}]}`)
payload, _ = sjson.SetBytes(payload, "messages.0.content.0.signature", signature)
out := ConvertClaudeRequestToCodex("grok-4.5", payload, false)
reasoning := gjson.GetBytes(out, "input.0")
if reasoning.Get("type").String() != "reasoning" {
t.Fatalf("input.0 type = %q, want reasoning; output=%s", reasoning.Get("type").String(), out)
}
if got := reasoning.Get("encrypted_content").String(); got != signature {
t.Fatalf("encrypted_content = %q, want Grok signature", got)
}
}
func TestConvertClaudeRequestToCodex_IgnoresGrokSignatureForNonGrokTargets(t *testing.T) {
signature := "HmlYdr2aCAqCYP/m9mr8PS6KOsdMs72FGDigmydR+Jsmuv8KX97yWPlbOwmXJgWn0CbHaCacdQD3+n5EvpgLfPNmafS3kdICBjRuDf4bzHy7uBiUhNVhqPtp/ee1y9q4imPE4LYgD1VZ4J+bp9mTeqA1+nC9Oue58CiNEMV9SVaGenCD+aBnVuSTzQhD32Y+68i6HLJW0Dx6ifaRfb8hxYtA/sPM+/FTvAMW11nRho5a2BBSkpnzfqqAz/e/vGJ77/bygpXM823QA9wL9i0X"
payload := []byte(`{"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"summary","signature":""},{"type":"text","text":"answer"}]},{"role":"user","content":"next"}]}`)
payload, _ = sjson.SetBytes(payload, "messages.0.content.0.signature", signature)
for _, modelName := range []string{"gpt-5.4", "claude-sonnet-4-6"} {
t.Run(modelName, func(t *testing.T) {
out := ConvertClaudeRequestToCodex(modelName, payload, false)
if got := countRequestInputItemsByType(out, "reasoning"); got != 0 {
t.Fatalf("got %d reasoning items for non-Grok target, want 0; output=%s", got, out)
}
})
}
}
func TestConvertClaudeRequestToCodex_IgnoresNonCodexThinkingSignatures(t *testing.T) {
tests := []struct {
name string