From 968745bd16ce03e9cd3c1ca0ead3eee34942abbf Mon Sep 17 00:00:00 2001 From: Zezae Oh Date: Tue, 30 Jun 2026 20:25:28 +0900 Subject: [PATCH] feat(providers): enable Codex prompt cache controls (#1310) --- internal/agent/loop_pipeline_callbacks.go | 41 +++++++++++++++- .../agent/loop_pipeline_callbacks_test.go | 48 +++++++++++++++++++ internal/providers/adapter_codex.go | 2 +- internal/providers/adapter_codex_test.go | 32 +++++++++++++ internal/providers/capabilities.go | 2 +- internal/providers/chatgpt_oauth_router.go | 21 ++++++++ internal/providers/codex.go | 2 +- internal/providers/codex_build.go | 7 +++ internal/providers/codex_test.go | 26 ++++++++++ 9 files changed, 176 insertions(+), 5 deletions(-) diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index 35eb36f6..bf7d1454 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -2,6 +2,8 @@ package agent import ( "context" + "crypto/sha256" + "encoding/hex" "fmt" "log/slog" "strings" @@ -328,8 +330,12 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c chatReq.Options[providers.OptPeerKind] = req.PeerKind chatReq.Options[providers.OptLocalKey] = req.LocalKey chatReq.Options[providers.OptWorkspace] = tools.ToolWorkspaceFromCtx(ctx) - if tid := store.TenantIDFromContext(ctx); tid != uuid.Nil { - chatReq.Options[providers.OptTenantID] = tid.String() + tenantID := store.TenantIDFromContext(ctx) + if tenantID != uuid.Nil { + chatReq.Options[providers.OptTenantID] = tenantID.String() + } + if supportsPromptCacheParams(provider) { + setDefaultPromptCacheOptions(chatReq.Options, tenantID, l.agentUUID, provider.Name(), req.SessionKey) } // Reasoning decision: resolve effort level for thinking models (o3, DeepSeek-R1, Kimi). @@ -697,3 +703,34 @@ func (l *Loop) reserveLLMUsageFor(ctx context.Context, req *RunRequest, iteratio MaxOutputTokens: l.maxOutputTokensFromRequest(chatReq), }) } + +func supportsPromptCacheParams(provider providers.Provider) bool { + switch provider.(type) { + case *providers.CodexProvider, *providers.ChatGPTOAuthRouter: + return true + default: + return false + } +} + +func setDefaultPromptCacheOptions(opts map[string]any, tenantID, agentID uuid.UUID, providerName, sessionKey string) { + if opts == nil { + return + } + if _, ok := opts[providers.OptPromptCacheKey]; !ok { + opts[providers.OptPromptCacheKey] = defaultPromptCacheKey(tenantID, agentID, providerName, sessionKey) + } + if _, ok := opts[providers.OptPromptCacheRetention]; !ok { + opts[providers.OptPromptCacheRetention] = "24h" + } +} + +func defaultPromptCacheKey(tenantID, agentID uuid.UUID, providerName, sessionKey string) string { + h := sha256.Sum256([]byte(strings.Join([]string{ + tenantID.String(), + agentID.String(), + providerName, + sessionKey, + }, "\x00"))) + return "goclaw/" + hex.EncodeToString(h[:16]) +} diff --git a/internal/agent/loop_pipeline_callbacks_test.go b/internal/agent/loop_pipeline_callbacks_test.go index 2e7617df..18d374b3 100644 --- a/internal/agent/loop_pipeline_callbacks_test.go +++ b/internal/agent/loop_pipeline_callbacks_test.go @@ -2,8 +2,11 @@ package agent import ( "context" + "strings" "testing" + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/pipeline" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/pkg/protocol" @@ -54,3 +57,48 @@ func TestMakeCallLLM_StreamsFinalThinkingWhenNoThinkingChunkArrives(t *testing.T t.Fatalf("thinking payload = %+v", thinking[0].Payload) } } + +func TestPromptCacheOptionsHelpers(t *testing.T) { + tenantID := uuid.MustParse("11111111-1111-1111-1111-111111111111") + agentID := uuid.MustParse("22222222-2222-2222-2222-222222222222") + + key1 := defaultPromptCacheKey(tenantID, agentID, "codex", "session-a") + key2 := defaultPromptCacheKey(tenantID, agentID, "codex", "session-a") + key3 := defaultPromptCacheKey(tenantID, agentID, "codex", "session-b") + if key1 != key2 { + t.Fatalf("defaultPromptCacheKey not stable: %q != %q", key1, key2) + } + if key1 == key3 { + t.Fatal("defaultPromptCacheKey should vary by session") + } + if !strings.HasPrefix(key1, "goclaw/") { + t.Fatalf("defaultPromptCacheKey = %q, want goclaw/ prefix", key1) + } + + opts := map[string]any{} + setDefaultPromptCacheOptions(opts, tenantID, agentID, "codex", "session-a") + if opts[providers.OptPromptCacheKey] != key1 { + t.Fatalf("prompt cache key = %v, want %s", opts[providers.OptPromptCacheKey], key1) + } + if opts[providers.OptPromptCacheRetention] != "24h" { + t.Fatalf("prompt cache retention = %v, want 24h", opts[providers.OptPromptCacheRetention]) + } + + opts = map[string]any{ + providers.OptPromptCacheKey: "custom-key", + providers.OptPromptCacheRetention: "in_memory", + } + setDefaultPromptCacheOptions(opts, tenantID, agentID, "codex", "session-a") + if opts[providers.OptPromptCacheKey] != "custom-key" || opts[providers.OptPromptCacheRetention] != "in_memory" { + t.Fatalf("custom prompt cache options were overwritten: %+v", opts) + } +} + +func TestSupportsPromptCacheParams(t *testing.T) { + if !supportsPromptCacheParams(providers.NewCodexProvider("codex", nil, "", "")) { + t.Fatal("CodexProvider should support prompt cache params") + } + if supportsPromptCacheParams(finalThinkingStreamProvider{}) { + t.Fatal("generic provider should not support prompt cache params") + } +} diff --git a/internal/providers/adapter_codex.go b/internal/providers/adapter_codex.go index a77ec01f..674b1a87 100644 --- a/internal/providers/adapter_codex.go +++ b/internal/providers/adapter_codex.go @@ -51,7 +51,7 @@ func (a *CodexAdapter) Capabilities() ProviderCapabilities { StreamWithTools: true, Thinking: true, Vision: true, - CacheControl: false, + CacheControl: true, ImageGeneration: true, // Codex (OpenAI Responses API) supports native image_generation tool MaxContextWindow: 1_050_000, TokenizerID: "o200k_base", diff --git a/internal/providers/adapter_codex_test.go b/internal/providers/adapter_codex_test.go index fbb5d3a9..88b7c6ba 100644 --- a/internal/providers/adapter_codex_test.go +++ b/internal/providers/adapter_codex_test.go @@ -386,3 +386,35 @@ func TestCodexAdapter_FromResponse_PromptTokensDetailsAlias(t *testing.T) { t.Fatal("PromptTokensIncludeCachedSegments = false, want true") } } + +func TestCodexAdapter_ToRequestPromptCacheControls(t *testing.T) { + a, _ := NewCodexAdapter(ProviderConfig{}) + body, _, err := a.ToRequest(ChatRequest{ + Messages: []Message{{Role: "user", Content: "hi"}}, + Options: map[string]any{ + OptPromptCacheKey: "tenant/agent/session", + OptPromptCacheRetention: "1h", + }, + }) + if err != nil { + t.Fatalf("ToRequest: %v", err) + } + + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("decode body: %v", err) + } + if got := payload["prompt_cache_key"]; got != "tenant/agent/session" { + t.Fatalf("prompt_cache_key = %v, want tenant/agent/session", got) + } + if got := payload["prompt_cache_retention"]; got != "1h" { + t.Fatalf("prompt_cache_retention = %v, want 1h", got) + } +} + +func TestCodexAdapter_CapabilitiesCacheControl(t *testing.T) { + a, _ := NewCodexAdapter(ProviderConfig{}) + if !a.Capabilities().CacheControl { + t.Fatal("CodexAdapter Capabilities().CacheControl = false, want true") + } +} diff --git a/internal/providers/capabilities.go b/internal/providers/capabilities.go index 2c6bebdf..8ae37790 100644 --- a/internal/providers/capabilities.go +++ b/internal/providers/capabilities.go @@ -10,7 +10,7 @@ type ProviderCapabilities struct { StreamWithTools bool // can stream while tool calls are in-flight Thinking bool // supports extended thinking / reasoning Vision bool // supports image inputs - CacheControl bool // supports cache_control blocks (Anthropic) + CacheControl bool // supports provider prompt cache controls ImageGeneration bool // supports native image_generation tool (Codex/OpenAI Responses API) MaxContextWindow int // default context window for default model TokenizerID string // for tokencount package mapping diff --git a/internal/providers/chatgpt_oauth_router.go b/internal/providers/chatgpt_oauth_router.go index 540777e1..cbda091e 100644 --- a/internal/providers/chatgpt_oauth_router.go +++ b/internal/providers/chatgpt_oauth_router.go @@ -70,6 +70,27 @@ func (p *ChatGPTOAuthRouter) DefaultModel() string { func (p *ChatGPTOAuthRouter) SupportsThinking() bool { return true } +// Capabilities reports the Codex capability set exposed by the routed OAuth pool. +func (p *ChatGPTOAuthRouter) Capabilities() ProviderCapabilities { + selection, err := p.orderedProviders(context.Background(), chatGPTOAuthModalityChat, false) + if err == nil && len(selection) > 0 { + if ca, ok := selection[0].(CapabilitiesAware); ok { + return ca.Capabilities() + } + } + return ProviderCapabilities{ + Streaming: true, + ToolCalling: true, + StreamWithTools: true, + Thinking: true, + Vision: true, + CacheControl: true, + ImageGeneration: true, + MaxContextWindow: 1_050_000, + TokenizerID: "o200k_base", + } +} + func (p *ChatGPTOAuthRouter) HasRegisteredProviders() bool { return len(p.registeredProviders()) > 0 } diff --git a/internal/providers/codex.go b/internal/providers/codex.go index 9e182a43..717a7128 100644 --- a/internal/providers/codex.go +++ b/internal/providers/codex.go @@ -79,7 +79,7 @@ func (p *CodexProvider) Capabilities() ProviderCapabilities { StreamWithTools: true, Thinking: true, Vision: true, - CacheControl: false, + CacheControl: true, ImageGeneration: true, // Codex (OpenAI Responses API) supports native image_generation tool MaxContextWindow: 1_050_000, TokenizerID: "o200k_base", diff --git a/internal/providers/codex_build.go b/internal/providers/codex_build.go index 4d9d4304..4c296802 100644 --- a/internal/providers/codex_build.go +++ b/internal/providers/codex_build.go @@ -140,6 +140,13 @@ func (p *CodexProvider) buildRequestBody(req ChatRequest, stream bool) map[strin body["reasoning"] = map[string]any{"effort": level} } + if cacheKey, ok := req.Options[OptPromptCacheKey]; ok { + body["prompt_cache_key"] = cacheKey + } + if retention, ok := req.Options[OptPromptCacheRetention]; ok { + body["prompt_cache_retention"] = retention + } + return body } diff --git a/internal/providers/codex_test.go b/internal/providers/codex_test.go index 54899ecb..8d69aa67 100644 --- a/internal/providers/codex_test.go +++ b/internal/providers/codex_test.go @@ -1128,3 +1128,29 @@ func TestCodexBuildRequestBody_NilFunction_HandlesGracefully(t *testing.T) { t.Errorf("expected function tool 'web_search', got: %v", tool) } } + +func TestCodexBuildRequestBodyPromptCacheControls(t *testing.T) { + p := NewCodexProvider("test", &staticTokenSource{token: "tok"}, "", "gpt-4o") + + body := p.buildRequestBody(ChatRequest{ + Messages: []Message{{Role: "user", Content: "Hello"}}, + Options: map[string]any{ + OptPromptCacheKey: "agent/session/provider", + OptPromptCacheRetention: "24h", + }, + }, true) + + if got := body["prompt_cache_key"]; got != "agent/session/provider" { + t.Fatalf("prompt_cache_key = %v, want agent/session/provider", got) + } + if got := body["prompt_cache_retention"]; got != "24h" { + t.Fatalf("prompt_cache_retention = %v, want 24h", got) + } +} + +func TestCodexProviderCapabilitiesCacheControl(t *testing.T) { + p := NewCodexProvider("test", &staticTokenSource{token: "tok"}, "", "gpt-4o") + if !p.Capabilities().CacheControl { + t.Fatal("CodexProvider Capabilities().CacheControl = false, want true") + } +}