From f1e9347f512ad716058da376a5d928c0b3870c07 Mon Sep 17 00:00:00 2001 From: sususu Date: Sun, 12 Jul 2026 23:19:36 +0800 Subject: [PATCH] fix(xai): preserve tool-call-only replay batches --- internal/cache/xai_reasoning_replay_cache.go | 9 ++-- .../cache/xai_reasoning_replay_cache_test.go | 50 +++++++++++++++++++ .../runtime/executor/xai_executor_test.go | 41 +++++++++++++++ 3 files changed, 96 insertions(+), 4 deletions(-) diff --git a/internal/cache/xai_reasoning_replay_cache.go b/internal/cache/xai_reasoning_replay_cache.go index 57c60ac2..156bbd4f 100644 --- a/internal/cache/xai_reasoning_replay_cache.go +++ b/internal/cache/xai_reasoning_replay_cache.go @@ -243,17 +243,18 @@ func xaiReasoningReplayKVKey(modelName, sessionKey string) string { func normalizeXAIReasoningReplayItems(items [][]byte) ([][]byte, bool) { normalized := make([][]byte, 0, len(items)) - hasReasoning := false + hasReplayAnchor := 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 + switch strings.TrimSpace(gjson.GetBytes(normalizedItem, "type").String()) { + case "reasoning", "function_call", "custom_tool_call": + hasReplayAnchor = true } } } - return normalized, hasReasoning + return normalized, hasReplayAnchor } func normalizeXAIReasoningReplayItem(item []byte) ([]byte, bool) { diff --git a/internal/cache/xai_reasoning_replay_cache_test.go b/internal/cache/xai_reasoning_replay_cache_test.go index de141c12..2945c1c9 100644 --- a/internal/cache/xai_reasoning_replay_cache_test.go +++ b/internal/cache/xai_reasoning_replay_cache_test.go @@ -172,6 +172,56 @@ func TestXAIReasoningReplayCacheRejectsAssistantMessageWithoutReasoning(t *testi } } +func TestXAIReasoningReplayCacheStoresToolCallWithoutReasoning(t *testing.T) { + ClearXAIReasoningReplayCache() + t.Cleanup(ClearXAIReasoningReplayCache) + + tests := []struct { + name string + sessionKey string + item []byte + wantType string + wantPayload string + }{ + { + name: "function call", + sessionKey: "prompt-cache:function-call-only", + item: []byte(`{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"q\":\"weather\"}"}`), + wantType: "function_call", + wantPayload: `{"q":"weather"}`, + }, + { + name: "custom tool call", + sessionKey: "prompt-cache:custom-tool-call-only", + item: []byte(`{"type":"custom_tool_call","call_id":"call_2","name":"shell","input":"pwd"}`), + wantType: "custom_tool_call", + wantPayload: "pwd", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if !CacheXAIReasoningReplayItems("grok-4.3", tt.sessionKey, [][]byte{tt.item}) { + t.Fatal("tool-call-only replay batch must be cached") + } + items, ok := GetXAIReasoningReplayItems("grok-4.3", tt.sessionKey) + if !ok || len(items) != 1 { + t.Fatalf("cached items = %q, %v, want one item", items, ok) + } + if got := gjson.GetBytes(items[0], "type").String(); got != tt.wantType { + t.Fatalf("cached type = %q, want %q; item=%s", got, tt.wantType, items[0]) + } + payloadPath := "arguments" + if tt.wantType == "custom_tool_call" { + payloadPath = "input" + } + if got := gjson.GetBytes(items[0], payloadPath).String(); got != tt.wantPayload { + t.Fatalf("cached %s = %q, want %q; item=%s", payloadPath, got, tt.wantPayload, items[0]) + } + }) + } +} + func TestXAIReasoningReplayRequiredHomeExpireFailureReturnsItems(t *testing.T) { ClearXAIReasoningReplayCache() t.Cleanup(ClearXAIReasoningReplayCache) diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 5f0b7c2c..ce955eba 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -1877,6 +1877,47 @@ func TestApplyXAIReasoningReplayCacheFallsBackWhenReadFails(t *testing.T) { } } +func TestXAIReasoningReplayCacheReplaysFunctionCallWithoutReasoning(t *testing.T) { + internalcache.ClearXAIReasoningReplayCache() + t.Cleanup(internalcache.ClearXAIReasoningReplayCache) + + const executionSessionID = "xai-tool-call-only" + cacheXAIReasoningReplayFromCompleted(context.Background(), xaiReasoningReplayScope{ + modelName: "grok-4.3", + sessionKey: "execution:" + executionSessionID, + }, []byte(`{"response":{"output":[{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"q\":\"weather\"}"}]}}`)) + + body := []byte(`{"model":"grok-4.3","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"call lookup"}]},{"type":"function_call_output","call_id":"call_1","output":"sunny"}]}`) + updated, scope, errReplay := applyXAIReasoningReplayCacheRequired(context.Background(), sdktranslator.FormatClaude, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: body, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID, + }, + }, body) + if errReplay != nil { + t.Fatalf("applyXAIReasoningReplayCacheRequired() error = %v", errReplay) + } + if !scope.valid() { + t.Fatal("tool-call-only replay scope must remain valid") + } + input := gjson.GetBytes(updated, "input").Array() + if len(input) != 3 { + t.Fatalf("input length = %d, want 3; body=%s", len(input), updated) + } + wantTypes := []string{"message", "function_call", "function_call_output"} + for i, wantType := range wantTypes { + if got := input[i].Get("type").String(); got != wantType { + t.Fatalf("input.%d.type = %q, want %q; body=%s", i, got, wantType, updated) + } + } + if got := input[1].Get("call_id").String(); got != "call_1" { + t.Fatalf("replayed call_id = %q, want call_1; body=%s", got, updated) + } +} + func TestXAIExecutorReasoningReplayCacheReplaysFunctionCallForClaudeToolResult(t *testing.T) { internalcache.ClearXAIReasoningReplayCache() t.Cleanup(internalcache.ClearXAIReasoningReplayCache) -- 2.51.2