From 7068b0a1db9d3589768e4bc71d6a553d2b91c30b Mon Sep 17 00:00:00 2001 From: Alexander Petric Date: Wed, 29 Oct 2025 17:03:36 -0400 Subject: [PATCH] feat: Add AI_HTTP_HEADERS environment variable for custom AI request headers (#6994) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit adds support for setting custom HTTP headers for all AI API requests via the AI_HTTP_HEADERS environment variable. Usage: AI_HTTP_HEADERS="customheader1: hello, customheader2: world" The environment variable accepts a comma-separated list of header:value pairs. These headers will be applied to all AI requests made through both the worker (AI agent jobs) and the API (AI proxy requests). Changes: - backend/windmill-worker/src/ai_executor.rs: Parse and apply custom headers - backend/windmill-api/src/ai.rs: Parse and apply custom headers 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: windmill-internal-app[bot] --- backend/windmill-api/src/ai.rs | 44 ++++++++++++++++++++-- backend/windmill-worker/src/ai_executor.rs | 32 ++++++++++++++++ 2 files changed, 73 insertions(+), 3 deletions(-) diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 64f075be01..e4d07cfdc6 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -22,6 +22,33 @@ lazy_static::lazy_static! { static ref OPENAI_AZURE_BASE_PATH: Option = std::env::var("OPENAI_AZURE_BASE_PATH").ok(); pub static ref AI_REQUEST_CACHE: Cache<(String, AIProvider), ExpiringAIRequestConfig> = Cache::new(500); + + /// Parse AI_HTTP_HEADERS environment variable into a vector of (header_name, header_value) tuples + /// Format: "header1: value1, header2: value2" + static ref AI_HTTP_HEADERS: Vec<(String, String)> = { + std::env::var("AI_HTTP_HEADERS") + .ok() + .map(|headers_str| { + headers_str + .split(',') + .filter_map(|header| { + let parts: Vec<&str> = header.splitn(2, ':').collect(); + if parts.len() == 2 { + let name = parts[0].trim().to_string(); + let value = parts[1].trim().to_string(); + if !name.is_empty() && !value.is_empty() { + Some((name, value)) + } else { + None + } + } else { + None + } + }) + .collect() + }) + .unwrap_or_default() + }; } #[derive(Deserialize, Debug)] @@ -201,6 +228,11 @@ impl AIRequestConfig { request = request.header("OpenAI-Organization", org_id); } + // Apply custom headers from AI_HTTP_HEADERS environment variable + for (header_name, header_value) in AI_HTTP_HEADERS.iter() { + request = request.header(header_name.as_str(), header_value.as_str()); + } + Ok(request) } @@ -288,11 +320,17 @@ async fn global_proxy( let url = format!("{}/{}", base_url, ai_path); - let request = HTTP_CLIENT + let mut request = HTTP_CLIENT .request(method, url) .header("content-type", "application/json") - .header("Authorization", format!("Bearer {}", api_key)) - .body(body); + .header("Authorization", format!("Bearer {}", api_key)); + + // Apply custom headers from AI_HTTP_HEADERS environment variable + for (header_name, header_value) in AI_HTTP_HEADERS.iter() { + request = request.header(header_name.as_str(), header_value.as_str()); + } + + let request = request.body(body); let response = request.send().await.map_err(to_anyhow)?; diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index a7c3bf12a7..f9207e23a5 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -46,6 +46,33 @@ use crate::{ lazy_static::lazy_static! { static ref TOOL_NAME_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_]+$").unwrap(); + + /// Parse AI_HTTP_HEADERS environment variable into a vector of (header_name, header_value) tuples + /// Format: "header1: value1, header2: value2" + static ref AI_HTTP_HEADERS: Vec<(String, String)> = { + std::env::var("AI_HTTP_HEADERS") + .ok() + .map(|headers_str| { + headers_str + .split(',') + .filter_map(|header| { + let parts: Vec<&str> = header.splitn(2, ':').collect(); + if parts.len() == 2 { + let name = parts[0].trim().to_string(); + let value = parts[1].trim().to_string(); + if !name.is_empty() && !value.is_empty() { + Some((name, value)) + } else { + None + } + } else { + None + } + }) + .collect() + }) + .unwrap_or_default() + }; } const MAX_AGENT_ITERATIONS: usize = 10; @@ -560,6 +587,11 @@ pub async fn run_agent( request = request.header(*header_name, header_value.clone()); } + // Apply custom headers from AI_HTTP_HEADERS environment variable + for (header_name, header_value) in AI_HTTP_HEADERS.iter() { + request = request.header(header_name.as_str(), header_value.as_str()); + } + if args.provider.kind.is_azure_openai(&base_url) { request = request.query(&[("api-version", AZURE_API_VERSION)]) }