diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 3cc83774..6583fb0d 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -173,6 +173,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) Messages: copyMessages(messages), Tools: exposed, ReasoningEffort: options.ReasoningEffort, + PromptCacheKey: options.SessionID, } // Report the per-category context budget for this turn so a surface can @@ -208,6 +209,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) Messages: copyMessages(messages), Tools: exposed, ReasoningEffort: options.ReasoningEffort, + PromptCacheKey: options.SessionID, } // Pre-content connect after a context-limit compaction: route through the // reconnect helper so a transient upstream hiccup here doesn't fail the @@ -268,6 +270,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) Messages: copyMessages(messages), Tools: exposed, ReasoningEffort: options.ReasoningEffort, + PromptCacheKey: options.SessionID, } retryStream, retryStreamErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options)) if retryStreamErr != nil { @@ -326,6 +329,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) Messages: copyMessages(messages), Tools: exposed, ReasoningEffort: options.ReasoningEffort, + PromptCacheKey: options.SessionID, } retryStream, retryErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options)) if retryErr != nil { @@ -705,6 +709,7 @@ func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, messages [ stream, err := streamWithReconnect(ctx, provider, zeroruntime.CompletionRequest{ Messages: copyMessages(finalMessages), ReasoningEffort: options.ReasoningEffort, + PromptCacheKey: options.SessionID, }, reconnectNoticeFor(options)) if err != nil { return "", messages, "" diff --git a/internal/providers/openai/provider.go b/internal/providers/openai/provider.go index bc738bb2..06420e7a 100644 --- a/internal/providers/openai/provider.go +++ b/internal/providers/openai/provider.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "net/url" + "os" "strings" "time" @@ -421,6 +422,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() { + mapped.PromptCacheKey = key + } if len(request.Tools) > 0 { mapped.Tools = make([]toolDefinition, 0, len(request.Tools)) for _, tool := range request.Tools { @@ -437,6 +444,15 @@ func (provider *Provider) openAIRequest(request zeroruntime.CompletionRequest) c return mapped } +// promptCacheKeyDisabled reports whether the ZERO_DISABLE_PROMPT_CACHE_KEY +// kill switch is set to a truthy value. "0" and "false" (any case) are +// no-ops, matching how ZERO_FORMAT_ON_WRITE parses boolean flags, so an +// explicitly-disabled toggle never flips the behavior it names. +func promptCacheKeyDisabled() bool { + value := strings.TrimSpace(os.Getenv("ZERO_DISABLE_PROMPT_CACHE_KEY")) + return value != "" && value != "0" && !strings.EqualFold(value, "false") +} + // openAIReasoningEffort normalizes a requested effort to a value the OpenAI chat // completions API accepts, or "" to omit the field. "none" (and anything else) // is dropped rather than risking a 400 on an unrecognized enum. diff --git a/internal/providers/openai/provider_test.go b/internal/providers/openai/provider_test.go index 5c76a8d8..0fc98055 100644 --- a/internal/providers/openai/provider_test.go +++ b/internal/providers/openai/provider_test.go @@ -1259,3 +1259,61 @@ func TestOpenAIRequestEmptyContentHandling(t *testing.T) { t.Fatalf("assistant tool calls dropped: %#v", req.Messages[1]) } } + +// 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. +func TestOpenAIRequestPromptCacheKey(t *testing.T) { + provider, err := New(Options{Model: "gpt-test"}) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + messages := []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hi"}} + + req := provider.openAIRequest(zeroruntime.CompletionRequest{ + Messages: messages, + PromptCacheKey: "sess_123", + }) + if req.PromptCacheKey != "sess_123" { + t.Fatalf("PromptCacheKey = %q, want sess_123", req.PromptCacheKey) + } + data, err := json.Marshal(req) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if !strings.Contains(string(data), `"prompt_cache_key":"sess_123"`) { + t.Fatalf("prompt_cache_key not serialized: %s", data) + } + + req = provider.openAIRequest(zeroruntime.CompletionRequest{Messages: messages}) + if data, err = json.Marshal(req); err != nil { + t.Fatalf("marshal: %v", err) + } + if strings.Contains(string(data), "prompt_cache_key") { + t.Fatalf("keyless request must omit prompt_cache_key: %s", data) + } + + t.Setenv("ZERO_DISABLE_PROMPT_CACHE_KEY", "1") + req = provider.openAIRequest(zeroruntime.CompletionRequest{ + Messages: messages, + PromptCacheKey: "sess_123", + }) + if req.PromptCacheKey != "" { + t.Fatalf("kill switch ignored; PromptCacheKey = %q", req.PromptCacheKey) + } + + // Explicitly-falsy kill switch values must NOT disable forwarding — only + // truthy values flip the toggle (same parsing as ZERO_FORMAT_ON_WRITE). + for _, value := range []string{"0", "false", "FALSE"} { + t.Setenv("ZERO_DISABLE_PROMPT_CACHE_KEY", value) + req = provider.openAIRequest(zeroruntime.CompletionRequest{ + Messages: messages, + PromptCacheKey: "sess_123", + }) + if req.PromptCacheKey != "sess_123" { + t.Fatalf("ZERO_DISABLE_PROMPT_CACHE_KEY=%q must be a no-op; PromptCacheKey = %q", value, req.PromptCacheKey) + } + } +} diff --git a/internal/providers/openai/types.go b/internal/providers/openai/types.go index 54511a6f..4035c1d3 100644 --- a/internal/providers/openai/types.go +++ b/internal/providers/openai/types.go @@ -8,6 +8,11 @@ type chatCompletionRequest struct { ReasoningEffort string `json:"reasoning_effort,omitempty"` Stream bool `json:"stream"` StreamOptions *streamOptions `json:"stream_options,omitempty"` + // 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. + PromptCacheKey string `json:"prompt_cache_key,omitempty"` } // streamOptions requests the final usage chunk on a streaming response. Without diff --git a/internal/zeroruntime/types.go b/internal/zeroruntime/types.go index b2487b68..514f9791 100644 --- a/internal/zeroruntime/types.go +++ b/internal/zeroruntime/types.go @@ -203,6 +203,13 @@ type CompletionRequest struct { // Anthropic/Gemini thinking budgets) and ignores it for models that do not // support reasoning. Empty means "let the provider decide". ReasoningEffort string + // PromptCacheKey, when non-empty, is an opaque stable identifier for the + // conversation (the session ID). Providers with server-side prefix-cache + // routing forward it — OpenAI `prompt_cache_key` — so consecutive requests + // land on a replica that already holds the cached prompt prefix instead of + // re-billing the full prefix each turn. Providers without an equivalent + // ignore it. + PromptCacheKey string } // Provider streams normalized completion events for one request.