switchyard_runner/
eligibility.rs1use 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
39pub 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
112pub(crate) struct TargetRequirements {
114 required_labels: BTreeSet<String>,
115}
116
117impl TargetRequirements {
118 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
130fn 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}