Skip to main content

edgequake_llm/
application_context.rs

1//! Application identity and attribution context for LLM provider calls.
2//!
3//! Single semantic type at the SDK boundary; provider-specific HTTP headers and
4//! body fields are resolved via [`crate::http::attribution`].
5
6use std::collections::HashMap;
7
8use crate::error::{LlmError, Result};
9
10/// Provider family for attribution resolution.
11#[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/// How strictly attribution must propagate to the upstream provider.
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
34pub enum AttributionPolicy {
35    /// Propagate what the provider supports; emit warnings for gaps.
36    #[default]
37    BestEffort,
38    /// Fail when `app_id` is set but cannot reach the provider.
39    RequireAppId,
40    /// Disable attribution injection (testing).
41    Disabled,
42}
43
44/// Non-fatal attribution issues surfaced to callers and tracing spans.
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum AttributionWarning {
47    /// OpenRouter requires `HTTP-Referer` for app leaderboard attribution.
48    OpenRouterMissingReferer,
49    /// Localhost referer without a display title.
50    OpenRouterLocalhostMissingTitle,
51    /// Custom header name or value was invalid and skipped.
52    InvalidHeader { name: String, reason: String },
53    /// Provider does not support application attribution headers.
54    ProviderUnsupported { provider: String },
55    /// Approaching Azure custom header passthrough limit (~10).
56    AzureHeaderBudget { count: usize },
57}
58
59/// Caller-supplied identity for attribution and request correlation.
60#[derive(Debug, Clone, Default, PartialEq, Eq)]
61pub struct ApplicationContext {
62    /// Stable application identifier (slug). Primary attribution key.
63    pub app_id: Option<String>,
64    /// Human-readable name for provider dashboards.
65    pub app_name: Option<String>,
66    /// Public URL of the application (OpenRouter `HTTP-Referer`).
67    pub app_url: Option<String>,
68    /// Multi-tenant partition.
69    pub tenant_id: Option<String>,
70    /// Per-request correlation ID (OpenAI `X-Client-Request-Id`, etc.).
71    pub request_id: Option<String>,
72    /// End-user identifier for provider safety/billing fields (OpenAI `user`, Azure `end_user_id`).
73    pub end_user_id: Option<String>,
74    /// Additional passthrough headers (`traceparent`, vendor-specific).
75    pub extra_headers: HashMap<String, String>,
76}
77
78impl ApplicationContext {
79    /// Create an empty context.
80    pub fn new() -> Self {
81        Self::default()
82    }
83
84    /// Parse canonical ingress headers into a validated context.
85    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    /// Load defaults from environment (`EDGEQUAKE_APP_*`).
115    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    /// Merge `other` into `self`. Per-field and extra_headers from `other` win.
141    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    /// Whether any attribution or correlation field is set.
164    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    /// Whether application-level attribution (not just request_id) is present.
175    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    /// Build from optional app id (Python SDK convenience).
180    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/// Builder for `ApplicationContext` with validation on build.
189#[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
255// ---------------------------------------------------------------------------
256// Sanitizers
257// ---------------------------------------------------------------------------
258
259/// Validate and normalize an application ID slug.
260pub 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}