Skip to content

feat(drivers): support custom HTTP headers for OpenAI-compatible providers #1294

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lifaon74 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(drivers): support custom HTTP headers for OpenAI-compatible providers #1294

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions crates/openfang-api/src/routes.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7780,6 +7780,9 @@ pub async fn set_provider_key(
api_key_env: env_var.clone(),
base_url: None,
subprocess_timeout_secs: None,
// Headers are provider-specific — the old provider's
// headers must not leak onto the new one.
headers: None,
};
let mut guard = state
.kernel
Expand Down Expand Up @@ -7938,6 +7941,21 @@ pub async fn test_provider(

// Attempt a lightweight connectivity test
let start = std::time::Instant::now();
// Attach provider-configured extra headers ([default_model].headers or
// [[fallback_providers]].headers) so the connectivity test exercises the
// same request shape as real traffic (e.g. x-opencode-session for
// OpenCode Go).
let headers = if state.kernel.config.default_model.provider == name {
state.kernel.config.default_model.headers.clone()
} else {
state
.kernel
.config
.fallback_providers
.iter()
.find(|fb| fb.provider == name)
.and_then(|fb| fb.headers.clone())
};
let driver_config = openfang_runtime::llm_driver::DriverConfig {
provider: name.clone(),
api_key,
Expand All @@ -7948,6 +7966,7 @@ pub async fn test_provider(
},
skip_permissions: true,
subprocess_timeout_secs: None,
headers,
};

match openfang_runtime::drivers::create_driver(&driver_config) {
Expand Down
24 changes: 24 additions & 0 deletions crates/openfang-kernel/src/kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -696,6 +696,7 @@ impl OpenFangKernel {
}),
skip_permissions: true,
subprocess_timeout_secs: config.default_model.subprocess_timeout_secs,
headers: config.default_model.headers.clone(),
};
// Primary driver failure is non-fatal: the dashboard should remain accessible
// even if the LLM provider is misconfigured. Users can fix config via dashboard.
Expand All @@ -722,6 +723,10 @@ impl OpenFangKernel {
// Inherit operator's default-model timeout intent: auto-detect
// is replacing the *provider*, not the timeout policy.
subprocess_timeout_secs: config.default_model.subprocess_timeout_secs,
// Headers are provider-specific (e.g. x-opencode-session
// is bound to the OpenCode Go endpoint) — do NOT carry
// default_model.headers onto a different provider.
headers: None,
};
match drivers::create_driver(&auto_config) {
Ok(d) => {
Expand Down Expand Up @@ -771,6 +776,7 @@ impl OpenFangKernel {
.or_else(|| config.provider_urls.get(&fb.provider).cloned()),
skip_permissions: true,
subprocess_timeout_secs: fb.subprocess_timeout_secs,
headers: fb.headers.clone(),
};
match drivers::create_driver(&fb_config) {
Ok(d) => {
Expand Down Expand Up @@ -5647,12 +5653,24 @@ impl OpenFangKernel {
.and_then(|fb| fb.subprocess_timeout_secs)
};

// Same per-provider resolution for extra HTTP headers
// ([default_model].headers / [[fallback_providers]].headers).
let primary_headers = if agent_provider == default_provider {
effective_default.headers.clone()
} else {
effective_fallbacks
.iter()
.find(|fb| &fb.provider == agent_provider)
.and_then(|fb| fb.headers.clone())
};

let driver_config = DriverConfig {
provider: agent_provider.clone(),
api_key,
base_url,
skip_permissions: true,
subprocess_timeout_secs: primary_timeout,
headers: primary_headers,
};

match drivers::create_driver(&driver_config) {
Expand Down Expand Up @@ -5738,6 +5756,11 @@ impl OpenFangKernel {
} else {
None
},
headers: if resolved_to_default {
dm.headers.clone()
} else {
None
},
};
match drivers::create_driver(&config) {
Ok(d) => chain.push((d, strip_provider_prefix(&fb_model_name, &fb_provider))),
Expand Down Expand Up @@ -5773,6 +5796,7 @@ impl OpenFangKernel {
.or_else(|| self.lookup_provider_url(&fb.provider)),
skip_permissions: true,
subprocess_timeout_secs: fb.subprocess_timeout_secs,
headers: fb.headers.clone(),
};
match drivers::create_driver(&fb_config) {
Ok(d) => {
Expand Down
8 changes: 8 additions & 0 deletions crates/openfang-runtime/src/agent_loop.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1267,6 +1267,10 @@ async fn call_with_retry(
base_url: fb.base_url.clone(),
skip_permissions: true,
subprocess_timeout_secs: None,
// Manifest fallbacks carry no per-provider header
// config; global [[fallback_providers]] headers
// are handled by the kernel-built driver chain.
headers: None,
};
let fb_driver = match crate::drivers::create_driver(&fb_config) {
Ok(d) => d,
Expand Down Expand Up @@ -1451,6 +1455,10 @@ async fn stream_with_retry(
base_url: fb.base_url.clone(),
skip_permissions: true,
subprocess_timeout_secs: None,
// Manifest fallbacks carry no per-provider header
// config; global [[fallback_providers]] headers
// are handled by the kernel-built driver chain.
headers: None,
};
let fb_driver = match crate::drivers::create_driver(&fb_config) {
Ok(d) => d,
Expand Down
87 changes: 81 additions & 6 deletions crates/openfang-runtime/src/drivers/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,27 @@ fn provider_defaults(provider: &str) -> Option<ProviderDefaults> {
}
}

/// Attach user-configured extra HTTP headers (`config.headers`) to an
/// OpenAI-compatible driver. No-op when none are set.
///
/// Headers come from `[default_model].headers` / `[[fallback_providers]].headers`
/// in `config.toml` (e.g. the `x-opencode-session` header required by the
/// OpenCode Go endpoint).
fn apply_extra_headers(
driver: openai::OpenAIDriver,
config: &DriverConfig,
) -> openai::OpenAIDriver {
match &config.headers {
Some(headers) if !headers.is_empty() => driver.with_extra_headers(
headers
.iter()
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
),
_ => driver,
}
}

/// Create an LLM driver based on provider name and configuration.
///
/// Supported providers:
Expand Down Expand Up @@ -382,7 +403,10 @@ pub fn create_driver(config: &DriverConfig) -> Result<Arc<dyn LlmDriver>, LlmErr
.base_url
.clone()
.unwrap_or_else(|| OPENAI_BASE_URL.to_string());
return Ok(Arc::new(openai::OpenAIDriver::new(api_key, base_url)));
return Ok(Arc::new(apply_extra_headers(
openai::OpenAIDriver::new(api_key, base_url),
config,
)));
}

// Claude Code CLI — subprocess-based, no API key needed
Expand Down Expand Up @@ -457,7 +481,10 @@ pub fn create_driver(config: &DriverConfig) -> Result<Arc<dyn LlmDriver>, LlmErr
https://{resource}.openai.azure.com/openai/deployments"
.to_string(),
})?;
return Ok(Arc::new(openai::OpenAIDriver::new_azure(api_key, base_url)));
return Ok(Arc::new(apply_extra_headers(
openai::OpenAIDriver::new_azure(api_key, base_url),
config,
)));
}

// Vertex AI — uses Google Cloud OAuth with service account credentials.
Expand Down Expand Up @@ -546,7 +573,10 @@ pub fn create_driver(config: &DriverConfig) -> Result<Arc<dyn LlmDriver>, LlmErr
.or_else(|| local_provider_url_from_env(provider))
.unwrap_or_else(|| defaults.base_url.to_string());

return Ok(Arc::new(openai::OpenAIDriver::new(api_key, base_url)));
return Ok(Arc::new(apply_extra_headers(
openai::OpenAIDriver::new(api_key, base_url),
config,
)));
}

// Unknown provider — if base_url is set, treat as custom OpenAI-compatible.
Expand All @@ -558,9 +588,9 @@ pub fn create_driver(config: &DriverConfig) -> Result<Arc<dyn LlmDriver>, LlmErr
let env_var = format!("{}_API_KEY", provider.to_uppercase().replace('-', "_"));
std::env::var(&env_var).unwrap_or_default()
});
return Ok(Arc::new(openai::OpenAIDriver::new(
api_key,
base_url.clone(),
return Ok(Arc::new(apply_extra_headers(
openai::OpenAIDriver::new(api_key, base_url.clone()),
config,
)));
}

Expand Down Expand Up @@ -779,11 +809,38 @@ mod tests {
base_url: Some("http://localhost:9999/v1".to_string()),
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(driver.is_ok());
}

#[test]
fn test_custom_provider_with_extra_headers() {
// Configured extra headers (e.g. `x-opencode-session` for the OpenCode
// Go endpoint) must flow into the OpenAI-compatible driver. We can't
// introspect the driver's headers from here, but the construction path
// through apply_extra_headers must succeed; the field itself is
// covered by openai.rs's test_with_extra_headers_stores_headers.
let config = DriverConfig {
provider: "my-custom-llm".to_string(),
api_key: Some("test".to_string()),
base_url: Some("http://localhost:9999/v1".to_string()),
skip_permissions: true,
subprocess_timeout_secs: None,
headers: Some(
[("x-opencode-session".to_string(), "oc-123".to_string())]
.into_iter()
.collect(),
),
};
let driver = create_driver(&config);
assert!(
driver.is_ok(),
"custom provider with headers should construct"
);
}

#[test]
fn test_unknown_provider_no_url_errors() {
let config = DriverConfig {
Expand All @@ -792,6 +849,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(driver.is_err());
Expand Down Expand Up @@ -913,6 +971,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(
Expand All @@ -931,6 +990,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(driver.is_err());
Expand All @@ -948,6 +1008,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(
Expand All @@ -967,6 +1028,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(driver.is_err());
Expand All @@ -984,6 +1046,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let result = create_driver(&config);
assert!(result.is_err());
Expand Down Expand Up @@ -1012,6 +1075,7 @@ mod tests {
base_url: Some("https://api.example.com/v1".to_string()),
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(driver.is_ok());
Expand Down Expand Up @@ -1040,6 +1104,7 @@ mod tests {
base_url: Some("https://myresource.openai.azure.com/openai/deployments".to_string()),
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(driver.is_ok(), "Azure driver with key + URL should succeed");
Expand All @@ -1053,6 +1118,7 @@ mod tests {
base_url: Some("https://myresource.openai.azure.com/openai/deployments".to_string()),
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let result = create_driver(&config);
assert!(result.is_err(), "Azure driver without key should error");
Expand All @@ -1072,6 +1138,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let result = create_driver(&config);
assert!(result.is_err(), "Azure driver without URL should error");
Expand All @@ -1091,6 +1158,7 @@ mod tests {
base_url: Some("https://myresource.openai.azure.com/openai/deployments".to_string()),
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(
Expand All @@ -1116,6 +1184,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
// Should succeed because api_key is provided
let driver = create_driver(&config);
Expand All @@ -1135,6 +1204,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(driver.is_ok(), "claude-code driver should construct");
Expand All @@ -1150,6 +1220,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: Some(480),
headers: None,
};
let driver = create_driver(&config);
assert!(
Expand All @@ -1170,6 +1241,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: Some(120),
headers: None,
};
let driver = create_driver(&config);
std::env::remove_var("OPENFANG_SUBPROCESS_TIMEOUT_SECS");
Expand All @@ -1189,6 +1261,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: Some(420),
headers: None,
};
let driver = create_driver(&config);
std::env::remove_var("OPENFANG_SUBPROCESS_TIMEOUT_SECS");
Expand Down Expand Up @@ -1279,6 +1352,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(
Expand All @@ -1303,6 +1377,7 @@ mod tests {
base_url: None,
skip_permissions: true,
subprocess_timeout_secs: None,
headers: None,
};
let driver = create_driver(&config);
assert!(driver.is_ok(), "lmstudio default should construct");
Expand Down
Loading