[Frontend] Add cache salt and tracing to structured decisions (#60786)

Signed-off-by: Taewoong Kim <ktw2172@gmail.com>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>

docs/serving/online_serving/structured_decisions.md

@@ -135,6 +135,11 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a
 | `questions` | question id to `{type, instructions, criteria}`, asked in this order |
 | `instructions` | optional context placed ahead of the questions |
 | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` |
+| `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question |
 
 A question with any field other than `type`, `instructions` and `criteria` is
 rejected with a 400.
+
+When [tracing](../../../examples/observability/opentelemetry/README.md) is enabled
+with `--otlp-traces-endpoint`, the HTTP `traceparent` and `tracestate` headers are
+forwarded to every question's generation request.

docs/usage/security.md

@@ -572,7 +572,7 @@ An attacker sharing the same backend can measure differences in Time to First To
 
 vLLM accepts an optional `cache_salt` parameter on requests. The salt is mixed into the hash of the first KV cache block, so only requests carrying the same salt can share cached prefix blocks. See [Automatic Prefix Caching](../design/prefix_caching.md) for the implementation details.
 
-`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, and by the Anthropic `/v1/messages` endpoint.
+`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, the Anthropic `/v1/messages` endpoint, and the structured decisions `/v1/systemone` endpoint.
 
 #### Usage with the OpenAI Python client
 

tests/entrypoints/generate/structured_decisions/test_strategies.py

@@ -2,24 +2,46 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
 from typing import Any
+from unittest.mock import MagicMock
 
 import pytest
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.testclient import TestClient
 
+from vllm.config import ModelConfig
+from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.generate.structured_decisions.api_router import (
     register_structured_decisions_api_router,
 )
+from vllm.entrypoints.generate.structured_decisions.protocol import (
+    StructuredDecisionRequest,
+    StructuredDecisionResponse,
+)
 from vllm.entrypoints.generate.structured_decisions.question_types import LABELS
-from vllm.entrypoints.generate.structured_decisions.serving import state_text
+from vllm.entrypoints.generate.structured_decisions.serving import (
+    ServingStructuredDecisions,
+    state_text,
+)
 from vllm.entrypoints.generate.structured_decisions.strategies import (
     NextTokenStrategy,
+    ReadContext,
     reply_label_ids,
     select_read_strategy,
 )
+from vllm.entrypoints.openai.models.protocol import BaseModelPath
+from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.serve.exception_handling.register import (
     init_exception_handler,
 )
+from vllm.logprobs import Logprob
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.renderers.hf import HfRenderer
+from vllm.renderers.online_renderer import OnlineRenderer
+
+TRACE_HEADERS = {
+    "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01",
+    "tracestate": "vendor=value",
+}
 
 
 def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any:
@@ -69,9 +91,9 @@ def test_structured_state_keeps_unicode_in_prompt():
 
 @pytest.fixture(scope="module")
 def qwen():
-    from transformers import AutoTokenizer
+    from vllm.tokenizers import get_tokenizer
 
-    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
+    tokenizer = get_tokenizer("Qwen/Qwen3-0.6B")
     prompt_ids = tokenizer.apply_chat_template(
         [{"role": "user", "content": "state"}],
         add_generation_prompt=True,
@@ -93,3 +115,108 @@ def test_labels_start_the_reply(qwen):
     # After a colon, Qwen writes ":A" as one token, so "A" is not one token.
     with pytest.raises(ValueError, match="not one distinct token"):
         reply_label_ids(tokenizer, tokenizer.encode("team:"))
+
+
+@pytest.fixture
+def serving_and_engine(qwen):
+    tokenizer, _ = qwen
+    model_config = ModelConfig(model="Qwen/Qwen3-0.6B", dtype="float32")
+    renderer = HfRenderer(
+        SimpleNamespace(
+            model_config=model_config,
+            parallel_config=SimpleNamespace(_api_process_rank=0),
+        ),
+        tokenizer,
+    )
+    engine = MagicMock(spec=EngineClient)
+    engine.errored = False
+    engine.model_config = model_config
+    engine.renderer = renderer
+    engine.input_processor = MagicMock()
+    online_renderer = OnlineRenderer(
+        model_config,
+        renderer,
+        request_logger=None,
+        chat_template=None,
+        chat_template_content_format="auto",
+    )
+
+    async def generate(prompt, params, request_id, **kwargs):
+        label_ids = params.logprob_token_ids
+        yield RequestOutput(
+            request_id=request_id,
+            prompt=None,
+            prompt_token_ids=prompt["prompt_token_ids"],
+            prompt_logprobs=None,
+            outputs=[
+                CompletionOutput(
+                    index=0,
+                    text="A",
+                    token_ids=[label_ids[0]],
+                    cumulative_logprob=None,
+                    logprobs=[{token_id: Logprob(-1.0) for token_id in label_ids}],
+                    finish_reason="length",
+                )
+            ],
+            finished=True,
+        )
+
+    engine.generate.side_effect = generate
+    strategy = NextTokenStrategy(ReadContext(engine, online_renderer, None, "auto", {}))
+    models = OpenAIServingModels(
+        engine, [BaseModelPath(name=model_config.model, model_path=model_config.model)]
+    )
+    yield ServingStructuredDecisions(models, strategy), engine
+    renderer.shutdown()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "cache_salt,tracing_enabled,request_headers",
+    [
+        ("tenant-a", True, {"authorization": "Bearer test", **TRACE_HEADERS}),
+        ("tenant-b", False, {"authorization": "Bearer test", **TRACE_HEADERS}),
+        (None, True, {"authorization": "Bearer test", **TRACE_HEADERS}),
+        (None, True, {}),
+        ("tenant-direct", True, None),
+    ],
+)
+async def test_request_metadata_reaches_every_question(
+    serving_and_engine, cache_salt, tracing_enabled, request_headers
+):
+    serving, engine = serving_and_engine
+    engine.is_tracing_enabled.return_value = tracing_enabled
+    raw_request = None
+    if request_headers is not None:
+        raw_request = Request(
+            {
+                "type": "http",
+                "headers": [
+                    (k.encode(), v.encode()) for k, v in request_headers.items()
+                ],
+            }
+        )
+    request = StructuredDecisionRequest(
+        state={"team": "billing", "language": "English"},
+        questions={
+            "team": {
+                "type": "choice",
+                "instructions": "Which team does the message name?",
+                "criteria": {"billing": "payments", "shipping": "deliveries"},
+            },
+            "lang": {
+                "type": "choice",
+                "instructions": "Which language does the message name?",
+                "criteria": {"English": "English text", "French": "French text"},
+            },
+        },
+        cache_salt=cache_salt,
+    )
+    response = await serving.create_decision(request, raw_request)
+    assert isinstance(response, StructuredDecisionResponse)
+    assert engine.generate.call_count == 2
+    for call in engine.generate.call_args_list:
+        assert call.args[0].get("cache_salt") == cache_salt
+        assert call.kwargs["trace_headers"] == (
+            TRACE_HEADERS if tracing_enabled and request_headers else None
+        )

tests/entrypoints/unit_tests/test_non_object_body_validation.py

@@ -13,6 +13,9 @@ request models whose before-validators were missing the same guard.
 import pytest
 from pydantic import ValidationError
 
+from vllm.entrypoints.generate.structured_decisions.protocol import (
+    StructuredDecisionRequest,
+)
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
@@ -74,6 +77,18 @@ def test_tokenize_chat_request_still_validates_dict_bodies():
 
 
 CACHE_SALT_MODELS = [
+    (
+        StructuredDecisionRequest,
+        {
+            "state": "hello",
+            "questions": {
+                "lang": {
+                    "type": "choice",
+                    "criteria": {"English": "Text written in English."},
+                }
+            },
+        },
+    ),
     (CompletionRequest, {"prompt": "hello"}),
     (ChatCompletionRequest, {"messages": [{"role": "user", "content": "hello"}]}),
     (ResponsesRequest, {"input": "hello"}),

vllm/entrypoints/generate/structured_decisions/protocol.py

@@ -5,10 +5,11 @@
 import time
 from typing import Any, Literal
 
-from pydantic import Field
+from pydantic import Field, model_validator
 
 from vllm.config import ModelConfig
 from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
+from vllm.entrypoints.generate.base.protocol import validate_cache_salt
 from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel
 from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
 from vllm.utils import random_uuid
@@ -30,9 +31,27 @@ class StructuredDecisionRequest(OpenAIBaseModel):
         default=None, description="Context placed ahead of the questions."
     )
     chat_template_kwargs: dict[str, Any] | None = None
+    cache_salt: str | None = Field(
+        default=None,
+        description=(
+            "If specified, the prefix cache will be salted with the provided "
+            "string to prevent an attacker to guess prompts in multi-user "
+            "environments. The salt should be random, protected from "
+            "access by 3rd parties, and long enough to be "
+            "unpredictable (e.g., 43 characters base64-encoded, corresponding "
+            "to 256 bit)."
+        ),
+    )
     priority: int = Field(default=0, ge=-(2**63), le=2**63 - 1)
     request_id: str = Field(default_factory=random_uuid)
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_cache_salt_support(cls, data: Any) -> Any:
+        if isinstance(data, dict):
+            validate_cache_salt(data.get("cache_salt"))
+        return data
+
 
 class DecisionUsage(OpenAIBaseModel):
     input_tokens: int
@@ -65,6 +84,7 @@ class ReadPromptRequest(OpenAIBaseModel):
     would follow a thought rather than start the reply."""
 
     chat_template_kwargs: dict[str, Any] | None = None
+    cache_salt: str | None = None
 
     def build_chat_params(
         self,

vllm/entrypoints/generate/structured_decisions/serving.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
 import json
+from collections.abc import Mapping
 from typing import Any
 
 from fastapi import Request
@@ -11,6 +12,11 @@ from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.engine.serving import BaseServing
 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
 from vllm.logger import init_logger
+from vllm.tracing import (
+    contains_trace_headers,
+    extract_trace_headers,
+    log_tracing_disabled_warning,
+)
 
 from .protocol import (
     DecisionUsage,
@@ -71,6 +77,16 @@ class ServingStructuredDecisions(BaseServing):
         self.strategy = strategy
         self.limits = strategy.limits()
 
+    async def _get_trace_headers(
+        self, headers: Mapping[str, str]
+    ) -> Mapping[str, str] | None:
+        if not contains_trace_headers(headers):
+            return None
+        if not await self.strategy.context.engine_client.is_tracing_enabled():
+            log_tracing_disabled_warning()
+            return None
+        return extract_trace_headers(headers)
+
     async def create_decision(
         self,
         request: StructuredDecisionRequest,
@@ -88,6 +104,11 @@ class ServingStructuredDecisions(BaseServing):
             questions = parse_questions(request, self.limits)
             lora_request = self._maybe_get_adapters(request)  # type: ignore[arg-type]
             engine_client.check_admission(len(questions))
+            trace_headers = (
+                None
+                if raw_request is None
+                else await self._get_trace_headers(raw_request.headers)
+            )
             reads = await self.strategy.read(
                 questions,
                 request.instructions,
@@ -96,6 +117,8 @@ class ServingStructuredDecisions(BaseServing):
                 chat_template_kwargs=request.chat_template_kwargs,
                 lora_request=lora_request,
                 priority=request.priority,
+                cache_salt=request.cache_salt,
+                trace_headers=trace_headers,
             )
         except StructuredDecisionError as e:
             return self.create_error_response(e)

vllm/entrypoints/generate/structured_decisions/strategies.py

@@ -7,7 +7,7 @@ Only models in NEXT_TOKEN_ARCHITECTURES are currently supported.
 
 import math
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any
 
@@ -74,6 +74,8 @@ class ReadStrategy(ABC):
         chat_template_kwargs: dict[str, Any] | None,
         lora_request: LoRARequest | None,
         priority: int,
+        cache_salt: str | None = None,
+        trace_headers: Mapping[str, str] | None = None,
     ) -> list[QuestionRead]:
         """One read per question, in order."""
 
@@ -136,13 +138,15 @@ class NextTokenStrategy(ReadStrategy):
         return DecisionLimits(max_questions=64, max_options=len(LABELS))
 
     def _read_request(
-        self, chat_template_kwargs: dict[str, Any] | None
+        self, chat_template_kwargs: dict[str, Any] | None, cache_salt: str | None = None
     ) -> ReadPromptRequest:
         if (chat_template_kwargs or {}).get("enable_thinking"):
             raise StructuredDecisionError(
                 "a read is the reply's first token, so thinking must be off"
             )
-        return ReadPromptRequest(chat_template_kwargs=chat_template_kwargs)
+        return ReadPromptRequest(
+            chat_template_kwargs=chat_template_kwargs, cache_salt=cache_salt
+        )
 
     async def _render(
         self, read_request: ReadPromptRequest, messages: list[dict[str, Any]]
@@ -170,9 +174,11 @@ class NextTokenStrategy(ReadStrategy):
         chat_template_kwargs: dict[str, Any] | None,
         lora_request: LoRARequest | None,
         priority: int,
+        cache_salt: str | None = None,
+        trace_headers: Mapping[str, str] | None = None,
     ) -> list[QuestionRead]:
         ctx = self.context
-        read_request = self._read_request(chat_template_kwargs)
+        read_request = self._read_request(chat_template_kwargs, cache_salt)
 
         slots, engine_inputs = [], []
         for q in questions:
@@ -197,6 +203,7 @@ class NextTokenStrategy(ReadStrategy):
             ],
             request_id,
             lora_request=lora_request,
+            trace_headers=trace_headers,
             priority=priority,
         )