Compare commits

...
Author SHA1 Message Date
KRATOS 346776be01 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
2026-07-10 15:58:10 +05:30
5 changed files with 141 additions and 49 deletions

No files matched your search

+18 -12
View File
@@ -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{
+52
View File
@@ -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) {
+45 -34
View File
@@ -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 {
+23 -2
View File
@@ -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,
+3 -1
View File
@@ -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"`
}