fix(openai): enable Astra messages prompt caching

This commit is contained in:
wucm667
2026-09-05 10:58:05 +08:00
parent b1748c4ea9
commit 4aaf5cb73b
3 changed files with 94 additions and 1 deletions
@@ -494,6 +494,66 @@ func TestForwardAsAnthropic_AutoDerivesPromptCacheKeyWhenMessagesDispatchHasNoSe
require.Equal(t, generateSessionUUID(isolateOpenAISessionID(0, cacheKey)), upstream.lastReq.Header.Get("session_id"))
}
func TestForwardAsAnthropic_GPT6AstraPromptCacheIdentityStableAcrossAppendedTurns(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
for _, mappedModel := range []string{"gpt-6-astra", "gpt-6"} {
mappedModel := mappedModel
t.Run(mappedModel, func(t *testing.T) {
t.Parallel()
upstream := &httpUpstreamRecorder{responses: []*http.Response{
openAICompatSSECompletedResponse("resp_gpt6_first", mappedModel),
openAICompatSSECompletedResponse("resp_gpt6_second", mappedModel),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
}
account := &Account{
ID: 6615,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://api.openai.com/v1",
},
}
bodies := [][]byte{
[]byte(`{"model":"claude-sonnet-4-5","max_tokens":16,"system":"You are helpful.","messages":[{"role":"user","content":"Inspect the repository"}],"stream":false}`),
[]byte(`{"model":"claude-sonnet-4-5","max_tokens":16,"system":"You are helpful.","messages":[{"role":"user","content":"Inspect the repository"},{"role":"assistant","content":"Done."},{"role":"user","content":"Run the tests"}],"stream":false}`),
}
for _, body := range bodies {
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", mappedModel)
require.NoError(t, err)
require.NotNil(t, result)
}
require.Len(t, upstream.bodies, 2)
require.Len(t, upstream.requests, 2)
require.Equal(t, mappedModel, gjson.GetBytes(upstream.bodies[0], "model").String())
require.Equal(t, mappedModel, gjson.GetBytes(upstream.bodies[1], "model").String())
firstCacheKey := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
secondCacheKey := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
require.NotEmpty(t, firstCacheKey)
require.Equal(t, firstCacheKey, secondCacheKey)
firstSessionID := upstream.requests[0].Header.Get("session_id")
secondSessionID := upstream.requests[1].Header.Get("session_id")
require.NotEmpty(t, firstSessionID)
require.Equal(t, firstSessionID, secondSessionID)
})
}
}
func TestForwardAsAnthropic_DoesNotAutoDerivePromptCacheKeyForNonCodexModel(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
@@ -13,6 +13,12 @@ const compatPromptCacheKeyPrefix = "compat_cc_"
func shouldAutoInjectPromptCacheKeyForCompat(model string) bool {
trimmed := strings.TrimSpace(strings.ToLower(model))
canonical := canonicalizeOpenAIModelAliasSpelling(trimmed)
// GPT-6 is the public alias for Astra. Keep this deliberately scoped to
// Astra so other GPT-6 families do not inherit Messages compatibility state.
if canonical == "gpt-6" || canonical == "gpt-6-astra" {
return true
}
// 仅对 Responses 兼容路径支持的 GPT-5 族开启自动注入,避免 normalizeCodexModel
// 的默认兜底把任意模型(如 gpt-4o、claude-*)误判为 gpt-5.4。
if !strings.Contains(trimmed, "gpt-5") && !strings.Contains(trimmed, "codex") {
@@ -21,7 +27,6 @@ func shouldAutoInjectPromptCacheKeyForCompat(model string) bool {
normalized := strings.TrimSpace(strings.ToLower(normalizeCodexModel(trimmed)))
return strings.HasPrefix(normalized, "gpt-5") || strings.Contains(normalized, "codex")
}
func deriveCompatPromptCacheKey(req *apicompat.ChatCompletionsRequest, mappedModel string) string {
if req == nil {
return ""
@@ -26,6 +26,34 @@ func TestShouldAutoInjectPromptCacheKeyForCompat(t *testing.T) {
require.False(t, shouldAutoInjectPromptCacheKeyForCompat("gpt-4o"))
}
func TestShouldAutoInjectPromptCacheKeyForCompat_GPT6AstraForms(t *testing.T) {
for _, model := range []string{
"gpt-6",
"gpt-6-astra",
"openai/gpt-6",
"openai/gpt-6-astra",
"OPENAI/GPT-6_ASTRA",
"provider/gpt-6-astra",
} {
require.True(t, shouldAutoInjectPromptCacheKeyForCompat(model), model)
}
for _, model := range []string{
"gpt-6-terra",
"gpt-6.1",
"gpt-6-astra-preview",
"gpt-6-astra-unrelated-name",
"gpt-6-astra-2026-09-01",
"gpt-6-astra-2026-09-01-extra",
"gpt-6-astra-v",
"gpt-6-astra-v1.beta",
"claude-sonnet-4-5",
"gpt-4o",
} {
require.False(t, shouldAutoInjectPromptCacheKeyForCompat(model), model)
}
}
func TestDeriveCompatPromptCacheKey_StableAcrossLaterTurns(t *testing.T) {
base := &apicompat.ChatCompletionsRequest{
Model: "gpt-5.4",