From 556921549e23c5a05a037fef96222a821171f3b2 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:28:34 -0700 Subject: [PATCH] [#20367][fix] Accept prompt_cache_key in Chat Completions requests ChatCompletionRequest forbids unknown fields, so requests carrying OpenAI's standard prompt_cache_key were rejected with 400. Accept the field and ignore it, like user and parallel_tool_calls. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- tensorrt_llm/serve/openai_protocol.py | 4 + .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../references/trtllm_serve_api.yaml | 6 ++ .../test_openai_protocol_prompt_cache_key.py | 73 +++++++++++++++++++ 4 files changed, 84 insertions(+) create mode 100644 tests/unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 7c2ab0812bf2..a280593d426f 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1102,6 +1102,10 @@ class ChatCompletionRequest(OpenAIBaseModel): # emits per turn, so parallel emission remains model behavior either way. parallel_tool_calls: Optional[bool] = None user: Optional[str] = None + # Standard OpenAI field, accepted for compatibility and otherwise ignored: + # KV cache reuse needs no client hint, and cache_salt is the field that + # isolates cached prefixes between users. + prompt_cache_key: Optional[str] = None reasoning_effort: Optional[ReasoningEffort | Literal[ "low", "medium", "high", "max", "none"]] = Field( default=ReasoningEffort.LOW, diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 0f1239eef75a..7f1b9512b8db 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -169,6 +169,7 @@ l0_cpu: - unittest/llmapi/apps/test_harmony_parsing.py::TestStripIncompleteMessagesReporting - unittest/llmapi/apps/test_kimi_serve_extensions.py - unittest/llmapi/apps/test_openai_chat_disagg_multimodal.py + - unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py - unittest/llmapi/apps/test_reasoning_prompt_resolution.py # test_request_metrics.py as a whole is not registered on any list, so only the # two cases added with the KV-transfer fix are selected here. Registering the diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index e25a699f6c97..2e71ac35d59a 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -640,6 +640,12 @@ models: default: null status: stable required: false + prompt_cache_key: + kind: openai + type: Optional[str] + default: null + status: stable + required: false reasoning_effort: kind: openai type: Optional[ReasoningEffort | Literal['low', 'medium', 'high', 'max', 'none']] diff --git a/tests/unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py b/tests/unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py new file mode 100644 index 000000000000..22a4c43ba618 --- /dev/null +++ b/tests/unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ChatCompletionRequest must accept the standard `prompt_cache_key` field. + +The request schema is extra="forbid", so the field must exist for requests +from OpenAI clients that set it (for example agent frameworks that pass a +per-session key) to validate. The value is not used. +""" + +import pytest +from pydantic import ValidationError + +from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest + +pytestmark = pytest.mark.cpu_only + + +def _request(**extra) -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + **extra, + ) + + +@pytest.mark.parametrize("value", ["session-1234", ""]) +def test_prompt_cache_key_accepted(value): + req = _request(prompt_cache_key=value) + assert req.prompt_cache_key == value + + +def test_prompt_cache_key_defaults_to_none(): + assert _request().prompt_cache_key is None + + +def test_prompt_cache_key_accepted_from_json(): + req = ChatCompletionRequest.model_validate_json( + '{"model": "test-model", "messages": [{"role": "user", "content": "hi"}],' + ' "prompt_cache_key": "session-1234"}' + ) + assert req.prompt_cache_key == "session-1234" + + +@pytest.mark.parametrize("value", [1, ["session"], {"key": "session"}]) +def test_prompt_cache_key_must_be_a_string(value): + with pytest.raises(ValidationError): + _request(prompt_cache_key=value) + + +def test_prompt_cache_key_leaves_other_fields_unchanged(): + plain = _request(cache_salt="tenant") + keyed = _request(cache_salt="tenant", prompt_cache_key="session-1234") + assert keyed.cache_salt == "tenant" + assert keyed.model_dump(exclude={"prompt_cache_key"}) == plain.model_dump( + exclude={"prompt_cache_key"} + ) + + +def test_unknown_fields_remain_rejected(): + with pytest.raises(ValidationError): + _request(prompt_cache_keys="typo")