diff --git a/DESIGN.md b/DESIGN.md index fff5158..90da7ac 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -1021,8 +1021,10 @@ caching := interceptors.NewOpenAIPromptCachingWithOrgExtractor( `NewXAIPromptCaching(convID)` — Enables xAI/Grok prompt caching: -- **Automatic prefix caching:** xAI caches from the start of the messages array automatically -- **Cache routing:** Adds `x-grok-conv-id` HTTP header to route requests to the same server where cache lives +- **Automatic prefix caching:** xAI caches from the start of the messages/input array automatically +- **Chat Completions cache routing:** Adds `x-grok-conv-id` HTTP header to route requests to the same server where cache lives +- **Responses API cache routing:** Adds `prompt_cache_key` to the request body (functionally identical sticky routing) +- **Key preference:** Preserves client-supplied `x-grok-conv-id` / `prompt_cache_key`, mirrors between header and body when needed, then falls back to configured/extracted keys - **Conversation ID:** Use a stable value (conversation ID, session ID, or deterministic hash of static content) - **Key rule:** Never reorder or modify earlier messages — only append diff --git a/README.md b/README.md index 0d5f2d3..ea8d4c5 100644 --- a/README.md +++ b/README.md @@ -339,7 +339,7 @@ llmproxy.WithInterceptor(interceptors.NewOpenAIPromptCaching(interceptors.CacheR // OpenAI prompt caching with auto-derived key and tenant namespace llmproxy.WithInterceptor(interceptors.NewOpenAIPromptCachingAuto("tenant-123", interceptors.CacheRetentionDefault)) -// xAI/Grok prompt caching (uses x-grok-conv-id header) +// xAI/Grok prompt caching (x-grok-conv-id for Chat Completions, prompt_cache_key for Responses) llmproxy.WithInterceptor(interceptors.NewXAIPromptCaching("conv-abc123")) // Fireworks prompt caching (uses x-session-affinity and x-prompt-cache-isolation-key headers) diff --git a/aigateway_regression_live_test.go b/aigateway_regression_live_test.go index 25e4df4..a0f612d 100644 --- a/aigateway_regression_live_test.go +++ b/aigateway_regression_live_test.go @@ -5,15 +5,19 @@ package llmproxy_test import ( "bytes" "context" + "encoding/json" "io" "net/http" "net/http/httptest" "os" + "sort" + "strconv" "strings" "testing" "time" "github.com/agentuity/llmproxy" + "github.com/agentuity/llmproxy/providers/anthropic" "github.com/agentuity/llmproxy/providers/openai_compatible" ) @@ -93,6 +97,206 @@ func TestLiveAIGatewayRegressionModels(t *testing.T) { } } +func TestLiveAnthropicMessagesStreamCompletes(t *testing.T) { + if os.Getenv("LLMPROXY_LIVE_AIGATEWAY_REGRESSION") != "1" { + t.Skip("set LLMPROXY_LIVE_AIGATEWAY_REGRESSION=1 to run live AI Gateway regression checks") + } + + apiKey := firstNonEmptyLiveEnv("ANTHROPIC_API_KEY", "GATEWAY_ANTHROPIC_API_KEY") + if apiKey == "" { + t.Skip("set ANTHROPIC_API_KEY or GATEWAY_ANTHROPIC_API_KEY to run this live regression") + } + + model := os.Getenv("LLMPROXY_LIVE_ANTHROPIC_MODEL") + if model == "" { + model = "anthropic/claude-haiku-4-5-20251001" + } + + provider, err := anthropic.New(apiKey) + if err != nil { + t.Fatalf("create anthropic provider: %v", err) + } + + router := llmproxy.NewAutoRouter( + llmproxy.WithAutoRouterDetector(llmproxy.ProviderDetectorFunc(func(hint llmproxy.ProviderHint) string { + return "anthropic" + })), + llmproxy.WithAutoRouterHTTPClient(&http.Client{Timeout: 60 * time.Second}), + ) + router.RegisterProvider(provider) + + rawBody, err := anthropicLiveStreamRequestBody(model) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(rawBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + rec := httptest.NewRecorder() + + router.ServeHTTP(rec, req) + readAndAssertLiveStream(t, rec.Result()) +} + +func TestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes(t *testing.T) { + if os.Getenv("LLMPROXY_LIVE_AIGATEWAY_REGRESSION") != "1" { + t.Skip("set LLMPROXY_LIVE_AIGATEWAY_REGRESSION=1 to run live AI Gateway regression checks") + } + + apiKey := firstNonEmptyLiveEnv( + "AGENTUITY_AIGATEWAY_KEY", + "AIGATEWAY_API_KEY", + "AGENTUITY_SDK_KEY", + "AGENTUITY_CODER_API_KEY", + "AGENTUITY_CLI_API_KEY", + "AGENTUITY_CLI_KEY", + ) + if apiKey == "" { + t.Skip("set an Agentuity AI Gateway or SDK API key to run this live regression") + } + + baseURL := firstNonEmptyLiveEnv("AGENTUITY_AIGATEWAY_URL", "AIGATEWAY_URL") + if baseURL == "" { + baseURL = "https://aigateway-usc.agentuity.cloud" + } + model := os.Getenv("LLMPROXY_LIVE_ANTHROPIC_MODEL") + if model == "" { + model = "anthropic/claude-haiku-4-5-20251001" + } + + rawBody, err := anthropicLiveStreamRequestBody(model) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + req, err := http.NewRequest(http.MethodPost, strings.TrimRight(baseURL, "/")+"/v1/messages", bytes.NewReader(rawBody)) + if err != nil { + t.Fatalf("create request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("Authorization", "Bearer "+apiKey) + req.Header.Set("x-api-key", apiKey) + if orgID := firstNonEmptyLiveEnv("AGENTUITY_AIGATEWAY_ORGID", "AGENTUITY_ORG_ID", "AGENTUITY_CLOUD_ORG_ID"); orgID != "" { + req.Header.Set("x-agentuity-orgid", orgID) + } + + resp, err := (&http.Client{Timeout: 60 * time.Second}).Do(req) + if err != nil { + t.Fatalf("send request: %v", err) + } + readAndAssertLiveStream(t, resp) +} + +func anthropicLiveStreamRequestBody(model string) ([]byte, error) { + body := map[string]any{ + "model": model, + "stream": true, + "max_tokens": 512, + "thinking": map[string]any{ + "type": "disabled", + }, + "messages": []map[string]any{ + { + "role": "user", + "content": []map[string]any{ + {"type": "text", "text": "Reply with GENESIS_DRIVER_SMOKE_OK and nothing else."}, + }, + }, + }, + } + return json.Marshal(body) +} + +func readAndAssertLiveStream(t *testing.T, resp *http.Response) { + t.Helper() + defer resp.Body.Close() + raw, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read response: %v", err) + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + t.Fatalf("status %d: %s", resp.StatusCode, truncateLiveRegressionBody(raw)) + } + assertAnthropicLiveStream(t, raw) +} + +func assertAnthropicLiveStream(t *testing.T, raw []byte) { + t.Helper() + summary := summarizeAnthropicLiveStream(raw) + if summary.EventTypes["message_stop"] == 0 { + t.Fatalf("stream did not include message_stop: %s\n%s", summary.String(), truncateLiveRegressionBody(raw)) + } + if summary.TextDeltas == 0 { + t.Fatalf("stream did not include any text deltas: %s\n%s", summary.String(), truncateLiveRegressionBody(raw)) + } + if !strings.Contains(summary.Text, "GENESIS_DRIVER_SMOKE_OK") { + t.Fatalf("stream text did not include expected sentinel: %s\ntext=%q\n%s", summary.String(), summary.Text, truncateLiveRegressionBody(raw)) + } + if summary.ThinkingDeltas > 0 { + t.Fatalf("stream included thinking_delta despite thinking disabled: %s\n%s", summary.String(), truncateLiveRegressionBody(raw)) + } +} + +type anthropicLiveStreamSummary struct { + EventTypes map[string]int + TextDeltas int + ThinkingDeltas int + Text string +} + +func (s anthropicLiveStreamSummary) String() string { + eventTypes := make([]string, 0, len(s.EventTypes)) + for eventType, count := range s.EventTypes { + eventTypes = append(eventTypes, eventType+"="+fmtInt(count)) + } + sort.Strings(eventTypes) + return "events=[" + strings.Join(eventTypes, ",") + "] text_deltas=" + fmtInt(s.TextDeltas) + " thinking_deltas=" + fmtInt(s.ThinkingDeltas) +} + +func summarizeAnthropicLiveStream(raw []byte) anthropicLiveStreamSummary { + summary := anthropicLiveStreamSummary{ + EventTypes: make(map[string]int), + } + for _, block := range strings.Split(string(raw), "\n\n") { + var data string + for _, line := range strings.Split(block, "\n") { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "data:") { + data = strings.TrimSpace(strings.TrimPrefix(line, "data:")) + } + } + if data == "" { + continue + } + var payload struct { + Type string `json:"type"` + Delta struct { + Type string `json:"type"` + Text string `json:"text"` + Thinking string `json:"thinking"` + } `json:"delta"` + } + if err := json.Unmarshal([]byte(data), &payload); err != nil { + continue + } + if payload.Type != "" { + summary.EventTypes[payload.Type]++ + } + switch payload.Delta.Type { + case "text_delta": + summary.TextDeltas++ + summary.Text += payload.Delta.Text + case "thinking_delta": + summary.ThinkingDeltas++ + } + } + return summary +} + +func fmtInt(value int) string { + return strconv.Itoa(value) +} + func truncateLiveRegressionBody(body []byte) string { value := strings.TrimSpace(string(body)) if len(value) <= 500 { @@ -100,3 +304,12 @@ func truncateLiveRegressionBody(body []byte) string { } return value[:500] + "..." } + +func firstNonEmptyLiveEnv(names ...string) string { + for _, name := range names { + if value := strings.TrimSpace(os.Getenv(name)); value != "" { + return value + } + } + return "" +} diff --git a/interceptors/promptcaching.go b/interceptors/promptcaching.go index bd4c49b..5d6ac61 100644 --- a/interceptors/promptcaching.go +++ b/interceptors/promptcaching.go @@ -118,16 +118,48 @@ func (i *PromptCachingInterceptor) Intercept(req *http.Request, meta llmproxy.Bo } func (i *PromptCachingInterceptor) interceptXAI(req *http.Request, meta llmproxy.BodyMetadata, rawBody []byte, next llmproxy.RoundTripFunc) (*http.Response, llmproxy.ResponseMetadata, []byte, error) { - if req.Header.Get("x-grok-conv-id") != "" { - return next(req) + fields, fieldsOK := parseJSONObjectRaw(rawBody) + apiType := xaiAPIType(req, meta, fields, fieldsOK) + headerKey := strings.TrimSpace(req.Header.Get("x-grok-conv-id")) + bodyKey := "" + if fieldsOK { + bodyKey = xaiBodyPromptCacheKeyFromFields(fields) + } + cacheKey := headerKey + if cacheKey == "" { + cacheKey = bodyKey + } + if cacheKey == "" { + cacheKey = i.resolveDynamicCacheKey(req, meta, rawBody) } - - cacheKey := i.resolveDynamicCacheKey(req, meta, rawBody) if cacheKey == "" { return next(req) } - req.Header.Set("x-grok-conv-id", cacheKey) + modifiedBody := rawBody + bodyChanged := false + + // Chat Completions sticky routing uses the x-grok-conv-id header. + // Responses sticky routing uses prompt_cache_key in the body (functionally + // identical per xAI docs). Mirror a client-supplied key across both forms + // so either API path gets sticky routing. + if headerKey == "" { + req.Header.Set("x-grok-conv-id", cacheKey) + } + if apiType == llmproxy.APITypeResponses && bodyKey == "" && fieldsOK { + var err error + modifiedBody, bodyChanged, err = setXAIPromptCacheKeyFromFields(fields, cacheKey) + if err != nil { + return next(req) + } + } + + if bodyChanged { + if req.Body != nil { + req.Body.Close() + } + req = cloneRequestWithBody(req, modifiedBody) + } resp, respMeta, rawRespBody, err := next(req) if err != nil { @@ -143,6 +175,80 @@ func (i *PromptCachingInterceptor) interceptXAI(req *http.Request, meta llmproxy return resp, respMeta, rawRespBody, err } +func parseJSONObjectRaw(rawBody []byte) (map[string]json.RawMessage, bool) { + var fields map[string]json.RawMessage + if err := json.Unmarshal(rawBody, &fields); err != nil || fields == nil { + return nil, false + } + return fields, true +} + +func xaiAPIType(req *http.Request, meta llmproxy.BodyMetadata, fields map[string]json.RawMessage, fieldsOK bool) llmproxy.APIType { + if req != nil && req.URL != nil { + if apiType := llmproxy.DetectAPITypeFromPath(req.URL.Path); apiType != "" { + return apiType + } + } + if meta.Custom != nil { + switch v := meta.Custom["api_type"].(type) { + case llmproxy.APIType: + if v != "" { + return v + } + case string: + if v != "" { + return llmproxy.APIType(v) + } + } + } + if !fieldsOK { + return llmproxy.APITypeChatCompletions + } + _, hasInput := fields["input"] + _, hasMessages := fields["messages"] + if hasInput && !hasMessages { + return llmproxy.APITypeResponses + } + _, hasPrompt := fields["prompt"] + if hasPrompt && !hasMessages { + return llmproxy.APITypeCompletions + } + return llmproxy.APITypeChatCompletions +} + +func xaiBodyPromptCacheKeyFromFields(fields map[string]json.RawMessage) string { + raw, ok := fields["prompt_cache_key"] + if !ok { + return "" + } + var key string + if err := json.Unmarshal(raw, &key); err != nil { + return "" + } + return strings.TrimSpace(key) +} + +func setXAIPromptCacheKeyFromFields(fields map[string]json.RawMessage, cacheKey string) ([]byte, bool, error) { + if existing := xaiBodyPromptCacheKeyFromFields(fields); existing != "" { + return nil, false, nil + } + keyRaw, err := json.Marshal(cacheKey) + if err != nil { + return nil, false, err + } + // Copy so callers can keep using the original parsed map unchanged. + next := make(map[string]json.RawMessage, len(fields)+1) + for k, v := range fields { + next[k] = v + } + next["prompt_cache_key"] = json.RawMessage(keyRaw) + modified, err := json.Marshal(next) + if err != nil { + return nil, false, err + } + return modified, true, nil +} + func (i *PromptCachingInterceptor) interceptFireworks(req *http.Request, meta llmproxy.BodyMetadata, rawBody []byte, next llmproxy.RoundTripFunc) (*http.Response, llmproxy.ResponseMetadata, []byte, error) { orgID := i.extractOrgID(req, meta) @@ -607,8 +713,10 @@ func (i *PromptCachingInterceptor) resolveDynamicCacheKey(req *http.Request, met func DeriveCacheKeyFromPrefix(meta llmproxy.BodyMetadata, rawBody []byte) string { var body struct { - System interface{} `json:"system"` - Messages []struct { + System interface{} `json:"system"` + Instructions interface{} `json:"instructions"` + Input interface{} `json:"input"` + Messages []struct { Role string `json:"role"` Content interface{} `json:"content"` } `json:"messages"` @@ -621,6 +729,10 @@ func DeriveCacheKeyFromPrefix(meta llmproxy.BodyMetadata, rawBody []byte) string sysBytes, _ := json.Marshal(body.System) prefix.Write(sysBytes) } + if body.Instructions != nil { + instructionsBytes, _ := json.Marshal(body.Instructions) + prefix.Write(instructionsBytes) + } if body.Tools != nil { toolsBytes, _ := json.Marshal(body.Tools) prefix.Write(toolsBytes) @@ -631,6 +743,18 @@ func DeriveCacheKeyFromPrefix(meta llmproxy.BodyMetadata, rawBody []byte) string prefix.Write(msgBytes) } } + // Responses API: hash prior input turns the same way as chat messages. + if len(body.Messages) == 0 { + switch input := body.Input.(type) { + case []interface{}: + for i, item := range input { + if i < len(input)-1 { + itemBytes, _ := json.Marshal(item) + prefix.Write(itemBytes) + } + } + } + } if prefix.Len() == 0 { return "" diff --git a/interceptors/promptcaching_test.go b/interceptors/promptcaching_test.go index e26a93b..83376bc 100644 --- a/interceptors/promptcaching_test.go +++ b/interceptors/promptcaching_test.go @@ -2088,6 +2088,206 @@ func TestPromptCachingInterceptor_XAICacheKeyHeader(t *testing.T) { } } +func TestPromptCachingInterceptor_XAIResponsesAddsPromptCacheKey(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("x-grok-conv-id") != "conv-responses" { + t.Errorf("x-grok-conv-id header = %q, want conv-responses", r.Header.Get("x-grok-conv-id")) + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + var body map[string]any + if err := json.Unmarshal(raw, &body); err != nil { + t.Fatalf("invalid json: %v", err) + } + if body["prompt_cache_key"] != "conv-responses" { + t.Errorf("prompt_cache_key = %v, want conv-responses", body["prompt_cache_key"]) + } + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer upstream.Close() + + caching := NewXAIPromptCaching("conv-responses") + reqBody := []byte(`{"model":"grok-4.5","input":"What is prompt caching?"}`) + req, _ := http.NewRequest("POST", upstream.URL+"/v1/responses", bytes.NewReader(reqBody)) + meta := llmproxy.BodyMetadata{ + Model: "grok-4.5", + Custom: map[string]any{"api_type": llmproxy.APITypeResponses}, + } + + next := func(req *http.Request) (*http.Response, llmproxy.ResponseMetadata, []byte, error) { + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, llmproxy.ResponseMetadata{}, nil, err + } + body, _ := io.ReadAll(resp.Body) + return resp, llmproxy.ResponseMetadata{}, body, nil + } + + if _, _, _, err := caching.Intercept(req, meta, reqBody, next); err != nil { + t.Fatalf("Intercept returned error: %v", err) + } +} + +func TestPromptCachingInterceptor_XAIResponsesPreservesExistingPromptCacheKey(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("x-grok-conv-id") != "existing-body-key" { + t.Errorf("x-grok-conv-id header = %q, want existing-body-key", r.Header.Get("x-grok-conv-id")) + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + var body map[string]any + if err := json.Unmarshal(raw, &body); err != nil { + t.Fatalf("invalid json: %v", err) + } + if body["prompt_cache_key"] != "existing-body-key" { + t.Errorf("prompt_cache_key = %v, want existing-body-key", body["prompt_cache_key"]) + } + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer upstream.Close() + + caching := NewXAIPromptCaching("config-key") + reqBody := []byte(`{"model":"grok-4.5","input":"hi","prompt_cache_key":"existing-body-key"}`) + req, _ := http.NewRequest("POST", upstream.URL+"/v1/responses", bytes.NewReader(reqBody)) + meta := llmproxy.BodyMetadata{Model: "grok-4.5"} + + next := func(req *http.Request) (*http.Response, llmproxy.ResponseMetadata, []byte, error) { + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, llmproxy.ResponseMetadata{}, nil, err + } + body, _ := io.ReadAll(resp.Body) + return resp, llmproxy.ResponseMetadata{}, body, nil + } + + if _, _, _, err := caching.Intercept(req, meta, reqBody, next); err != nil { + t.Fatalf("Intercept returned error: %v", err) + } +} + +func TestPromptCachingInterceptor_XAIResponsesMirrorsHeaderToBody(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("x-grok-conv-id") != "header-conv" { + t.Errorf("x-grok-conv-id header = %q, want header-conv", r.Header.Get("x-grok-conv-id")) + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + var body map[string]any + if err := json.Unmarshal(raw, &body); err != nil { + t.Fatalf("invalid json: %v", err) + } + if body["prompt_cache_key"] != "header-conv" { + t.Errorf("prompt_cache_key = %v, want header-conv", body["prompt_cache_key"]) + } + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer upstream.Close() + + caching := NewXAIPromptCaching("config-key") + reqBody := []byte(`{"model":"grok-4.5","input":"hi"}`) + req, _ := http.NewRequest("POST", upstream.URL+"/v1/responses", bytes.NewReader(reqBody)) + req.Header.Set("x-grok-conv-id", "header-conv") + meta := llmproxy.BodyMetadata{Model: "grok-4.5"} + + next := func(req *http.Request) (*http.Response, llmproxy.ResponseMetadata, []byte, error) { + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, llmproxy.ResponseMetadata{}, nil, err + } + body, _ := io.ReadAll(resp.Body) + return resp, llmproxy.ResponseMetadata{}, body, nil + } + + if _, _, _, err := caching.Intercept(req, meta, reqBody, next); err != nil { + t.Fatalf("Intercept returned error: %v", err) + } +} + +func TestPromptCachingInterceptor_XAIChatMirrorsBodyKeyToHeader(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("x-grok-conv-id") != "body-conv" { + t.Errorf("x-grok-conv-id header = %q, want body-conv", r.Header.Get("x-grok-conv-id")) + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + var body map[string]any + if err := json.Unmarshal(raw, &body); err != nil { + t.Fatalf("invalid json: %v", err) + } + if body["prompt_cache_key"] != "body-conv" { + t.Errorf("prompt_cache_key = %v, want body-conv (chat path should leave body key intact)", body["prompt_cache_key"]) + } + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer upstream.Close() + + caching := NewXAIPromptCaching("config-key") + reqBody := []byte(`{"model":"grok-4.5","messages":[{"role":"user","content":"hi"}],"prompt_cache_key":"body-conv"}`) + req, _ := http.NewRequest("POST", upstream.URL+"/v1/chat/completions", bytes.NewReader(reqBody)) + meta := llmproxy.BodyMetadata{Model: "grok-4.5"} + + next := func(req *http.Request) (*http.Response, llmproxy.ResponseMetadata, []byte, error) { + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, llmproxy.ResponseMetadata{}, nil, err + } + body, _ := io.ReadAll(resp.Body) + return resp, llmproxy.ResponseMetadata{}, body, nil + } + + if _, _, _, err := caching.Intercept(req, meta, reqBody, next); err != nil { + t.Fatalf("Intercept returned error: %v", err) + } +} + +func TestDeriveCacheKeyFromPrefix_ResponsesInstructions(t *testing.T) { + key1 := DeriveCacheKeyFromPrefix(llmproxy.BodyMetadata{}, []byte(`{"model":"grok-4.5","instructions":"Be helpful.","input":[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi"},{"role":"user","content":"Next"}]}`)) + key2 := DeriveCacheKeyFromPrefix(llmproxy.BodyMetadata{}, []byte(`{"model":"grok-4.5","instructions":"Be helpful.","input":[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi"},{"role":"user","content":"Other"}]}`)) + key3 := DeriveCacheKeyFromPrefix(llmproxy.BodyMetadata{}, []byte(`{"model":"grok-4.5","instructions":"Different.","input":[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi"},{"role":"user","content":"Next"}]}`)) + if key1 == "" { + t.Fatal("expected non-empty responses prefix key") + } + if key1 != key2 { + t.Fatalf("expected stable key across last-turn changes, got %q vs %q", key1, key2) + } + if key1 == key3 { + t.Fatal("expected different key when instructions change") + } +} + +func TestSetXAIPromptCacheKeyFromFieldsPreservesLargeIntegers(t *testing.T) { + raw := []byte(`{"model":"grok-4.5","input":"hi","max_tokens":9007199254740993}`) + fields, ok := parseJSONObjectRaw(raw) + if !ok { + t.Fatal("expected parseable body") + } + modified, changed, err := setXAIPromptCacheKeyFromFields(fields, "conv-1") + if err != nil { + t.Fatalf("setXAIPromptCacheKeyFromFields error = %v", err) + } + if !changed { + t.Fatal("expected body change") + } + if !bytes.Contains(modified, []byte(`"max_tokens":9007199254740993`)) { + t.Fatalf("expected large integer to be preserved, got %s", modified) + } + if !bytes.Contains(modified, []byte(`"prompt_cache_key":"conv-1"`)) { + t.Fatalf("expected prompt_cache_key injection, got %s", modified) + } +} + func TestPromptCachingInterceptor_FireworksCacheKeyHeader(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { sessionID := r.Header.Get(HeaderFireworksSessionAffinity)