Skip to content

Commit c2b714c

Browse files
committed
feat(openai): forward prompt_cache_key for server-side prefix cache routing
Send the session ID as OpenAI's prompt_cache_key on every completion request so consecutive turns route to a backend replica that already holds the conversation prefix in its prompt cache, instead of re-billing the full ~11k-token system+tools prefix each turn on OpenAI-compatible endpoints (which get no cache_control breakpoints, unlike Anthropic). Keyless requests serialize byte-identically to before, and ZERO_DISABLE_PROMPT_CACHE_KEY=1 suppresses the field for strict endpoints that reject unknown parameters. The Codex (Responses API) flavor is deliberately untouched: the ChatGPT backend is a separate, stricter surface.
1 parent 3c81fea commit c2b714c

5 files changed

Lines changed: 70 additions & 0 deletions

File tree

‎internal/agent/loop.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
173173
Messages: copyMessages(messages),
174174
Tools: exposed,
175175
ReasoningEffort: options.ReasoningEffort,
176+
PromptCacheKey: options.SessionID,
176177
}
177178

178179
// 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)
208209
Messages: copyMessages(messages),
209210
Tools: exposed,
210211
ReasoningEffort: options.ReasoningEffort,
212+
PromptCacheKey: options.SessionID,
211213
}
212214
// Pre-content connect after a context-limit compaction: route through the
213215
// 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)
268270
Messages: copyMessages(messages),
269271
Tools: exposed,
270272
ReasoningEffort: options.ReasoningEffort,
273+
PromptCacheKey: options.SessionID,
271274
}
272275
retryStream, retryStreamErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options))
273276
if retryStreamErr != nil {
@@ -326,6 +329,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
326329
Messages: copyMessages(messages),
327330
Tools: exposed,
328331
ReasoningEffort: options.ReasoningEffort,
332+
PromptCacheKey: options.SessionID,
329333
}
330334
retryStream, retryErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options))
331335
if retryErr != nil {
@@ -705,6 +709,7 @@ func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, messages [
705709
stream, err := streamWithReconnect(ctx, provider, zeroruntime.CompletionRequest{
706710
Messages: copyMessages(finalMessages),
707711
ReasoningEffort: options.ReasoningEffort,
712+
PromptCacheKey: options.SessionID,
708713
}, reconnectNoticeFor(options))
709714
if err != nil {
710715
return "", messages, ""

‎internal/providers/openai/provider.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"io"
1010
"net/http"
1111
"net/url"
12+
"os"
1213
"strings"
1314
"time"
1415

@@ -421,6 +422,13 @@ func (provider *Provider) openAIRequest(request zeroruntime.CompletionRequest) c
421422
if effort := openAIReasoningEffort(request.ReasoningEffort); effort != "" {
422423
mapped.ReasoningEffort = effort
423424
}
425+
// prompt_cache_key is a documented OpenAI parameter; compatible servers
426+
// ignore unknown fields, but a strict endpoint that rejects it can be
427+
// accommodated with ZERO_DISABLE_PROMPT_CACHE_KEY=1.
428+
if key := strings.TrimSpace(request.PromptCacheKey); key != "" &&
429+
os.Getenv("ZERO_DISABLE_PROMPT_CACHE_KEY") == "" {
430+
mapped.PromptCacheKey = key
431+
}
424432
if len(request.Tools) > 0 {
425433
mapped.Tools = make([]toolDefinition, 0, len(request.Tools))
426434
for _, tool := range request.Tools {

‎internal/providers/openai/provider_test.go‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1259,3 +1259,48 @@ func TestOpenAIRequestEmptyContentHandling(t *testing.T) {
12591259
t.Fatalf("assistant tool calls dropped: %#v", req.Messages[1])
12601260
}
12611261
}
1262+
1263+
// TestOpenAIRequestPromptCacheKey locks in prompt_cache_key forwarding: a
1264+
// session-carrying request serializes the key so the backend can route to a
1265+
// replica holding the cached prefix, a keyless request omits the field
1266+
// entirely (strict servers see byte-identical requests to before), and
1267+
// ZERO_DISABLE_PROMPT_CACHE_KEY suppresses it for endpoints that reject it.
1268+
func TestOpenAIRequestPromptCacheKey(t *testing.T) {
1269+
provider, err := New(Options{Model: "gpt-test"})
1270+
if err != nil {
1271+
t.Fatalf("New returned error: %v", err)
1272+
}
1273+
messages := []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hi"}}
1274+
1275+
req := provider.openAIRequest(zeroruntime.CompletionRequest{
1276+
Messages: messages,
1277+
PromptCacheKey: "sess_123",
1278+
})
1279+
if req.PromptCacheKey != "sess_123" {
1280+
t.Fatalf("PromptCacheKey = %q, want sess_123", req.PromptCacheKey)
1281+
}
1282+
data, err := json.Marshal(req)
1283+
if err != nil {
1284+
t.Fatalf("marshal: %v", err)
1285+
}
1286+
if !strings.Contains(string(data), `"prompt_cache_key":"sess_123"`) {
1287+
t.Fatalf("prompt_cache_key not serialized: %s", data)
1288+
}
1289+
1290+
req = provider.openAIRequest(zeroruntime.CompletionRequest{Messages: messages})
1291+
if data, err = json.Marshal(req); err != nil {
1292+
t.Fatalf("marshal: %v", err)
1293+
}
1294+
if strings.Contains(string(data), "prompt_cache_key") {
1295+
t.Fatalf("keyless request must omit prompt_cache_key: %s", data)
1296+
}
1297+
1298+
t.Setenv("ZERO_DISABLE_PROMPT_CACHE_KEY", "1")
1299+
req = provider.openAIRequest(zeroruntime.CompletionRequest{
1300+
Messages: messages,
1301+
PromptCacheKey: "sess_123",
1302+
})
1303+
if req.PromptCacheKey != "" {
1304+
t.Fatalf("kill switch ignored; PromptCacheKey = %q", req.PromptCacheKey)
1305+
}
1306+
}

‎internal/providers/openai/types.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,11 @@ type chatCompletionRequest struct {
88
ReasoningEffort string `json:"reasoning_effort,omitempty"`
99
Stream bool `json:"stream"`
1010
StreamOptions *streamOptions `json:"stream_options,omitempty"`
11+
// PromptCacheKey asks the backend to route the request to a replica that
12+
// already holds this conversation's prefix in its prompt cache (the OpenAI
13+
// `prompt_cache_key` parameter). Omitted when the caller carries no session
14+
// identity or when ZERO_DISABLE_PROMPT_CACHE_KEY is set.
15+
PromptCacheKey string `json:"prompt_cache_key,omitempty"`
1116
}
1217

1318
// streamOptions requests the final usage chunk on a streaming response. Without

‎internal/zeroruntime/types.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,13 @@ type CompletionRequest struct {
203203
// Anthropic/Gemini thinking budgets) and ignores it for models that do not
204204
// support reasoning. Empty means "let the provider decide".
205205
ReasoningEffort string
206+
// PromptCacheKey, when non-empty, is an opaque stable identifier for the
207+
// conversation (the session ID). Providers with server-side prefix-cache
208+
// routing forward it — OpenAI `prompt_cache_key` — so consecutive requests
209+
// land on a replica that already holds the cached prompt prefix instead of
210+
// re-billing the full prefix each turn. Providers without an equivalent
211+
// ignore it.
212+
PromptCacheKey string
206213
}
207214

208215
// Provider streams normalized completion events for one request.

0 commit comments

Comments
 (0)