From a8f31d2bb84e178daaefcf79e0e5c845e5b89c89 Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Mon, 13 Jul 2026 13:05:01 -0500 Subject: [PATCH 1/4] Add xAI Responses prompt_cache_key sticky routing. Mirror conversation affinity keys across Chat Completions headers and Responses body fields so Grok cache hits follow xAI's documented routing. Co-authored-by: Cursor --- DESIGN.md | 6 +- README.md | 2 +- interceptors/promptcaching.go | 107 +++++++++++++++-- interceptors/promptcaching_test.go | 179 +++++++++++++++++++++++++++++ 4 files changed, 284 insertions(+), 10 deletions(-) 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/interceptors/promptcaching.go b/interceptors/promptcaching.go index bd4c49b..83326ac 100644 --- a/interceptors/promptcaching.go +++ b/interceptors/promptcaching.go @@ -118,16 +118,44 @@ 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) + apiType := xaiAPIType(req, meta, rawBody) + headerKey := strings.TrimSpace(req.Header.Get("x-grok-conv-id")) + bodyKey := xaiBodyPromptCacheKey(rawBody) + 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 == "" { + var err error + modifiedBody, bodyChanged, err = setXAIPromptCacheKey(rawBody, 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 +171,53 @@ func (i *PromptCachingInterceptor) interceptXAI(req *http.Request, meta llmproxy return resp, respMeta, rawRespBody, err } +func xaiAPIType(req *http.Request, meta llmproxy.BodyMetadata, rawBody []byte) 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) + } + } + } + return llmproxy.DetectAPIType(rawBody) +} + +func xaiBodyPromptCacheKey(rawBody []byte) string { + var body struct { + PromptCacheKey string `json:"prompt_cache_key"` + } + if err := json.Unmarshal(rawBody, &body); err != nil { + return "" + } + return strings.TrimSpace(body.PromptCacheKey) +} + +func setXAIPromptCacheKey(rawBody []byte, cacheKey string) ([]byte, bool, error) { + var body map[string]interface{} + if err := json.Unmarshal(rawBody, &body); err != nil { + return rawBody, false, err + } + if existing, ok := body["prompt_cache_key"].(string); ok && strings.TrimSpace(existing) != "" { + return rawBody, false, nil + } + body["prompt_cache_key"] = cacheKey + modified, err := json.Marshal(body) + if err != nil { + return rawBody, 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 +682,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 +698,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 +712,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..f45fc71 100644 --- a/interceptors/promptcaching_test.go +++ b/interceptors/promptcaching_test.go @@ -2088,6 +2088,185 @@ 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 TestPromptCachingInterceptor_FireworksCacheKeyHeader(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { sessionID := r.Header.Get(HeaderFireworksSessionAffinity) From be6d6df584d53703aa825c60e75b2bcce71210be Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Tue, 26 May 2026 08:29:31 -0500 Subject: [PATCH 2/4] Add Anthropic AI Gateway stream regression tests --- aigateway_regression_live_test.go | 220 ++++++++++++++++++++++++++++++ 1 file changed, 220 insertions(+) diff --git a/aigateway_regression_live_test.go b/aigateway_regression_live_test.go index 25e4df4..828da7f 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,213 @@ 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) + + body := `{"model":"` + model + `","stream":true,"max_tokens":64000,"thinking":{"type":"disabled"},"messages":[{"role":"user","content":[{"type":"text","text":"Reply with GENESIS_DRIVER_SMOKE_OK and nothing else."}]}]}` + req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader([]byte(body))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + rec := httptest.NewRecorder() + + router.ServeHTTP(rec, req) + + resp := rec.Result() + 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 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" + } + + body := map[string]any{ + "model": model, + "stream": true, + "max_tokens": 64_000, + "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."}, + }, + }, + }, + } + rawBody, err := json.Marshal(body) + 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) + } + 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 eventType string + var data string + for _, line := range strings.Split(block, "\n") { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "event:") { + eventType = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + } + if strings.HasPrefix(line, "data:") { + data = strings.TrimSpace(strings.TrimPrefix(line, "data:")) + } + } + if eventType != "" { + summary.EventTypes[eventType]++ + } + 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 +311,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 "" +} From 2110cd36cf3e9c205f50a0efb08aaa53373f18da Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Mon, 13 Jul 2026 13:21:00 -0500 Subject: [PATCH 3/4] Address review feedback for xAI caching and Anthropic live tests. Preserve request JSON numbers via RawMessage injection, parse the body once in the xAI caching path, and fold Anthropic AI Gateway stream regressions with corrected event counting and lower max_tokens. Co-authored-by: Cursor --- aigateway_regression_live_test.go | 11 +---- interceptors/promptcaching.go | 71 +++++++++++++++++++++--------- interceptors/promptcaching_test.go | 21 +++++++++ 3 files changed, 74 insertions(+), 29 deletions(-) diff --git a/aigateway_regression_live_test.go b/aigateway_regression_live_test.go index 828da7f..066a6bd 100644 --- a/aigateway_regression_live_test.go +++ b/aigateway_regression_live_test.go @@ -125,7 +125,7 @@ func TestLiveAnthropicMessagesStreamCompletes(t *testing.T) { ) router.RegisterProvider(provider) - body := `{"model":"` + model + `","stream":true,"max_tokens":64000,"thinking":{"type":"disabled"},"messages":[{"role":"user","content":[{"type":"text","text":"Reply with GENESIS_DRIVER_SMOKE_OK and nothing else."}]}]}` + body := `{"model":"` + model + `","stream":true,"max_tokens":512,"thinking":{"type":"disabled"},"messages":[{"role":"user","content":[{"type":"text","text":"Reply with GENESIS_DRIVER_SMOKE_OK and nothing else."}]}]}` req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader([]byte(body))) req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "text/event-stream") @@ -175,7 +175,7 @@ func TestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes(t *testing.T) { body := map[string]any{ "model": model, "stream": true, - "max_tokens": 64_000, + "max_tokens": 512, "thinking": map[string]any{ "type": "disabled", }, @@ -258,20 +258,13 @@ func summarizeAnthropicLiveStream(raw []byte) anthropicLiveStreamSummary { EventTypes: make(map[string]int), } for _, block := range strings.Split(string(raw), "\n\n") { - var eventType string var data string for _, line := range strings.Split(block, "\n") { line = strings.TrimSpace(line) - if strings.HasPrefix(line, "event:") { - eventType = strings.TrimSpace(strings.TrimPrefix(line, "event:")) - } if strings.HasPrefix(line, "data:") { data = strings.TrimSpace(strings.TrimPrefix(line, "data:")) } } - if eventType != "" { - summary.EventTypes[eventType]++ - } if data == "" { continue } diff --git a/interceptors/promptcaching.go b/interceptors/promptcaching.go index 83326ac..5d6ac61 100644 --- a/interceptors/promptcaching.go +++ b/interceptors/promptcaching.go @@ -118,9 +118,13 @@ 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) { - apiType := xaiAPIType(req, meta, rawBody) + fields, fieldsOK := parseJSONObjectRaw(rawBody) + apiType := xaiAPIType(req, meta, fields, fieldsOK) headerKey := strings.TrimSpace(req.Header.Get("x-grok-conv-id")) - bodyKey := xaiBodyPromptCacheKey(rawBody) + bodyKey := "" + if fieldsOK { + bodyKey = xaiBodyPromptCacheKeyFromFields(fields) + } cacheKey := headerKey if cacheKey == "" { cacheKey = bodyKey @@ -142,9 +146,9 @@ func (i *PromptCachingInterceptor) interceptXAI(req *http.Request, meta llmproxy if headerKey == "" { req.Header.Set("x-grok-conv-id", cacheKey) } - if apiType == llmproxy.APITypeResponses && bodyKey == "" { + if apiType == llmproxy.APITypeResponses && bodyKey == "" && fieldsOK { var err error - modifiedBody, bodyChanged, err = setXAIPromptCacheKey(rawBody, cacheKey) + modifiedBody, bodyChanged, err = setXAIPromptCacheKeyFromFields(fields, cacheKey) if err != nil { return next(req) } @@ -171,7 +175,15 @@ func (i *PromptCachingInterceptor) interceptXAI(req *http.Request, meta llmproxy return resp, respMeta, rawRespBody, err } -func xaiAPIType(req *http.Request, meta llmproxy.BodyMetadata, rawBody []byte) llmproxy.APIType { +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 @@ -189,31 +201,50 @@ func xaiAPIType(req *http.Request, meta llmproxy.BodyMetadata, rawBody []byte) l } } } - return llmproxy.DetectAPIType(rawBody) + 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 xaiBodyPromptCacheKey(rawBody []byte) string { - var body struct { - PromptCacheKey string `json:"prompt_cache_key"` +func xaiBodyPromptCacheKeyFromFields(fields map[string]json.RawMessage) string { + raw, ok := fields["prompt_cache_key"] + if !ok { + return "" } - if err := json.Unmarshal(rawBody, &body); err != nil { + var key string + if err := json.Unmarshal(raw, &key); err != nil { return "" } - return strings.TrimSpace(body.PromptCacheKey) + return strings.TrimSpace(key) } -func setXAIPromptCacheKey(rawBody []byte, cacheKey string) ([]byte, bool, error) { - var body map[string]interface{} - if err := json.Unmarshal(rawBody, &body); err != nil { - return rawBody, false, err +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 } - if existing, ok := body["prompt_cache_key"].(string); ok && strings.TrimSpace(existing) != "" { - return rawBody, false, nil + // 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 } - body["prompt_cache_key"] = cacheKey - modified, err := json.Marshal(body) + next["prompt_cache_key"] = json.RawMessage(keyRaw) + modified, err := json.Marshal(next) if err != nil { - return rawBody, false, err + return nil, false, err } return modified, true, nil } diff --git a/interceptors/promptcaching_test.go b/interceptors/promptcaching_test.go index f45fc71..83376bc 100644 --- a/interceptors/promptcaching_test.go +++ b/interceptors/promptcaching_test.go @@ -2267,6 +2267,27 @@ func TestDeriveCacheKeyFromPrefix_ResponsesInstructions(t *testing.T) { } } +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) From f0f3d856f2498c88a08d048bd23519113a2fb7c8 Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Mon, 13 Jul 2026 13:27:28 -0500 Subject: [PATCH 4/4] Deduplicate Anthropic live stream test helpers. Share request body marshaling and response read/assert logic across the local router and deployed AI Gateway Anthropic stream regressions. Co-authored-by: Cursor --- aigateway_regression_live_test.go | 64 +++++++++++++++---------------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/aigateway_regression_live_test.go b/aigateway_regression_live_test.go index 066a6bd..a0f612d 100644 --- a/aigateway_regression_live_test.go +++ b/aigateway_regression_live_test.go @@ -125,25 +125,17 @@ func TestLiveAnthropicMessagesStreamCompletes(t *testing.T) { ) router.RegisterProvider(provider) - body := `{"model":"` + model + `","stream":true,"max_tokens":512,"thinking":{"type":"disabled"},"messages":[{"role":"user","content":[{"type":"text","text":"Reply with GENESIS_DRIVER_SMOKE_OK and nothing else."}]}]}` - req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader([]byte(body))) + 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) - - resp := rec.Result() - 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) + readAndAssertLiveStream(t, rec.Result()) } func TestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes(t *testing.T) { @@ -172,23 +164,7 @@ func TestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes(t *testing.T) { model = "anthropic/claude-haiku-4-5-20251001" } - 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."}, - }, - }, - }, - } - rawBody, err := json.Marshal(body) + rawBody, err := anthropicLiveStreamRequestBody(model) if err != nil { t.Fatalf("marshal request: %v", err) } @@ -208,6 +184,31 @@ func TestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes(t *testing.T) { 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 { @@ -216,7 +217,6 @@ func TestLiveAgentuityAIGatewayAnthropicMessagesStreamCompletes(t *testing.T) { if resp.StatusCode < 200 || resp.StatusCode >= 300 { t.Fatalf("status %d: %s", resp.StatusCode, truncateLiveRegressionBody(raw)) } - assertAnthropicLiveStream(t, raw) }