diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index ddc717f2f7..e6d438baea 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -187,6 +187,22 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions( } } + if account.Type == AccountTypeAPIKey { + if trimmedKey := strings.TrimSpace(promptCacheKey); trimmedKey != "" { + var reqBody map[string]any + if err := json.Unmarshal(responsesBody, &reqBody); err != nil { + return nil, fmt.Errorf("unmarshal for prompt cache key injection: %w", err) + } + if existing, ok := reqBody["prompt_cache_key"].(string); !ok || strings.TrimSpace(existing) == "" { + reqBody["prompt_cache_key"] = trimmedKey + responsesBody, err = json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("remarshal after prompt cache key injection: %w", err) + } + } + } + } + // 4b. Apply OpenAI fast policy (may filter service_tier or block the request). updatedBody, policyErr := s.applyOpenAIFastPolicyToBody(ctx, account, upstreamModel, responsesBody) if policyErr != nil { @@ -214,7 +230,8 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions( } if promptCacheKey != "" { - upstreamReq.Header.Set("session_id", generateSessionUUID(promptCacheKey)) + apiKeyID := getAPIKeyIDFromContext(c) + upstreamReq.Header.Set("session_id", generateSessionUUID(isolateOpenAISessionID(apiKeyID, promptCacheKey))) } // 7. Send request diff --git a/backend/internal/service/openai_gateway_chat_completions_test.go b/backend/internal/service/openai_gateway_chat_completions_test.go index 58aea1d716..779565037d 100644 --- a/backend/internal/service/openai_gateway_chat_completions_test.go +++ b/backend/internal/service/openai_gateway_chat_completions_test.go @@ -11,6 +11,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -112,7 +113,10 @@ func TestForwardAsChatCompletions_UnknownModelDoesNotUseDefaultMappedModel(t *te Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"model not found"}}`)), }} - svc := &OpenAIGatewayService{httpUpstream: upstream} + svc := &OpenAIGatewayService{ + cfg: &config.Config{}, + httpUpstream: upstream, + } account := &Account{ ID: 1, Name: "openai-oauth", @@ -133,6 +137,50 @@ func TestForwardAsChatCompletions_UnknownModelDoesNotUseDefaultMappedModel(t *te require.Equal(t, http.StatusBadRequest, rec.Code) } +func TestForwardAsChatCompletions_APIKeyPropagatesPromptCacheKeyInResponsesBody(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":false}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set("api_key", &APIKey{ID: 99}) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_chat_prompt_cache"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop before response parsing"}}`)), + }} + + svc := &OpenAIGatewayService{ + cfg: &config.Config{}, + httpUpstream: upstream, + } + account := &Account{ + ID: 2, + Name: "openai-compatible", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-compatible", + }, + Extra: map[string]any{ + "openai_responses_supported": true, + }, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "cache-key-123", "gpt-5.4") + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, "cache-key-123", gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()) + require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "https://api.openai.com/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer sk-compatible", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, generateSessionUUID(isolateOpenAISessionID(99, "cache-key-123")), upstream.lastReq.Header.Get("session_id")) +} + func TestForwardAsChatCompletions_ClientDisconnectDrainsUpstreamUsage(t *testing.T) { gin.SetMode(gin.TestMode)