From ae676d8cfbbae663901756b2ce0c5a9e80befece Mon Sep 17 00:00:00 2001 From: aimlapi Date: Wed, 9 Sep 2026 21:02:16 +0500 Subject: [PATCH] feat(providers): let a provider send configured extra headers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `OpenAiTransport::build_projected_request` writes exactly two headers — authorization and content type — and there is no way to add a third. Several gateways ask callers to identify themselves on the wire (an app name, a referer, a partner id) and cannot attribute traffic without it, so anything routed through this crate arrives anonymous no matter what the caller configured further up. Adds `extra_headers` to `TransportCompat`, applied by the OpenAI and Anthropic transports. Nothing here names a provider: which headers to send is a per-deployment question, so it belongs in configuration rather than in a match arm on a host name. Three details that are deliberate: - **Reserved headers win.** The map is written first and the protocol headers after it, so a config typo cannot unauthenticate a request, change the wire format, or downgrade `anthropic-version`. No blocklist to keep in sync — the ordering is the rule, and tests pin it for both transports. - **Merging is per key.** `TransportCompat::merge` extends rather than replaces, so setting one header in user config does not drop the preset's others the way `.or()` semantics would. - **Bedrock and Vertex are not covered.** Bedrock signs its headers with SigV4; a header added outside the signing step invalidates the signature instead of being attributed. Better to leave those transports alone than to ship something that fails confusingly. Backward compatible on disk: the field is `#[serde(default)]` with `skip_serializing_if`, so an empty map never appears in a rendered config and existing files parse unchanged. A test pins that too, since it is the part a release would break quietly. Verified: aion-config and aion-providers suites pass (240 in the providers lib alone). Reverting the two `insert_extra_headers` calls while keeping the tests fails exactly the three that assert the new behaviour, and leaves the reserved-headers-win test passing — which is what it should do, since that one holds either way. --- crates/aion-config/src/compat.rs | 33 +++++++++ crates/aion-config/src/compat_test.rs | 42 +++++++++++ crates/aion-providers/src/transport.rs | 26 ++++++- crates/aion-providers/src/transport_test.rs | 79 +++++++++++++++++++++ 4 files changed, 178 insertions(+), 2 deletions(-) diff --git a/crates/aion-config/src/compat.rs b/crates/aion-config/src/compat.rs index 851299a6..2ce49975 100644 --- a/crates/aion-config/src/compat.rs +++ b/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, + /// 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, + /// Maximum serialized provider request body size in bytes. /// Default: None (no local preflight limit). pub max_request_body_bytes: Option, @@ -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 { + &self.transport.extra_headers + } + pub fn openai_api_path(&self) -> &str { match self.openai_api_mode() { OpenAiApiMode::ChatCompletions => self.api_path(), diff --git a/crates/aion-config/src/compat_test.rs b/crates/aion-config/src/compat_test.rs index 4f76d3b1..edd47d11 100644 --- a/crates/aion-config/src/compat_test.rs +++ b/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), diff --git a/crates/aion-providers/src/transport.rs b/crates/aion-providers/src/transport.rs index 9ad87f90..7dd19ee3 100644 --- a/crates/aion-providers/src/transport.rs +++ b/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 { 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 { 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 { 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), diff --git a/crates/aion-providers/src/transport_test.rs b/crates/aion-providers/src/transport_test.rs index 8887308c..e3f03510 100644 --- a/crates/aion-providers/src/transport_test.rs +++ b/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/");