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 and bounded task retention.
5
6use std::collections::{BTreeSet, HashMap};
7
8use parking_lot::Mutex;
9use serde_json::Value;
10use switchyard_protocol::{LlmClientError, Request, WireFormat};
11
12const REQUIREMENTS_KEY: &str = "switchyard.internal.required_target_labels";
13const MAX_LABELS: usize = 64;
14const MAX_LABEL_BYTES: usize = 128;
15const MAX_REQUIREMENTS_BYTES: usize = MAX_LABELS * (MAX_LABEL_BYTES * 6 + 4) + 2;
16const MAX_TASK_ID_BYTES: usize = 512;
17const MAX_TASKS: usize = 4_096;
18const RESPONSES_STATE_FIELDS: [&str; 2] = ["previous_response_id", "conversation"];
19
20pub(crate) struct EligibilityPolicy {
21    require_explicit_requirements: bool,
22    tasks: Option<Mutex<HashMap<String, BTreeSet<String>>>>,
23}
24
25impl EligibilityPolicy {
26    pub(crate) fn new(require_explicit_requirements: bool, task_scope: bool) -> Self {
27        Self {
28            require_explicit_requirements,
29            tasks: task_scope.then(Mutex::default),
30        }
31    }
32
33    // Retain valid requirements before checking provider-held context for this turn.
34    pub(crate) fn admit(&self, request: &Request) -> Result<TargetRequirements, LlmClientError> {
35        let supplied = target_requirements(request)?;
36        if self.require_explicit_requirements && supplied.is_none() {
37            return Err(invalid("route requires explicit target requirements"));
38        }
39        let mut required_labels = supplied.unwrap_or_default();
40        if let Some(tasks) = &self.tasks {
41            let task_id = request
42                .metadata
43                .as_ref()
44                .and_then(|metadata| metadata.task_id.as_deref())
45                .filter(|id| !id.is_empty() && id.len() <= MAX_TASK_ID_BYTES && id.trim() == *id)
46                .ok_or_else(|| invalid("target eligibility needs a valid task id"))?;
47            // Validate and merge under one lock so concurrent turns cannot lose requirements.
48            let mut tasks = tasks.lock();
49            if let Some(retained) = tasks.get_mut(task_id) {
50                if retained.len() + required_labels.difference(retained).count() > MAX_LABELS {
51                    return Err(invalid("too many target requirement labels"));
52                }
53                retained.append(&mut required_labels);
54                required_labels.clone_from(retained);
55            } else if !required_labels.is_empty() {
56                if tasks.len() >= MAX_TASKS {
57                    return Err(invalid("target requirement task capacity exhausted"));
58                }
59                tasks.insert(task_id.to_owned(), required_labels.clone());
60            }
61        }
62        validate_request_context(request)?;
63        Ok(TargetRequirements { required_labels })
64    }
65}
66
67/// Adds trusted target requirements to an in-process request without removing earlier labels.
68/// An empty collection explicitly supplies no new requirements. The host owns classification and approvals.
69pub fn require_target_labels(
70    request: &mut Request,
71    labels: impl IntoIterator<Item = impl AsRef<str>>,
72) -> Result<(), LlmClientError> {
73    let mut required = target_requirements(request)?.unwrap_or_default();
74    for label in labels {
75        let label = label.as_ref();
76        validate_label(label)?;
77        required.insert(label.to_owned());
78        if required.len() > MAX_LABELS {
79            return Err(invalid("too many target requirement labels"));
80        }
81    }
82    let serialized = serde_json::to_string(&required)
83        .map_err(|_| invalid("invalid target requirement labels"))?;
84    request
85        .metadata
86        .get_or_insert_default()
87        .extra_metadata
88        .get_or_insert_default()
89        .insert(REQUIREMENTS_KEY.to_owned(), serialized);
90    Ok(())
91}
92
93pub(crate) fn has_target_requirements(request: &Request) -> bool {
94    request
95        .metadata
96        .as_ref()
97        .and_then(|metadata| metadata.extra_metadata.as_ref())
98        .is_some_and(|metadata| metadata.contains_key(REQUIREMENTS_KEY))
99}
100
101pub(crate) fn validate_labels(labels: &BTreeSet<String>) -> Result<(), LlmClientError> {
102    if labels.len() > MAX_LABELS {
103        return Err(invalid("too many target requirement labels"));
104    }
105    labels.iter().try_for_each(|label| validate_label(label))
106}
107
108fn validate_label(label: &str) -> Result<(), LlmClientError> {
109    if label.is_empty() || label.len() > MAX_LABEL_BYTES || label.trim() != label {
110        return Err(invalid(
111            "target labels must be nonblank, trimmed, and at most 128 bytes",
112        ));
113    }
114    Ok(())
115}
116
117fn target_requirements(request: &Request) -> Result<Option<BTreeSet<String>>, LlmClientError> {
118    let Some(metadata) = request
119        .metadata
120        .as_ref()
121        .and_then(|metadata| metadata.extra_metadata.as_ref())
122    else {
123        return Ok(None);
124    };
125    let labels: Option<BTreeSet<String>> = metadata
126        .get(REQUIREMENTS_KEY)
127        .map(|encoded| {
128            if encoded.len() > MAX_REQUIREMENTS_BYTES {
129                return Err(invalid("target requirement labels exceed the size limit"));
130            }
131            serde_json::from_str(encoded).map_err(|_| invalid("invalid target requirement labels"))
132        })
133        .transpose()?;
134    if let Some(labels) = &labels {
135        validate_labels(labels)?;
136    }
137    Ok(labels)
138}
139
140/// Effective target requirements after combining caller requirements and retained task state.
141pub(crate) struct TargetRequirements {
142    required_labels: BTreeSet<String>,
143}
144
145impl TargetRequirements {
146    /// Every returned label must be present on each target receiving the request.
147    pub(crate) fn required_labels(&self) -> &BTreeSet<String> {
148        &self.required_labels
149    }
150}
151
152fn invalid(message: &'static str) -> LlmClientError {
153    LlmClientError::InvalidRequest {
154        message: message.to_owned(),
155    }
156}
157
158// Provider-held history may contain data absent from the request being admitted.
159fn validate_request_context(request: &Request) -> Result<(), LlmClientError> {
160    let has_state = |body: &Value| {
161        RESPONSES_STATE_FIELDS
162            .iter()
163            .any(|field| body.get(*field).is_some_and(|value| !value.is_null()))
164    };
165    if RESPONSES_STATE_FIELDS.iter().any(|field| {
166        request
167            .llm_request
168            .extensions
169            .fields
170            .get(*field)
171            .is_some_and(|value| !value.is_null())
172    }) || request
173        .llm_request
174        .preservation
175        .requests
176        .iter()
177        .any(|(format, body)| {
178            format.as_str() == WireFormat::OpenAiResponses.as_str() && has_state(body)
179        })
180    {
181        return Err(invalid(
182            "eligibility routes do not support OpenAI Responses continuation state",
183        ));
184    }
185    Ok(())
186}
187
188#[cfg(test)]
189mod tests {
190    use std::collections::BTreeMap;
191
192    use super::*;
193
194    fn task(id: &str, labels: &[&str]) -> Request {
195        let mut request = Request::default();
196        request.metadata.get_or_insert_default().task_id = Some(id.into());
197        require_target_labels(&mut request, labels).expect("trusted requirements");
198        request
199    }
200
201    fn labels(values: &[&str]) -> BTreeSet<String> {
202        values.iter().map(|label| (*label).to_owned()).collect()
203    }
204
205    fn metadata(request: &mut Request) -> &mut BTreeMap<String, String> {
206        request
207            .metadata
208            .get_or_insert_default()
209            .extra_metadata
210            .get_or_insert_default()
211    }
212
213    fn assert_labels(policy: &EligibilityPolicy, request: &Request, expected: &[&str]) {
214        assert_eq!(
215            policy.admit(request).unwrap().required_labels(),
216            &labels(expected)
217        );
218    }
219
220    #[test]
221    fn requirements_are_normalized_and_only_accumulate() {
222        let policy = EligibilityPolicy::new(false, false);
223        let mut request = task("one", &["customer_data", "private", "private"]);
224        require_target_labels(&mut request, ["approved"]).expect("additional requirement");
225        require_target_labels(&mut request, [] as [&str; 0]).expect("cannot clear requirements");
226        assert_labels(&policy, &request, &["approved", "customer_data", "private"]);
227        assert_labels(&policy, &task("one", &[]), &[]);
228    }
229
230    #[test]
231    fn invalid_requirements_do_not_replace_valid_requirements() {
232        let mut request = task("one", &["private"]);
233        for invalid_label in [
234            "",
235            " ",
236            " private",
237            "private ",
238            &"x".repeat(MAX_LABEL_BYTES + 1),
239        ] {
240            assert!(require_target_labels(&mut request, [invalid_label]).is_err());
241        }
242        assert_eq!(
243            target_requirements(&request).unwrap(),
244            Some(labels(&["private"]))
245        );
246        for encoded in ["null", "{}", "[1]", "[\"\"]", "["] {
247            metadata(&mut request).insert(REQUIREMENTS_KEY.into(), encoded.into());
248            assert!(has_target_requirements(&request));
249            assert!(
250                EligibilityPolicy::new(false, false)
251                    .admit(&request)
252                    .is_err()
253            );
254            assert!(require_target_labels(&mut request, ["private"]).is_err());
255        }
256    }
257
258    #[test]
259    fn required_labels_distinguish_missing_from_empty() {
260        let policy = EligibilityPolicy::new(true, false);
261        assert!(policy.admit(&Request::default()).is_err());
262        let decision = policy
263            .admit(&task("one", &[]))
264            .expect("explicit empty requirements");
265        assert!(decision.required_labels().is_empty());
266    }
267
268    #[test]
269    fn task_requirements_accumulate_without_affecting_other_tasks() {
270        let policy = EligibilityPolicy::new(false, true);
271        policy
272            .admit(&task("one", &["private"]))
273            .expect("first requirement");
274        let decision = policy
275            .admit(&task("one", &["customer_data"]))
276            .expect("union");
277        assert_eq!(
278            decision.required_labels(),
279            &labels(&["private", "customer_data"])
280        );
281        assert_labels(&policy, &task("one", &[]), &["private", "customer_data"]);
282        assert_labels(&policy, &task("two", &[]), &[]);
283        let mut absent = Request::default();
284        absent.metadata.get_or_insert_default().task_id = Some("one".into());
285        assert_labels(&policy, &absent, &["private", "customer_data"]);
286
287        let explicit = EligibilityPolicy::new(true, true);
288        assert_labels(&explicit, &task("one", &["private"]), &["private"]);
289        assert!(explicit.admit(&absent).is_err());
290        assert_labels(&explicit, &task("one", &[]), &["private"]);
291    }
292
293    // Invalid labels and task identities must not alter retained requirements.
294    #[test]
295    fn task_requirements_validate_signals_and_ids_before_retaining() {
296        let policy = EligibilityPolicy::new(false, true);
297        assert_labels(&policy, &task("one", &["private"]), &["private"]);
298        for encoded in ["null", "{}", "[1]", "[\"\"]", "["] {
299            let mut request = task("one", &[]);
300            metadata(&mut request).insert(REQUIREMENTS_KEY.into(), encoded.into());
301            assert!(policy.admit(&request).is_err());
302            assert_labels(&policy, &task("one", &[]), &["private"]);
303        }
304        assert!(policy.admit(&Request::default()).is_err());
305        for id in ["", " ", " one", "one ", &"x".repeat(MAX_TASK_ID_BYTES + 1)] {
306            assert!(policy.admit(&task(id, &["customer_data"])).is_err());
307        }
308        assert_eq!(policy.tasks.as_ref().unwrap().lock().len(), 1);
309        assert_labels(
310            &policy,
311            &task(&"x".repeat(MAX_TASK_ID_BYTES), &["private"]),
312            &["private"],
313        );
314    }
315
316    #[test]
317    fn rejected_provider_history_does_not_erase_valid_requirements() {
318        let policy = EligibilityPolicy::new(false, true);
319        for field in RESPONSES_STATE_FIELDS {
320            let mut request = task(field, &["private"]);
321            request
322                .llm_request
323                .extensions
324                .fields
325                .insert(field.into(), Value::String("hidden".into()));
326            assert!(policy.admit(&request).is_err());
327            assert_labels(&policy, &task(field, &[]), &["private"]);
328            let mut preserved = task(field, &["customer_data"]);
329            preserved.llm_request.preservation.requests.insert(
330                WireFormat::OpenAiResponses.as_str().into(),
331                serde_json::json!({field: "hidden"}),
332            );
333            assert!(policy.admit(&preserved).is_err());
334            assert_labels(&policy, &task(field, &[]), &["private", "customer_data"]);
335        }
336    }
337
338    #[test]
339    fn label_and_task_limits_reject_without_partially_retaining() {
340        let mut request = task("one", &["private"]);
341        assert!(
342            require_target_labels(&mut request, (0..MAX_LABELS).map(|index| index.to_string()))
343                .is_err()
344        );
345        assert_eq!(
346            target_requirements(&request).unwrap(),
347            Some(labels(&["private"]))
348        );
349        let policy = EligibilityPolicy::new(false, true);
350        let retained = (0..MAX_LABELS).map(|index| index.to_string()).collect();
351        policy.tasks.as_ref().unwrap().lock().extend(
352            (0..MAX_TASKS).map(|index| (index.to_string(), BTreeSet::from(["private".into()]))),
353        );
354        assert!(policy.admit(&task("new", &["private"])).is_err());
355        assert_labels(&policy, &task("new", &[]), &[]);
356        assert_eq!(policy.tasks.as_ref().unwrap().lock().len(), MAX_TASKS);
357        policy
358            .tasks
359            .as_ref()
360            .unwrap()
361            .lock()
362            .insert("0".into(), retained);
363        assert!(policy.admit(&task("0", &["customer_data"])).is_err());
364        assert_eq!(
365            policy
366                .admit(&task("0", &[]))
367                .unwrap()
368                .required_labels()
369                .len(),
370            MAX_LABELS
371        );
372    }
373}
switchyard_runner/eligibility.rs

switchyard_runner/
eligibility.rs
