Skip to content

Commit fe68905

Browse files
david-strejcclaudeautofix-ci[bot]tusharmath
authored
feat: add custom_headers for providers + built-in Kimi Coding provider (#2576)
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Tushar Mathur <tusharmath@gmail.com>
1 parent 29db91a commit fe68905

13 files changed

Lines changed: 131 additions & 1 deletion

File tree

‎crates/forge_app/src/command_generator.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,7 @@ mod tests {
211211
auth_details: AuthDetails::ApiKey("test-key".to_string().into()),
212212
url_params: Default::default(),
213213
}),
214+
custom_headers: None,
214215
})
215216
}
216217

‎crates/forge_app/src/dto/openai/transformers/pipeline.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@ mod tests {
126126
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
127127
url_params: vec![],
128128
credential: make_credential(ProviderId::FORGE, key),
129+
custom_headers: None,
129130
models: Some(ModelSource::Url(
130131
Url::parse("https://antinomy.ai/api/v1/models").unwrap(),
131132
)),
@@ -141,6 +142,7 @@ mod tests {
141142
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
142143
url_params: vec![],
143144
credential: make_credential(ProviderId::ZAI, key),
145+
custom_headers: None,
144146
models: Some(ModelSource::Url(
145147
Url::parse("https://api.z.ai/api/paas/v4/models").unwrap(),
146148
)),
@@ -156,6 +158,7 @@ mod tests {
156158
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
157159
url_params: vec![],
158160
credential: make_credential(ProviderId::ZAI_CODING, key),
161+
custom_headers: None,
159162
models: Some(ModelSource::Url(
160163
Url::parse("https://api.z.ai/api/paas/v4/models").unwrap(),
161164
)),
@@ -171,6 +174,7 @@ mod tests {
171174
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
172175
url_params: vec![],
173176
credential: make_credential(ProviderId::OPENAI, key),
177+
custom_headers: None,
174178
models: Some(ModelSource::Url(
175179
Url::parse("https://api.openai.com/v1/models").unwrap(),
176180
)),
@@ -186,6 +190,7 @@ mod tests {
186190
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
187191
url_params: vec![],
188192
credential: make_credential(ProviderId::XAI, key),
193+
custom_headers: None,
189194
models: Some(ModelSource::Url(
190195
Url::parse("https://api.x.ai/v1/models").unwrap(),
191196
)),
@@ -201,6 +206,7 @@ mod tests {
201206
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
202207
url_params: vec![],
203208
credential: make_credential(ProviderId::REQUESTY, key),
209+
custom_headers: None,
204210
models: Some(ModelSource::Url(
205211
Url::parse("https://api.requesty.ai/v1/models").unwrap(),
206212
)),
@@ -216,6 +222,7 @@ mod tests {
216222
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
217223
url_params: vec![],
218224
credential: make_credential(ProviderId::OPEN_ROUTER, key),
225+
custom_headers: None,
219226
models: Some(ModelSource::Url(
220227
Url::parse("https://openrouter.ai/api/v1/models").unwrap(),
221228
)),
@@ -231,6 +238,7 @@ mod tests {
231238
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
232239
url_params: vec![],
233240
credential: make_credential(ProviderId::ANTHROPIC, key),
241+
custom_headers: None,
234242
models: Some(ModelSource::Url(
235243
Url::parse("https://api.anthropic.com/v1/models").unwrap(),
236244
)),

‎crates/forge_domain/src/provider.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,9 @@ pub struct Provider<T> {
206206
#[serde(default)]
207207
pub url_params: Vec<crate::URLParam>,
208208
pub credential: Option<AuthCredential>,
209+
/// Custom HTTP headers to include in API requests for this provider.
210+
#[serde(default, skip_serializing_if = "Option::is_none")]
211+
pub custom_headers: Option<std::collections::HashMap<String, String>>,
209212
}
210213

211214
/// Type alias for a provider with template URLs (not yet rendered)
@@ -346,6 +349,7 @@ mod test_helpers {
346349
auth_methods: vec![crate::AuthMethod::ApiKey],
347350
url_params: vec![],
348351
credential: make_credential(ProviderId::ZAI, key),
352+
custom_headers: None,
349353
models: Some(ModelSource::Url(
350354
Url::parse("https://api.z.ai/api/paas/v4/models").unwrap(),
351355
)),
@@ -362,6 +366,7 @@ mod test_helpers {
362366
auth_methods: vec![crate::AuthMethod::ApiKey],
363367
url_params: vec![],
364368
credential: make_credential(ProviderId::ZAI_CODING, key),
369+
custom_headers: None,
365370
models: Some(ModelSource::Url(
366371
Url::parse("https://api.z.ai/api/paas/v4/models").unwrap(),
367372
)),
@@ -378,6 +383,7 @@ mod test_helpers {
378383
auth_methods: vec![crate::AuthMethod::ApiKey],
379384
url_params: vec![],
380385
credential: make_credential(ProviderId::OPENAI, key),
386+
custom_headers: None,
381387
models: Some(ModelSource::Url(
382388
Url::parse("https://api.openai.com/v1/models").unwrap(),
383389
)),
@@ -394,6 +400,7 @@ mod test_helpers {
394400
auth_methods: vec![crate::AuthMethod::ApiKey],
395401
url_params: vec![],
396402
credential: make_credential(ProviderId::XAI, key),
403+
custom_headers: None,
397404
models: Some(ModelSource::Url(
398405
Url::parse("https://api.x.ai/v1/models").unwrap(),
399406
)),
@@ -436,6 +443,7 @@ mod test_helpers {
436443
.map(|&s| s.to_string().into())
437444
.collect(),
438445
credential: make_credential(ProviderId::VERTEX_AI, key),
446+
custom_headers: None,
439447
models: Some(ModelSource::Url(Url::parse(&model_url).unwrap())),
440448
}
441449
}
@@ -451,6 +459,7 @@ mod test_helpers {
451459
auth_methods: vec![crate::AuthMethod::ApiKey],
452460
url_params: vec![],
453461
credential: make_credential(ProviderId::IO_INTELLIGENCE, key),
462+
custom_headers: None,
454463
models: Some(ModelSource::Url(
455464
Url::parse("https://api.intelligence.io.solutions/api/v1/models").unwrap(),
456465
)),
@@ -484,6 +493,7 @@ mod test_helpers {
484493
.map(|&s| s.to_string().into())
485494
.collect(),
486495
credential: make_credential(ProviderId::AZURE, key),
496+
custom_headers: None,
487497
models: Some(ModelSource::Url(Url::parse(&model_url).unwrap())),
488498
}
489499
}
@@ -558,6 +568,7 @@ mod tests {
558568
models: Some(ModelSource::Url(
559569
Url::from_str("https://api.intelligence.io.solutions/api/v1/models").unwrap(),
560570
)),
571+
custom_headers: None,
561572
};
562573
assert_eq!(actual, expected);
563574
}
@@ -581,6 +592,7 @@ mod tests {
581592
models: Some(ModelSource::Url(
582593
Url::from_str("https://api.x.ai/v1/models").unwrap(),
583594
)),
595+
custom_headers: None,
584596
};
585597
assert_eq!(actual, expected);
586598
}

‎crates/forge_infra/src/http.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,10 @@ impl<F: forge_app::FileWriterInfra + 'static> ForgeHttpInfra<F> {
167167
// - `X-Title`: Sets/modifies your app's title
168168
fn headers(&self, headers: Option<HeaderMap>) -> HeaderMap {
169169
let mut headers = headers.unwrap_or_default();
170-
headers.insert("User-Agent", HeaderValue::from_static("Forge"));
170+
// Only set User-Agent if the provider hasn't already set one
171+
if !headers.contains_key("User-Agent") {
172+
headers.insert("User-Agent", HeaderValue::from_static("Forge"));
173+
}
171174
headers.insert("X-Title", HeaderValue::from_static("forge"));
172175
headers.insert(
173176
"x-app-version",

‎crates/forge_main/src/model.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1028,6 +1028,7 @@ mod tests {
10281028
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
10291029
url_params: vec![],
10301030
credential: None,
1031+
custom_headers: None,
10311032
models: Some(ModelSource::Url(
10321033
Url::parse("https://api.openai.com/v1/models").unwrap(),
10331034
)),
@@ -1048,6 +1049,7 @@ mod tests {
10481049
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
10491050
url_params: vec![],
10501051
credential: None,
1052+
custom_headers: None,
10511053
models: Some(ModelSource::Url(
10521054
Url::parse("https://openrouter.ai/api/v1/models").unwrap(),
10531055
)),
@@ -1068,6 +1070,7 @@ mod tests {
10681070
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
10691071
url_params: vec![],
10701072
credential: None,
1073+
custom_headers: None,
10711074
models: Some(ModelSource::Url(
10721075
Url::parse("http://localhost:8080/models").unwrap(),
10731076
)),
@@ -1088,6 +1091,7 @@ mod tests {
10881091
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
10891092
url_params: vec![],
10901093
credential: None,
1094+
custom_headers: None,
10911095
models: Some(ModelSource::Url(Template::new(
10921096
"https://api.anthropic.com/v1/models",
10931097
))),
@@ -1108,6 +1112,7 @@ mod tests {
11081112
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
11091113
url_params: vec![],
11101114
credential: None,
1115+
custom_headers: None,
11111116
models: Some(ModelSource::Url(
11121117
Url::parse("http://192.168.1.1:8080/models").unwrap(),
11131118
)),

‎crates/forge_repo/src/provider/anthropic.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,7 @@ mod tests {
284284
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
285285
url_params: vec![],
286286
models: Some(forge_domain::ModelSource::Url(model_url)),
287+
custom_headers: None,
287288
};
288289

289290
Ok(Anthropic::new(
@@ -351,6 +352,7 @@ mod tests {
351352
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
352353
url_params: vec![],
353354
models: Some(forge_domain::ModelSource::Url(model_url.clone())),
355+
custom_headers: None,
354356
};
355357

356358
let anthropic = Anthropic::new(
@@ -490,6 +492,7 @@ mod tests {
490492
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
491493
url_params: vec![],
492494
models: Some(forge_domain::ModelSource::Url(model_url)),
495+
custom_headers: None,
493496
};
494497

495498
let fixture = Anthropic::new(
@@ -568,6 +571,7 @@ mod tests {
568571
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
569572
url_params: vec![],
570573
models: Some(forge_domain::ModelSource::Url(model_url)),
574+
custom_headers: None,
571575
};
572576

573577
let fixture = Anthropic::new(
@@ -644,6 +648,7 @@ mod tests {
644648
auth_methods: vec![forge_domain::AuthMethod::GoogleAdc],
645649
url_params: vec![],
646650
models: Some(forge_domain::ModelSource::Hardcoded(vec![])),
651+
custom_headers: None,
647652
};
648653

649654
let _anthropic = Anthropic::new(

‎crates/forge_repo/src/provider/bedrock.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -988,6 +988,7 @@ mod tests {
988988
auth_details: AuthDetails::ApiKey(ApiKey::from(token.to_string())),
989989
url_params,
990990
}),
991+
custom_headers: None,
991992
}
992993
}
993994

‎crates/forge_repo/src/provider/openai.rs‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,12 @@ impl<H: HttpInfra> OpenAIProvider<H> {
9494
}
9595
forge_domain::AuthMethod::GoogleAdc => {}
9696
});
97+
// Append provider-level custom headers (from provider.json config)
98+
if let Some(custom_headers) = &self.provider.custom_headers {
99+
for (k, v) in custom_headers {
100+
headers.push((k.clone(), v.clone()));
101+
}
102+
}
97103
headers
98104
}
99105

@@ -274,6 +280,7 @@ mod tests {
274280
response: Some(ProviderResponse::OpenAI),
275281
url: Url::parse("https://api.openai.com/v1/chat/completions").unwrap(),
276282
credential: make_credential(ProviderId::OPENAI, key),
283+
custom_headers: None,
277284
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
278285
url_params: vec![],
279286
models: Some(forge_domain::ModelSource::Url(
@@ -289,6 +296,7 @@ mod tests {
289296
response: Some(ProviderResponse::OpenAI),
290297
url: Url::parse("https://api.z.ai/api/paas/v4/chat/completions").unwrap(),
291298
credential: make_credential(ProviderId::ZAI, key),
299+
custom_headers: None,
292300
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
293301
url_params: vec![],
294302
models: Some(forge_domain::ModelSource::Url(
@@ -304,6 +312,7 @@ mod tests {
304312
response: Some(ProviderResponse::OpenAI),
305313
url: Url::parse("https://api.z.ai/api/coding/paas/v4/chat/completions").unwrap(),
306314
credential: make_credential(ProviderId::ZAI_CODING, key),
315+
custom_headers: None,
307316
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
308317
url_params: vec![],
309318
models: Some(forge_domain::ModelSource::Url(
@@ -319,6 +328,7 @@ mod tests {
319328
response: Some(ProviderResponse::Anthropic),
320329
url: Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
321330
credential: make_credential(ProviderId::ANTHROPIC, key),
331+
custom_headers: None,
322332
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
323333
url_params: vec![],
324334
models: Some(forge_domain::ModelSource::Url(
@@ -383,6 +393,7 @@ mod tests {
383393
response: Some(ProviderResponse::OpenAI),
384394
url: reqwest::Url::parse(base_url)?,
385395
credential: make_credential(ProviderId::OPENAI, "test-api-key"),
396+
custom_headers: None,
386397
auth_methods: vec![forge_domain::AuthMethod::ApiKey],
387398
url_params: vec![],
388399
models: Some(forge_domain::ModelSource::Url(
@@ -691,6 +702,51 @@ mod tests {
691702
let error_string = format!("{:#}", actual);
692703
insta::assert_snapshot!(error_string);
693704
}
705+
706+
#[test]
707+
fn test_get_headers_includes_custom_headers() {
708+
let mut provider = openai("test-key");
709+
let mut custom = std::collections::HashMap::new();
710+
custom.insert("User-Agent".to_string(), "KimiCLI/1.0.0".to_string());
711+
custom.insert("X-Custom".to_string(), "custom-value".to_string());
712+
provider.custom_headers = Some(custom);
713+
714+
let http_client = Arc::new(MockHttpClient::new());
715+
let openai_provider = OpenAIProvider::new(provider, http_client);
716+
let headers = openai_provider.get_headers();
717+
718+
assert!(
719+
headers
720+
.iter()
721+
.any(|(k, v)| k == "User-Agent" && v == "KimiCLI/1.0.0")
722+
);
723+
assert!(
724+
headers
725+
.iter()
726+
.any(|(k, v)| k == "X-Custom" && v == "custom-value")
727+
);
728+
assert!(
729+
headers
730+
.iter()
731+
.any(|(k, v)| k == "authorization" && v == "Bearer test-key")
732+
);
733+
}
734+
735+
#[test]
736+
fn test_get_headers_no_custom_headers() {
737+
let provider = openai("test-key");
738+
let http_client = Arc::new(MockHttpClient::new());
739+
let openai_provider = OpenAIProvider::new(provider, http_client);
740+
let headers = openai_provider.get_headers();
741+
742+
// Only authorization header should be present
743+
assert_eq!(headers.len(), 1);
744+
assert!(
745+
headers
746+
.iter()
747+
.any(|(k, v)| k == "authorization" && v == "Bearer test-key")
748+
);
749+
}
694750
}
695751

696752
/// Repository for OpenAI-compatible provider responses

0 commit comments

Comments
 (0)