1use std::collections::HashMap;
7
8use crate::error::{LlmError, Result};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum AttributionProviderKind {
13 OpenAI,
14 AzureOpenAI,
15 Anthropic,
16 Gemini,
17 VertexAI,
18 OpenRouter,
19 OpenAICompatible,
20 Mistral,
21 Nvidia,
22 Cohere,
23 Bedrock,
24 XAI,
25 HuggingFace,
26 LMStudio,
27 Ollama,
28 VsCodeCopilot,
29 Mock,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
34pub enum AttributionPolicy {
35 #[default]
37 BestEffort,
38 RequireAppId,
40 Disabled,
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum AttributionWarning {
47 OpenRouterMissingReferer,
49 OpenRouterLocalhostMissingTitle,
51 InvalidHeader { name: String, reason: String },
53 ProviderUnsupported { provider: String },
55 AzureHeaderBudget { count: usize },
57}
58
59#[derive(Debug, Clone, Default, PartialEq, Eq)]
61pub struct ApplicationContext {
62 pub app_id: Option<String>,
64 pub app_name: Option<String>,
66 pub app_url: Option<String>,
68 pub tenant_id: Option<String>,
70 pub request_id: Option<String>,
72 pub end_user_id: Option<String>,
74 pub extra_headers: HashMap<String, String>,
76}
77
78impl ApplicationContext {
79 pub fn new() -> Self {
81 Self::default()
82 }
83
84 pub fn from_ingress_headers(headers: &HashMap<String, String>) -> Result<Self> {
86 let mut ctx = Self::new();
87 if let Some(v) = headers.get("x-edgequake-app-id") {
88 ctx.app_id = Some(sanitize_app_id(v)?);
89 }
90 if let Some(v) = headers.get("x-edgequake-app-name") {
91 ctx.app_name = Some(sanitize_app_name(v)?);
92 }
93 if let Some(v) = headers.get("x-edgequake-app-url") {
94 ctx.app_url = Some(sanitize_app_url(v)?);
95 }
96 if let Some(v) = headers.get("x-edgequake-tenant-id") {
97 ctx.tenant_id = Some(sanitize_tenant_id(v)?);
98 }
99 if let Some(v) = headers.get("x-edgequake-request-id") {
100 ctx.request_id = Some(sanitize_request_id(v)?);
101 }
102 for (k, v) in headers {
103 let lower = k.to_ascii_lowercase();
104 if lower.starts_with("x-edgequake-") {
105 continue;
106 }
107 if lower == "traceparent" || lower == "tracestate" || lower == "baggage" {
108 ctx.extra_headers.insert(k.clone(), v.clone());
109 }
110 }
111 Ok(ctx)
112 }
113
114 pub fn from_env() -> Self {
116 let mut ctx = Self::new();
117 if let Ok(v) = std::env::var("EDGEQUAKE_APP_ID") {
118 if let Ok(id) = sanitize_app_id(&v) {
119 ctx.app_id = Some(id);
120 }
121 }
122 if let Ok(v) = std::env::var("EDGEQUAKE_APP_NAME") {
123 if let Ok(name) = sanitize_app_name(&v) {
124 ctx.app_name = Some(name);
125 }
126 }
127 if let Ok(v) = std::env::var("EDGEQUAKE_APP_URL") {
128 if let Ok(url) = sanitize_app_url(&v) {
129 ctx.app_url = Some(url);
130 }
131 }
132 if let Ok(v) = std::env::var("EDGEQUAKE_TENANT_ID") {
133 if let Ok(tid) = sanitize_tenant_id(&v) {
134 ctx.tenant_id = Some(tid);
135 }
136 }
137 ctx
138 }
139
140 pub fn merge(&mut self, other: ApplicationContext) {
142 if other.app_id.is_some() {
143 self.app_id = other.app_id;
144 }
145 if other.app_name.is_some() {
146 self.app_name = other.app_name;
147 }
148 if other.app_url.is_some() {
149 self.app_url = other.app_url;
150 }
151 if other.tenant_id.is_some() {
152 self.tenant_id = other.tenant_id;
153 }
154 if other.request_id.is_some() {
155 self.request_id = other.request_id;
156 }
157 if other.end_user_id.is_some() {
158 self.end_user_id = other.end_user_id;
159 }
160 self.extra_headers.extend(other.extra_headers);
161 }
162
163 pub fn is_empty(&self) -> bool {
165 self.app_id.is_none()
166 && self.app_name.is_none()
167 && self.app_url.is_none()
168 && self.tenant_id.is_none()
169 && self.request_id.is_none()
170 && self.end_user_id.is_none()
171 && self.extra_headers.is_empty()
172 }
173
174 pub fn has_app_attribution(&self) -> bool {
176 self.app_id.is_some() || self.app_name.is_some() || self.app_url.is_some()
177 }
178
179 pub fn with_app_id(app_id: impl Into<String>) -> Result<Self> {
181 Ok(Self {
182 app_id: Some(sanitize_app_id(&app_id.into())?),
183 ..Default::default()
184 })
185 }
186}
187
188#[derive(Debug, Default)]
190pub struct ApplicationContextBuilder {
191 inner: ApplicationContext,
192}
193
194impl ApplicationContextBuilder {
195 pub fn new() -> Self {
196 Self::default()
197 }
198
199 pub fn app_id(mut self, value: impl Into<String>) -> Self {
200 self.inner.app_id = Some(value.into());
201 self
202 }
203
204 pub fn app_name(mut self, value: impl Into<String>) -> Self {
205 self.inner.app_name = Some(value.into());
206 self
207 }
208
209 pub fn app_url(mut self, value: impl Into<String>) -> Self {
210 self.inner.app_url = Some(value.into());
211 self
212 }
213
214 pub fn tenant_id(mut self, value: impl Into<String>) -> Self {
215 self.inner.tenant_id = Some(value.into());
216 self
217 }
218
219 pub fn request_id(mut self, value: impl Into<String>) -> Self {
220 self.inner.request_id = Some(value.into());
221 self
222 }
223
224 pub fn end_user_id(mut self, value: impl Into<String>) -> Self {
225 self.inner.end_user_id = Some(value.into());
226 self
227 }
228
229 pub fn extra_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
230 self.inner.extra_headers.insert(key.into(), value.into());
231 self
232 }
233
234 pub fn build(self) -> Result<ApplicationContext> {
235 let mut ctx = self.inner;
236 if let Some(ref id) = ctx.app_id {
237 ctx.app_id = Some(sanitize_app_id(id)?);
238 }
239 if let Some(ref name) = ctx.app_name {
240 ctx.app_name = Some(sanitize_app_name(name)?);
241 }
242 if let Some(ref url) = ctx.app_url {
243 ctx.app_url = Some(sanitize_app_url(url)?);
244 }
245 if let Some(ref tid) = ctx.tenant_id {
246 ctx.tenant_id = Some(sanitize_tenant_id(tid)?);
247 }
248 if let Some(ref rid) = ctx.request_id {
249 ctx.request_id = Some(sanitize_request_id(rid)?);
250 }
251 Ok(ctx)
252 }
253}
254
255pub fn sanitize_app_id(value: &str) -> Result<String> {
261 let trimmed = value.trim();
262 if trimmed.is_empty() {
263 return Err(LlmError::InvalidRequest("app_id must not be empty".into()));
264 }
265 if trimmed.len() > 128 {
266 return Err(LlmError::InvalidRequest(
267 "app_id must be at most 128 characters".into(),
268 ));
269 }
270 let lower = trimmed.to_ascii_lowercase();
271 if lower.starts_with("sk-") || lower.starts_with("bearer") {
272 return Err(LlmError::InvalidRequest(
273 "app_id must not resemble an API key or credential".into(),
274 ));
275 }
276 if !trimmed
277 .chars()
278 .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-')
279 {
280 return Err(LlmError::InvalidRequest(
281 "app_id may only contain ASCII letters, digits, '.', '_', and '-'".into(),
282 ));
283 }
284 if !trimmed
285 .chars()
286 .next()
287 .is_some_and(|c| c.is_ascii_alphanumeric())
288 {
289 return Err(LlmError::InvalidRequest(
290 "app_id must start with an ASCII letter or digit".into(),
291 ));
292 }
293 Ok(trimmed.to_string())
294}
295
296pub fn sanitize_app_name(value: &str) -> Result<String> {
297 let trimmed = value.trim();
298 if trimmed.is_empty() {
299 return Err(LlmError::InvalidRequest(
300 "app_name must not be empty".into(),
301 ));
302 }
303 if trimmed.len() > 256 {
304 return Err(LlmError::InvalidRequest(
305 "app_name must be at most 256 characters".into(),
306 ));
307 }
308 if !trimmed.is_ascii() {
309 return Err(LlmError::InvalidRequest(
310 "app_name must be ASCII only".into(),
311 ));
312 }
313 Ok(trimmed.to_string())
314}
315
316pub fn sanitize_app_url(value: &str) -> Result<String> {
317 let trimmed = value.trim();
318 if trimmed.is_empty() {
319 return Err(LlmError::InvalidRequest("app_url must not be empty".into()));
320 }
321 if trimmed.len() > 2048 {
322 return Err(LlmError::InvalidRequest(
323 "app_url must be at most 2048 characters".into(),
324 ));
325 }
326 if !trimmed.starts_with("http://") && !trimmed.starts_with("https://") {
327 return Err(LlmError::InvalidRequest(
328 "app_url must start with http:// or https://".into(),
329 ));
330 }
331 Ok(trimmed.to_string())
332}
333
334pub fn sanitize_tenant_id(value: &str) -> Result<String> {
335 sanitize_app_id(value)
336}
337
338pub fn sanitize_request_id(value: &str) -> Result<String> {
339 let trimmed = value.trim();
340 if trimmed.is_empty() {
341 return Err(LlmError::InvalidRequest(
342 "request_id must not be empty".into(),
343 ));
344 }
345 if trimmed.len() > 512 {
346 return Err(LlmError::InvalidRequest(
347 "request_id must be at most 512 characters (OpenAI limit)".into(),
348 ));
349 }
350 if !trimmed.is_ascii() {
351 return Err(LlmError::InvalidRequest(
352 "request_id must be ASCII only".into(),
353 ));
354 }
355 Ok(trimmed.to_string())
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361 use serial_test::serial;
362
363 #[test]
364 fn sanitize_app_id_accepts_valid_slug() {
365 assert_eq!(
366 sanitize_app_id("my-app_v2.backend").unwrap(),
367 "my-app_v2.backend"
368 );
369 }
370
371 #[test]
372 fn sanitize_app_id_rejects_api_key_like() {
373 assert!(sanitize_app_id("sk-abc123").is_err());
374 assert!(sanitize_app_id("Bearer-token").is_err());
375 }
376
377 #[test]
378 fn sanitize_app_id_rejects_invalid_chars() {
379 assert!(sanitize_app_id("app with spaces").is_err());
380 }
381
382 #[test]
383 fn sanitize_request_id_rejects_non_ascii() {
384 assert!(sanitize_request_id("req-日本語").is_err());
385 }
386
387 #[test]
388 fn sanitize_request_id_truncates_policy_at_512() {
389 let long = "a".repeat(513);
390 assert!(sanitize_request_id(&long).is_err());
391 assert!(sanitize_request_id(&"a".repeat(512)).is_ok());
392 }
393
394 #[test]
395 fn builder_validates_on_build() {
396 let ctx = ApplicationContextBuilder::new()
397 .app_id("valid-app")
398 .request_id("req-1")
399 .build()
400 .unwrap();
401 assert_eq!(ctx.app_id.as_deref(), Some("valid-app"));
402 }
403
404 #[test]
405 fn merge_prefers_other_values() {
406 let mut base = ApplicationContext {
407 app_id: Some("old".into()),
408 ..Default::default()
409 };
410 base.merge(ApplicationContext {
411 app_id: Some("new".into()),
412 request_id: Some("r1".into()),
413 ..Default::default()
414 });
415 assert_eq!(base.app_id.as_deref(), Some("new"));
416 assert_eq!(base.request_id.as_deref(), Some("r1"));
417 }
418
419 #[test]
420 fn from_ingress_headers_parses_canonical_keys() {
421 let mut headers = HashMap::new();
422 headers.insert("x-edgequake-app-id".into(), "workspace-api".into());
423 headers.insert("traceparent".into(), "00-abc-def-01".into());
424 let ctx = ApplicationContext::from_ingress_headers(&headers).unwrap();
425 assert_eq!(ctx.app_id.as_deref(), Some("workspace-api"));
426 assert_eq!(
427 ctx.extra_headers.get("traceparent").map(String::as_str),
428 Some("00-abc-def-01")
429 );
430 }
431
432 #[test]
433 fn has_app_attribution() {
434 assert!(!ApplicationContext::default().has_app_attribution());
435 assert!(ApplicationContext {
436 app_id: Some("x".into()),
437 ..Default::default()
438 }
439 .has_app_attribution());
440 }
441
442 #[test]
443 fn sanitize_app_id_rejects_empty() {
444 assert!(sanitize_app_id("").is_err());
445 assert!(sanitize_app_id(" ").is_err());
446 }
447
448 #[test]
449 fn sanitize_app_id_rejects_leading_invalid() {
450 assert!(sanitize_app_id("-bad").is_err());
451 assert!(sanitize_app_id(".bad").is_err());
452 }
453
454 #[test]
455 fn sanitize_app_url_requires_parseable_url() {
456 assert!(sanitize_app_url("https://app.example.com").is_ok());
457 assert!(sanitize_app_url("not-a-url").is_err());
458 }
459
460 #[test]
461 fn sanitize_tenant_id_accepts_alphanumeric() {
462 assert_eq!(
463 sanitize_tenant_id("tenant-1_prod").unwrap(),
464 "tenant-1_prod"
465 );
466 }
467
468 #[test]
469 #[serial]
470 fn from_env_reads_edgequake_app_id() {
471 std::env::set_var("EDGEQUAKE_APP_ID", "env-app");
472 let ctx = ApplicationContext::from_env();
473 assert_eq!(ctx.app_id.as_deref(), Some("env-app"));
474 std::env::remove_var("EDGEQUAKE_APP_ID");
475 }
476
477 #[test]
478 #[serial]
479 fn from_env_ignores_invalid_app_id() {
480 std::env::remove_var("EDGEQUAKE_APP_ID");
481 std::env::set_var("EDGEQUAKE_APP_ID", "sk-secret");
482 let ctx = ApplicationContext::from_env();
483 assert!(ctx.app_id.is_none());
484 std::env::remove_var("EDGEQUAKE_APP_ID");
485 }
486
487 #[test]
488 fn is_empty_context() {
489 assert!(ApplicationContext::default().is_empty());
490 assert!(!ApplicationContext {
491 request_id: Some("r".into()),
492 ..Default::default()
493 }
494 .is_empty());
495 }
496
497 #[test]
498 fn builder_rejects_empty_app_id() {
499 let err = ApplicationContextBuilder::new().app_id("").build();
500 assert!(err.is_err());
501 }
502
503 #[test]
504 fn merge_ingress_preserves_traceparent() {
505 let mut headers = HashMap::new();
506 headers.insert("x-edgequake-app-id".into(), "app".into());
507 headers.insert("traceparent".into(), "00-abcd".into());
508 headers.insert("Authorization".into(), "Bearer x".into());
509 let ctx = ApplicationContext::from_ingress_headers(&headers).unwrap();
510 assert_eq!(ctx.app_id.as_deref(), Some("app"));
511 assert_eq!(
512 ctx.extra_headers.get("traceparent").map(String::as_str),
513 Some("00-abcd")
514 );
515 assert!(!ctx.extra_headers.contains_key("Authorization"));
516 }
517}