feat(providers): let a provider send configured extra headers

crates/aion-config/src/compat.rs

@@ -1,6 +1,8 @@
 // Configuration-driven provider compatibility layer.
 // Each provider type has default presets; users can override any field via config.
 
+use std::collections::BTreeMap;
+
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
@@ -56,6 +58,23 @@ pub struct TransportCompat {
     /// Default: "/chat/completions" for OpenAI-compatible providers.
     pub api_path: Option<String>,
 
+    /// Extra headers sent with every request to this provider.
+    ///
+    /// Some gateways ask callers to identify themselves on the wire — an app
+    /// name, a referer, a partner id — and refuse to attribute traffic without
+    /// it. Those headers are per-deployment, not per-provider-family, so they
+    /// belong in configuration rather than in a match arm here.
+    ///
+    /// Reserved headers cannot be replaced through this map: authorization,
+    /// content type and the provider's own protocol headers are written after
+    /// it and win. That is deliberate — a config typo should not be able to
+    /// unauthenticate a request or change the wire format.
+    ///
+    /// `BTreeMap` so the order is stable and two equal configs serialize
+    /// identically.
+    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
+    pub extra_headers: BTreeMap<String, String>,
+
     /// Maximum serialized provider request body size in bytes.
     /// Default: None (no local preflight limit).
     pub max_request_body_bytes: Option<usize>,
@@ -177,6 +196,14 @@ impl TransportCompat {
             api_path: user.api_path.or(defaults.api_path),
             max_request_body_bytes: user.max_request_body_bytes.or(defaults.max_request_body_bytes),
             include_stream_options: user.include_stream_options.or(defaults.include_stream_options),
+            extra_headers: {
+                // Per-key override rather than whole-map replacement: setting
+                // one header in user config should not drop the preset's
+                // others, which is what `.or()` semantics would do here.
+                let mut merged = defaults.extra_headers;
+                merged.extend(user.extra_headers);
+                merged
+            },
         }
     }
 }
@@ -371,6 +398,12 @@ impl ProviderCompat {
     /// The historical `/chat/completions` preset remains the default for Chat
     /// Completions. When Responses is selected, that inherited preset is
     /// replaced with `/responses`; non-default custom paths remain honored.
+    /// Headers this provider adds to every request. See
+    /// [`TransportCompat::extra_headers`]; reserved headers still win.
+    pub fn extra_headers(&self) -> &BTreeMap<String, String> {
+        &self.transport.extra_headers
+    }
+
     pub fn openai_api_path(&self) -> &str {
         match self.openai_api_mode() {
             OpenAiApiMode::ChatCompletions => self.api_path(),

crates/aion-config/src/compat_test.rs

@@ -75,6 +75,45 @@ max_tokens = 64000
         );
     }
 
+    #[test]
+    fn extra_headers_round_trip_through_toml_and_stay_absent_when_empty() {
+        let empty = ProviderCompat::default();
+        let rendered = toml::to_string(&empty).expect("serialize");
+        assert!(
+            !rendered.contains("extra_headers"),
+            "an empty map must not appear on disk, or every existing config file changes: {rendered}"
+        );
+
+        let parsed: ProviderCompat = toml::from_str(
+            r#"
+            [extra_headers]
+            "X-Partner-Id" = "part_abc123"
+            "HTTP-Referer" = "https://example.test"
+            "#,
+        )
+        .expect("parse");
+        assert_eq!(
+            parsed.extra_headers().get("X-Partner-Id").map(String::as_str),
+            Some("part_abc123")
+        );
+        assert_eq!(parsed.extra_headers().len(), 2);
+    }
+
+    #[test]
+    fn user_extra_headers_override_per_key_rather_than_replacing_the_preset_map() {
+        let mut defaults = TransportCompat::default();
+        defaults.extra_headers.insert("X-Kept".into(), "preset".into());
+        defaults.extra_headers.insert("X-Replaced".into(), "preset".into());
+
+        let mut user = TransportCompat::default();
+        user.extra_headers.insert("X-Replaced".into(), "user".into());
+
+        let merged = TransportCompat::merge(defaults, user);
+
+        assert_eq!(merged.extra_headers.get("X-Kept").map(String::as_str), Some("preset"));
+        assert_eq!(merged.extra_headers.get("X-Replaced").map(String::as_str), Some("user"));
+    }
+
     #[test]
     fn test_flattened_compat_serializes_to_legacy_toml_keys() {
         let compat = ProviderCompat {
@@ -90,6 +129,9 @@ max_tokens = 64000
                 api_path: Some("/chat/completions".to_string()),
                 max_request_body_bytes: Some(1_048_576),
                 include_stream_options: Some(false),
+                // Empty and `skip_serializing_if`, so the legacy key set this
+                // test pins is unchanged by the field existing.
+                extra_headers: Default::default(),
             },
             messages: MessageCompat {
                 merge_assistant_messages: Some(true),

crates/aion-providers/src/transport.rs

@@ -2,7 +2,7 @@ use aion_config::compat::{OpenAiApiMode, ProviderCompat};
 use aion_types::llm::LlmRequest;
 use futures::StreamExt;
 use reqwest::ResponseBuilderExt;
-use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderValue};
+use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue};
 use serde_json::Value;
 
 use crate::bedrock::BedrockTransportState;
@@ -69,6 +69,22 @@ pub(crate) struct ProjectedHttpRequest {
     pub tool_wire_shape: ResolvedToolWireShape,
 }
 
+/// Apply [`ProviderCompat::extra_headers`] to a request under construction.
+///
+/// Bedrock and Vertex are deliberately not covered: Bedrock signs its headers
+/// with SigV4, so an extra header added outside the signing step invalidates
+/// the signature rather than being attributed.
+fn insert_extra_headers(headers: &mut HeaderMap, compat: &ProviderCompat) -> Result<(), ProviderError> {
+    for (name, value) in compat.extra_headers() {
+        let header_name = HeaderName::from_bytes(name.as_bytes())
+            .map_err(|error| ProviderError::Connection(format!("Invalid extra header name '{name}': {error}")))?;
+        let header_value = HeaderValue::from_str(value)
+            .map_err(|error| ProviderError::Connection(format!("Invalid extra header value for '{name}': {error}")))?;
+        headers.insert(header_name, header_value);
+    }
+    Ok(())
+}
+
 impl OpenAiTransport {
     pub(crate) fn new(api_key: &str, base_url: &str) -> Self {
         Self {
@@ -85,6 +101,9 @@ impl OpenAiTransport {
         tool_wire_shape: ResolvedToolWireShape,
     ) -> Result<ProjectedHttpRequest, ProviderError> {
         let mut headers = HeaderMap::new();
+        // Written first so the reserved headers below overwrite anything the
+        // configuration tries to set for them.
+        insert_extra_headers(&mut headers, compat)?;
         let bearer = format!("Bearer {}", self.api_key);
         let auth = HeaderValue::from_str(&bearer)
             .map_err(|error| ProviderError::Connection(format!("Invalid authorization header: {error}")))?;
@@ -118,9 +137,12 @@ impl AnthropicTransport {
     pub(crate) fn build_projected_request(
         &self,
         body: Value,
+        compat: &ProviderCompat,
         tool_wire_shape: ResolvedToolWireShape,
     ) -> Result<ProjectedHttpRequest, ProviderError> {
         let mut headers = HeaderMap::new();
+        // See the OpenAI transport: extras first, protocol headers after.
+        insert_extra_headers(&mut headers, compat)?;
         let api_key = HeaderValue::from_str(&self.api_key)
             .map_err(|error| ProviderError::Connection(format!("Invalid x-api-key header: {error}")))?;
         headers.insert("x-api-key", api_key);
@@ -232,7 +254,7 @@ impl ProviderTransport {
     ) -> Result<ProjectedHttpRequest, ProviderError> {
         match self {
             Self::OpenAi(transport) => transport.build_projected_request(body, compat, tool_wire_shape),
-            Self::Anthropic(transport) => transport.build_projected_request(body, tool_wire_shape),
+            Self::Anthropic(transport) => transport.build_projected_request(body, compat, tool_wire_shape),
             Self::Vertex(transport) => transport
                 .inner
                 .build_projected_request(model, body, compat, tool_wire_shape),

crates/aion-providers/src/transport_test.rs

@@ -71,6 +71,85 @@ mod tests {
         assert_eq!(transport.decoder(&compat), StreamDecoder::OpenAiResponsesSse);
     }
 
+    fn compat_with_headers(pairs: &[(&str, &str)]) -> ProviderCompat {
+        let mut compat = ProviderCompat::openai_defaults();
+        compat.transport.extra_headers = pairs
+            .iter()
+            .map(|(name, value)| ((*name).to_owned(), (*value).to_owned()))
+            .collect();
+        compat
+    }
+
+    #[test]
+    fn openai_transport_sends_configured_extra_headers() {
+        let transport = OpenAiTransport::new("test-key", "https://api.example.test/v1");
+        let compat = compat_with_headers(&[
+            ("X-Partner-Id", "part_abc123"),
+            ("HTTP-Referer", "https://example.test"),
+        ]);
+
+        let request = transport
+            .build_projected_request(json!({ "model": "m" }), &compat, ResolvedToolWireShape::OpenAiFunction)
+            .expect("request projection should succeed");
+
+        assert_eq!(request.headers.get("x-partner-id").unwrap(), "part_abc123");
+        assert_eq!(request.headers.get("http-referer").unwrap(), "https://example.test");
+        // The transport still authenticates.
+        assert_eq!(request.headers.get(AUTHORIZATION).unwrap(), "Bearer test-key");
+    }
+
+    #[test]
+    fn openai_transport_extra_headers_cannot_replace_authorization_or_content_type() {
+        let transport = OpenAiTransport::new("real-key", "https://api.example.test/v1");
+        let compat = compat_with_headers(&[("authorization", "Bearer stolen"), ("content-type", "text/plain")]);
+
+        let request = transport
+            .build_projected_request(json!({ "model": "m" }), &compat, ResolvedToolWireShape::OpenAiFunction)
+            .expect("request projection should succeed");
+
+        // Reserved headers are written after the map, so configuration cannot
+        // unauthenticate a request or change the wire format.
+        assert_eq!(request.headers.get(AUTHORIZATION).unwrap(), "Bearer real-key");
+        assert_eq!(request.headers.get(CONTENT_TYPE).unwrap(), "application/json");
+    }
+
+    #[test]
+    fn anthropic_transport_sends_extra_headers_but_keeps_its_protocol_headers() {
+        let transport = AnthropicTransport::new("test-key", "https://api.example.test", false);
+        let compat = compat_with_headers(&[
+            ("X-Partner-Id", "part_abc123"),
+            ("x-api-key", "stolen"),
+            ("anthropic-version", "1999-01-01"),
+        ]);
+
+        let request = transport
+            .build_projected_request(
+                json!({ "model": "m" }),
+                &compat,
+                ResolvedToolWireShape::AnthropicInputSchema,
+            )
+            .expect("request projection should succeed");
+
+        assert_eq!(request.headers.get("x-partner-id").unwrap(), "part_abc123");
+        assert_eq!(request.headers.get("x-api-key").unwrap(), "test-key");
+        assert_eq!(request.headers.get("anthropic-version").unwrap(), "2023-06-01");
+    }
+
+    #[test]
+    fn extra_header_with_an_invalid_name_is_reported_rather_than_dropped() {
+        let transport = OpenAiTransport::new("test-key", "https://api.example.test/v1");
+        let compat = compat_with_headers(&[("bad header", "value")]);
+
+        let error = transport
+            .build_projected_request(json!({ "model": "m" }), &compat, ResolvedToolWireShape::OpenAiFunction)
+            .expect_err("an unusable header name should not be silently skipped");
+
+        assert!(
+            format!("{error}").contains("bad header"),
+            "error should name the header: {error}"
+        );
+    }
+
     #[test]
     fn openai_transport_appends_chat_completions_to_configured_base_url() {
         let transport = OpenAiTransport::new("test-key", "https://open.bigmodel.cn/api/paas/v4/");