Skip to main content

switchyard_runner/
eligibility.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Trusted target requirements for one request.
5
6use std::collections::BTreeSet;
7
8use serde_json::Value;
9use switchyard_protocol::{LlmClientError, Request, WireFormat};
10
11const REQUIREMENTS_KEY: &str = "switchyard.internal.required_target_labels";
12const MAX_LABELS: usize = 64;
13const MAX_LABEL_BYTES: usize = 128;
14const MAX_REQUIREMENTS_BYTES: usize = MAX_LABELS * (MAX_LABEL_BYTES * 6 + 4) + 2;
15const RESPONSES_STATE_FIELDS: [&str; 2] = ["previous_response_id", "conversation"];
16
17pub(crate) struct EligibilityPolicy {
18    require_explicit_requirements: bool,
19}
20
21impl EligibilityPolicy {
22    pub(crate) fn new(require_explicit_requirements: bool) -> Self {
23        Self {
24            require_explicit_requirements,
25        }
26    }
27
28    pub(crate) fn admit(&self, request: &Request) -> Result<TargetRequirements, LlmClientError> {
29        let supplied = target_requirements(request)?;
30        if self.require_explicit_requirements && supplied.is_none() {
31            return Err(invalid("route requires explicit target requirements"));
32        }
33        let required_labels = supplied.unwrap_or_default();
34        validate_request_context(request)?;
35        Ok(TargetRequirements { required_labels })
36    }
37}
38
39/// Adds trusted target requirements to an in-process request without removing earlier labels.
40/// An empty collection explicitly supplies no new requirements. The host owns classification and approvals.
41pub fn require_target_labels(
42    request: &mut Request,
43    labels: impl IntoIterator<Item = impl AsRef<str>>,
44) -> Result<(), LlmClientError> {
45    let mut required = target_requirements(request)?.unwrap_or_default();
46    for label in labels {
47        let label = label.as_ref();
48        validate_label(label)?;
49        required.insert(label.to_owned());
50        if required.len() > MAX_LABELS {
51            return Err(invalid("too many target requirement labels"));
52        }
53    }
54    let serialized = serde_json::to_string(&required)
55        .map_err(|_| invalid("invalid target requirement labels"))?;
56    request
57        .metadata
58        .get_or_insert_default()
59        .extra_metadata
60        .get_or_insert_default()
61        .insert(REQUIREMENTS_KEY.to_owned(), serialized);
62    Ok(())
63}
64
65pub(crate) fn has_target_requirements(request: &Request) -> bool {
66    request
67        .metadata
68        .as_ref()
69        .and_then(|metadata| metadata.extra_metadata.as_ref())
70        .is_some_and(|metadata| metadata.contains_key(REQUIREMENTS_KEY))
71}
72
73pub(crate) fn validate_labels(labels: &BTreeSet<String>) -> Result<(), LlmClientError> {
74    if labels.len() > MAX_LABELS {
75        return Err(invalid("too many target requirement labels"));
76    }
77    labels.iter().try_for_each(|label| validate_label(label))
78}
79
80fn validate_label(label: &str) -> Result<(), LlmClientError> {
81    if label.is_empty() || label.len() > MAX_LABEL_BYTES || label.trim() != label {
82        return Err(invalid(
83            "target labels must be nonblank, trimmed, and at most 128 bytes",
84        ));
85    }
86    Ok(())
87}
88
89fn target_requirements(request: &Request) -> Result<Option<BTreeSet<String>>, LlmClientError> {
90    let Some(metadata) = request
91        .metadata
92        .as_ref()
93        .and_then(|metadata| metadata.extra_metadata.as_ref())
94    else {
95        return Ok(None);
96    };
97    let labels: Option<BTreeSet<String>> = metadata
98        .get(REQUIREMENTS_KEY)
99        .map(|encoded| {
100            if encoded.len() > MAX_REQUIREMENTS_BYTES {
101                return Err(invalid("target requirement labels exceed the size limit"));
102            }
103            serde_json::from_str(encoded).map_err(|_| invalid("invalid target requirement labels"))
104        })
105        .transpose()?;
106    if let Some(labels) = &labels {
107        validate_labels(labels)?;
108    }
109    Ok(labels)
110}
111
112/// Validated target requirements supplied by the caller.
113pub(crate) struct TargetRequirements {
114    required_labels: BTreeSet<String>,
115}
116
117impl TargetRequirements {
118    /// Every returned label must be present on each target receiving the request.
119    pub(crate) fn required_labels(&self) -> &BTreeSet<String> {
120        &self.required_labels
121    }
122}
123
124fn invalid(message: &'static str) -> LlmClientError {
125    LlmClientError::InvalidRequest {
126        message: message.to_owned(),
127    }
128}
129
130// Provider-held history may contain data absent from the request being admitted.
131fn validate_request_context(request: &Request) -> Result<(), LlmClientError> {
132    let has_state = |body: &Value| {
133        RESPONSES_STATE_FIELDS
134            .iter()
135            .any(|field| body.get(*field).is_some_and(|value| !value.is_null()))
136    };
137    if RESPONSES_STATE_FIELDS.iter().any(|field| {
138        request
139            .llm_request
140            .extensions
141            .fields
142            .get(*field)
143            .is_some_and(|value| !value.is_null())
144    }) || request
145        .llm_request
146        .preservation
147        .requests
148        .iter()
149        .any(|(format, body)| {
150            format.as_str() == WireFormat::OpenAiResponses.as_str() && has_state(body)
151        })
152    {
153        return Err(invalid(
154            "eligibility routes do not support OpenAI Responses continuation state",
155        ));
156    }
157    Ok(())
158}
159
160#[cfg(test)]
161mod tests {
162    use std::collections::BTreeMap;
163
164    use super::*;
165
166    fn request(labels: &[&str]) -> Request {
167        let mut request = Request::default();
168        require_target_labels(&mut request, labels).expect("trusted requirements");
169        request
170    }
171
172    fn labels(values: &[&str]) -> BTreeSet<String> {
173        values.iter().map(|label| (*label).to_owned()).collect()
174    }
175
176    fn metadata(request: &mut Request) -> &mut BTreeMap<String, String> {
177        request
178            .metadata
179            .get_or_insert_default()
180            .extra_metadata
181            .get_or_insert_default()
182    }
183
184    fn assert_labels(policy: &EligibilityPolicy, request: &Request, expected: &[&str]) {
185        assert_eq!(
186            policy.admit(request).unwrap().required_labels(),
187            &labels(expected)
188        );
189    }
190
191    #[test]
192    fn requirements_are_normalized_and_only_accumulate() {
193        let policy = EligibilityPolicy::new(false);
194        let mut request = request(&["customer_data", "private", "private"]);
195        require_target_labels(&mut request, ["approved"]).expect("additional requirement");
196        require_target_labels(&mut request, [] as [&str; 0]).expect("cannot clear requirements");
197        assert_labels(&policy, &request, &["approved", "customer_data", "private"]);
198    }
199
200    #[test]
201    fn invalid_requirements_do_not_replace_valid_requirements() {
202        let mut request = request(&["private"]);
203        for invalid_label in [
204            "",
205            " ",
206            " private",
207            "private ",
208            &"x".repeat(MAX_LABEL_BYTES + 1),
209        ] {
210            assert!(require_target_labels(&mut request, [invalid_label]).is_err());
211        }
212        assert_eq!(
213            target_requirements(&request).unwrap(),
214            Some(labels(&["private"]))
215        );
216        for encoded in ["null", "{}", "[1]", "[\"\"]", "["] {
217            metadata(&mut request).insert(REQUIREMENTS_KEY.into(), encoded.into());
218            assert!(has_target_requirements(&request));
219            assert!(EligibilityPolicy::new(false).admit(&request).is_err());
220            assert!(require_target_labels(&mut request, ["private"]).is_err());
221        }
222    }
223
224    #[test]
225    fn required_labels_distinguish_missing_from_empty() {
226        let policy = EligibilityPolicy::new(true);
227        assert!(policy.admit(&Request::default()).is_err());
228        let decision = policy
229            .admit(&request(&[]))
230            .expect("explicit empty requirements");
231        assert!(decision.required_labels().is_empty());
232    }
233
234    #[test]
235    fn rejected_provider_history_does_not_erase_valid_requirements() {
236        let policy = EligibilityPolicy::new(false);
237        for field in RESPONSES_STATE_FIELDS {
238            let mut request = request(&["private"]);
239            request
240                .llm_request
241                .extensions
242                .fields
243                .insert(field.into(), Value::String("hidden".into()));
244            assert!(policy.admit(&request).is_err());
245            request.llm_request.extensions.fields.remove(field);
246            assert_labels(&policy, &request, &["private"]);
247            let mut preserved = self::request(&[]);
248            preserved.llm_request.preservation.requests.insert(
249                WireFormat::OpenAiResponses.as_str().into(),
250                serde_json::json!({field: "hidden"}),
251            );
252            assert!(policy.admit(&preserved).is_err());
253        }
254    }
255
256    #[test]
257    fn label_limits_reject_without_replacing_valid_requirements() {
258        let mut request = request(&["private"]);
259        assert!(
260            require_target_labels(&mut request, (0..MAX_LABELS).map(|index| index.to_string()))
261                .is_err()
262        );
263        assert_eq!(
264            target_requirements(&request).unwrap(),
265            Some(labels(&["private"]))
266        );
267    }
268}