tenant.openagents/omega

1use collections::HashMap;
2use http_client::CustomHeaders;
3use http_client::http::{HeaderName, HeaderValue};
4
5pub mod anthropic;
6pub mod anthropic_compatible;
7pub mod api_compatible;
8pub mod bedrock;
9pub mod cloud;
10pub mod copilot_chat;
11pub mod deepseek;
12pub mod google;
13pub mod llama_cpp;
14pub mod lmstudio;
15pub mod mistral;
16pub mod ollama;
17pub mod open_ai;
18pub mod open_ai_compatible;
19pub mod open_router;
20pub mod openai_subscribed;
21pub mod opencode;
22
23pub mod vercel_ai_gateway;
24pub mod x_ai;
25
26const COMMON_RESERVED_HEADER_NAMES: &[&str] = &["Authorization", "Content-Type", "Accept"];
27
28/// Validate the user-supplied custom-headers map once at settings load time,
29/// dropping reserved or malformed entries (each with a `log::warn!`) and
30/// returning a typed `CustomHeaders` ready to be appended to outgoing requests.
31pub(crate) fn resolve_custom_headers(
32    provider_name: &str,
33    settings: &HashMap<String, String>,
34    reserved_header_names: &[&str],
35) -> CustomHeaders {
36    let headers = settings
37        .iter()
38        .filter_map(|(name, value)| {
39            if COMMON_RESERVED_HEADER_NAMES
40                .iter()
41                .chain(reserved_header_names)
42                .any(|reserved| reserved.eq_ignore_ascii_case(name))
43            {
44                log::warn!(
45                    "ignoring custom {provider_name} header `{name}`: managed by Omega and cannot be overridden"
46                );
47                return None;
48            }
49            let header_name = match name.parse::<HeaderName>() {
50                Ok(header_name) => header_name,
51                Err(err) => {
52                    log::warn!("ignoring custom {provider_name} header `{name}`: invalid header name ({err})");
53                    return None;
54                }
55            };
56            let header_value = match HeaderValue::from_str(value) {
57                Ok(header_value) => header_value,
58                Err(err) => {
59                    log::warn!(
60                        "ignoring custom {provider_name} header `{name}`: invalid header value ({err})"
61                    );
62                    return None;
63                }
64            };
65            Some((header_name, header_value))
66        })
67        .collect();
68    CustomHeaders::new(headers)
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    fn map(pairs: &[(&str, &str)]) -> HashMap<String, String> {
76        pairs
77            .iter()
78            .map(|(key, value)| ((*key).to_string(), (*value).to_string()))
79            .collect()
80    }
81
82    fn names(headers: &CustomHeaders) -> Vec<String> {
83        let mut names: Vec<String> = headers
84            .iter()
85            .map(|(name, _)| name.as_str().to_owned())
86            .collect();
87        names.sort();
88        names
89    }
90
91    #[test]
92    fn drops_common_and_provider_reserved_headers() {
93        let settings = map(&[
94            ("Authorization", "Bearer leak"),
95            ("Content-Type", "text/plain"),
96            ("Accept", "text/plain"),
97            ("X-Api-Key", "leak"),
98            ("X-Allowed", "yes"),
99        ]);
100        let merged = resolve_custom_headers("Test", &settings, &["X-Api-Key"]);
101        assert_eq!(names(&merged), vec!["x-allowed".to_string()]);
102    }
103
104    #[test]
105    fn reserved_header_match_is_case_insensitive() {
106        let settings = map(&[
107            ("authorization", "Bearer leak"),
108            ("CONTENT-TYPE", "text/plain"),
109            ("x-api-key", "leak"),
110            ("X-Allowed", "yes"),
111        ]);
112        let merged = resolve_custom_headers("Test", &settings, &["X-Api-Key"]);
113        assert_eq!(names(&merged), vec!["x-allowed".to_string()]);
114    }
115
116    #[test]
117    fn headers_with_reserved_prefix_are_kept() {
118        let settings = map(&[("Authorization-Forwarded", "ok"), ("X-Api-Key-Trace", "ok")]);
119        let merged = resolve_custom_headers("Test", &settings, &["X-Api-Key"]);
120        assert_eq!(
121            names(&merged),
122            vec![
123                "authorization-forwarded".to_string(),
124                "x-api-key-trace".to_string(),
125            ]
126        );
127    }
128
129    #[test]
130    fn drops_invalid_header_name_and_value() {
131        let settings = map(&[
132            ("Bad Name", "ok"),
133            ("X-Bad-Value", "line1\nline2"),
134            ("X-Allowed", "yes"),
135        ]);
136        let merged = resolve_custom_headers("Test", &settings, &[]);
137        assert_eq!(names(&merged), vec!["x-allowed".to_string()]);
138    }
139}
140