From 346776be01e73a5b768b24f9a280fc4ed50569b0 Mon Sep 17 00:00:00 2001 From: KRATOS Date: Fri, 10 Jul 2026 15:58:10 +0530 Subject: [PATCH] fix(openai): omit prompt_cache_key for openai-compatible providers Interactive TUI always has a session ID, so every completion request forwarded OpenAI's prompt_cache_key. Strict openai-compatible gateways (e.g. NVIDIA NIM) reject unknown fields with a 400, while plain zero exec usually has no session and omits the field. Disable prompt_cache_key for ProviderKindOpenAICompatible; keep it for official OpenAI. ZERO_DISABLE_PROMPT_CACHE_KEY remains a global kill switch. Fixes #624 --- internal/providers/factory.go | 30 ++++---- internal/providers/factory_test.go | 52 ++++++++++++++ internal/providers/openai/provider.go | 79 ++++++++++++---------- internal/providers/openai/provider_test.go | 25 ++++++- internal/providers/openai/types.go | 4 +- 5 files changed, 141 insertions(+), 49 deletions(-) diff --git a/internal/providers/factory.go b/internal/providers/factory.go index 4c3f0c5bd..7aa4b4d8f 100644 --- a/internal/providers/factory.go +++ b/internal/providers/factory.go @@ -55,19 +55,25 @@ func New(profile config.ProviderProfile, options Options) (zeroruntime.Provider, switch resolved.providerKind { case config.ProviderKindOpenAI, config.ProviderKindOpenAICompatible: + // prompt_cache_key is an OpenAI-only chat-completions field. Strict + // openai-compatible gateways (NVIDIA NIM, etc.) reject it with a 400 + // instead of ignoring unknown parameters — so omit it for every + // openai-compatible profile. Official OpenAI keeps the field so + // multi-turn sessions still route to a cached-prefix replica. return openai.New(openai.Options{ - APIKey: profile.APIKey, - BaseURL: resolved.baseURL, - Model: resolved.apiModel, - AuthHeader: profile.AuthHeader, - AuthScheme: profile.AuthScheme, - AuthHeaderValue: profile.AuthHeaderValue, - CustomHeaders: profile.CustomHeaders, - OAuthResolver: options.OAuthResolver, - MaxTokens: resolved.maxOutputTokens, - HTTPClient: options.HTTPClient, - UserAgent: options.UserAgent, - ParseThinkTags: parseThinkTagsForProfile(profile, resolved), + APIKey: profile.APIKey, + BaseURL: resolved.baseURL, + Model: resolved.apiModel, + AuthHeader: profile.AuthHeader, + AuthScheme: profile.AuthScheme, + AuthHeaderValue: profile.AuthHeaderValue, + CustomHeaders: profile.CustomHeaders, + OAuthResolver: options.OAuthResolver, + MaxTokens: resolved.maxOutputTokens, + HTTPClient: options.HTTPClient, + UserAgent: options.UserAgent, + ParseThinkTags: parseThinkTagsForProfile(profile, resolved), + DisablePromptCacheKey: resolved.providerKind == config.ProviderKindOpenAICompatible, }) case config.ProviderKindAnthropic, config.ProviderKindAnthropicCompat: return anthropic.New(anthropic.Options{ diff --git a/internal/providers/factory_test.go b/internal/providers/factory_test.go index b9ef0d6e5..9fb9fed87 100644 --- a/internal/providers/factory_test.go +++ b/internal/providers/factory_test.go @@ -159,6 +159,58 @@ func TestNewSupportsOpenAIProviderKind(t *testing.T) { } } +// TestPromptCacheKeyOnlyOnOfficialOpenAI locks in #624: session-backed TUI +// turns always carry a PromptCacheKey, but openai-compatible gateways (NVIDIA +// NIM, strict local proxies) reject the OpenAI-only prompt_cache_key field. +// The factory must omit it for openai-compatible profiles while still +// forwarding it for official OpenAI so multi-turn cache routing stays intact. +func TestPromptCacheKeyOnlyOnOfficialOpenAI(t *testing.T) { + requestWithSession := zeroruntime.CompletionRequest{ + Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hello"}}, + PromptCacheKey: "sess_tui_123", + } + + for _, tc := range []struct { + name string + kind config.ProviderKind + wantCacheKey bool + }{ + {name: "openai", kind: config.ProviderKindOpenAI, wantCacheKey: true}, + {name: "openai-compatible", kind: config.ProviderKindOpenAICompatible, wantCacheKey: false}, + } { + t.Run(tc.name, func(t *testing.T) { + transport := &captureTransport{responseBody: "data: [DONE]\n\n"} + provider, err := New(config.ProviderProfile{ + Name: "test", + ProviderKind: tc.kind, + BaseURL: "https://provider.example/v1", + APIKey: "sk-test", + Model: "test-model", + }, Options{HTTPClient: &http.Client{Transport: transport}}) + if err != nil { + t.Fatalf("New() error = %v", err) + } + stream, err := provider.StreamCompletion(context.Background(), requestWithSession) + if err != nil { + t.Fatalf("StreamCompletion() error = %v", err) + } + for range stream { + } + var body map[string]any + if err := json.NewDecoder(transport.body()).Decode(&body); err != nil { + t.Fatalf("decode request body: %v", err) + } + _, hasKey := body["prompt_cache_key"] + if hasKey != tc.wantCacheKey { + t.Fatalf("prompt_cache_key present = %v, want %v; body = %#v", hasKey, tc.wantCacheKey, body) + } + if tc.wantCacheKey && body["prompt_cache_key"] != "sess_tui_123" { + t.Fatalf("prompt_cache_key = %#v, want sess_tui_123", body["prompt_cache_key"]) + } + }) + } +} + func TestParseThinkTagsForProfileUsesConservativeDefaultsAndOverride(t *testing.T) { openAICompatible := resolvedProfile{providerKind: config.ProviderKindOpenAICompatible, apiModel: "qwen3-coder:480b"} if !parseThinkTagsForProfile(config.ProviderProfile{}, openAICompatible) { diff --git a/internal/providers/openai/provider.go b/internal/providers/openai/provider.go index 1e19eab01..8bf2bc5ba 100644 --- a/internal/providers/openai/provider.go +++ b/internal/providers/openai/provider.go @@ -58,25 +58,33 @@ type Options struct { // "originator" value. It is also called on the 401-refresh retry, so any // per-request state must be re-derivable from the live request. SetRequestExtra func(*http.Request) + // DisablePromptCacheKey omits OpenAI's prompt_cache_key even when the + // caller supplies a session identity. Used for openai-compatible gateways + // (NVIDIA NIM, strict local proxies, …) that validate and reject unknown + // request fields instead of ignoring them. Official OpenAI keeps the field + // enabled; the ZERO_DISABLE_PROMPT_CACHE_KEY env kill switch still applies + // on top for any endpoint. + DisablePromptCacheKey bool } // Provider streams completions from an OpenAI-compatible chat completions API. type Provider struct { - apiKey string - baseURL string - endpoint string - model string - authHeader string - authScheme string - authHeaderValue string - customHeaders map[string]string - oauthResolver providerio.TokenResolver - maxTokens int - httpClient *http.Client - userAgent string - streamIdleTimeout time.Duration - parseThinkTags bool - setRequestExtra func(*http.Request) + apiKey string + baseURL string + endpoint string + model string + authHeader string + authScheme string + authHeaderValue string + customHeaders map[string]string + oauthResolver providerio.TokenResolver + maxTokens int + httpClient *http.Client + userAgent string + streamIdleTimeout time.Duration + parseThinkTags bool + setRequestExtra func(*http.Request) + disablePromptCacheKey bool } // New creates an OpenAI-compatible provider. @@ -116,21 +124,22 @@ func New(options Options) (*Provider, error) { } return &Provider{ - apiKey: options.APIKey, - baseURL: baseURL, - endpoint: endpoint, - model: model, - authHeader: strings.TrimSpace(options.AuthHeader), - authScheme: strings.TrimSpace(options.AuthScheme), - authHeaderValue: strings.TrimSpace(options.AuthHeaderValue), - customHeaders: providerio.CopyHeaders(options.CustomHeaders), - oauthResolver: options.OAuthResolver, - maxTokens: maxTokens, - httpClient: httpClient, - userAgent: options.UserAgent, - streamIdleTimeout: providerio.ResolveStreamIdleTimeout(options.StreamIdleTimeout), - parseThinkTags: options.ParseThinkTags, - setRequestExtra: options.SetRequestExtra, + apiKey: options.APIKey, + baseURL: baseURL, + endpoint: endpoint, + model: model, + authHeader: strings.TrimSpace(options.AuthHeader), + authScheme: strings.TrimSpace(options.AuthScheme), + authHeaderValue: strings.TrimSpace(options.AuthHeaderValue), + customHeaders: providerio.CopyHeaders(options.CustomHeaders), + oauthResolver: options.OAuthResolver, + maxTokens: maxTokens, + httpClient: httpClient, + userAgent: options.UserAgent, + streamIdleTimeout: providerio.ResolveStreamIdleTimeout(options.StreamIdleTimeout), + parseThinkTags: options.ParseThinkTags, + setRequestExtra: options.SetRequestExtra, + disablePromptCacheKey: options.DisablePromptCacheKey, }, nil } @@ -458,10 +467,12 @@ func (provider *Provider) openAIRequest(request zeroruntime.CompletionRequest) c if effort := openAIReasoningEffort(request.ReasoningEffort); effort != "" { mapped.ReasoningEffort = effort } - // prompt_cache_key is a documented OpenAI parameter; compatible servers - // ignore unknown fields, but a strict endpoint that rejects it can be - // accommodated with ZERO_DISABLE_PROMPT_CACHE_KEY=1. - if key := strings.TrimSpace(request.PromptCacheKey); key != "" && !promptCacheKeyDisabled() { + // prompt_cache_key is a documented OpenAI parameter for server-side prefix + // cache routing. Official OpenAI accepts it; many openai-compatible + // gateways (NVIDIA NIM, strict local proxies) reject unknown fields with a + // 400. Those providers are constructed with DisablePromptCacheKey, and any + // endpoint can still force-omit via ZERO_DISABLE_PROMPT_CACHE_KEY=1. + if key := strings.TrimSpace(request.PromptCacheKey); key != "" && !provider.disablePromptCacheKey && !promptCacheKeyDisabled() { mapped.PromptCacheKey = key } if len(request.Tools) > 0 { diff --git a/internal/providers/openai/provider_test.go b/internal/providers/openai/provider_test.go index 2a473f51a..948b95aaf 100644 --- a/internal/providers/openai/provider_test.go +++ b/internal/providers/openai/provider_test.go @@ -1300,8 +1300,9 @@ func TestOpenAIRequestEmptyContentHandling(t *testing.T) { // TestOpenAIRequestPromptCacheKey locks in prompt_cache_key forwarding: a // session-carrying request serializes the key so the backend can route to a // replica holding the cached prefix, a keyless request omits the field -// entirely (strict servers see byte-identical requests to before), and -// ZERO_DISABLE_PROMPT_CACHE_KEY suppresses it for endpoints that reject it. +// entirely (strict servers see byte-identical requests to before), +// DisablePromptCacheKey (used for openai-compatible gateways) suppresses it, +// and ZERO_DISABLE_PROMPT_CACHE_KEY is the env kill switch for any endpoint. func TestOpenAIRequestPromptCacheKey(t *testing.T) { provider, err := New(Options{Model: "gpt-test"}) if err != nil { @@ -1332,6 +1333,26 @@ func TestOpenAIRequestPromptCacheKey(t *testing.T) { t.Fatalf("keyless request must omit prompt_cache_key: %s", data) } + // openai-compatible providers are constructed with DisablePromptCacheKey so + // strict gateways (NVIDIA NIM, …) never see the OpenAI-only field. + compat, err := New(Options{Model: "gpt-test", DisablePromptCacheKey: true}) + if err != nil { + t.Fatalf("New(DisablePromptCacheKey) returned error: %v", err) + } + req = compat.openAIRequest(zeroruntime.CompletionRequest{ + Messages: messages, + PromptCacheKey: "sess_123", + }) + if req.PromptCacheKey != "" { + t.Fatalf("DisablePromptCacheKey ignored; PromptCacheKey = %q", req.PromptCacheKey) + } + if data, err = json.Marshal(req); err != nil { + t.Fatalf("marshal: %v", err) + } + if strings.Contains(string(data), "prompt_cache_key") { + t.Fatalf("compatible provider must omit prompt_cache_key: %s", data) + } + t.Setenv("ZERO_DISABLE_PROMPT_CACHE_KEY", "1") req = provider.openAIRequest(zeroruntime.CompletionRequest{ Messages: messages, diff --git a/internal/providers/openai/types.go b/internal/providers/openai/types.go index 4035c1d35..568a3439f 100644 --- a/internal/providers/openai/types.go +++ b/internal/providers/openai/types.go @@ -11,7 +11,9 @@ type chatCompletionRequest struct { // PromptCacheKey asks the backend to route the request to a replica that // already holds this conversation's prefix in its prompt cache (the OpenAI // `prompt_cache_key` parameter). Omitted when the caller carries no session - // identity or when ZERO_DISABLE_PROMPT_CACHE_KEY is set. + // identity, when the provider was constructed with DisablePromptCacheKey + // (openai-compatible gateways that reject unknown fields), or when + // ZERO_DISABLE_PROMPT_CACHE_KEY is set. PromptCacheKey string `json:"prompt_cache_key,omitempty"` }