language_models: Support custom HTTP headers for provider requests

crates/agent_ui/src/agent_configuration/add_llm_provider_modal.rs

@@ -279,6 +279,7 @@ fn save_provider_to_settings(
                         OpenAiCompatibleSettingsContent {
                             api_url,
                             available_models: models,
+                            custom_headers: None,
                         },
                     );
             });

crates/anthropic/src/anthropic.rs

@@ -6,7 +6,10 @@ use anyhow::{Context as _, Result};
 use chrono::{DateTime, Utc};
 use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};
 use http_client::http::{self, HeaderMap, HeaderValue};
-use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest, StatusCode};
+use http_client::{
+    AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt,
+    StatusCode,
+};
 use serde::{Deserialize, Serialize};
 use strum::EnumString;
 use thiserror::Error;
@@ -202,10 +205,18 @@ pub async fn stream_completion(
     api_key: &str,
     request: Request,
     beta_headers: Option<String>,
+    extra_headers: &CustomHeaders,
 ) -> Result<BoxStream<'static, Result<Event, AnthropicError>>, AnthropicError> {
-    stream_completion_with_rate_limit_info(client, api_url, api_key, request, beta_headers)
-        .await
-        .map(|output| output.0)
+    stream_completion_with_rate_limit_info(
+        client,
+        api_url,
+        api_key,
+        request,
+        beta_headers,
+        extra_headers,
+    )
+    .await
+    .map(|output| output.0)
 }
 
 /// A raw model entry returned by the Anthropic models listing endpoint.
@@ -233,6 +244,7 @@ pub async fn list_models(
     client: &dyn HttpClient,
     api_url: &str,
     api_key: &str,
+    extra_headers: &CustomHeaders,
 ) -> Result<Vec<Model>> {
     let uri = format!("{api_url}/v1/models?limit=1000");
 
@@ -242,6 +254,7 @@ pub async fn list_models(
         .header("Anthropic-Version", "2023-06-01")
         .header("X-Api-Key", api_key.trim())
         .header("Accept", "application/json")
+        .extra_headers(extra_headers)
         .body(AsyncBody::default())
         .context("failed to build Anthropic models list request")?;
 
@@ -282,9 +295,17 @@ pub async fn non_streaming_completion(
     api_key: &str,
     request: Request,
     beta_headers: Option<String>,
+    extra_headers: &CustomHeaders,
 ) -> Result<Response, AnthropicError> {
-    let (mut response, rate_limits) =
-        send_request(client, api_url, api_key, &request, beta_headers).await?;
+    let (mut response, rate_limits) = send_request(
+        client,
+        api_url,
+        api_key,
+        &request,
+        beta_headers,
+        extra_headers,
+    )
+    .await?;
 
     if response.status().is_success() {
         let mut body = String::new();
@@ -306,6 +327,7 @@ async fn send_request(
     api_key: &str,
     request: impl Serialize,
     beta_headers: Option<String>,
+    extra_headers: &CustomHeaders,
 ) -> Result<(http::Response<AsyncBody>, RateLimitInfo), AnthropicError> {
     let uri = format!("{api_url}/v1/messages");
 
@@ -323,6 +345,7 @@ async fn send_request(
     let serialized_request =
         serde_json::to_string(&request).map_err(AnthropicError::SerializeRequest)?;
     let request = request_builder
+        .extra_headers(extra_headers)
         .body(AsyncBody::from(serialized_request))
         .map_err(AnthropicError::BuildRequestBody)?;
 
@@ -462,6 +485,7 @@ pub async fn stream_completion_with_rate_limit_info(
     api_key: &str,
     request: Request,
     beta_headers: Option<String>,
+    extra_headers: &CustomHeaders,
 ) -> Result<
     (
         BoxStream<'static, Result<Event, AnthropicError>>,
@@ -474,8 +498,15 @@ pub async fn stream_completion_with_rate_limit_info(
         stream: true,
     };
 
-    let (response, rate_limits) =
-        send_request(client, api_url, api_key, &request, beta_headers).await?;
+    let (response, rate_limits) = send_request(
+        client,
+        api_url,
+        api_key,
+        &request,
+        beta_headers,
+        extra_headers,
+    )
+    .await?;
 
     if response.status().is_success() {
         let reader = BufReader::new(response.into_body());

crates/bedrock/Cargo.toml

@@ -20,6 +20,7 @@ anyhow.workspace = true
 aws-sdk-bedrockruntime = { workspace = true, features = ["behavior-version-latest"] }
 aws-smithy-types = {workspace = true}
 futures.workspace = true
+http_client.workspace = true
 schemars = { workspace = true, optional = true }
 serde.workspace = true
 serde_json.workspace = true

crates/bedrock/src/bedrock.rs

@@ -35,6 +35,7 @@ pub use crate::models::*;
 pub async fn stream_completion(
     client: bedrock::Client,
     request: Request,
+    extra_headers: http_client::CustomHeaders,
 ) -> Result<BoxStream<'static, Result<BedrockStreamingResponse, anyhow::Error>>, BedrockError> {
     let mut response = bedrock::Client::converse_stream(&client)
         .model_id(request.model.clone())
@@ -99,30 +100,45 @@ pub async fn stream_completion(
         );
     }
 
-    let output = response.send().await.map_err(|err| match err {
-        bedrock::error::SdkError::ServiceError(ctx) => {
-            use bedrock::operation::converse_stream::ConverseStreamError;
-            let err = ctx.into_err();
-            match &err {
-                ConverseStreamError::ValidationException(e) => {
-                    BedrockError::Validation(e.message().unwrap_or("validation error").to_string())
-                }
-                ConverseStreamError::ThrottlingException(_) => BedrockError::RateLimited,
-                ConverseStreamError::ServiceUnavailableException(_)
-                | ConverseStreamError::ModelNotReadyException(_) => {
-                    BedrockError::ServiceUnavailable
-                }
-                ConverseStreamError::AccessDeniedException(e) => {
-                    BedrockError::AccessDenied(e.message().unwrap_or("access denied").to_string())
+    let output = response
+        .customize()
+        .mutate_request(move |http_request| {
+            let headers = http_request.headers_mut();
+            for (name, value) in extra_headers.iter() {
+                headers.insert(
+                    name.as_str().to_owned(),
+                    value.to_str().unwrap_or("").to_owned(),
+                );
+            }
+        })
+        .send()
+        .await
+        .map_err(|err| match err {
+            bedrock::error::SdkError::ServiceError(ctx) => {
+                use bedrock::operation::converse_stream::ConverseStreamError;
+                let err = ctx.into_err();
+                match &err {
+                    ConverseStreamError::ValidationException(e) => BedrockError::Validation(
+                        e.message().unwrap_or("validation error").to_string(),
+                    ),
+                    ConverseStreamError::ThrottlingException(_) => BedrockError::RateLimited,
+                    ConverseStreamError::ServiceUnavailableException(_)
+                    | ConverseStreamError::ModelNotReadyException(_) => {
+                        BedrockError::ServiceUnavailable
+                    }
+                    ConverseStreamError::AccessDeniedException(e) => BedrockError::AccessDenied(
+                        e.message().unwrap_or("access denied").to_string(),
+                    ),
+                    ConverseStreamError::InternalServerException(e) => {
+                        BedrockError::InternalServer(
+                            e.message().unwrap_or("internal server error").to_string(),
+                        )
+                    }
+                    _ => BedrockError::Other(err.into()),
                 }
-                ConverseStreamError::InternalServerException(e) => BedrockError::InternalServer(
-                    e.message().unwrap_or("internal server error").to_string(),
-                ),
-                _ => BedrockError::Other(err.into()),
             }
-        }
-        other => BedrockError::Other(other.into()),
-    });
+            other => BedrockError::Other(other.into()),
+        });
 
     let stream = Box::pin(stream::unfold(
         output?.stream,

crates/deepseek/src/deepseek.rs

@@ -4,7 +4,9 @@ use futures::{
     io::BufReader,
     stream::{BoxStream, StreamExt},
 };
-use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest};
+use http_client::{
+    AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt,
+};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::convert::TryFrom;
@@ -297,15 +299,16 @@ pub async fn stream_completion(
     api_url: &str,
     api_key: &str,
     request: Request,
+    extra_headers: &CustomHeaders,
 ) -> Result<BoxStream<'static, Result<StreamResponse>>> {
     let uri = format!("{api_url}/chat/completions");
-    let request_builder = HttpRequest::builder()
+    let request = HttpRequest::builder()
         .method(Method::POST)
         .uri(uri)
         .header("Content-Type", "application/json")
-        .header("Authorization", format!("Bearer {}", api_key.trim()));
-
-    let request = request_builder.body(AsyncBody::from(serde_json::to_string(&request)?))?;
+        .header("Authorization", format!("Bearer {}", api_key.trim()))
+        .extra_headers(extra_headers)
+        .body(AsyncBody::from(serde_json::to_string(&request)?))?;
     let mut response = client.send(request).await?;
 
     if response.status().is_success() {

crates/edit_prediction_cli/src/anthropic_client.rs

@@ -63,6 +63,7 @@ impl PlainLlmClient {
             &self.api_key,
             request,
             None,
+            &http_client::CustomHeaders::default(),
         )
         .await
         .map_err(|e| anyhow::anyhow!("{:?}", e))?;
@@ -104,6 +105,7 @@ impl PlainLlmClient {
             &self.api_key,
             request,
             None,
+            &http_client::CustomHeaders::default(),
         )
         .await
         .map_err(|e| anyhow::anyhow!("{:?}", e))?;

crates/google_ai/src/google_ai.rs

@@ -2,7 +2,9 @@ use std::mem;
 
 use anyhow::{Result, anyhow, bail};
 use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};
-use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest};
+use http_client::{
+    AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt,
+};
 pub use language_model_core::ModelMode as GoogleModelMode;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 pub mod completion;
@@ -14,6 +16,7 @@ pub async fn stream_generate_content(
     api_url: &str,
     api_key: &str,
     mut request: GenerateContentRequest,
+    extra_headers: &CustomHeaders,
 ) -> Result<BoxStream<'static, Result<GenerateContentResponse>>> {
     let api_key = api_key.trim();
     validate_generate_content_request(&request)?;
@@ -24,12 +27,12 @@ pub async fn stream_generate_content(
     let uri =
         format!("{api_url}/v1beta/models/{model_id}:streamGenerateContent?alt=sse&key={api_key}",);
 
-    let request_builder = HttpRequest::builder()
+    let request = HttpRequest::builder()
         .method(Method::POST)
         .uri(uri)
-        .header("Content-Type", "application/json");
-
-    let request = request_builder.body(AsyncBody::from(serde_json::to_string(&request)?))?;
+        .header("Content-Type", "application/json")
+        .extra_headers(extra_headers)
+        .body(AsyncBody::from(serde_json::to_string(&request)?))?;
     let mut response = client.send(request).await?;
     if response.status().is_success() {
         let reader = BufReader::new(response.into_body());

crates/http_client/src/http_client.rs

@@ -7,8 +7,8 @@ pub mod github_download;
 pub use anyhow::{Result, anyhow};
 pub use async_body::{AsyncBody, Inner, Json};
 use derive_more::Deref;
-use http::HeaderValue;
 pub use http::{self, Method, Request, Response, StatusCode, Uri, request::Builder};
+use http::{HeaderName, HeaderValue};
 
 use futures::future::BoxFuture;
 use parking_lot::Mutex;
@@ -57,6 +57,58 @@ impl HttpRequestExt for http::request::Builder {
     }
 }
 
+/// A set of pre-validated user-supplied HTTP headers.
+///
+/// Construction (and the per-name validation that goes with it) happens once
+/// at settings load time. Cloning is `Arc`-cheap, so providers can hand a copy
+/// to each outgoing request without re-parsing or re-allocating.
+#[derive(Default, Clone, Debug)]
+pub struct CustomHeaders(Arc<[(HeaderName, HeaderValue)]>);
+
+impl CustomHeaders {
+    pub fn new(headers: Vec<(HeaderName, HeaderValue)>) -> Self {
+        Self(headers.into())
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    pub fn iter(&self) -> impl ExactSizeIterator<Item = (&HeaderName, &HeaderValue)> {
+        self.0.iter().map(|(n, v)| (n, v))
+    }
+}
+
+impl PartialEq for CustomHeaders {
+    fn eq(&self, other: &Self) -> bool {
+        self.0.len() == other.0.len()
+            && self
+                .0
+                .iter()
+                .zip(other.0.iter())
+                .all(|(a, b)| a.0 == b.0 && a.1 == b.1)
+    }
+}
+
+pub trait RequestBuilderExt {
+    /// Append every header in `headers` to the request being built.
+    fn extra_headers(self, headers: &CustomHeaders) -> Self;
+}
+
+impl RequestBuilderExt for http::request::Builder {
+    fn extra_headers(mut self, headers: &CustomHeaders) -> Self {
+        if headers.is_empty() {
+            return self;
+        }
+        if let Some(map) = self.headers_mut() {
+            for (name, value) in headers.iter() {
+                map.append(name.clone(), value.clone());
+            }
+        }
+        self
+    }
+}
+
 pub trait HttpClient: 'static + Send + Sync {
     fn user_agent(&self) -> Option<&HeaderValue>;